-- | The Vectorisation monad.
module VectMonad (
- Scope(..),
VM,
noV, traceNoV, ensureV, traceEnsureV, tryV, maybeV, traceMaybeV, orElseV,
combinePDVar, scalarZip, closureCtrFun,
builtin, builtins,
- GlobalEnv(..),
setFamInstEnv,
readGEnv, setGEnv, updGEnv,
- LocalEnv(..),
readLEnv, setLEnv, updLEnv,
getBindName, inBind,
#include "HsVersions.h"
import VectBuiltIn
+import Vectorise.Env
import HscTypes hiding ( MonadThings(..) )
import Module ( PackageId )
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 ----------------------------------------------------
import VectMonad
import VectUtils
import VectCore
+import Vectorise.Env
import HscTypes ( TypeEnv, extendTypeEnvList, typeEnvTyCons )
import BasicTypes
import VectCore
import VectMonad
+import Vectorise.Env
import MkCore ( mkCoreTup, mkWildCase )
import CoreSyn
import BasicTypes ( Boxity(..), Arity )
import Literal ( Literal, mkMachInt )
+
import Outputable
import FastString
import VectCore
import VectMonad
import VectType
+import Vectorise.Env
import CoreSyn
import Type
import Var
import VectVar
import VectType
import VectCore
+import Vectorise.Env
import HscTypes hiding ( MonadThings(..) )
--- /dev/null
+
+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]]
+