5 -- * Local Environments
9 -- * Global Environments
12 extendImportedVarsEnv,
40 -- | Indicates what scope something (a variable) is in.
46 -- LocalEnv -------------------------------------------------------------------
47 -- | The local environment.
50 -- Mapping from local variables to their vectorised and lifted versions.
51 local_vars :: VarEnv (Var, Var)
53 -- In-scope type variables.
54 , local_tyvars :: [TyVar]
56 -- Mapping from tyvars to their PA dictionaries.
57 , local_tyvar_pa :: VarEnv CoreExpr
59 -- Local binding name.
60 , local_bind_name :: FastString
64 -- | Create an empty local environment.
65 emptyLocalEnv :: LocalEnv
66 emptyLocalEnv = LocalEnv {
67 local_vars = emptyVarEnv
69 , local_tyvar_pa = emptyVarEnv
70 , local_bind_name = fsLit "fn"
74 -- GlobalEnv ------------------------------------------------------------------
76 -- |The global environment: entities that exist at top-level.
80 { global_vars :: VarEnv Var
81 -- ^Mapping from global variables to their vectorised versions — aka the /vectorisation
84 , global_vect_decls :: VarEnv (Type, CoreExpr)
85 -- ^Mapping from global variables that have a vectorisation declaration to the right-hand
86 -- side of that declaration and its type. This mapping only applies to non-scalar
87 -- vectorisation declarations. All variables with a scalar vectorisation declaration are
88 -- mentioned in 'global_scalars_vars'.
90 , global_scalar_vars :: VarSet
91 -- ^Purely scalar variables. Code which mentions only these variables doesn't have to be
92 -- lifted. This includes variables from the current module that have a scalar
93 -- vectorisation declaration and those that the vectoriser determines to be scalar.
95 , global_scalar_tycons :: NameSet
96 -- ^Type constructors whose values can only contain scalar data. Scalar code may only
97 -- operate on such data.
99 , global_exported_vars :: VarEnv (Var, Var)
100 -- ^Exported variables which have a vectorised version.
102 , global_tycons :: NameEnv TyCon
103 -- ^Mapping from TyCons to their vectorised versions.
104 -- TyCons which do not have to be vectorised are mapped to themselves.
106 , global_datacons :: NameEnv DataCon
107 -- ^Mapping from DataCons to their vectorised versions.
109 , global_pa_funs :: NameEnv Var
110 -- ^Mapping from TyCons to their PA dfuns.
112 , global_pr_funs :: NameEnv Var
113 -- ^Mapping from TyCons to their PR dfuns.
115 , global_boxed_tycons :: NameEnv TyCon
116 -- ^Mapping from unboxed TyCons to their boxed versions.
118 , global_inst_env :: (InstEnv, InstEnv)
119 -- ^External package inst-env & home-package inst-env for class instances.
121 , global_fam_inst_env :: FamInstEnvs
122 -- ^External package inst-env & home-package inst-env for family instances.
124 , global_bindings :: [(Var, CoreExpr)]
125 -- ^Hoisted bindings.
128 -- |Create an initial global environment.
130 initGlobalEnv :: VectInfo -> [CoreVect] -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
131 initGlobalEnv info vectDecls instEnvs famInstEnvs
133 { global_vars = mapVarEnv snd $ vectInfoVar info
134 , global_vect_decls = mkVarEnv vects
135 , global_scalar_vars = vectInfoScalarVars info `extendVarSetList` scalars
136 , global_scalar_tycons = vectInfoScalarTyCons info
137 , global_exported_vars = emptyVarEnv
138 , global_tycons = mapNameEnv snd $ vectInfoTyCon info
139 , global_datacons = mapNameEnv snd $ vectInfoDataCon info
140 , global_pa_funs = mapNameEnv snd $ vectInfoPADFun info
141 , global_pr_funs = emptyNameEnv
142 , global_boxed_tycons = emptyNameEnv
143 , global_inst_env = instEnvs
144 , global_fam_inst_env = famInstEnvs
145 , global_bindings = []
148 vects = [(var, (varType var, exp)) | Vect var (Just exp) <- vectDecls]
149 scalars = [var | Vect var Nothing <- vectDecls]
152 -- Operators on Global Environments -------------------------------------------
154 -- |Extend the list of global variables in an environment.
156 extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
157 extendImportedVarsEnv ps genv
158 = genv { global_vars = extendVarEnvList (global_vars genv) ps }
160 -- |Extend the set of scalar variables in an environment.
162 extendScalars :: [Var] -> GlobalEnv -> GlobalEnv
163 extendScalars vs genv
164 = genv { global_scalar_vars = extendVarSetList (global_scalar_vars genv) vs }
166 -- |Set the list of type family instances in an environment.
168 setFamEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
169 setFamEnv l_fam_inst genv
170 = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
171 where (g_fam_inst, _) = global_fam_inst_env genv
173 -- |Extend the list of type family instances.
175 extendFamEnv :: [FamInst] -> GlobalEnv -> GlobalEnv
176 extendFamEnv new genv
177 = genv { global_fam_inst_env = (g_fam_inst, extendFamInstEnvList l_fam_inst new) }
178 where (g_fam_inst, l_fam_inst) = global_fam_inst_env genv
180 -- |Extend the list of type constructors in an environment.
182 extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
183 extendTyConsEnv ps genv
184 = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
186 -- |Extend the list of data constructors in an environment.
188 extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
189 extendDataConsEnv ps genv
190 = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
192 -- |Extend the list of PA functions in an environment.
194 extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
195 extendPAFunsEnv ps genv
196 = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
198 -- |Set the list of PR functions in an environment.
200 setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
202 = genv { global_pr_funs = mkNameEnv ps }
204 -- |Set the list of boxed type constructor in an environment.
206 setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
207 setBoxedTyConsEnv ps genv
208 = genv { global_boxed_tycons = mkNameEnv ps }
210 -- |Compute vectorisation information that goes into 'ModGuts' (and is stored in interface files).
211 -- The incoming 'vectInfo' is that from the 'HscEnv' and 'EPS'. The outgoing one contains only the
212 -- definitions for the currently compiled module.
214 modVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
215 modVectInfo env tyenv info
217 { vectInfoVar = global_exported_vars env
218 , vectInfoTyCon = mk_env typeEnvTyCons global_tycons
219 , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
220 , vectInfoPADFun = mk_env typeEnvTyCons global_pa_funs
221 , vectInfoScalarVars = global_scalar_vars env `minusVarSet` vectInfoScalarVars info
222 , vectInfoScalarTyCons = global_scalar_tycons env `minusNameSet` vectInfoScalarTyCons info
225 mk_env from_tyenv from_env
226 = mkNameEnv [(name, (from,to))
227 | from <- from_tyenv tyenv
228 , let name = getName from
229 , Just to <- [lookupNameEnv (from_env env) name]]