    module GHC.Stg.Subst where

    import GHC.Prelude

    import GHC.Types.Id
    import GHC.Types.Var.Env
    import GHC.Utils.Monad.State.Strict

    import GHC.Utils.Outputable
    import GHC.Utils.Misc
    import GHC.Utils.Panic

    -- TODO: This code might make folly of the work done in CorePrep where
    -- we clone local ids in order to ensure *all* local binders are unique.
    -- It's my understanding that here we use "the rapier"/uniqAway which makes up
    -- uniques based on the ids in scope. Which can give the same unique to different
    -- binders as long as they are in different scopes. A guarantee which isn't
    -- strong enough for code generation in general. See Note [CorePrep Overview].

    -- | 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 -> State Subst Id) -> f Id -> State Subst (f Id)
    forall (t :: * -> *) (f :: * -> *) a b.
    (Traversable t, Applicative f) =>
    (a -> f b) -> t a -> f (t b)
    forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> f a -> f (f b)
    traverse ((Subst -> (Id, Subst)) -> State Subst Id
    forall s a. (s -> (a, s)) -> State s a
    state ((Subst -> (Id, Subst)) -> State Subst Id)
    -> (Id -> Subst -> (Id, Subst)) -> Id -> State Subst 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 = Bool -> String -> SDoc -> Id -> Id
    forall a. HasCallStack => Bool -> String -> SDoc -> a -> a
    warnPprTrace Bool
    True String
    "StgSubst.lookupIdSubst" (Id -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr Id
    id SDoc -> SDoc -> SDoc
    forall doc. IsDoc doc => doc -> doc -> doc
    $$ InScopeSet -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr InScopeSet
    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)
      = Bool -> SDoc -> Subst -> Subst
    forall a. HasCallStack => Bool -> SDoc -> a -> a
    assertPpr (Id
    new_id Id -> InScopeSet -> Bool
    `elemInScopeSet` InScopeSet
    in_scope) (Id -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr Id
    id SDoc -> SDoc -> SDoc
    forall doc. IsLine doc => doc -> doc -> doc
    <+> Id -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr Id
    new_id SDoc -> SDoc -> SDoc
    forall doc. IsDoc doc => doc -> doc -> doc
    $$ InScopeSet -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr InScopeSet
    in_scope) (Subst -> Subst) -> Subst -> Subst
    forall a b. (a -> b) -> a -> b
    $
        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)