Split out vectoriser environments into own module
[ghc-hetmet.git] / compiler / vectorise / Vectorise / Env.hs
diff --git a/compiler/vectorise/Vectorise/Env.hs b/compiler/vectorise/Vectorise/Env.hs
new file mode 100644 (file)
index 0000000..51d45a4
--- /dev/null
@@ -0,0 +1,197 @@
+
+module Vectorise.Env (
+       Scope(..),
+
+       -- * Local Environments
+       LocalEnv(..),
+       emptyLocalEnv,
+       
+       -- * Global Environments
+       GlobalEnv(..),
+       initGlobalEnv,
+       extendImportedVarsEnv,
+       extendScalars,
+       setFamInstEnv,
+       extendTyConsEnv,
+       extendDataConsEnv,
+       extendPAFunsEnv,
+       setPRFunsEnv,
+       setBoxedTyConsEnv,
+       updVectInfo
+) where
+import HscTypes
+import InstEnv
+import FamInstEnv
+import CoreSyn
+import TyCon
+import DataCon
+import Type
+import VarEnv
+import VarSet
+import Var
+import Name
+import NameEnv
+import FastString
+
+
+-- | Indicates what scope something (a variable) is in.
+data Scope a b 
+       = Global a 
+       | Local  b
+
+
+-- LocalEnv -------------------------------------------------------------------
+-- | 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 empty local environment.
+emptyLocalEnv :: LocalEnv
+emptyLocalEnv = LocalEnv {
+                   local_vars     = emptyVarEnv
+                 , local_tyvars   = []
+                 , local_tyvar_pa = emptyVarEnv
+                 , local_bind_name  = fsLit "fn"
+                 }
+
+
+-- GlobalEnv ------------------------------------------------------------------
+-- | The global environment.
+--     These are things the exist at top-level.
+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)]
+        }
+
+
+-- | 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 -------------------------------------------
+-- | Extend the list of global variables in an environment.
+extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
+extendImportedVarsEnv ps genv
+  = genv { global_vars  = extendVarEnvList (global_vars genv) ps }
+
+
+-- | Extend the set of scalar variables in an environment.
+extendScalars :: [Var] -> GlobalEnv -> GlobalEnv
+extendScalars vs genv
+  = genv { global_scalars = extendVarSetList (global_scalars genv) vs }
+
+
+-- | Set the list of type family instances in an environment.
+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
+
+
+-- | Extend the list of type constructors in an environment.
+extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
+extendTyConsEnv ps genv
+  = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
+
+
+-- | Extend the list of data constructors in an environment.
+extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
+extendDataConsEnv ps genv
+  = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
+
+
+-- | Extend the list of PA functions in an environment.
+extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
+extendPAFunsEnv ps genv
+  = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
+
+
+-- | Set the list of PR functions in an environment.
+setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
+setPRFunsEnv ps genv
+  = genv { global_pr_funs = mkNameEnv ps }
+
+
+-- | Set the list of boxed type constructor in an environment.
+setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
+setBoxedTyConsEnv ps genv
+  = genv { global_boxed_tycons = mkNameEnv ps }
+
+
+-- | TODO: What is this for?
+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 from_tyenv from_env 
+       = mkNameEnv [(name, (from,to))
+                        | from     <- from_tyenv tyenv
+                        , let name =  getName from
+                        , Just to  <- [lookupNameEnv (from_env env) name]]
+