2 module Vectorise.Monad.Global (
30 import Vectorise.Monad.Base
40 -- Global Environment ---------------------------------------------------------
41 -- | Project something from the global environment.
42 readGEnv :: (GlobalEnv -> a) -> VM a
43 readGEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
46 -- | Set the value of the global environment.
47 setGEnv :: GlobalEnv -> VM ()
48 setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ())
51 -- | Update the global environment using the provided function.
52 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
53 updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
56 -- Vars -----------------------------------------------------------------------
57 -- | Add a mapping between a global var and its vectorised version to the state.
58 defGlobalVar :: Var -> Var -> VM ()
59 defGlobalVar v v' = updGEnv $ \env ->
60 env { global_vars = extendVarEnv (global_vars env) v v'
61 , global_exported_vars = upd (global_exported_vars env)
64 upd env | isExportedId v = extendVarEnv env v (v, v')
68 -- Scalars --------------------------------------------------------------------
69 -- | Get the set of global scalar variables.
70 globalScalars :: VM VarSet
72 = readGEnv global_scalars
75 -- TyCons ---------------------------------------------------------------------
76 -- | Lookup the vectorised version of a `TyCon` from the global environment.
77 lookupTyCon :: TyCon -> VM (Maybe TyCon)
79 | isUnLiftedTyCon tc || isTupleTyCon tc
83 = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
86 -- | Lookup the vectorised version of a boxed `TyCon` from the global environment.
87 lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
89 = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
93 -- | Add a mapping between plain and vectorised `TyCon`s to the global environment.
94 defTyCon :: TyCon -> TyCon -> VM ()
95 defTyCon tc tc' = updGEnv $ \env ->
96 env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
99 -- DataCons -------------------------------------------------------------------
100 -- | Lookup the vectorised version of a `DataCon` from the global environment.
101 lookupDataCon :: DataCon -> VM (Maybe DataCon)
103 | isTupleTyCon (dataConTyCon dc)
107 = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
110 -- | Add the mapping between plain and vectorised `DataCon`s to the global environment.
111 defDataCon :: DataCon -> DataCon -> VM ()
112 defDataCon dc dc' = updGEnv $ \env ->
113 env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
116 -- PA dictionaries ------------------------------------------------------------
117 -- | Lookup a PA `TyCon` from the global environment.
118 lookupTyConPA :: TyCon -> VM (Maybe Var)
120 = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
123 -- | Add a mapping between a PA TyCon and is vectorised version to the global environment.
124 defTyConPA :: TyCon -> Var -> VM ()
125 defTyConPA tc pa = updGEnv $ \env ->
126 env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
129 -- | Add several mapping between PA TyCons and their vectorised versions to the global environment.
130 defTyConPAs :: [(TyCon, Var)] -> VM ()
131 defTyConPAs ps = updGEnv $ \env ->
132 env { global_pa_funs = extendNameEnvList (global_pa_funs env)
133 [(tyConName tc, pa) | (tc, pa) <- ps] }
136 -- PR Dictionaries ------------------------------------------------------------
137 lookupTyConPR :: TyCon -> VM (Maybe Var)
138 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)