Fix binder-swap bug

This patch fixes #21229 properly, by avoiding doing a
binder-swap on dictionary Ids.  This is pretty subtle, and explained
in Note [Care with binder-swap on dictionaries].

Test is already in simplCore/should_run/T21229

This allows us to restore a feature to the specialiser that we had
to revert: see Note [Specialising polymorphic dictionaries].
(This is done in a separate patch.)

I also modularised things, using a new function scrutBinderSwap_maybe
in all the places where we are (effectively) doing a binder-swap,
notably

* Simplify.Iteration.addAltUnfoldings
* SpecConstr.extendCaseBndrs

In Simplify.Iteration.addAltUnfoldings I also eliminated a guard
    Many <- idMult case_bndr
because we concluded, in #22123, that it was doing no good.

7 files changed, 162 insertions, 106 deletions

diff --git a/compiler/GHC/Core/Opt/OccurAnal.hs b/compiler/GHC/Core/Opt/OccurAnal.hs
index 5e59e149a9..c60dbfc39d 100644
--- a/compiler/GHC/Core/Opt/OccurAnal.hs
+++ b/compiler/GHC/Core/Opt/OccurAnal.hs
@@ -19,7 +19,7 @@ core expression with (hopefully) improved usage information.
 module GHC.Core.Opt.OccurAnal (
     occurAnalysePgm,
     occurAnalyseExpr,
-    zapLambdaBndrs
+    zapLambdaBndrs, scrutBinderSwap_maybe
   ) where
 
 import GHC.Prelude hiding ( head, init, last, tail )
@@ -27,11 +27,12 @@ import GHC.Prelude hiding ( head, init, last, tail )
 import GHC.Core
 import GHC.Core.FVs
 import GHC.Core.Utils   ( exprIsTrivial, isDefaultAlt, isExpandableApp,
-                          stripTicksTopE, mkTicks )
+                          mkCastMCo, mkTicks )
 import GHC.Core.Opt.Arity   ( joinRhsArity, isOneShotBndr )
 import GHC.Core.Coercion
+import GHC.Core.Predicate   ( isDictId )
 import GHC.Core.Type
-import GHC.Core.TyCo.FVs( tyCoVarsOfMCo )
+import GHC.Core.TyCo.FVs    ( tyCoVarsOfMCo )
 
 import GHC.Data.Maybe( isJust, orElse )
 import GHC.Data.Graph.Directed ( SCC(..), Node(..)
@@ -2464,8 +2465,8 @@ 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)
+           -- then please replace x by (y |> mco)
+           -- Invariant of course: idType x = exprType (y |> mco)
            , occ_bs_env  :: !(VarEnv (OutId, MCoercion))
            , occ_bs_rng  :: !VarSet   -- Vars free in the range of occ_bs_env
                    -- Domain is Global and Local Ids
@@ -2671,7 +2672,7 @@ The binder-swap is implemented by the occ_bs_env field of OccEnv.
 There are two main pieces:
 
 * Given    case x |> co of b { alts }
-  we add [x :-> (b, co)] to the occ_bs_env environment; this is
+  we add [x :-> (b, sym co)] to the occ_bs_env environment; this is
   done by addBndrSwap.
 
 * Then, at an occurrence of a variable, we look up in the occ_bs_env
@@ -2739,30 +2740,8 @@ Some tricky corners:
 (BS5) We have to apply the occ_bs_env substitution uniformly,
       including to (local) rules and unfoldings.
 
-Historical note
----------------
-We used to do the binder-swap transformation by introducing
-a proxy let-binding, thus;
-
-   case x of b { pi -> ri }
-      ==>
-   case x of b { pi -> let x = b in ri }
-
-But that had two problems:
-
-1. If 'x' is an imported GlobalId, we'd end up with a GlobalId
-   on the LHS of a let-binding which isn't allowed.  We worked
-   around this for a while by "localising" x, but it turned
-   out to be very painful #16296,
-
-2. In CorePrep we use the occurrence analyser to do dead-code
-   elimination (see Note [Dead code in CorePrep]).  But that
-   occasionally led to an unlifted let-binding
-       case x of b { DEFAULT -> let x::Int# = b in ... }
-   which disobeys one of CorePrep's output invariants (no unlifted
-   let-bindings) -- see #5433.
-
-Doing a substitution (via occ_bs_env) is much better.
+(BS6) We must be very careful with dictionaries.
+      See Note [Care with binder-swap on dictionaries]
 
 Note [Case of cast]
 ~~~~~~~~~~~~~~~~~~~
