Run StgCse after unarise, fixes #15300

Given two unboxed sum terms:

    (# 1 | #) :: (# Int | Int# #)
    (# 1 | #) :: (# Int | Int  #)

These two terms are not equal as they unarise to different unboxed
tuples. However StgCse was thinking that these are equal, and replacing
one of these with a binder to the other.

To not deal with unboxed sums in StgCse we now do it after unarise. For
StgCse to maintain post-unarise invariants we factor-out case binder
in-scopeness check to `stgCaseBndrInScope` and use it in StgCse.

Also did some refactoring in SimplStg.

Another way to fix this would be adding a special case in StgCse to not
bring unboxed sum binders in scope:

    diff --git a/compiler/simplStg/StgCse.hs
b/compiler/simplStg/StgCse.hs
    index 6c740ca4cb..93a0f8f6ad 100644
    --- a/compiler/simplStg/StgCse.hs
    +++ b/compiler/simplStg/StgCse.hs
    @@ -332,7 +332,11 @@ stgCseExpr env (StgLetNoEscape binds body)
     stgCseAlt :: CseEnv -> OutId -> InStgAlt -> OutStgAlt
     stgCseAlt env case_bndr (DataAlt dataCon, args, rhs)
         = let (env1, args') = substBndrs env args
    -          env2 = addDataCon case_bndr dataCon (map StgVarArg
args') env1
    +          env2
    +            | isUnboxedSumCon dataCon
    +            = env1
    +            | otherwise
    +            = addDataCon case_bndr dataCon (map StgVarArg args')
env1
                 -- see note [Case 2: CSEing case binders]
               rhs' = stgCseExpr env2 rhs
           in (DataAlt dataCon, args', rhs')

I think this patch seems better in that it doesn't add a special case to
StgCse.

Test Plan:
Validate.

I tried to come up with a minimal example but failed. I thought a simple
program like

    data T = T (# Int | Int #) (# Int# | Int #)

    case T (# 1 | #) (# 1 | #) of ...

should be enough to trigger this bug, but for some reason StgCse
doesn't do
anything on this program.

Reviewers: simonpj, bgamari

Reviewed By: simonpj

Subscribers: rwbarton, thomie, carter

GHC Trac Issues: #15300

Differential Revision: https://phabricator.haskell.org/D4962

5 files changed, 38 insertions, 35 deletions

diff --git a/compiler/simplStg/SimplStg.hs b/compiler/simplStg/SimplStg.hs
index 854bb92258..36bf5101d6 100644
--- a/compiler/simplStg/SimplStg.hs
+++ b/compiler/simplStg/SimplStg.hs
@@ -21,7 +21,6 @@ import StgCse           ( stgCse )
 
 import DynFlags
 import ErrUtils
-import SrcLoc
 import UniqSupply       ( mkSplitUniqSupply )
 import Outputable
 import Control.Monad
@@ -34,27 +33,19 @@ stg2stg dflags binds
   = do  { showPass dflags "Stg2Stg"
         ; us <- mkSplitUniqSupply 'g'
 
-        ; when (dopt Opt_D_verbose_stg2stg dflags)
-               (putLogMsg dflags NoReason SevDump noSrcSpan
-                  (defaultDumpStyle dflags) (text "VERBOSE STG-TO-STG:"))
-
-        ; binds' <- end_pass "Stg2Stg" binds
-
                 -- Do the main business!
-        ; processed_binds <- foldM do_stg_pass binds' (getStgToDo dflags)
-
         ; dumpIfSet_dyn dflags Opt_D_dump_stg "Pre unarise:"
-                        (pprStgTopBindings processed_binds)
-
-        ; let un_binds = unarise us processed_binds
+                        (pprStgTopBindings binds)
 
+        ; stg_linter False "Pre-unarise" binds
+        ; let un_binds = unarise us binds
         ; stg_linter True "Unarise" un_binds
 
         ; dumpIfSet_dyn dflags Opt_D_dump_stg "STG syntax:"
                         (pprStgTopBindings un_binds)
 
-        ; return un_binds
-   }
+        ; foldM do_stg_pass un_binds (getStgToDo dflags)
+        }
 
   where
     stg_linter unarised
@@ -65,8 +56,7 @@ stg2stg dflags binds
     do_stg_pass binds to_do
       = case to_do of
           D_stg_stats ->
-             trace (showStgStats binds)
-             end_pass "StgStats" binds
+             trace (showStgStats binds) (return binds)
 
           StgCSE ->
              {-# SCC "StgCse" #-}
@@ -78,8 +68,8 @@ stg2stg dflags binds
     end_pass what binds2
       = do -- report verbosely, if required
            dumpIfSet_dyn dflags Opt_D_verbose_stg2stg what
-              (vcat (map ppr binds2))
-           stg_linter False what binds2
+              (pprStgTopBindings binds2)
+           stg_linter True what binds2
            return binds2
 
 -- -----------------------------------------------------------------------------
diff --git a/compiler/simplStg/StgCse.hs b/compiler/simplStg/StgCse.hs
index 6c740ca4cb..1ae1213960 100644
--- a/compiler/simplStg/StgCse.hs
+++ b/compiler/simplStg/StgCse.hs
@@ -302,7 +302,7 @@ stgCseExpr env (StgCase scrut bndr ty alts)
     env2 | StgApp trivial_scrut [] <- scrut' = addTrivCaseBndr bndr trivial_scrut env1
                  -- See Note [Trivial case scrutinee]
          | otherwise                         = env1
-    alts' = map (stgCseAlt env2 bndr') alts
+    alts' = map (stgCseAlt env2 ty bndr') alts
 
 
 -- A constructor application.
@@ -329,14 +329,24 @@ stgCseExpr env (StgLetNoEscape binds body)
 
 -- Case alternatives
 -- Extend the CSE environment
-stgCseAlt :: CseEnv -> OutId -> InStgAlt -> OutStgAlt
-stgCseAlt env case_bndr (DataAlt dataCon, args, rhs)
+stgCseAlt :: CseEnv -> AltType -> OutId -> InStgAlt -> OutStgAlt
+stgCseAlt env ty case_bndr (DataAlt dataCon, args, rhs)
     = let (env1, args') = substBndrs env args
-          env2 = addDataCon case_bndr dataCon (map StgVarArg args') env1
+          env2
+            -- To avoid dealing with unboxed sums StgCse runs after unarise and
+            -- should maintain invariants listed in Note [Post-unarisation
+            -- invariants]. One of the invariants is that some binders are not
+            -- used (unboxed tuple case binders) which is what we check with
+            -- `stgCaseBndrInScope` here. If the case binder is not in scope we
+            -- don't add it to the CSE env. See also #15300.
+            | stgCaseBndrInScope ty True -- CSE runs after unarise
+            = addDataCon case_bndr dataCon (map StgVarArg args') env1
+            | otherwise
+            = env1
             -- see note [Case 2: CSEing case binders]
           rhs' = stgCseExpr env2 rhs
       in (DataAlt dataCon, args', rhs')
-stgCseAlt env _ (altCon, args, rhs)
+stgCseAlt env _ _ (altCon, args, rhs)
     = let (env1, args') = substBndrs env args
           rhs' = stgCseExpr env1 rhs
       in (altCon, args', rhs')
diff --git a/compiler/stgSyn/StgLint.hs b/compiler/stgSyn/StgLint.hs
index bb2064ab48..58f14a1b3f 100644
--- a/compiler/stgSyn/StgLint.hs
+++ b/compiler/stgSyn/StgLint.hs
@@ -172,15 +172,7 @@ lintStgExpr (StgCase scrut bndr alts_type alts) = do
     lintStgExpr scrut
 
     lf <- getLintFlags
-    let in_scope =
-          case alts_type of
-             AlgAlt _      -> True
-             PrimAlt _     -> True
-             -- Case binders of unboxed tuple or unboxed sum type always dead
-             -- after the unariser has run.
-             -- See Note [Post-unarisation invariants].
-             MultiValAlt _ -> not (lf_unarised lf)
-             PolyAlt       -> True
+    let in_scope = stgCaseBndrInScope alts_type (lf_unarised lf)
 
     addInScopeVars [bndr | in_scope] (mapM_ lintAlt alts)
 
diff --git a/compiler/stgSyn/StgSyn.hs b/compiler/stgSyn/StgSyn.hs
index 608a028673..eb905f7456 100644
--- a/compiler/stgSyn/StgSyn.hs
+++ b/compiler/stgSyn/StgSyn.hs
@@ -39,6 +39,7 @@ module StgSyn (
         isDllConApp,
         stgArgType,
         stripStgTicksTop,
+        stgCaseBndrInScope,
 
         pprStgBinding, pprStgTopBindings
     ) where
@@ -155,6 +156,18 @@ stripStgTicksTop p = go []
    where go ts (StgTick t e) | p t = go (t:ts) e
          go ts other               = (reverse ts, other)
 
+-- | Given an alt type and whether the program is unarised, return whether the
+-- case binder is in scope.
+--
+-- Case binders of unboxed tuple or unboxed sum type always dead after the
+-- unariser has run. See Note [Post-unarisation invariants].
+stgCaseBndrInScope :: AltType -> Bool {- ^ unarised? -} -> Bool
+stgCaseBndrInScope alt_ty unarised =
+    case alt_ty of
+      AlgAlt _      -> True
+      PrimAlt _     -> True
+      MultiValAlt _ -> not unarised
+      PolyAlt       -> True
 
 {-
 ************************************************************************
diff --git a/testsuite/tests/simplStg/should_compile/all.T b/testsuite/tests/simplStg/should_compile/all.T
index 19fa5134b4..2cb89749c4 100644
--- a/testsuite/tests/simplStg/should_compile/all.T
+++ b/testsuite/tests/simplStg/should_compile/all.T
@@ -17,6 +17,4 @@ def checkStgString(needle):
             return "%s not contained in -ddump-simpl\n" % needle
     return normalise_errmsg_fun(norm)
 
-
-
-test('T13588', [ checkStgString('case') ] , compile, ['-ddump-stg'])
+test('T13588', [ checkStgString('case') ] , compile, ['-dverbose-stg2stg'])