Split out vectoriser environments into own module

compiler/vectorise/VectMonad.hs

@@ -2,7 +2,6 @@
 
 -- | The Vectorisation monad.
 module VectMonad (
-  Scope(..),
   VM,
 
   noV, traceNoV, ensureV, traceEnsureV, tryV, maybeV, traceMaybeV, orElseV,
@@ -17,11 +16,9 @@ module VectMonad (
   combinePDVar, scalarZip, closureCtrFun,
   builtin, builtins,
 
-  GlobalEnv(..),
   setFamInstEnv,
   readGEnv, setGEnv, updGEnv,
 
-  LocalEnv(..),
   readLEnv, setLEnv, updLEnv,
 
   getBindName, inBind,
@@ -41,6 +38,7 @@ module VectMonad (
 #include "HsVersions.h"
 
 import VectBuiltIn
+import Vectorise.Env
 
 import HscTypes hiding  ( MonadThings(..) )
 import Module           ( PackageId )
@@ -67,155 +65,6 @@ import SrcLoc        ( noSrcSpan )
 
 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 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 ----------------------------------------------------
 

compiler/vectorise/VectType.hs

@@ -9,6 +9,7 @@ where
 import VectMonad
 import VectUtils
 import VectCore
+import Vectorise.Env
 
 import HscTypes          ( TypeEnv, extendTypeEnvList, typeEnvTyCons )
 import BasicTypes

compiler/vectorise/VectUtils.hs

@@ -24,6 +24,7 @@ module VectUtils (
 
 import VectCore
 import VectMonad
+import Vectorise.Env
 
 import MkCore ( mkCoreTup, mkWildCase )
 import CoreSyn
@@ -41,6 +42,7 @@ import TysWiredIn
 import BasicTypes         ( Boxity(..), Arity )
 import Literal            ( Literal, mkMachInt )
 
+
 import Outputable
 import FastString
 

compiler/vectorise/VectVar.hs

@@ -14,6 +14,7 @@ import VectUtils
 import VectCore
 import VectMonad
 import VectType
+import Vectorise.Env
 import CoreSyn
 import Type
 import Var

compiler/vectorise/Vectorise.hs

@@ -8,6 +8,7 @@ import VectUtils
 import VectVar
 import VectType
 import VectCore
+import Vectorise.Env
 
 import HscTypes hiding      ( MonadThings(..) )
 

compiler/vectorise/Vectorise/Env.hs

+
+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]]
+