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_novect_vars :: VarSet
100 -- ^Variables that are not vectorised. (They may be referenced in the right-hand sides
101 -- of vectorisation declarations, though.)
103 , global_exported_vars :: VarEnv (Var, Var)
104 -- ^Exported variables which have a vectorised version.
106 , global_tycons :: NameEnv TyCon
107 -- ^Mapping from TyCons to their vectorised versions.
108 -- TyCons which do not have to be vectorised are mapped to themselves.
110 , global_datacons :: NameEnv DataCon
111 -- ^Mapping from DataCons to their vectorised versions.
113 , global_pa_funs :: NameEnv Var
114 -- ^Mapping from TyCons to their PA dfuns.
116 , global_pr_funs :: NameEnv Var
117 -- ^Mapping from TyCons to their PR dfuns.
119 , global_boxed_tycons :: NameEnv TyCon
120 -- ^Mapping from unboxed TyCons to their boxed versions.
122 , global_inst_env :: (InstEnv, InstEnv)
123 -- ^External package inst-env & home-package inst-env for class instances.
125 , global_fam_inst_env :: FamInstEnvs
126 -- ^External package inst-env & home-package inst-env for family instances.
128 , global_bindings :: [(Var, CoreExpr)]
129 -- ^Hoisted bindings.
132 -- |Create an initial global environment.
134 initGlobalEnv :: VectInfo -> [CoreVect] -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
135 initGlobalEnv info vectDecls instEnvs famInstEnvs
137 { global_vars = mapVarEnv snd $ vectInfoVar info
138 , global_vect_decls = mkVarEnv vects
139 , global_scalar_vars = vectInfoScalarVars info `extendVarSetList` scalars
140 , global_scalar_tycons = vectInfoScalarTyCons info
141 , global_novect_vars = mkVarSet novects
142 , global_exported_vars = emptyVarEnv
143 , global_tycons = mapNameEnv snd $ vectInfoTyCon info
144 , global_datacons = mapNameEnv snd $ vectInfoDataCon info
145 , global_pa_funs = mapNameEnv snd $ vectInfoPADFun info
146 , global_pr_funs = emptyNameEnv
147 , global_boxed_tycons = emptyNameEnv
148 , global_inst_env = instEnvs
149 , global_fam_inst_env = famInstEnvs
150 , global_bindings = []
153 vects = [(var, (varType var, exp)) | Vect var (Just exp) <- vectDecls]
154 scalars = [var | Vect var Nothing <- vectDecls]
155 novects = [var | NoVect var <- vectDecls]
158 -- Operators on Global Environments -------------------------------------------
160 -- |Extend the list of global variables in an environment.
162 extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
163 extendImportedVarsEnv ps genv
164 = genv { global_vars = extendVarEnvList (global_vars genv) ps }
166 -- |Extend the set of scalar variables in an environment.
168 extendScalars :: [Var] -> GlobalEnv -> GlobalEnv
169 extendScalars vs genv
170 = genv { global_scalar_vars = extendVarSetList (global_scalar_vars genv) vs }
172 -- |Set the list of type family instances in an environment.
174 setFamEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
175 setFamEnv l_fam_inst genv
176 = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
177 where (g_fam_inst, _) = global_fam_inst_env genv
179 -- |Extend the list of type family instances.
181 extendFamEnv :: [FamInst] -> GlobalEnv -> GlobalEnv
182 extendFamEnv new genv
183 = genv { global_fam_inst_env = (g_fam_inst, extendFamInstEnvList l_fam_inst new) }
184 where (g_fam_inst, l_fam_inst) = global_fam_inst_env genv
186 -- |Extend the list of type constructors in an environment.
188 extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
189 extendTyConsEnv ps genv
190 = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
192 -- |Extend the list of data constructors in an environment.
194 extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
195 extendDataConsEnv ps genv
196 = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
198 -- |Extend the list of PA functions in an environment.
200 extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
201 extendPAFunsEnv ps genv
202 = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
204 -- |Set the list of PR functions in an environment.
206 setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
208 = genv { global_pr_funs = mkNameEnv ps }
210 -- |Set the list of boxed type constructor in an environment.
212 setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
213 setBoxedTyConsEnv ps genv
214 = genv { global_boxed_tycons = mkNameEnv ps }
216 -- |Compute vectorisation information that goes into 'ModGuts' (and is stored in interface files).
217 -- The incoming 'vectInfo' is that from the 'HscEnv' and 'EPS'. The outgoing one contains only the
218 -- definitions for the currently compiled module.
220 modVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
221 modVectInfo env tyenv info
223 { vectInfoVar = global_exported_vars env
224 , vectInfoTyCon = mk_env typeEnvTyCons global_tycons
225 , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
226 , vectInfoPADFun = mk_env typeEnvTyCons global_pa_funs
227 , vectInfoScalarVars = global_scalar_vars env `minusVarSet` vectInfoScalarVars info
228 , vectInfoScalarTyCons = global_scalar_tycons env `minusNameSet` vectInfoScalarTyCons info
231 mk_env from_tyenv from_env
232 = mkNameEnv [(name, (from,to))
233 | from <- from_tyenv tyenv
234 , let name = getName from
235 , Just to <- [lookupNameEnv (from_env env) name]]