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