vectorise/Vectorise/Monad/Local.hs

    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 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 environment 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)