compiler/vectorise/Vectorise/Monad/Local.hs

   1
   2 module Vectorise.Monad.Local (
   3         readLEnv,
   4         setLEnv,
   5         updLEnv,
   6         localV,
   7         closedV,
   8         getBindName,
   9         inBind,
  10         lookupTyVarPA,
  11         defLocalTyVar,
  12         defLocalTyVarWithPA,
  13         localTyVars
  14 ) where
  15 import Vectorise.Monad.Base
  16 import Vectorise.Env
  17
  18 import CoreSyn
  19 import Name
  20 import VarEnv
  21 import Var
  22 import FastString
  23
  24 -- Local Environment ----------------------------------------------------------
  25 -- | Project something from the local environment.
  26 readLEnv :: (LocalEnv -> a) -> VM a
  27 readLEnv f      = VM $ \_ genv lenv -> return (Yes genv lenv (f lenv))
  28
  29
  30 -- | Set the local environment.
  31 setLEnv :: LocalEnv -> VM ()
  32 setLEnv lenv    = VM $ \_ genv _ -> return (Yes genv lenv ())
  33
  34
  35 -- | Update the enviroment using the provided function.
  36 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
  37 updLEnv f       = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
  38
  39
  40 -- | Perform a computation in its own local environment.
  41 --      This does not alter the environment of the current state.
  42 localV :: VM a -> VM a
  43 localV p
  44  = do   env <- readLEnv id
  45         x   <- p
  46         setLEnv env
  47         return x
  48
  49
  50 -- | Perform a computation in an empty local environment.
  51 closedV :: VM a -> VM a
  52 closedV p
  53  = do   env <- readLEnv id
  54         setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
  55         x   <- p
  56         setLEnv env
  57         return x
  58
  59
  60 -- | Get the name of the local binding currently being vectorised.
  61 getBindName :: VM FastString
  62 getBindName = readLEnv local_bind_name
  63
  64
  65 -- | Run a vectorisation computation in a local environment,
  66 --   with this id set as the current binding.
  67 inBind :: Id -> VM a -> VM a
  68 inBind id p
  69   = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
  70        p
  71
  72
  73 -- | Lookup a PA tyvars from the local environment.
  74 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
  75 lookupTyVarPA tv
  76         = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
  77
  78
  79 -- | Add a tyvar to the local environment.
  80 defLocalTyVar :: TyVar -> VM ()
  81 defLocalTyVar tv = updLEnv $ \env ->
  82   env { local_tyvars   = tv : local_tyvars env
  83       , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
  84       }
  85
  86 -- | Add mapping between a tyvar and pa dictionary to the local environment.
  87 defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
  88 defLocalTyVarWithPA tv pa = updLEnv $ \env ->
  89   env { local_tyvars   = tv : local_tyvars env
  90       , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
  91       }
  92
  93
  94 -- | Get the set of tyvars from the local environment.
  95 localTyVars :: VM [TyVar]
  96 localTyVars = readLEnv (reverse . local_tyvars)
  97
  98