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 :: HasDebugCallStack => Id -> Subst -> Id
lookupIdSubst :: HasDebugCallStack => 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 :: HasDebugCallStack => Id -> Subst -> Id
noWarnLookupIdSubst :: HasDebugCallStack => 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)