{-# LANGUAGE DataKinds #-}

module Data.Record.Internal.Record.Resolution.Internal (
    getRecordInfo
  , putRecordInfo
  ) where

import Data.Map (Map)
import Data.Maybe (fromMaybe)
import Language.Haskell.TH.Syntax

import qualified Data.Map as Map

import Data.Record.Internal.Naming
import Data.Record.Internal.Record

import qualified Data.Record.Internal.TH.Name as N

{-------------------------------------------------------------------------------
  Internal state

  As keys we use the names of the internal constructor, because when we do name
  resolution, that is what we would normally use to query ghc. We use /global/
  names, which uniquely identify a name (qualified by package and module).
-------------------------------------------------------------------------------}

newtype TypeEnv = WrapTypeEnv {
      TypeEnv -> Map (Name 'DataName 'Global) (Record ())
unwrapTypeEnv :: Map (N.Name 'DataName 'N.Global) (Record ())
    }

getTypeEnv :: Quasi m => m TypeEnv
getTypeEnv :: m TypeEnv
getTypeEnv = TypeEnv -> Maybe TypeEnv -> TypeEnv
forall a. a -> Maybe a -> a
fromMaybe (Map (Name 'DataName 'Global) (Record ()) -> TypeEnv
WrapTypeEnv Map (Name 'DataName 'Global) (Record ())
forall k a. Map k a
Map.empty) (Maybe TypeEnv -> TypeEnv) -> m (Maybe TypeEnv) -> m TypeEnv
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> m (Maybe TypeEnv)
forall (m :: Type -> Type) a. (Quasi m, Typeable a) => m (Maybe a)
qGetQ

putTypeEnv :: Quasi m => TypeEnv -> m ()
putTypeEnv :: TypeEnv -> m ()
putTypeEnv = TypeEnv -> m ()
forall (m :: Type -> Type) a. (Quasi m, Typeable a) => a -> m ()
qPutQ

{-------------------------------------------------------------------------------
  Accessing the internal state
-------------------------------------------------------------------------------}

getRecordInfo ::
     Quasi m
  => N.Name 'DataName 'N.Global  -- ^ Name of the internal constructor
  -> m (Maybe (Record ()))
getRecordInfo :: Name 'DataName 'Global -> m (Maybe (Record ()))
getRecordInfo Name 'DataName 'Global
internalConstr =
    Name 'DataName 'Global
-> Map (Name 'DataName 'Global) (Record ()) -> Maybe (Record ())
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Name 'DataName 'Global
internalConstr (Map (Name 'DataName 'Global) (Record ()) -> Maybe (Record ()))
-> (TypeEnv -> Map (Name 'DataName 'Global) (Record ()))
-> TypeEnv
-> Maybe (Record ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeEnv -> Map (Name 'DataName 'Global) (Record ())
unwrapTypeEnv (TypeEnv -> Maybe (Record ()))
-> m TypeEnv -> m (Maybe (Record ()))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> m TypeEnv
forall (m :: Type -> Type). Quasi m => m TypeEnv
getTypeEnv

-- | Add 'RecordInfo' to the environment
--
-- NOTE: Must be called whilst processing the module in which the record is
-- defined.
putRecordInfo :: Quasi m => Record () -> m ()
putRecordInfo :: Record () -> m ()
putRecordInfo Record ()
info = do
    Map (Name 'DataName 'Global) (Record ())
env <- TypeEnv -> Map (Name 'DataName 'Global) (Record ())
unwrapTypeEnv (TypeEnv -> Map (Name 'DataName 'Global) (Record ()))
-> m TypeEnv -> m (Map (Name 'DataName 'Global) (Record ()))
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> m TypeEnv
forall (m :: Type -> Type). Quasi m => m TypeEnv
getTypeEnv

    -- In order to be able to resolve the record info later, we need to properly
    -- quantify the record name. We do this by requesting the /current/ TH
    -- location. This is justified by the precondition to the function.

    Loc
loc <- Q Loc -> m Loc
forall (m :: Type -> Type) a. Quasi m => Q a -> m a
runQ Q Loc
location
    let internalConstr :: N.Name 'DataName 'N.Global
        internalConstr :: Name 'DataName 'Global
internalConstr =
          OccName -> NameFlavour 'Global -> Name 'DataName 'Global
forall (flavour :: Flavour) (ns :: NameSpace).
OccName -> NameFlavour flavour -> Name ns flavour
N.Name
            (String -> OccName
OccName (String -> String
nameRecordInternalConstr (Record () -> String
forall a. Record a -> String
recordConstr Record ()
info)))
            (NameSpace -> PkgName -> ModName -> NameFlavour 'Global
N.NameGlobal
              NameSpace
DataName
              (String -> PkgName
mkPkgName (Loc -> String
loc_package Loc
loc))
              (String -> ModName
mkModName (Loc -> String
loc_module  Loc
loc))
            )

    TypeEnv -> m ()
forall (m :: Type -> Type). Quasi m => TypeEnv -> m ()
putTypeEnv (TypeEnv -> m ()) -> TypeEnv -> m ()
forall a b. (a -> b) -> a -> b
$ Map (Name 'DataName 'Global) (Record ()) -> TypeEnv
WrapTypeEnv (Map (Name 'DataName 'Global) (Record ()) -> TypeEnv)
-> Map (Name 'DataName 'Global) (Record ()) -> TypeEnv
forall a b. (a -> b) -> a -> b
$ Name 'DataName 'Global
-> Record ()
-> Map (Name 'DataName 'Global) (Record ())
-> Map (Name 'DataName 'Global) (Record ())
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Name 'DataName 'Global
internalConstr Record ()
info Map (Name 'DataName 'Global) (Record ())
env