module Vectorise.Monad.Local
  ( readLEnv
  , setLEnv
  , updLEnv
  , localV
  , closedV
  , getBindName
  , inBind
  , lookupTyVarPA
  , defLocalTyVar
  , defLocalTyVarWithPA
  , localTyVars
  )
where

import GhcPrelude

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)