2 module Vectorise.Monad.Local (
15 import Vectorise.Monad.Base
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))
30 -- | Set the local environment.
31 setLEnv :: LocalEnv -> VM ()
32 setLEnv lenv = VM $ \_ genv _ -> return (Yes genv lenv ())
35 -- | Update the enviroment using the provided function.
36 updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
37 updLEnv f = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
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
44 = do env <- readLEnv id
50 -- | Perform a computation in an empty local environment.
51 closedV :: VM a -> VM a
53 = do env <- readLEnv id
54 setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
60 -- | Get the name of the local binding currently being vectorised.
61 getBindName :: VM FastString
62 getBindName = readLEnv local_bind_name
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
69 = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
73 -- | Lookup a PA tyvars from the local environment.
74 lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
76 = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
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
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
94 -- | Get the set of tyvars from the local environment.
95 localTyVars :: VM [TyVar]
96 localTyVars = readLEnv (reverse . local_tyvars)