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