module Vectorise.Monad.Global (
readGEnv,
setGEnv,
updGEnv,
defGlobalVar,
globalScalars,
lookupTyCon,
lookupBoxedTyCon,
defTyCon,
lookupDataCon,
defDataCon,
lookupTyConPA,
defTyConPA,
defTyConPAs,
lookupTyConPR
) where
import Vectorise.Monad.Base
import Vectorise.Env
import TyCon
import DataCon
import NameEnv
import Var
import VarEnv
import VarSet
readGEnv :: (GlobalEnv -> a) -> VM a
readGEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
setGEnv :: GlobalEnv -> VM ()
setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
defGlobalVar :: Var -> Var -> VM ()
defGlobalVar v v' = updGEnv $ \env ->
env { global_vars = extendVarEnv (global_vars env) v v'
, global_exported_vars = upd (global_exported_vars env)
}
where
upd env | isExportedId v = extendVarEnv env v (v, v')
| otherwise = env
globalScalars :: VM VarSet
globalScalars
= readGEnv global_scalars
lookupTyCon :: TyCon -> VM (Maybe TyCon)
lookupTyCon tc
| isUnLiftedTyCon tc || isTupleTyCon tc
= return (Just tc)
| otherwise
= readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
lookupBoxedTyCon tc
= readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
(tyConName tc)
defTyCon :: TyCon -> TyCon -> VM ()
defTyCon tc tc' = updGEnv $ \env ->
env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
lookupDataCon :: DataCon -> VM (Maybe DataCon)
lookupDataCon dc
| isTupleTyCon (dataConTyCon dc)
= return (Just dc)
| otherwise
= readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
defDataCon :: DataCon -> DataCon -> VM ()
defDataCon dc dc' = updGEnv $ \env ->
env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
lookupTyConPA :: TyCon -> VM (Maybe Var)
lookupTyConPA tc
= readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
defTyConPA :: TyCon -> Var -> VM ()
defTyConPA tc pa = updGEnv $ \env ->
env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
defTyConPAs :: [(TyCon, Var)] -> VM ()
defTyConPAs ps = updGEnv $ \env ->
env { global_pa_funs = extendNameEnvList (global_pa_funs env)
[(tyConName tc, pa) | (tc, pa) <- ps] }
lookupTyConPR :: TyCon -> VM (Maybe Var)
lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)