diff options
Diffstat (limited to 'compiler/vectorise/Vectorise/Monad/Global.hs')
| -rw-r--r-- | compiler/vectorise/Vectorise/Monad/Global.hs | 71 |
1 files changed, 41 insertions, 30 deletions
diff --git a/compiler/vectorise/Vectorise/Monad/Global.hs b/compiler/vectorise/Vectorise/Monad/Global.hs index e471ebbc03..96448fb26a 100644 --- a/compiler/vectorise/Vectorise/Monad/Global.hs +++ b/compiler/vectorise/Vectorise/Monad/Global.hs @@ -1,3 +1,4 @@ +-- Operations on the global state of the vectorisation monad. module Vectorise.Monad.Global ( readGEnv, @@ -11,12 +12,11 @@ module Vectorise.Monad.Global ( lookupVectDecl, noVectDecl, -- * Scalars - globalScalars, isGlobalScalar, + globalScalarVars, isGlobalScalar, globalScalarTyCons, -- * TyCons - lookupTyCon, - lookupBoxedTyCon, - defTyCon, + lookupTyCon, lookupBoxedTyCon, + defTyCon, globalVectTyCons, -- * Datacons lookupDataCon, @@ -24,7 +24,6 @@ module Vectorise.Monad.Global ( -- * PA Dictionaries lookupTyConPA, - defTyConPA, defTyConPAs, -- * PR Dictionaries @@ -39,6 +38,7 @@ import Type import TyCon import DataCon import NameEnv +import NameSet import Var import VarEnv import VarSet @@ -49,17 +49,17 @@ import VarSet -- |Project something from the global environment. -- readGEnv :: (GlobalEnv -> a) -> VM a -readGEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f genv)) +readGEnv f = VM $ \_ genv lenv -> return (Yes genv lenv (f genv)) -- |Set the value of the global environment. -- setGEnv :: GlobalEnv -> VM () -setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ()) +setGEnv genv = VM $ \_ _ lenv -> return (Yes genv lenv ()) -- |Update the global environment using the provided function. -- updGEnv :: (GlobalEnv -> GlobalEnv) -> VM () -updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ()) +updGEnv f = VM $ \_ genv lenv -> return (Yes (f genv) lenv ()) -- Vars ----------------------------------------------------------------------- @@ -93,13 +93,19 @@ noVectDecl var = readGEnv $ \env -> elemVarSet var (global_novect_vars env) -- |Get the set of global scalar variables. -- -globalScalars :: VM VarSet -globalScalars = readGEnv global_scalar_vars +globalScalarVars :: VM VarSet +globalScalarVars = readGEnv global_scalar_vars -- |Check whether a given variable is in the set of global scalar variables. -- isGlobalScalar :: Var -> VM Bool -isGlobalScalar var = readGEnv $ \env -> elemVarSet var (global_scalar_vars env) +isGlobalScalar var = readGEnv $ \env -> var `elemVarSet` global_scalar_vars env + +-- |Get the set of global scalar type constructors including both those scalar type constructors +-- declared in an imported module and those declared in the current module. +-- +globalScalarTyCons :: VM NameSet +globalScalarTyCons = readGEnv global_scalar_tycons -- TyCons --------------------------------------------------------------------- @@ -110,25 +116,32 @@ lookupTyCon :: TyCon -> VM (Maybe TyCon) lookupTyCon tc | isUnLiftedTyCon tc || isTupleTyCon tc = return (Just tc) - | otherwise = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc) --- | Lookup the vectorised version of a boxed `TyCon` from the global environment. +-- |Lookup the vectorised version of a boxed `TyCon` from the global environment. +-- lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon) lookupBoxedTyCon tc - = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env) - (tyConName tc) + = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env) + (tyConName tc) --- | Add a mapping between plain and vectorised `TyCon`s to the global environment. +-- |Add a mapping between plain and vectorised `TyCon`s to the global environment. +-- defTyCon :: TyCon -> TyCon -> VM () defTyCon tc tc' = updGEnv $ \env -> env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' } +-- |Get the set of all vectorised type constructors. +-- +globalVectTyCons :: VM (NameEnv TyCon) +globalVectTyCons = readGEnv global_tycons + -- DataCons ------------------------------------------------------------------- --- | Lookup the vectorised version of a `DataCon` from the global environment. +-- |Lookup the vectorised version of a `DataCon` from the global environment. +-- lookupDataCon :: DataCon -> VM (Maybe DataCon) lookupDataCon dc | isTupleTyCon (dataConTyCon dc) @@ -137,27 +150,24 @@ lookupDataCon dc | otherwise = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc) - --- | Add the mapping between plain and vectorised `DataCon`s to the global environment. +-- |Add the mapping between plain and vectorised `DataCon`s to the global environment. +-- defDataCon :: DataCon -> DataCon -> VM () defDataCon dc dc' = updGEnv $ \env -> env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' } --- PA dictionaries ------------------------------------------------------------ --- | Lookup a PA `TyCon` from the global environment. +-- 'PA' dictionaries ------------------------------------------------------------ + +-- |Lookup the 'PA' dfun of a vectorised type constructor in the global environment. +-- lookupTyConPA :: TyCon -> VM (Maybe Var) lookupTyConPA tc - = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc) - + = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc) --- | Add a mapping between a PA TyCon and is vectorised version to the global environment. -defTyConPA :: TyCon -> Var -> VM () -defTyConPA tc pa = updGEnv $ \env -> - env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa } - - --- | Add several mapping between PA TyCons and their vectorised versions to the global environment. +-- |Associate vectorised type constructors with the dfun of their 'PA' instances in the global +-- environment. +-- defTyConPAs :: [(TyCon, Var)] -> VM () defTyConPAs ps = updGEnv $ \env -> env { global_pa_funs = extendNameEnvList (global_pa_funs env) @@ -165,6 +175,7 @@ defTyConPAs ps = updGEnv $ \env -> -- PR Dictionaries ------------------------------------------------------------ + lookupTyConPR :: TyCon -> VM (Maybe Var) lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc) |