compiler/vectorise/Vectorise/Monad/Global.hs

   1
   2 module Vectorise.Monad.Global (
   3         readGEnv,
   4         setGEnv,
   5         updGEnv,
   6
   7         -- * Vars
   8         defGlobalVar,
   9
  10         -- * Scalars
  11         globalScalars,
  12
  13         -- * TyCons
  14         lookupTyCon,
  15         lookupBoxedTyCon,
  16         defTyCon,
  17
  18         -- * Datacons
  19         lookupDataCon,
  20         defDataCon,
  21
  22         -- * PA Dictionaries
  23         lookupTyConPA,
  24         defTyConPA,
  25         defTyConPAs,
  26
  27         -- * PR Dictionaries
  28         lookupTyConPR
  29 ) where
  30 import Vectorise.Monad.Base
  31 import Vectorise.Env
  32 import TyCon
  33 import DataCon
  34 import NameEnv
  35 import Var
  36 import VarEnv
  37 import VarSet
  38
  39
  40 -- Global Environment ---------------------------------------------------------
  41 -- | Project something from the global environment.
  42 readGEnv :: (GlobalEnv -> a) -> VM a
  43 readGEnv f      = VM $ \_ genv lenv -> return (Yes genv lenv (f genv))
  44
  45
  46 -- | Set the value of the global environment.
  47 setGEnv :: GlobalEnv -> VM ()
  48 setGEnv genv    = VM $ \_ _ lenv -> return (Yes genv lenv ())
  49
  50
  51 -- | Update the global environment using the provided function.
  52 updGEnv :: (GlobalEnv -> GlobalEnv) -> VM ()
  53 updGEnv f       = VM $ \_ genv lenv -> return (Yes (f genv) lenv ())
  54
  55
  56 -- Vars -----------------------------------------------------------------------
  57 -- | Add a mapping between a global var and its vectorised version to the state.
  58 defGlobalVar :: Var -> Var -> VM ()
  59 defGlobalVar v v' = updGEnv $ \env ->
  60   env { global_vars = extendVarEnv (global_vars env) v v'
  61       , global_exported_vars = upd (global_exported_vars env)
  62       }
  63   where
  64     upd env | isExportedId v = extendVarEnv env v (v, v')
  65             | otherwise      = env
  66
  67
  68 -- Scalars --------------------------------------------------------------------
  69 -- | Get the set of global scalar variables.
  70 globalScalars :: VM VarSet
  71 globalScalars
  72         = readGEnv global_scalars
  73
  74
  75 -- TyCons ---------------------------------------------------------------------
  76 -- | Lookup the vectorised version of a `TyCon` from the global environment.
  77 lookupTyCon :: TyCon -> VM (Maybe TyCon)
  78 lookupTyCon tc
  79   | isUnLiftedTyCon tc || isTupleTyCon tc
  80   = return (Just tc)
  81
  82   | otherwise
  83   = readGEnv $ \env -> lookupNameEnv (global_tycons env) (tyConName tc)
  84
  85
  86 -- | Lookup the vectorised version of a boxed `TyCon` from the global environment.
  87 lookupBoxedTyCon :: TyCon -> VM (Maybe TyCon)
  88 lookupBoxedTyCon tc
  89         = readGEnv $ \env -> lookupNameEnv (global_boxed_tycons env)
  90                                            (tyConName tc)
  91
  92
  93 -- | Add a mapping between plain and vectorised `TyCon`s to the global environment.
  94 defTyCon :: TyCon -> TyCon -> VM ()
  95 defTyCon tc tc' = updGEnv $ \env ->
  96   env { global_tycons = extendNameEnv (global_tycons env) (tyConName tc) tc' }
  97
  98
  99 -- DataCons -------------------------------------------------------------------
 100 -- | Lookup the vectorised version of a `DataCon` from the global environment.
 101 lookupDataCon :: DataCon -> VM (Maybe DataCon)
 102 lookupDataCon dc
 103   | isTupleTyCon (dataConTyCon dc)
 104   = return (Just dc)
 105
 106   | otherwise
 107   = readGEnv $ \env -> lookupNameEnv (global_datacons env) (dataConName dc)
 108
 109
 110 -- | Add the mapping between plain and vectorised `DataCon`s to the global environment.
 111 defDataCon :: DataCon -> DataCon -> VM ()
 112 defDataCon dc dc' = updGEnv $ \env ->
 113   env { global_datacons = extendNameEnv (global_datacons env) (dataConName dc) dc' }
 114
 115
 116 -- PA dictionaries ------------------------------------------------------------
 117 -- | Lookup a PA `TyCon` from the global environment.
 118 lookupTyConPA :: TyCon -> VM (Maybe Var)
 119 lookupTyConPA tc
 120         = readGEnv $ \env -> lookupNameEnv (global_pa_funs env) (tyConName tc)
 121
 122
 123 -- | Add a mapping between a PA TyCon and is vectorised version to the global environment.
 124 defTyConPA :: TyCon -> Var -> VM ()
 125 defTyConPA tc pa = updGEnv $ \env ->
 126   env { global_pa_funs = extendNameEnv (global_pa_funs env) (tyConName tc) pa }
 127
 128
 129 -- | Add several mapping between PA TyCons and their vectorised versions to the global environment.
 130 defTyConPAs :: [(TyCon, Var)] -> VM ()
 131 defTyConPAs ps = updGEnv $ \env ->
 132   env { global_pa_funs = extendNameEnvList (global_pa_funs env)
 133                                            [(tyConName tc, pa) | (tc, pa) <- ps] }
 134
 135
 136 -- PR Dictionaries ------------------------------------------------------------
 137 lookupTyConPR :: TyCon -> VM (Maybe Var)
 138 lookupTyConPR tc = readGEnv $ \env -> lookupNameEnv (global_pr_funs env) (tyConName tc)
 139
 140