Improve CSE.combineAlts

This patch improves the way that CSE combines identical
alternatives.  See #17901.

I'm still not happy about the duplication between CSE.combineAlts
and GHC.Core.Utils.combineIdenticalAlts; see the Notes with those
functions.  But this patch is a step forward.

Metric Decrease:
    T12425
    T5642

6 files changed, 140 insertions, 25 deletions

diff --git a/compiler/GHC/Core/Utils.hs b/compiler/GHC/Core/Utils.hs
index c846d2ac2e..82d920a58e 100644
--- a/compiler/GHC/Core/Utils.hs
+++ b/compiler/GHC/Core/Utils.hs
@@ -864,10 +864,27 @@ This gave rise to a horrible sequence of cases
 
 and similarly in cascade for all the join points!
 
-NB: it's important that all this is done in [InAlt], *before* we work
-on the alternatives themselves, because Simplify.simplAlt may zap the
-occurrence info on the binders in the alternatives, which in turn
-defeats combineIdenticalAlts (see #7360).
+Note [Combine identical alternatives: wrinkles]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+* It's important that we try to combine alternatives *before*
+  simplifying them, rather than after. Reason: because
+  Simplify.simplAlt may zap the occurrence info on the binders in the
+  alternatives, which in turn defeats combineIdenticalAlts use of
+  isDeadBinder (see #7360).
+
+  You can see this in the call to combineIdenticalAlts in
+  SimplUtils.prepareAlts.  Here the alternatives have type InAlt
+  (the "In" meaning input) rather than OutAlt.
+
+* combineIdenticalAlts does not work well for nullary constructors
+      case x of y
+         []    -> f []
+         (_:_) -> f y
+  Here we won't see that [] and y are the same.  Sigh! This problem
+  is solved in CSE, in CSE.combineAlts, which does a better version of
+  combineIdenticalAlts. But sadly it doesn't have the occurrence info
+  we have here.  See Note [Combine case alts: awkward corner] in CSE).
 
 Note [Care with impossible-constructors when combining alternatives]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/compiler/simplCore/CSE.hs b/compiler/simplCore/CSE.hs
index 637312c3ca..0f0d5e49d3 100644
--- a/compiler/simplCore/CSE.hs
+++ b/compiler/simplCore/CSE.hs
@@ -17,8 +17,8 @@ import GhcPrelude
 
 import GHC.Core.Subst
 import Var              ( Var )
-import VarEnv           ( elemInScopeSet, mkInScopeSet )
-import Id               ( Id, idType, isDeadBinder, idHasRules
+import VarEnv           ( mkInScopeSet )
+import Id               ( Id, idType, idHasRules
                         , idInlineActivation, setInlineActivation
                         , zapIdOccInfo, zapIdUsageInfo, idInlinePragma
                         , isJoinId, isJoinId_maybe )
@@ -31,7 +31,7 @@ import GHC.Core
 import Outputable
 import BasicTypes
 import GHC.Core.Map
-import Util             ( filterOut )
+import Util             ( filterOut, equalLength, debugIsOn )
 import Data.List        ( mapAccumL )
 
 {-
@@ -618,15 +618,8 @@ cseCase env scrut bndr ty alts
     arg_tys :: [OutType]
     arg_tys = tyConAppArgs (idType bndr3)
 
-    -- Given case x of { K y z -> ...K y z... }
-    -- CSE K y z into x...
+    -- See Note [CSE for case alternatives]
     cse_alt (DataAlt con, args, rhs)
-        | not (null args)
-                -- ... but don't try CSE if there are no args; it just increases the number
-                -- of live vars.  E.g.
-                --      case x of { True -> ....True.... }
-                -- Don't replace True by x!
-                -- Hence the 'null args', which also deal with literals and DEFAULT
         = (DataAlt con, args', tryForCSE new_env rhs)
         where
           (env', args') = addBinders alt_env args
@@ -638,21 +631,61 @@ cseCase env scrut bndr ty alts
         where
           (env', args') = addBinders alt_env args
 
-combineAlts :: CSEnv -> [InAlt] -> [InAlt]
+combineAlts :: CSEnv -> [OutAlt] -> [OutAlt]
 -- See Note [Combine case alternatives]
-combineAlts env ((_,bndrs1,rhs1) : rest_alts)
-  | all isDeadBinder bndrs1
-  = (DEFAULT, [], rhs1) : filtered_alts
+combineAlts env alts
+  | (Just alt1, rest_alts) <- find_bndr_free_alt alts
+  , (_,bndrs1,rhs1) <- alt1
+  , let filtered_alts = filterOut (identical_alt rhs1) rest_alts
+  , not (equalLength rest_alts filtered_alts)
+  = ASSERT2( null bndrs1, ppr alts )
+    (DEFAULT, [], rhs1) : filtered_alts
+
+  | otherwise
+  = alts
   where
     in_scope = substInScope (csEnvSubst env)
-    filtered_alts = filterOut identical rest_alts
-    identical (_con, bndrs, rhs) = all ok bndrs && eqExpr in_scope rhs1 rhs
-    ok bndr = isDeadBinder bndr || not (bndr `elemInScopeSet` in_scope)
-
-combineAlts _ alts = alts  -- Default case
 
-{- Note [Combine case alternatives]
+    find_bndr_free_alt :: [CoreAlt] -> (Maybe CoreAlt, [CoreAlt])
+       -- The (Just alt) is a binder-free alt
+       -- See Note [Combine case alts: awkward corner]
+    find_bndr_free_alt []
+      = (Nothing, [])
+    find_bndr_free_alt (alt@(_,bndrs,_) : alts)
+      | null bndrs = (Just alt, alts)
+      | otherwise  = case find_bndr_free_alt alts of
+                       (mb_bf, alts) -> (mb_bf, alt:alts)
+
+    identical_alt rhs1 (_,_,rhs) = eqExpr in_scope rhs1 rhs
+       -- Even if this alt has binders, they will have been cloned
+       -- If any of these binders are mentioned in 'rhs', then
+       -- 'rhs' won't compare equal to 'rhs1' (which is from an
+       -- alt with no binders).
+
+{- Note [CSE for case alternatives]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Consider   case e of x
+            K1 y -> ....(K1 y)...
+            K2   -> ....K2....
+
+We definitely want to CSE that (K1 y) into just x.
+
+But what about the lone K2?  At first you would think "no" because
+turning K2 into 'x' increases the number of live variables.  But
+
+* Turning K2 into x increases the chance of combining identical alts.
+  Example      case xs of
+                  (_:_) -> f xs
+                  []    -> f []
+  See #17901 and simplCore/should_compile/T17901 for more examples
+  of this kind.
+
+* The next run of the simplifier will turn 'x' back into K2, so we won't
+  permanently bloat the free-var count.
+
+
+Note [Combine case alternatives]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 combineAlts is just a more heavyweight version of the use of
 combineIdenticalAlts in SimplUtils.prepareAlts.  The basic idea is
 to transform
@@ -673,6 +706,26 @@ to be doing, which is why I put it here.
 I actually saw some examples in the wild, where some inlining made e1 too
 big for cheapEqExpr to catch it.
 
+Note [Combine case alts: awkward corner]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+We would really like to check isDeadBinder on the binders in the
+alternative.  But alas, the simplifer zaps occ-info on binders in case
+alternatives; see Note [Case alternative occ info] in Simplify.
+
+* One alternative (perhaps a good one) would be to do OccAnal
+  just before CSE.  Then perhaps we could get rid of combineIdenticalAlts
+  in the Simplifier, which might save work.
+
+* Another would be for CSE to return free vars as it goes.
+
+* But the current solution is to find a nullary alternative (including
+  the DEFAULT alt, if any). This will not catch
+      case x of
+        A y   -> blah
+        B z p -> blah
+  where no alternative is nullary or DEFAULT.  But the current
+  solution is at least cheap.
+
 
 ************************************************************************
 *                                                                      *
diff --git a/testsuite/tests/simplCore/should_compile/Makefile b/testsuite/tests/simplCore/should_compile/Makefile
index b40c02175b..396c375110 100644
--- a/testsuite/tests/simplCore/should_compile/Makefile
+++ b/testsuite/tests/simplCore/should_compile/Makefile
@@ -266,3 +266,8 @@ T15631:
 T17140:
 	$(RM) -f T17140*.hi T17140*.o
 	'$(TEST_HC)' $(TEST_HC_OPTS) --make -O T17140 T17140a -v0
+
+T17901:
+	$(RM) -f T17901.o T17901.hi
+	'$(TEST_HC)' $(TEST_HC_OPTS) -O -c -ddump-simpl -dsuppress-uniques T17901.hs | grep 'wombat'
+        # All three functions should get their case alternatives combined
diff --git a/testsuite/tests/simplCore/should_compile/T17901.hs b/testsuite/tests/simplCore/should_compile/T17901.hs
new file mode 100644
index 0000000000..b60c1b6e18
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T17901.hs
@@ -0,0 +1,21 @@
+module T17901 where
+
+data T = A | B | C
+
+f1 wombat1 x = case x of
+                A -> wombat1 A
+                B -> wombat1 B
+                C -> wombat1 C
+
+data S = SA Int | SB
+
+f2 wombat2 x = case x of
+                 SA _ -> wombat2 x
+                 SB   -> wombat2 x
+
+data W = WB | WA Int
+
+f3 wombat3 x = case x of
+                WA _ -> wombat3 x
+                WB   -> wombat3 x
+
diff --git a/testsuite/tests/simplCore/should_compile/T17901.stdout b/testsuite/tests/simplCore/should_compile/T17901.stdout
new file mode 100644
index 0000000000..406e81ef5f
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T17901.stdout
@@ -0,0 +1,16 @@
+                 (wombat1 [Occ=Once*!] :: T -> p)
+                   A -> wombat1 T17901.A;
+                   B -> wombat1 T17901.B;
+                   C -> wombat1 T17901.C
+  = \ (@p) (wombat1 :: T -> p) (x :: T) ->
+      case x of wild { __DEFAULT -> wombat1 wild }
+                 (wombat2 [Occ=Once*!] :: S -> p)
+                   SA _ [Occ=Dead] -> wombat2 wild;
+                   SB -> wombat2 T17901.SB
+  = \ (@p) (wombat2 :: S -> p) (x :: S) ->
+      case x of wild { __DEFAULT -> wombat2 wild }
+                 (wombat3 [Occ=Once*!] :: W -> p)
+                   WB -> wombat3 T17901.WB;
+                   WA _ [Occ=Dead] -> wombat3 wild
+  = \ (@p) (wombat3 :: W -> p) (x :: W) ->
+      case x of wild { __DEFAULT -> wombat3 wild }
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index 31b7989b43..ed89a35690 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -316,3 +316,6 @@ test('T17590', normal, compile, ['-dcore-lint -O2'])
 test('T17722', normal, multimod_compile, ['T17722B', '-dcore-lint -O2 -v0'])
 test('T17724', normal, compile, ['-dcore-lint -O2'])
 test('T17787',  [ grep_errmsg(r'foo') ], compile, ['-ddump-simpl -dsuppress-uniques'])
+test('T17901',
+     normal,
+     makefile_test, ['T17901'])