    {-# LANGUAGE CPP #-}

    module GHC.Stg.Subst where

    #include "HsVersions.h"

    import GHC.Prelude

    import GHC.Types.Id
    import GHC.Types.Var.Env
    import Control.Monad.Trans.State.Strict
    import GHC.Utils.Outputable
    import GHC.Utils.Misc

    -- | A renaming substitution from 'Id's to 'Id's. Like 'RnEnv2', but not
    -- maintaining pairs of substitutions. Like 'GHC.Core.Subst.Subst', but
    -- with the domain being 'Id's instead of entire 'CoreExpr'.
    data Subst = Subst InScopeSet IdSubstEnv

    type IdSubstEnv = IdEnv Id

    -- | @emptySubst = 'mkEmptySubst' 'emptyInScopeSet'@
    emptySubst :: Subst
    emptySubst :: Subst
    emptySubst = InScopeSet -> Subst
    mkEmptySubst InScopeSet
    emptyInScopeSet

    -- | Constructs a new 'Subst' assuming the variables in the given 'InScopeSet'
    -- are in scope.
    mkEmptySubst :: InScopeSet -> Subst
    mkEmptySubst :: InScopeSet -> Subst
    mkEmptySubst InScopeSet
    in_scope = InScopeSet -> IdSubstEnv -> Subst
    Subst InScopeSet
    in_scope IdSubstEnv
    forall a. VarEnv a
    emptyVarEnv

    -- | Substitutes an 'Id' for another one according to the 'Subst' given in a way
    -- that avoids shadowing the 'InScopeSet', returning the result and an updated
    -- 'Subst' that should be used by subsequent substitutions.
    substBndr :: Id -> Subst -> (Id, Subst)
    substBndr :: Id -> Subst -> (Id, Subst)
    substBndr Id
    id (Subst InScopeSet
    in_scope IdSubstEnv
    env)
      = (Id
    new_id, InScopeSet -> IdSubstEnv -> Subst
    Subst InScopeSet
    new_in_scope IdSubstEnv
    new_env)
      where
        new_id :: Id
    new_id = InScopeSet -> Id -> Id
    uniqAway InScopeSet
    in_scope Id
    id
        no_change :: Bool
    no_change = Id
    new_id Id -> Id -> Bool
    forall a. Eq a => a -> a -> Bool
    == Id
    id -- in case nothing shadowed
        new_in_scope :: InScopeSet
    new_in_scope = InScopeSet
    in_scope InScopeSet -> Id -> InScopeSet
    `extendInScopeSet` Id
    new_id
        new_env :: IdSubstEnv
    new_env
          | Bool
    no_change = IdSubstEnv -> Id -> IdSubstEnv
    forall a. VarEnv a -> Id -> VarEnv a
    delVarEnv IdSubstEnv
    env Id
    id
          | Bool
    otherwise = IdSubstEnv -> Id -> Id -> IdSubstEnv
    forall a. VarEnv a -> Id -> a -> VarEnv a
    extendVarEnv IdSubstEnv
    env Id
    id Id
    new_id

    -- | @substBndrs = runState . traverse (state . substBndr)@
    substBndrs :: Traversable f => f Id -> Subst -> (f Id, Subst)
    substBndrs :: forall (f :: * -> *).
    Traversable f =>
    f Id -> Subst -> (f Id, Subst)
    substBndrs = State Subst (f Id) -> Subst -> (f Id, Subst)
    forall s a. State s a -> s -> (a, s)
    runState (State Subst (f Id) -> Subst -> (f Id, Subst))
    -> (f Id -> State Subst (f Id)) -> f Id -> Subst -> (f Id, Subst)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Id -> StateT Subst Identity Id) -> f Id -> State Subst (f Id)
    forall (t :: * -> *) (f :: * -> *) a b.
    (Traversable t, Applicative f) =>
    (a -> f b) -> t a -> f (t b)
    traverse ((Subst -> (Id, Subst)) -> StateT Subst Identity Id
    forall (m :: * -> *) s a. Monad m => (s -> (a, s)) -> StateT s m a
    state ((Subst -> (Id, Subst)) -> StateT Subst Identity Id)
    -> (Id -> Subst -> (Id, Subst)) -> Id -> StateT Subst Identity Id
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Id -> Subst -> (Id, Subst)
    substBndr)

    -- | Substitutes an occurrence of an identifier for its counterpart recorded
    -- in the 'Subst'.
    lookupIdSubst :: HasCallStack => Id -> Subst -> Id
    lookupIdSubst :: HasCallStack => Id -> Subst -> Id
    lookupIdSubst Id
    id (Subst InScopeSet
    in_scope IdSubstEnv
    env)
      | Bool -> Bool
    not (Id -> Bool
    isLocalId Id
    id) = Id
    id
      | Just Id
    id' <- IdSubstEnv -> Id -> Maybe Id
    forall a. VarEnv a -> Id -> Maybe a
    lookupVarEnv IdSubstEnv
    env Id
    id = Id
    id'
      | Just Id
    id' <- InScopeSet -> Id -> Maybe Id
    lookupInScope InScopeSet
    in_scope Id
    id = Id
    id'
      | Bool
    otherwise = WARN( True, text "StgSubst.lookupIdSubst" <+> ppr id $$ ppr in_scope)
                    Id
    id

    -- | Substitutes an occurrence of an identifier for its counterpart recorded
    -- in the 'Subst'. Does not generate a debug warning if the identifier to
    -- to substitute wasn't in scope.
    noWarnLookupIdSubst :: HasCallStack => Id -> Subst -> Id
    noWarnLookupIdSubst :: HasCallStack => Id -> Subst -> Id
    noWarnLookupIdSubst Id
    id (Subst InScopeSet
    in_scope IdSubstEnv
    env)
      | Bool -> Bool
    not (Id -> Bool
    isLocalId Id
    id) = Id
    id
      | Just Id
    id' <- IdSubstEnv -> Id -> Maybe Id
    forall a. VarEnv a -> Id -> Maybe a
    lookupVarEnv IdSubstEnv
    env Id
    id = Id
    id'
      | Just Id
    id' <- InScopeSet -> Id -> Maybe Id
    lookupInScope InScopeSet
    in_scope Id
    id = Id
    id'
      | Bool
    otherwise = Id
    id

    -- | Add the 'Id' to the in-scope set and remove any existing substitutions for
    -- it.
    extendInScope :: Id -> Subst -> Subst
    extendInScope :: Id -> Subst -> Subst
    extendInScope Id
    id (Subst InScopeSet
    in_scope IdSubstEnv
    env) = InScopeSet -> IdSubstEnv -> Subst
    Subst (InScopeSet
    in_scope InScopeSet -> Id -> InScopeSet
    `extendInScopeSet` Id
    id) IdSubstEnv
    env

    -- | Add a substitution for an 'Id' to the 'Subst': you must ensure that the
    -- in-scope set is such that TyCoSubst Note [The substitution invariant]
    -- holds after extending the substitution like this.
    extendSubst :: Id -> Id -> Subst -> Subst
    extendSubst :: Id -> Id -> Subst -> Subst
    extendSubst Id
    id Id
    new_id (Subst InScopeSet
    in_scope IdSubstEnv
    env)
      = ASSERT2( new_id `elemInScopeSet` in_scope, ppr id <+> ppr new_id $$ ppr in_scope )
        InScopeSet -> IdSubstEnv -> Subst
    Subst InScopeSet
    in_scope (IdSubstEnv -> Id -> Id -> IdSubstEnv
    forall a. VarEnv a -> Id -> a -> VarEnv a
    extendVarEnv IdSubstEnv
    env Id
    id Id
    new_id)