5 -- * Local Environments
9 -- * Global Environments
12 extendImportedVarsEnv,
39 -- | Indicates what scope something (a variable) is in.
45 -- LocalEnv -------------------------------------------------------------------
46 -- | The local environment.
49 -- Mapping from local variables to their vectorised and lifted versions.
50 local_vars :: VarEnv (Var, Var)
52 -- In-scope type variables.
53 , local_tyvars :: [TyVar]
55 -- Mapping from tyvars to their PA dictionaries.
56 , local_tyvar_pa :: VarEnv CoreExpr
58 -- Local binding name.
59 , local_bind_name :: FastString
63 -- | Create an empty local environment.
64 emptyLocalEnv :: LocalEnv
65 emptyLocalEnv = LocalEnv {
66 local_vars = emptyVarEnv
68 , local_tyvar_pa = emptyVarEnv
69 , local_bind_name = fsLit "fn"
73 -- GlobalEnv ------------------------------------------------------------------
74 -- | The global environment.
75 -- These are things the exist at top-level.
78 -- | Mapping from global variables to their vectorised versions — aka the /vectorisation
80 global_vars :: VarEnv Var
82 -- | Mapping from global variables that have a vectorisation declaration to the right-hand
83 -- side of that declaration and its type. This mapping only applies to non-scalar
84 -- vectorisation declarations. All variables with a scalar vectorisation declaration are
85 -- mentioned in 'global_scalars'.
86 , global_vect_decls :: VarEnv (Type, CoreExpr)
88 -- | Purely scalar variables. Code which mentions only these variables doesn't have to be
89 -- lifted. This includes variables from the current module that have a scalar
90 -- vectorisation declaration and those that the vectoriser determines to be scalar.
91 , global_scalars :: VarSet
93 -- | Exported variables which have a vectorised version.
94 , global_exported_vars :: VarEnv (Var, Var)
96 -- | Mapping from TyCons to their vectorised versions.
97 -- TyCons which do not have to be vectorised are mapped to themselves.
98 , global_tycons :: NameEnv TyCon
100 -- | Mapping from DataCons to their vectorised versions.
101 , global_datacons :: NameEnv DataCon
103 -- | Mapping from TyCons to their PA dfuns.
104 , global_pa_funs :: NameEnv Var
106 -- | Mapping from TyCons to their PR dfuns.
107 , global_pr_funs :: NameEnv Var
109 -- | Mapping from unboxed TyCons to their boxed versions.
110 , global_boxed_tycons :: NameEnv TyCon
112 -- | External package inst-env & home-package inst-env for class instances.
113 , global_inst_env :: (InstEnv, InstEnv)
115 -- | External package inst-env & home-package inst-env for family instances.
116 , global_fam_inst_env :: FamInstEnvs
118 -- | Hoisted bindings.
119 , global_bindings :: [(Var, CoreExpr)]
122 -- | Create an initial global environment
123 initGlobalEnv :: VectInfo -> [CoreVect] -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
124 initGlobalEnv info vectDecls instEnvs famInstEnvs
126 { global_vars = mapVarEnv snd $ vectInfoVar info
127 , global_vect_decls = mkVarEnv vects
128 , global_scalars = mkVarSet scalars
129 , global_exported_vars = emptyVarEnv
130 , global_tycons = mapNameEnv snd $ vectInfoTyCon info
131 , global_datacons = mapNameEnv snd $ vectInfoDataCon info
132 , global_pa_funs = mapNameEnv snd $ vectInfoPADFun info
133 , global_pr_funs = emptyNameEnv
134 , global_boxed_tycons = emptyNameEnv
135 , global_inst_env = instEnvs
136 , global_fam_inst_env = famInstEnvs
137 , global_bindings = []
140 vects = [(var, (varType var, exp)) | Vect var (Just exp) <- vectDecls]
141 scalars = [var | Vect var Nothing <- vectDecls]
144 -- Operators on Global Environments -------------------------------------------
145 -- | Extend the list of global variables in an environment.
146 extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
147 extendImportedVarsEnv ps genv
148 = genv { global_vars = extendVarEnvList (global_vars genv) ps }
150 -- | Extend the set of scalar variables in an environment.
151 extendScalars :: [Var] -> GlobalEnv -> GlobalEnv
152 extendScalars vs genv
153 = genv { global_scalars = extendVarSetList (global_scalars genv) vs }
155 -- | Set the list of type family instances in an environment.
156 setFamEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
157 setFamEnv l_fam_inst genv
158 = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
159 where (g_fam_inst, _) = global_fam_inst_env genv
161 extendFamEnv :: [FamInst] -> GlobalEnv -> GlobalEnv
162 extendFamEnv new genv
163 = genv { global_fam_inst_env = (g_fam_inst, extendFamInstEnvList l_fam_inst new) }
164 where (g_fam_inst, l_fam_inst) = global_fam_inst_env genv
167 -- | Extend the list of type constructors in an environment.
168 extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
169 extendTyConsEnv ps genv
170 = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
173 -- | Extend the list of data constructors in an environment.
174 extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
175 extendDataConsEnv ps genv
176 = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
179 -- | Extend the list of PA functions in an environment.
180 extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
181 extendPAFunsEnv ps genv
182 = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
185 -- | Set the list of PR functions in an environment.
186 setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
188 = genv { global_pr_funs = mkNameEnv ps }
191 -- | Set the list of boxed type constructor in an environment.
192 setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
193 setBoxedTyConsEnv ps genv
194 = genv { global_boxed_tycons = mkNameEnv ps }
197 -- | TODO: What is this for?
198 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
199 updVectInfo env tyenv info
201 { vectInfoVar = global_exported_vars env
202 , vectInfoTyCon = mk_env typeEnvTyCons global_tycons
203 , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
204 , vectInfoPADFun = mk_env typeEnvTyCons global_pa_funs
207 mk_env from_tyenv from_env
208 = mkNameEnv [(name, (from,to))
209 | from <- from_tyenv tyenv
210 , let name = getName from
211 , Just to <- [lookupNameEnv (from_env env) name]]