Fix a shadowing issue in StgUnarise.wip/andreask/fix_unarise

For I assume performance reasons we don't record no-op replacements
during unarise. This lead to problems with code like this:

    f = \(Eta_B0 :: VoidType) x1 x2 ->
       ... let foo = \(Eta_B0 :: LiftedType) -> g x y Eta_B0
           in ...

Here we would record the outer Eta_B0 as void rep, but would not
shadow Eta_B0 inside `foo` because this arg is single-rep and so
doesn't need to replaced. But this means when looking at occurence
sites we would check the env and assume it's void rep based on the
entry we made for the (no longer in scope) outer `Eta_B0`.

Fixes #21396 and the ticket has a few more details.

1 files changed, 23 insertions, 1 deletions

diff --git a/compiler/GHC/Stg/Unarise.hs b/compiler/GHC/Stg/Unarise.hs
index d46719298e..23c2646f73 100644
--- a/compiler/GHC/Stg/Unarise.hs
+++ b/compiler/GHC/Stg/Unarise.hs
@@ -186,6 +186,18 @@ So we pass type arguments of the DataCon's TyCon in StgConApp to decide what
 layout to use. Note that unlifted values can't be let-bound, so we don't need
 types in StgRhsCon.
 
+Note [UnariseEnv]
+~~~~~~~~~~~~~~~~~~
+At any variable occurrence 'v',
+* If the UnariseEnv has a binding for 'v', the binding says what 'v' is bound to
+* If not, 'v' stands just for itself.
+
+Most variables are unaffected by unarisation, and (for efficiency) we don't put
+them in the UnariseEnv at all.  But NB: when we go under a binding for 'v' we must
+remember to delete 'v' from the UnariseEnv, lest occurrences of 'v' see the outer
+binding for the variable (#21396).
+
+
 Note [UnariseEnv can map to literals]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 To avoid redundant case expressions when unarising unboxed sums, UnariseEnv
@@ -305,6 +317,8 @@ instance Outputable UnariseVal where
   ppr (UnaryVal arg)   = text "UnaryVal" <+> ppr arg
 
 -- | Extend the environment, checking the UnariseEnv invariant.
+-- The id is mapped to one or more things.
+-- See Note [UnariseEnv]
 extendRho :: UnariseEnv -> Id -> UnariseVal -> UnariseEnv
 extendRho rho x (MultiVal args)
   = assert (all (isNvUnaryType . stgArgType) args)
@@ -312,6 +326,14 @@ extendRho rho x (MultiVal args)
 extendRho rho x (UnaryVal val)
   = assert (isNvUnaryType (stgArgType val))
     extendVarEnv rho x (UnaryVal val)
+-- Properly shadow things from an outer scope.
+-- See Note [UnariseEnv]
+
+-- The id stands for itself so we don't record a mapping.
+-- See Note [UnariseEnv]
+extendRhoWithoutValue :: UnariseEnv -> Id -> UnariseEnv
+extendRhoWithoutValue rho x = delVarEnv rho x
+
 
 --------------------------------------------------------------------------------
 
@@ -754,7 +776,7 @@ unariseArgBinder is_con_arg rho x =
       -> do x' <- mkId (mkFastString "us") (primRepToType rep)
             return (extendRho rho x (MultiVal [StgVarArg x']), [x'])
       | otherwise
-      -> return (rho, [x])
+      -> return (extendRhoWithoutValue rho x, [x])
 
     reps -> do
       xs <- mkIds (mkFastString "us") (map primRepToType reps)