Fix shadowing lacuna in OccurAnal

Issue #22623 demonstrated another lacuna in the implementation
of wrinkle (BS3) in Note [The binder-swap substitution] in
the occurrence analyser.

I was failing to add TyVar lambda binders using
addInScope/addOneInScope and that led to a totally bogus binder-swap
transformation.

Very easy to fix.

(cherry picked from commit e193e53790dd5886feea3cf4c9c17625d188291b)

4 files changed, 125 insertions, 23 deletions

diff --git a/compiler/GHC/Core/Opt/OccurAnal.hs b/compiler/GHC/Core/Opt/OccurAnal.hs
index 3f9510cb60..da9fedae96 100644
--- a/compiler/GHC/Core/Opt/OccurAnal.hs
+++ b/compiler/GHC/Core/Opt/OccurAnal.hs
@@ -1725,7 +1725,7 @@ occAnalRhs :: OccEnv -> RecFlag -> Maybe JoinArity
            -> CoreExpr   -- RHS
            -> (UsageDetails, CoreExpr)
 occAnalRhs env is_rec mb_join_arity rhs
-  = case occAnalLamOrRhs env bndrs body of { (body_usage, bndrs', body') ->
+  = case occAnalLamOrRhs env1 bndrs body of { (body_usage, bndrs', body') ->
     let final_bndrs | isRec is_rec = bndrs'
                     | otherwise    = markJoinOneShots mb_join_arity bndrs'
                -- For a /non-recursive/ join point we can mark all
@@ -1737,6 +1737,7 @@ occAnalRhs env is_rec mb_join_arity rhs
     in (rhs_usage, mkLams final_bndrs body') }
   where
     (bndrs, body) = collectBinders rhs
+    env1          = addInScope env bndrs
 
 occAnalUnfolding :: OccEnv
                  -> RecFlag
@@ -2005,7 +2006,7 @@ partially applying lambdas. See the calls to zapLamBndrs in
 
 occAnal env expr@(Lam _ _)
   = -- See Note [Occurrence analysis for lambda binders]
-    case occAnalLamOrRhs env bndrs body of { (usage, tagged_bndrs, body') ->
+    case occAnalLamOrRhs env1 bndrs body of { (usage, tagged_bndrs, body') ->
     let
         expr'       = mkLams tagged_bndrs body'
         usage1      = markAllNonTail usage
@@ -2015,6 +2016,7 @@ occAnal env expr@(Lam _ _)
     (final_usage, expr') }
   where
     (bndrs, body) = collectBinders expr
+    env1          = addInScope env bndrs
 
 occAnal env (Case scrut bndr ty alts)
   = case occAnal (scrutCtxt env alts) scrut of { (scrut_usage, scrut') ->
@@ -2284,12 +2286,13 @@ data OccEnv
 
            -- See Note [The binder-swap substitution]
            -- If  x :-> (y, co)  is in the env,
-           -- then please replace x by (y |> sym mco)
-           -- Invariant of course: idType x = exprType (y |> sym mco)
-           , occ_bs_env  :: VarEnv (OutId, MCoercion)
-           , occ_bs_rng  :: VarSet   -- Vars free in the range of occ_bs_env
+           -- then please replace x by (y |> mco)
+           -- Invariant of course: idType x = exprType (y |> mco)
+           , occ_bs_env  :: !(IdEnv (OutId, MCoercion))
                    -- Domain is Global and Local Ids
                    -- Range is just Local Ids
+           , occ_bs_rng  :: !VarSet
+                   -- Vars (TyVars and Ids) free in the range of occ_bs_env
     }
 
 
@@ -2578,25 +2581,29 @@ Some tricky corners:
 
 (BS3) We need care when shadowing.  Suppose [x :-> b] is in occ_bs_env,
       and we encounter:
-         - \x. blah
-           Here we want to delete the x-binding from occ_bs_env
-
-         - \b. blah
-           This is harder: we really want to delete all bindings that
-           have 'b' free in the range.  That is a bit tiresome to implement,
-           so we compromise.  We keep occ_bs_rng, which is the set of
-           free vars of rng(occc_bs_env).  If a binder shadows any of these
-           variables, we discard all of occ_bs_env.  Safe, if a bit
-           brutal.  NB, however: the simplifer de-shadows the code, so the
-           next time around this won't happen.
+         (i) \x. blah
+             Here we want to delete the x-binding from occ_bs_env
+
+         (ii) \b. blah
+              This is harder: we really want to delete all bindings that
+              have 'b' free in the range.  That is a bit tiresome to implement,
+              so we compromise.  We keep occ_bs_rng, which is the set of
+              free vars of rng(occc_bs_env).  If a binder shadows any of these
+              variables, we discard all of occ_bs_env.  Safe, if a bit
+              brutal.  NB, however: the simplifer de-shadows the code, so the
+              next time around this won't happen.
 
       These checks are implemented in addInScope.
-
-      The occurrence analyser itself does /not/ do cloning. It could, in
-      principle, but it'd make it a bit more complicated and there is no
-      great benefit. The simplifer uses cloning to get a no-shadowing
-      situation, the care-when-shadowing behaviour above isn't needed for
-      long.
+      (i) is needed only for Ids, but (ii) is needed for tyvars too (#22623)
+      because if occ_bs_env has [x :-> ...a...] where `a` is a tyvar, we
+      must not replace `x` by `...a...` under /\a. ...x..., or similarly
+      under a case pattern match that binds `a`.
+
+      An alternative would be for the occurrence analyser to do cloning as
+      it goes.  In principle it could do so, but it'd make it a bit more
+      complicated and there is no great benefit. The simplifer uses
+      cloning to get a no-shadowing situation, the care-when-shadowing
+      behaviour above isn't needed for long.
 
 (BS4) The domain of occ_bs_env can include GlobaIds.  Eg
          case M.foo of b { alts }
diff --git a/testsuite/tests/simplCore/should_compile/T22623.hs b/testsuite/tests/simplCore/should_compile/T22623.hs
new file mode 100644
index 0000000000..8cd1004ecf
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T22623.hs
@@ -0,0 +1,34 @@
+{-# LANGUAGE GHC2021 #-}
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+
+module T22623 where
+
+import T22623a
+
+type BindNonEmptyList :: NonEmpty -> NonEmpty -> [Q]
+type family BindNonEmptyList (x :: NonEmpty) (y :: NonEmpty) :: [Q] where
+  BindNonEmptyList ('(:|) a as) c = Tail c ++ Foldr2 a c as
+
+sBindNonEmptyList ::
+  forall (t :: NonEmpty)
+         (c :: NonEmpty). SNonEmpty t -> SNonEmpty c -> SList (BindNonEmptyList t c :: [Q])
+sBindNonEmptyList
+  ((:%|) (sA :: SQ a) (sAs :: SList as)) (sC :: SNonEmpty c)
+  = let
+      sMyHead :: SNonEmpty c -> SQ (MyHead a c)
+      sMyHead ((:%|) x _) = x
+
+      sFoldr :: forall t. SList t -> SList (Foldr2 a c t)
+      sFoldr SNil = SNil
+      sFoldr (SCons _ sYs) = SCons (sMyHead sC) (sFoldr sYs)
+
+      sF :: Id (SLambda (ConstSym1 c))
+      sF = SLambda (const sC)
+
+      sBs :: SList (Tail c)
+      _ :%| sBs = applySing sF sA
+    in
+      sBs %++ sFoldr sAs
diff --git a/testsuite/tests/simplCore/should_compile/T22623a.hs b/testsuite/tests/simplCore/should_compile/T22623a.hs
new file mode 100644
index 0000000000..c1568a7dad
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T22623a.hs
@@ -0,0 +1,60 @@
+{-# LANGUAGE GHC2021 #-}
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE TypeFamilies #-}
+module T22623a where
+
+import Data.Kind
+
+type Id :: Type -> Type
+type family Id x
+type instance Id x = x
+
+data Q
+data SQ (x :: Q)
+
+data NonEmpty where
+  (:|) :: Q -> [Q] -> NonEmpty
+
+type Tail :: NonEmpty -> [Q]
+type family Tail y where
+  Tail ('(:|) _ y) = y
+type MyHead :: Q -> NonEmpty -> Q
+type family MyHead x y where
+  MyHead _ ('(:|) c _) = c
+
+type SList :: [Q] -> Type
+data SList z where
+  SNil  :: SList '[]
+  SCons :: SQ x -> SList xs -> SList (x:xs)
+
+type SNonEmpty :: NonEmpty -> Type
+data SNonEmpty z where
+  (:%|) :: SQ x -> SList xs -> SNonEmpty (x :| xs)
+
+data TyFun
+type F = TyFun -> Type
+
+type Apply :: F -> Q -> NonEmpty
+type family Apply f x
+
+type ConstSym1 :: NonEmpty -> F
+data ConstSym1 (x :: NonEmpty) :: F
+type instance Apply (ConstSym1 x) _ = x
+
+type SLambda :: F -> Type
+newtype SLambda (f :: F) =
+  SLambda { applySing :: forall t. SQ t -> SNonEmpty (f `Apply` t) }
+
+type Foldr2 :: Q -> NonEmpty -> [Q] -> [Q]
+type family Foldr2 a c x where
+  Foldr2 _ _ '[] = '[]
+  Foldr2 a c (_:ys) = MyHead a c : Foldr2 a c ys
+
+type (++) :: [Q] -> [Q] -> [Q]
+type family (++) xs ys where
+  (++) '[] ys = ys
+  (++) ('(:) x xs) ys = '(:) x (xs ++ ys)
+
+(%++) :: forall (x :: [Q]) (y :: [Q]). SList x -> SList y -> SList (x ++ y)
+(%++) SNil sYs = sYs
+(%++) (SCons sX sXs) sYs = SCons sX (sXs %++ sYs)
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index 5659ccf9b7..db09b00853 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -367,3 +367,4 @@ test('T20200', normal, compile, [''])
 test('T20820',  normal, compile, ['-O0'])
 test('T22491', normal, compile, ['-O2'])
 test('T22662', normal, compile, [''])
+test('T22623', normal, multimod_compile, ['T22623', '-O -v0'])