Finish breaking up VectBuiltIn and VectMonad, and add comments

[ghc-hetmet.git] / compiler / vectorise / Vectorise / Monad / Local.hs

diff --git a/compiler/vectorise/Vectorise/Monad/Local.hs b/compiler/vectorise/Vectorise/Monad/Local.hs
new file mode 100644 (file)
index 0000000..7d8b493
--- /dev/null
+++ b/compiler/vectorise/Vectorise/Monad/Local.hs
@@ -0,0 +1,100 @@
+
+module Vectorise.Monad.Local ( 
+       readLEnv,
+       setLEnv,
+       updLEnv,
+       localV,
+       closedV,
+       getBindName,
+       inBind,
+       lookupTyVarPA,
+       defLocalTyVar,
+       defLocalTyVarWithPA,
+       localTyVars
+) where
+import Vectorise.Monad.Base
+import Vectorise.Env
+
+import CoreSyn
+import Id
+import OccName
+import Name
+import VarEnv
+import Var
+import FastString
+
+-- Local Environment ----------------------------------------------------------
+-- | Project something from the local environment.
+readLEnv :: (LocalEnv -> a) -> VM a
+readLEnv f     = VM $ \_ genv lenv -> return (Yes genv lenv (f lenv))
+
+
+-- | Set the local environment.
+setLEnv :: LocalEnv -> VM ()
+setLEnv lenv   = VM $ \_ genv _ -> return (Yes genv lenv ())
+
+
+-- | Update the enviroment using the provided function.
+updLEnv :: (LocalEnv -> LocalEnv) -> VM ()
+updLEnv f      = VM $ \_ genv lenv -> return (Yes genv (f lenv) ())
+
+
+-- | Perform a computation in its own local environment.
+--     This does not alter the environment of the current state.
+localV :: VM a -> VM a
+localV p 
+ = do  env <- readLEnv id
+       x   <- p
+       setLEnv env
+       return x
+
+
+-- | Perform a computation in an empty local environment.
+closedV :: VM a -> VM a
+closedV p 
+ = do  env <- readLEnv id
+       setLEnv (emptyLocalEnv { local_bind_name = local_bind_name env })
+       x   <- p
+       setLEnv env
+       return x
+
+
+-- | Get the name of the local binding currently being vectorised.
+getBindName :: VM FastString
+getBindName = readLEnv local_bind_name
+
+
+-- | Run a vectorisation computation in a local environment, 
+--   with this id set as the current binding.
+inBind :: Id -> VM a -> VM a
+inBind id p
+  = do updLEnv $ \env -> env { local_bind_name = occNameFS (getOccName id) }
+       p
+
+
+-- | Lookup a PA tyvars from the local environment.
+lookupTyVarPA :: Var -> VM (Maybe CoreExpr)
+lookupTyVarPA tv 
+       = readLEnv $ \env -> lookupVarEnv (local_tyvar_pa env) tv
+
+
+-- | Add a tyvar to the local environment.
+defLocalTyVar :: TyVar -> VM ()
+defLocalTyVar tv = updLEnv $ \env ->
+  env { local_tyvars   = tv : local_tyvars env
+      , local_tyvar_pa = local_tyvar_pa env `delVarEnv` tv
+      }
+
+-- | Add mapping between a tyvar and pa dictionary to the local environment.
+defLocalTyVarWithPA :: TyVar -> CoreExpr -> VM ()
+defLocalTyVarWithPA tv pa = updLEnv $ \env ->
+  env { local_tyvars   = tv : local_tyvars env
+      , local_tyvar_pa = extendVarEnv (local_tyvar_pa env) tv pa
+      }
+
+
+-- | Get the set of tyvars from the local environment.
+localTyVars :: VM [TyVar]
+localTyVars = readLEnv (reverse . local_tyvars)
+
+