@@ -2772,6 +2751,54 @@ We'd like to eliminate the inner case.  That is the motivation for
 equation (2) in Note [Binder swap].  When we get to the inner case, we
 inline x, cancel the casts, and away we go.
 
+Note [Care with binder-swap on dictionaries]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+This Note explains why we need isDictId in scrutBinderSwap_maybe.
+Consider this tricky example (#21229, #21470):
+
+  class Sing (b :: Bool) where sing :: Bool
+  instance Sing 'True  where sing = True
+  instance Sing 'False where sing = False
+
+  f :: forall a. Sing a => blah
+
+  h = \ @(a :: Bool) ($dSing :: Sing a)
+      let the_co =  Main.N:Sing[0] <a> :: Sing a ~R# Bool
+      case ($dSing |> the_co) of wild
+        True  -> f @'True (True |> sym the_co)
+        False -> f @a     dSing
+
+Now do a binder-swap on the case-expression:
+
+  h = \ @(a :: Bool) ($dSing :: Sing a)
+      let the_co =  Main.N:Sing[0] <a> :: Sing a ~R# Bool
+      case ($dSing |> the_co) of wild
+        True  -> f @'True (True |> sym the_co)
+        False -> f @a     (wild |> sym the_co)
+
+And now substitute `False` for `wild` (since wild=False in the False branch):
+
+  h = \ @(a :: Bool) ($dSing :: Sing a)
+      let the_co =  Main.N:Sing[0] <a> :: Sing a ~R# Bool
+      case ($dSing |> the_co) of wild
+        True  -> f @'True (True  |> sym the_co)
+        False -> f @a     (False |> sym the_co)
+
+And now we have a problem.  The specialiser will specialise (f @a d)a (for all
+vtypes a and dictionaries d!!) with the dictionary (False |> sym the_co), using
+Note [Specialising polymorphic dictionaries] in GHC.Core.Opt.Specialise.
+
+The real problem is the binder-swap.  It swaps a dictionary variable $dSing
+(of kind Constraint) for a term variable wild (of kind Type).  And that is
+dangerous: a dictionary is a /singleton/ type whereas a general term variable is
+not.  In this particular example, Bool is most certainly not a singleton type!
+
+Conclusion:
+  for a /dictionary variable/ do not perform
+  the clever cast version of the binder-swap
+
+Hence the subtle isDictId in scrutBinderSwap_maybe.
+
 Note [Zap case binders in proxy bindings]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 From the original
@@ -2786,8 +2813,87 @@ binding x = cb.  See #5028.
 NB: the OccInfo on /occurrences/ really doesn't matter much; the simplifier
 doesn't use it. So this is only to satisfy the perhaps-over-picky Lint.
 
+-}
+
+addBndrSwap :: OutExpr -> Id -> OccEnv -> OccEnv
+-- See Note [The binder-swap substitution]
+addBndrSwap scrut case_bndr
+            env@(OccEnv { occ_bs_env = swap_env, occ_bs_rng = rng_vars })
+  | Just (scrut_var, mco) <- scrutBinderSwap_maybe scrut
+  , scrut_var /= case_bndr
+      -- Consider: case x of x { ... }
+      -- Do not add [x :-> x] to occ_bs_env, else lookupBndrSwap will loop
+  = env { occ_bs_env = extendVarEnv swap_env scrut_var (case_bndr', mco)
+        , occ_bs_rng = rng_vars `extendVarSet` case_bndr'
+                       `unionVarSet` tyCoVarsOfMCo mco }
+
+  | otherwise
+  = env
+  where
+    case_bndr' = zapIdOccInfo case_bndr
+                 -- See Note [Zap case binders in proxy bindings]
+
+scrutBinderSwap_maybe :: OutExpr -> Maybe (OutVar, MCoercion)
+-- If (scrutBinderSwap_maybe e = Just (v, mco), then
+--    v = e |> mco
+-- See Note [Case of cast]
+-- See Note [Care with binder-swap on dictionaries]
+--
+-- We use this same function in SpecConstr, and Simplify.Iteration,
+-- when something binder-swap-like is happening
+scrutBinderSwap_maybe (Var v)    = Just (v, MRefl)
+scrutBinderSwap_maybe (Cast (Var v) co)
+  | not (isDictId v)             = Just (v, MCo (mkSymCo co))
+        -- Cast: see Note [Case of cast]
+        -- isDictId: see Note [Care with binder-swap on dictionaries]
+        -- The isDictId rejects a Constraint/Constraint binder-swap, perhaps
+        -- over-conservatively. But I have never seen one, so I'm leaving
+        -- the code as simple as possible. Losing the binder-swap in a
+        -- rare case probably has very low impact.
+scrutBinderSwap_maybe (Tick _ e) = scrutBinderSwap_maybe e  -- Drop ticks
+scrutBinderSwap_maybe _          = Nothing
+
+lookupBndrSwap :: OccEnv -> Id -> (CoreExpr, Id)
+-- See Note [The binder-swap substitution]
+-- Returns an expression of the same type as Id
+lookupBndrSwap env@(OccEnv { occ_bs_env = bs_env })  bndr
+  = case lookupVarEnv bs_env bndr of {
+       Nothing           -> (Var bndr, bndr) ;
+       Just (bndr1, mco) ->
+
+    -- Why do we iterate here?
+    -- See (BS2) in Note [The binder-swap substitution]
+    case lookupBndrSwap env bndr1 of
+      (fun, fun_id) -> (mkCastMCo fun mco, fun_id) }
+
+
+{- Historical note [Proxy let-bindings]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+We used to do the binder-swap transformation by introducing
+a proxy let-binding, thus;
+
+   case x of b { pi -> ri }
+      ==>
+   case x of b { pi -> let x = b in ri }
+
+But that had two problems:
+
+1. If 'x' is an imported GlobalId, we'd end up with a GlobalId
+   on the LHS of a let-binding which isn't allowed.  We worked
+   around this for a while by "localising" x, but it turned
+   out to be very painful #16296,
+
+2. In CorePrep we use the occurrence analyser to do dead-code
+   elimination (see Note [Dead code in CorePrep]).  But that
+   occasionally led to an unlifted let-binding
+       case x of b { DEFAULT -> let x::Int# = b in ... }
+   which disobeys one of CorePrep's output invariants (no unlifted
+   let-bindings) -- see #5433.
+
+Doing a substitution (via occ_bs_env) is much better.
+
 Historical Note [no-case-of-case]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 We *used* to suppress the binder-swap in case expressions when
 -fno-case-of-case is on.  Old remarks:
     "This happens in the first simplifier pass,
@@ -2846,53 +2952,8 @@ binder-swap in OccAnal:
 It's fixed by doing the binder-swap in OccAnal because we can do the
 binder-swap unconditionally and still get occurrence analysis
 information right.
--}
 
-addBndrSwap :: OutExpr -> Id -> OccEnv -> OccEnv
--- See Note [The binder-swap substitution]
-addBndrSwap scrut case_bndr
-            env@(OccEnv { occ_bs_env = swap_env, occ_bs_rng = rng_vars })
-  | Just (scrut_var, mco) <- get_scrut_var (stripTicksTopE (const True) scrut)
-  , scrut_var /= case_bndr
-      -- Consider: case x of x { ... }
-      -- Do not add [x :-> x] to occ_bs_env, else lookupBndrSwap will loop
-  = env { occ_bs_env = extendVarEnv swap_env scrut_var (case_bndr', mco)
-        , occ_bs_rng = rng_vars `extendVarSet` case_bndr'
-                       `unionVarSet` tyCoVarsOfMCo mco }
-
-  | otherwise
-  = env
-  where
-    get_scrut_var :: OutExpr -> Maybe (OutVar, MCoercion)
-    get_scrut_var (Var v)           = Just (v, MRefl)
-    get_scrut_var (Cast (Var v) co) = Just (v, MCo co) -- See Note [Case of cast]
-    get_scrut_var _                 = Nothing
-
-    case_bndr' = zapIdOccInfo case_bndr
-                 -- See Note [Zap case binders in proxy bindings]
 
-lookupBndrSwap :: OccEnv -> Id -> (CoreExpr, Id)
--- See Note [The binder-swap substitution]
--- Returns an expression of the same type as Id
-lookupBndrSwap env@(OccEnv { occ_bs_env = bs_env })  bndr
-  = case lookupVarEnv bs_env bndr of {
-       Nothing           -> (Var bndr, bndr) ;
-       Just (bndr1, mco) ->
-
-    -- Why do we iterate here?
-    -- See (BS2) in Note [The binder-swap substitution]
-    case lookupBndrSwap env bndr1 of
-      (fun, fun_id) -> (add_cast fun mco, fun_id) }
-
-  where
-    add_cast fun MRefl    = fun
-    add_cast fun (MCo co) = Cast fun (mkSymCo co)
-    -- We must switch that 'co' to 'sym co';
-    -- see the comment with occ_bs_env
-    -- No need to test for isReflCo, because 'co' came from
-    -- a (Cast e co) and hence is unlikely to be Refl
-
-{-
 ************************************************************************
 *                                                                      *
 \subsection[OccurAnal-types]{OccEnv}
diff --git a/compiler/GHC/Core/Opt/SetLevels.hs b/compiler/GHC/Core/Opt/SetLevels.hs
index 168e0a1dd3..5ac447883a 100644
--- a/compiler/GHC/Core/Opt/SetLevels.hs
+++ b/compiler/GHC/Core/Opt/SetLevels.hs
@@ -51,17 +51,6 @@
   The simplifier tries to get rid of occurrences of x, in favour of wild,
   in the hope that there will only be one remaining occurrence of x, namely
   the scrutinee of the case, and we can inline it.
-
-  This can only work if @wild@ is an unrestricted binder. Indeed, even with the
-  extended typing rule (in the linter) for case expressions, if
-       case x of wild % 1 { p -> e}
-  is well-typed, then
-       case x of wild % 1 { p -> e[wild\x] }
-  is only well-typed if @e[wild\x] = e@ (that is, if @wild@ is not used in @e@
-  at all). In which case, it is, of course, pointless to do the substitution
-  anyway. So for a linear binder (and really anything which isn't unrestricted),
-  doing this substitution would either produce ill-typed terms or be the
-  identity.
 -}
 
 module GHC.Core.Opt.SetLevels (
@@ -1602,7 +1591,9 @@ extendCaseBndrEnv :: LevelEnv
                   -> LevelEnv
 extendCaseBndrEnv le@(LE { le_subst = subst, le_env = id_env })
                   case_bndr (Var scrut_var)
-    | Many <- varMult case_bndr
+  -- We could use OccurAnal. scrutBinderSwap_maybe here, and perhaps
+  -- get a bit more floating.  But we didn't in the past and it's
+  -- an unforced change, so I'm leaving it.
   = le { le_subst   = extendSubstWithVar subst case_bndr scrut_var
        , le_env     = add_id id_env (case_bndr, scrut_var) }
 extendCaseBndrEnv env _ _ = env
diff --git a/compiler/GHC/Core/Opt/Simplify/Iteration.hs b/compiler/GHC/Core/Opt/Simplify/Iteration.hs
index f3fb5c2f0b..49707fecf4 100644
--- a/compiler/GHC/Core/Opt/Simplify/Iteration.hs
+++ b/compiler/GHC/Core/Opt/Simplify/Iteration.hs
@@ -22,7 +22,7 @@ import GHC.Core.Opt.Simplify.Monad
 import GHC.Core.Type hiding ( substTy, substTyVar, extendTvSubst, extendCvSubst )
 import GHC.Core.Opt.Simplify.Env
 import GHC.Core.Opt.Simplify.Utils
-import GHC.Core.Opt.OccurAnal ( occurAnalyseExpr, zapLambdaBndrs )
+import GHC.Core.Opt.OccurAnal ( occurAnalyseExpr, zapLambdaBndrs, scrutBinderSwap_maybe )
 import GHC.Core.Make       ( FloatBind, mkImpossibleExpr, castBottomExpr )
 import qualified GHC.Core.Make
 import GHC.Core.Coercion hiding ( substCo, substCoVar )
@@ -3286,19 +3286,21 @@ zapIdOccInfoAndSetEvald str v =
                          -- see Note [Case alternative occ info]
 
 addAltUnfoldings :: SimplEnv -> Maybe OutExpr -> OutId -> OutExpr -> SimplM SimplEnv
-addAltUnfoldings env scrut case_bndr con_app
+addAltUnfoldings env mb_scrut case_bndr con_app
   = do { let con_app_unf = mk_simple_unf con_app
              env1 = addBinderUnfolding env case_bndr con_app_unf
 
              -- See Note [Add unfolding for scrutinee]
-             env2 | Many <- idMult case_bndr = case scrut of
-                      Just (Var v)           -> addBinderUnfolding env1 v con_app_unf
-                      Just (Cast (Var v) co) -> addBinderUnfolding env1 v $
-                                                mk_simple_unf (Cast con_app (mkSymCo co))
-                      _                      -> env1
+             env2 | Just scrut <- mb_scrut
+                  , Just (v,mco) <- scrutBinderSwap_maybe scrut
+                  = addBinderUnfolding env1 v $
+                       if isReflMCo mco  -- isReflMCo: avoid calling mk_simple_unf
+                       then con_app_unf  --            twice in the common case
+                       else mk_simple_unf (mkCastMCo con_app mco)
+
                   | otherwise = env1
 
-       ; traceSmpl "addAltUnf" (vcat [ppr case_bndr <+> ppr scrut, ppr con_app])
+       ; traceSmpl "addAltUnf" (vcat [ppr case_bndr <+> ppr mb_scrut, ppr con_app])
        ; return env2 }
   where
     -- Force the opts, so that the whole SimplEnv isn't retained
@@ -3361,9 +3363,6 @@ it's also good for case-elimination -- suppose that 'f' was inlined
 and did multi-level case analysis, then we'd solve it in one
 simplifier sweep instead of two.
 
-Exactly the same issue arises in GHC.Core.Opt.SpecConstr;
-see Note [Add scrutinee to ValueEnv too] in GHC.Core.Opt.SpecConstr
-
 HOWEVER, given
   case x of y { Just a -> r1; Nothing -> r2 }
 we do not want to add the unfolding x -> y to 'x', which might seem cool,
@@ -3374,8 +3373,11 @@ piece of information.
 So instead we add the unfolding x -> Just a, and x -> Nothing in the
 respective RHSs.
 
-Since this transformation is tantamount to a binder swap, the same caveat as in
-Note [Suppressing binder-swaps on linear case] in OccurAnal apply.
+Since this transformation is tantamount to a binder swap, we use
+GHC.Core.Opt.OccurAnal.scrutBinderSwap_maybe to do the check.
+
+Exactly the same issue arises in GHC.Core.Opt.SpecConstr;
+see Note [Add scrutinee to ValueEnv too] in GHC.Core.Opt.SpecConstr
 
 
 ************************************************************************
diff --git a/compiler/GHC/Core/Opt/SpecConstr.hs b/compiler/GHC/Core/Opt/SpecConstr.hs
index 0dd84fdf99..12d9ceb0fa 100644
--- a/compiler/GHC/Core/Opt/SpecConstr.hs
+++ b/compiler/GHC/Core/Opt/SpecConstr.hs
@@ -35,6 +35,7 @@ import GHC.Core.Unfold
 import GHC.Core.FVs     ( exprsFreeVarsList, exprFreeVars )
 import GHC.Core.Opt.Monad
 import GHC.Core.Opt.WorkWrap.Utils
+import GHC.Core.Opt.OccurAnal( scrutBinderSwap_maybe )
 import GHC.Core.DataCon
 import GHC.Core.Class( classTyVars )
 import GHC.Core.Coercion hiding( substCo )
@@ -1072,8 +1073,8 @@ extendCaseBndrs env scrut case_bndr con alt_bndrs
    = (env2, alt_bndrs')
  where
    live_case_bndr = not (isDeadBinder case_bndr)
-   env1 | Var v <- stripTicksTopE (const True) scrut
-                         = extendValEnv env v cval
+   env1 | Just (v, mco) <- scrutBinderSwap_maybe scrut
+        , isReflMCo mco  = extendValEnv env v cval
         | otherwise      = env  -- See Note [Add scrutinee to ValueEnv too]
    env2 | live_case_bndr = extendValEnv env1 case_bndr cval
         | otherwise      = env1
@@ -1167,6 +1168,10 @@ though the simplifier has systematically replaced uses of 'x' with 'y'
 and 'b' with 'c' in the code.  The use of 'b' in the ValueEnv came
 from outside the case.  See #4908 for the live example.
 
+It's very like the binder-swap story, so we use scrutBinderSwap_maybe
+to identify suitable scrutinees -- but only if there is no cast
+(isReflMCo) because that's all that the ValueEnv allows.
+
 Note [Avoiding exponential blowup]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 The sc_count field of the ScEnv says how many times we are prepared to
diff --git a/compiler/GHC/Core/Subst.hs b/compiler/GHC/Core/Subst.hs
index 232f7cb9cf..9bfca6184e 100644
--- a/compiler/GHC/Core/Subst.hs
+++ b/compiler/GHC/Core/Subst.hs
@@ -41,7 +41,6 @@ import GHC.Core
 import GHC.Core.FVs
 import GHC.Core.Seq
 import GHC.Core.Utils
-import GHC.Core.TyCo.Subst ( substCo )
 
         -- We are defining local versions
 import GHC.Core.Type hiding ( substTy )
diff --git a/testsuite/tests/linters/notes.stdout b/testsuite/tests/linters/notes.stdout
index ac7a453445..b50da6b3a2 100644
--- a/testsuite/tests/linters/notes.stdout
+++ b/testsuite/tests/linters/notes.stdout
@@ -2,7 +2,6 @@ ref    compiler/GHC/Core/Coercion/Axiom.hs:461:2:     Note [RoughMap and rm_empt
 ref    compiler/GHC/Core/Opt/OccurAnal.hs:857:15:     Note [Loop breaking]
 ref    compiler/GHC/Core/Opt/SetLevels.hs:1580:30:     Note [Top level scope]
 ref    compiler/GHC/Core/Opt/Simplify/Iteration.hs:2675:13:     Note [Case binder next]
-ref    compiler/GHC/Core/Opt/Simplify/Iteration.hs:3303:0:     Note [Suppressing binder-swaps on linear case]
 ref    compiler/GHC/Core/Opt/Simplify/Iteration.hs:3854:8:     Note [Lambda-bound unfoldings]
 ref    compiler/GHC/Core/Opt/Simplify/Utils.hs:1257:37:     Note [Gentle mode]
 ref    compiler/GHC/Core/Opt/Specialise.hs:1623:28:     Note [Arity decrease]
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index fc708ef9f0..9cbd330d0b 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -359,7 +359,6 @@ test('T19586', normal, compile, [''])
 
 test('T19599', normal, compile, ['-O -ddump-rules'])
 test('T19599a', normal, compile, ['-O -ddump-rules'])
-test('T13873',  [expect_broken(21229), grep_errmsg(r'SPEC') ], compile, ['-O -ddump-rules'])
 
 # Look for a specialisation rule for wimwam
 test('T19672', normal, compile, ['-O2 -ddump-rules'])