summaryrefslogtreecommitdiff
path: root/compiler/GHC/Stg/Subst.hs
blob: 487b8b8adc11ba666074d9b79402d4254c78dc84 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80


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
import GHC.Utils.Trace

-- | 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 = mkEmptySubst emptyInScopeSet

-- | Constructs a new 'Subst' assuming the variables in the given 'InScopeSet'
-- are in scope.
mkEmptySubst :: InScopeSet -> Subst
mkEmptySubst in_scope = Subst in_scope 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 in_scope env)
  = (new_id, Subst new_in_scope new_env)
  where
    new_id = uniqAway in_scope id
    no_change = new_id == id -- in case nothing shadowed
    new_in_scope = in_scope `extendInScopeSet` new_id
    new_env
      | no_change = delVarEnv env id
      | otherwise = extendVarEnv env id new_id

-- | @substBndrs = runState . traverse (state . substBndr)@
substBndrs :: Traversable f => f Id -> Subst -> (f Id, Subst)
substBndrs = runState . traverse (state . substBndr)

-- | Substitutes an occurrence of an identifier for its counterpart recorded
-- in the 'Subst'.
lookupIdSubst :: HasCallStack => Id -> Subst -> Id
lookupIdSubst id (Subst in_scope env)
  | not (isLocalId id) = id
  | Just id' <- lookupVarEnv env id = id'
  | Just id' <- lookupInScope in_scope id = id'
  | otherwise = warnPprTrace True (text "StgSubst.lookupIdSubst" <+> ppr id $$ ppr in_scope) 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 id (Subst in_scope env)
  | not (isLocalId id) = id
  | Just id' <- lookupVarEnv env id = id'
  | Just id' <- lookupInScope in_scope id = id'
  | otherwise = id

-- | Add the 'Id' to the in-scope set and remove any existing substitutions for
-- it.
extendInScope :: Id -> Subst -> Subst
extendInScope id (Subst in_scope env) = Subst (in_scope `extendInScopeSet` id) 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 new_id (Subst in_scope env)
  = assertPpr (new_id `elemInScopeSet` in_scope) (ppr id <+> ppr new_id $$ ppr in_scope) $
    Subst in_scope (extendVarEnv env id new_id)