compiler/vectorise/Vectorise/Env.hs

   1
   2 module Vectorise.Env (
   3         Scope(..),
   4
   5         -- * Local Environments
   6         LocalEnv(..),
   7         emptyLocalEnv,
   8
   9         -- * Global Environments
  10         GlobalEnv(..),
  11         initGlobalEnv,
  12         extendImportedVarsEnv,
  13         extendScalars,
  14         setFamInstEnv,
  15         extendTyConsEnv,
  16         extendDataConsEnv,
  17         extendPAFunsEnv,
  18         setPRFunsEnv,
  19         setBoxedTyConsEnv,
  20         updVectInfo
  21 ) where
  22 import HscTypes
  23 import InstEnv
  24 import FamInstEnv
  25 import CoreSyn
  26 import TyCon
  27 import DataCon
  28 import Type
  29 import VarEnv
  30 import VarSet
  31 import Var
  32 import Name
  33 import NameEnv
  34 import FastString
  35
  36
  37 -- | Indicates what scope something (a variable) is in.
  38 data Scope a b
  39         = Global a
  40         | Local  b
  41
  42
  43 -- LocalEnv -------------------------------------------------------------------
  44 -- | The local environment.
  45 data LocalEnv
  46         = LocalEnv {
  47         -- Mapping from local variables to their vectorised and lifted versions.
  48             local_vars          :: VarEnv (Var, Var)
  49
  50         -- In-scope type variables.
  51         , local_tyvars          :: [TyVar]
  52
  53         -- Mapping from tyvars to their PA dictionaries.
  54         , local_tyvar_pa        :: VarEnv CoreExpr
  55
  56         -- Local binding name.
  57         , local_bind_name       :: FastString
  58         }
  59
  60
  61 -- | Create an empty local environment.
  62 emptyLocalEnv :: LocalEnv
  63 emptyLocalEnv = LocalEnv {
  64                    local_vars     = emptyVarEnv
  65                  , local_tyvars   = []
  66                  , local_tyvar_pa = emptyVarEnv
  67                  , local_bind_name  = fsLit "fn"
  68                  }
  69
  70
  71 -- GlobalEnv ------------------------------------------------------------------
  72 -- | The global environment.
  73 --      These are things the exist at top-level.
  74 data GlobalEnv
  75         = GlobalEnv {
  76         -- | Mapping from global variables to their vectorised versions.
  77           global_vars           :: VarEnv Var
  78
  79         -- | Purely scalar variables. Code which mentions only these
  80         --   variables doesn't have to be lifted.
  81         , global_scalars        :: VarSet
  82
  83         -- | Exported variables which have a vectorised version.
  84         , global_exported_vars  :: VarEnv (Var, Var)
  85
  86         -- | Mapping from TyCons to their vectorised versions.
  87         --   TyCons which do not have to be vectorised are mapped to themselves.
  88         , global_tycons         :: NameEnv TyCon
  89
  90         -- | Mapping from DataCons to their vectorised versions.
  91         , global_datacons       :: NameEnv DataCon
  92
  93         -- | Mapping from TyCons to their PA dfuns.
  94         , global_pa_funs        :: NameEnv Var
  95
  96         -- | Mapping from TyCons to their PR dfuns.
  97         , global_pr_funs        :: NameEnv Var
  98
  99         -- | Mapping from unboxed TyCons to their boxed versions.
 100         , global_boxed_tycons   :: NameEnv TyCon
 101
 102         -- | External package inst-env & home-package inst-env for class instances.
 103         , global_inst_env       :: (InstEnv, InstEnv)
 104
 105         -- | External package inst-env & home-package inst-env for family instances.
 106         , global_fam_inst_env   :: FamInstEnvs
 107
 108         -- | Hoisted bindings.
 109         , global_bindings       :: [(Var, CoreExpr)]
 110         }
 111
 112
 113 -- | Create an initial global environment
 114 initGlobalEnv :: VectInfo -> (InstEnv, InstEnv) -> FamInstEnvs -> GlobalEnv
 115 initGlobalEnv info instEnvs famInstEnvs
 116         = GlobalEnv
 117         { global_vars          = mapVarEnv snd $ vectInfoVar info
 118         , global_scalars       = emptyVarSet
 119         , global_exported_vars = emptyVarEnv
 120         , global_tycons        = mapNameEnv snd $ vectInfoTyCon info
 121         , global_datacons      = mapNameEnv snd $ vectInfoDataCon info
 122         , global_pa_funs       = mapNameEnv snd $ vectInfoPADFun info
 123         , global_pr_funs       = emptyNameEnv
 124         , global_boxed_tycons  = emptyNameEnv
 125         , global_inst_env      = instEnvs
 126         , global_fam_inst_env  = famInstEnvs
 127         , global_bindings      = []
 128         }
 129
 130
 131
 132 -- Operators on Global Environments -------------------------------------------
 133 -- | Extend the list of global variables in an environment.
 134 extendImportedVarsEnv :: [(Var, Var)] -> GlobalEnv -> GlobalEnv
 135 extendImportedVarsEnv ps genv
 136   = genv { global_vars   = extendVarEnvList (global_vars genv) ps }
 137
 138
 139 -- | Extend the set of scalar variables in an environment.
 140 extendScalars :: [Var] -> GlobalEnv -> GlobalEnv
 141 extendScalars vs genv
 142   = genv { global_scalars = extendVarSetList (global_scalars genv) vs }
 143
 144
 145 -- | Set the list of type family instances in an environment.
 146 setFamInstEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
 147 setFamInstEnv l_fam_inst genv
 148   = genv { global_fam_inst_env = (g_fam_inst, l_fam_inst) }
 149   where (g_fam_inst, _) = global_fam_inst_env genv
 150
 151
 152 -- | Extend the list of type constructors in an environment.
 153 extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
 154 extendTyConsEnv ps genv
 155   = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
 156
 157
 158 -- | Extend the list of data constructors in an environment.
 159 extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
 160 extendDataConsEnv ps genv
 161   = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
 162
 163
 164 -- | Extend the list of PA functions in an environment.
 165 extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
 166 extendPAFunsEnv ps genv
 167   = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
 168
 169
 170 -- | Set the list of PR functions in an environment.
 171 setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
 172 setPRFunsEnv ps genv
 173   = genv { global_pr_funs = mkNameEnv ps }
 174
 175
 176 -- | Set the list of boxed type constructor in an environment.
 177 setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
 178 setBoxedTyConsEnv ps genv
 179   = genv { global_boxed_tycons = mkNameEnv ps }
 180
 181
 182 -- | TODO: What is this for?
 183 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
 184 updVectInfo env tyenv info
 185   = info
 186     { vectInfoVar     = global_exported_vars env
 187     , vectInfoTyCon   = mk_env typeEnvTyCons global_tycons
 188     , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
 189     , vectInfoPADFun  = mk_env typeEnvTyCons global_pa_funs
 190     }
 191   where
 192     mk_env from_tyenv from_env
 193         = mkNameEnv [(name, (from,to))
 194                         | from     <- from_tyenv tyenv
 195                         , let name =  getName from
 196                         , Just to  <- [lookupNameEnv (from_env env) name]]
 197