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         setFamEnv,
  15         extendFamEnv,
  16         extendTyConsEnv,
  17         extendDataConsEnv,
  18         extendPAFunsEnv,
  19         setPRFunsEnv,
  20         setBoxedTyConsEnv,
  21         updVectInfo
  22 ) where
  23 import HscTypes
  24 import InstEnv
  25 import FamInstEnv
  26 import CoreSyn
  27 import TyCon
  28 import DataCon
  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 setFamEnv :: FamInstEnv -> GlobalEnv -> GlobalEnv
 147 setFamEnv 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 extendFamEnv :: [FamInst] -> GlobalEnv -> GlobalEnv
 152 extendFamEnv new genv
 153   = genv { global_fam_inst_env = (g_fam_inst, extendFamInstEnvList l_fam_inst new) }
 154   where (g_fam_inst, l_fam_inst) = global_fam_inst_env genv
 155
 156
 157 -- | Extend the list of type constructors in an environment.
 158 extendTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
 159 extendTyConsEnv ps genv
 160   = genv { global_tycons = extendNameEnvList (global_tycons genv) ps }
 161
 162
 163 -- | Extend the list of data constructors in an environment.
 164 extendDataConsEnv :: [(Name, DataCon)] -> GlobalEnv -> GlobalEnv
 165 extendDataConsEnv ps genv
 166   = genv { global_datacons = extendNameEnvList (global_datacons genv) ps }
 167
 168
 169 -- | Extend the list of PA functions in an environment.
 170 extendPAFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
 171 extendPAFunsEnv ps genv
 172   = genv { global_pa_funs = extendNameEnvList (global_pa_funs genv) ps }
 173
 174
 175 -- | Set the list of PR functions in an environment.
 176 setPRFunsEnv :: [(Name, Var)] -> GlobalEnv -> GlobalEnv
 177 setPRFunsEnv ps genv
 178   = genv { global_pr_funs = mkNameEnv ps }
 179
 180
 181 -- | Set the list of boxed type constructor in an environment.
 182 setBoxedTyConsEnv :: [(Name, TyCon)] -> GlobalEnv -> GlobalEnv
 183 setBoxedTyConsEnv ps genv
 184   = genv { global_boxed_tycons = mkNameEnv ps }
 185
 186
 187 -- | TODO: What is this for?
 188 updVectInfo :: GlobalEnv -> TypeEnv -> VectInfo -> VectInfo
 189 updVectInfo env tyenv info
 190   = info
 191     { vectInfoVar     = global_exported_vars env
 192     , vectInfoTyCon   = mk_env typeEnvTyCons global_tycons
 193     , vectInfoDataCon = mk_env typeEnvDataCons global_datacons
 194     , vectInfoPADFun  = mk_env typeEnvTyCons global_pa_funs
 195     }
 196   where
 197     mk_env from_tyenv from_env
 198         = mkNameEnv [(name, (from,to))
 199                         | from     <- from_tyenv tyenv
 200                         , let name =  getName from
 201                         , Just to  <- [lookupNameEnv (from_env env) name]]
 202