import Control.Monad
+-- | Indicates what scope something (a variable) is in.
+data Scope a b = Global a | Local b
+
+
+-- | The global environment.
+data GlobalEnv = GlobalEnv {
+ -- | Mapping from global variables to their vectorised versions.
+ --
+ global_vars :: VarEnv Var
+
+ -- | Purely scalar variables. Code which mentions only these
+ -- variables doesn't have to be lifted.
+ , global_scalars :: VarSet
+
+ -- | Exported variables which have a vectorised version
+ --
+ , global_exported_vars :: VarEnv (Var, Var)
+
+ -- | Mapping from TyCons to their vectorised versions.
+ -- TyCons which do not have to be vectorised are mapped to
+ -- themselves.
+ --
+ , global_tycons :: NameEnv TyCon
+
+ -- | Mapping from DataCons to their vectorised versions
+ --
+ , global_datacons :: NameEnv DataCon
+
+ -- | Mapping from TyCons to their PA dfuns
+ --
+ , global_pa_funs :: NameEnv Var
+
+ -- | Mapping from TyCons to their PR dfuns
+ , global_pr_funs :: NameEnv Var
+
+ -- | Mapping from unboxed TyCons to their boxed versions
+ , global_boxed_tycons :: NameEnv TyCon
+
+ -- | External package inst-env & home-package inst-env for class
+ -- instances
+ --
+ , global_inst_env :: (InstEnv, InstEnv)
+
+ -- | External package inst-env & home-package inst-env for family
+ -- instances
+ --
+ , global_fam_inst_env :: FamInstEnvs
+
+ -- | Hoisted bindings
+ , global_bindings :: [(Var, CoreExpr)]
+ }
+
+-- | The local environment.
+data LocalEnv = LocalEnv {
+ -- Mapping from local variables to their vectorised and
+ -- lifted versions
+ --
+ local_vars :: VarEnv (Var, Var)
+
+ -- In-scope type variables
+ --
+ , local_tyvars :: [TyVar]
+
+ -- Mapping from tyvars to their PA dictionaries
+ , local_tyvar_pa :: VarEnv CoreExpr
+
+ -- Local binding name
+ , local_bind_name :: FastString
+ }
+
+
+-- | Create an initial global environment
+initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
+initGlobalEnv info instEnvs famInstEnvs
+ = GlobalEnv {
+ global_vars = mapVarEnv snd $ vectInfoVar info
+ , global_scalars = emptyVarSet
+ , global_exported_vars = emptyVarEnv
+ , global_tycons = mapNameEnv snd $ vectInfoTyCon info
+ , global_datacons = mapNameEnv snd $ vectInfoDataCon info
+ , global_pa_funs = mapNameEnv snd $ vectInfoPADFun info
+ , global_pr_funs = emptyNameEnv
+ , global_boxed_tycons = emptyNameEnv
+ , global_inst_env = instEnvs
+ , global_fam_inst_env = famInstEnvs
+ , global_bindings = []
+ }
+
+
+-- Operators on Global Environments -------------------------------------------
+extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
+extendImportedVarsEnv ps genv
+ = genv { global_vars = extendVarEnvList (global_vars genv) ps }
+
+extendScalars :: [Var] -> GlobalEnv -> GlobalEnv
+extendScalars vs genv
+ = genv { global_scalars = extendVarSetList (global_scalars genv) vs }
+
+setFamInstEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
+setFamInstEnv l_fam_inst genv
+ = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
+ where
+ (g_fam_inst, _) = global_fam_inst_env genv
+
+extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
+extendTyConsEnv ps genv
+ = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
+
+extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
+extendDataConsEnv ps genv
+ = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
+
+extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
+extendPAFunsEnv ps genv
+ = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
+
+setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
+setPRFunsEnv ps genv
+ = genv { global_pr_funs = mkNameEnv ps }
+
+setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
+setBoxedTyConsEnv ps genv
+ = genv { global_boxed_tycons = mkNameEnv ps }
+
+
+-- | Create an empty local environment.
+emptyLocalEnv :: LocalEnv
+emptyLocalEnv = LocalEnv {
+ local_vars = emptyVarEnv
+ , local_tyvars = []
+ , local_tyvar_pa = emptyVarEnv
+ , local_bind_name = fsLit "fn"
+ }
+
+-- FIXME
+updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
+updVectInfo env tyenv info
+ = info {
+ vectInfoVar = global_exported_vars env
+ , vectInfoTyCon = mk_env typeEnvTyCons global_tycons
+ , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
+ , vectInfoPADFun = mk_env typeEnvTyCons global_pa_funs
+ }
+ where
+ mk_env :: NamedThing from =>
+ (TypeEnv -> [from])
+ -> (GlobalEnv -> NameEnv to)
+ -> NameEnv (from,to)
+ mk_env from_tyenv from_env = mkNameEnv [(name, (from,to))
+ | from <- from_tyenv tyenv
+ , let name = getName from
+ , Just to <- [lookupNameEnv (from_env env) name]]
-- The Vectorisation Monad ----------------------------------------------------
--
lookupFamInst :: TyCon -> [Type] -> VM (TyCon, [Type])
lookupFamInst tycon tys
- = ASSERT( isOpenTyCon tycon )
+ = ASSERT( isFamilyTyCon tycon )
do { instEnv <- getFamInstEnv
; case lookupFamInstEnv instEnv tycon tys of
[(fam_inst, rep_tys)] -> return (famInstTyCon fam_inst, rep_tys)
projects / ghc-hetmet.git / blobdiff