Be a bit more selective about floating bottoming expressionswip/T22494

This MR arranges to float a bottoming expression to the top
only if it escapes a value lambda.

See #22494 and Note [Floating to the top] in SetLevels.

This has a generally beneficial effect in nofib

+-------------------------------++----------+
|                               ||tsv (rel) |
+===============================++==========+
|           imaginary/paraffins ||   -0.93% |
|                imaginary/rfib ||   -0.05% |
|                      real/fem ||   -0.03% |
|                    real/fluid ||   -0.01% |
|                   real/fulsom ||   +0.05% |
|                   real/gamteb ||   -0.27% |
|                       real/gg ||   -0.10% |
|                   real/hidden ||   -0.01% |
|                      real/hpg ||   -0.03% |
|                      real/scs ||  -11.13% |
|         shootout/k-nucleotide ||   -0.01% |
|               shootout/n-body ||   -0.08% |
|   shootout/reverse-complement ||   -0.00% |
|        shootout/spectral-norm ||   -0.02% |
|             spectral/fibheaps ||   -0.20% |
|           spectral/hartel/fft ||   -1.04% |
|         spectral/hartel/solid ||   +0.33% |
|     spectral/hartel/wave4main ||   -0.35% |
|                 spectral/mate ||   +0.76% |
+===============================++==========+
|                     geom mean ||   -0.12% |

The effect on compile time is generally slightly beneficial

Metrics: compile_time/bytes allocated
----------------------------------------------
MultiLayerModulesTH_OneShot(normal)  +0.3%
                  PmSeriesG(normal)  -0.2%
                  PmSeriesT(normal)  -0.1%
                     T10421(normal)  -0.1%
                    T10421a(normal)  -0.1%
                     T10858(normal)  -0.1%
                     T11276(normal)  -0.1%
                    T11303b(normal)  -0.2%
                     T11545(normal)  -0.1%
                     T11822(normal)  -0.1%
                     T12150(optasm)  -0.1%
                     T12234(optasm)  -0.3%
                     T13035(normal)  -0.2%
                     T16190(normal)  -0.1%
                     T16875(normal)  -0.4%
                    T17836b(normal)  -0.2%
                     T17977(normal)  -0.2%
                    T17977b(normal)  -0.2%
                     T18140(normal)  -0.1%
                     T18282(normal)  -0.1%
                     T18304(normal)  -0.2%
                    T18698a(normal)  -0.1%
                     T18923(normal)  -0.1%
                     T20049(normal)  -0.1%
                    T21839r(normal)  -0.1%
                      T5837(normal)  -0.4%
                      T6048(optasm)  +3.2% BAD
                      T9198(normal)  -0.2%
                      T9630(normal)  -0.1%
       TcPlugin_RewritePerf(normal)  -0.4%
             hard_hole_fits(normal)  -0.1%

                          geo. mean  -0.0%
                          minimum    -0.4%
                          maximum    +3.2%

The T6048 outlier is hard to pin down, but it may be the effect of
reading in more interface files definitions. It's a small program for
which compile time is very short, so I'm not bothered about it.

Metric Increase:
    T6048

4 files changed, 196 insertions, 43 deletions

diff --git a/compiler/GHC/Core/Opt/SetLevels.hs b/compiler/GHC/Core/Opt/SetLevels.hs
index 2fdd5ba362..95084cf7b6 100644
--- a/compiler/GHC/Core/Opt/SetLevels.hs
+++ b/compiler/GHC/Core/Opt/SetLevels.hs
@@ -685,15 +685,16 @@ lvlMFE env strict_ctxt ann_expr
     expr_ty      = exprType expr
     fvs          = freeVarsOf ann_expr
     fvs_ty       = tyCoVarsOfType expr_ty
-    is_bot       = isBottomThunk mb_bot_str
-    is_bot_lam   = isJust mb_bot_str
+    is_bot_lam   = isJust mb_bot_str   -- True of bottoming thunks too!
     is_function  = isFunction ann_expr
     mb_bot_str   = exprBotStrictness_maybe expr
                            -- See Note [Bottoming floats]
                            -- esp Bottoming floats (2)
     expr_ok_for_spec = exprOkForSpeculation expr
-    dest_lvl     = destLevel env fvs fvs_ty is_function is_bot False
-    abs_vars     = abstractVars dest_lvl env fvs
+    abs_vars = abstractVars dest_lvl env fvs
+    dest_lvl = destLevel env fvs fvs_ty is_function is_bot_lam False
+               -- NB: is_bot_lam not is_bot; see (3) in
+               --     Note [Bottoming floats]
 
     -- float_is_new_lam: the floated thing will be a new value lambda
     -- replacing, say (g (x+4)) by (lvl x).  No work is saved, nor is
@@ -725,7 +726,9 @@ lvlMFE env strict_ctxt ann_expr
     -- See Note [Floating to the top]
     saves_alloc =  isTopLvl dest_lvl
                 && floatConsts env
-                && (not strict_ctxt || is_bot || exprIsHNF expr)
+                && (   not strict_ctxt                     -- (a)
+                    || exprIsHNF expr                      -- (b)
+                    || (is_bot_lam && escapes_value_lam))  -- (c)
 
 hasFreeJoin :: LevelEnv -> DVarSet -> Bool
 -- Has a free join point which is not being floated to top level.
@@ -735,55 +738,63 @@ hasFreeJoin :: LevelEnv -> DVarSet -> Bool
 hasFreeJoin env fvs
   = not (maxFvLevel isJoinId env fvs == tOP_LEVEL)
 
-isBottomThunk :: Maybe (Arity, DmdSig, CprSig) -> Bool
--- See Note [Bottoming floats] (2)
-isBottomThunk (Just (0, _, _)) = True   -- Zero arity
-isBottomThunk _                = False
-
 {- Note [Floating to the top]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-We are keen to float something to the top level, even if it does not
-escape a value lambda (and hence save work), for two reasons:
-
-  * Doing so makes the function smaller, by floating out
-    bottoming expressions, or integer or string literals.  That in
-    turn makes it easier to inline, with less duplication.
-
-  * (Minor) Doing so may turn a dynamic allocation (done by machine
-    instructions) into a static one. Minor because we are assuming
-    we are not escaping a value lambda.
-
-But do not do so if (saves_alloc):
-     - the context is strict, and
-     - the expression is not a HNF, and
-     - the expression is not bottoming
+Suppose saves_work is False, i.e.
+ - 'e' does not escape a value lambda (escapes_value_lam), or
+ - 'e' would have added value lambdas if floated (float_is_new_lam)
+Then we may still be keen to float a sub-expression 'e' to the top level,
+for two reasons:
+
+ (i) Doing so makes the function smaller, by floating out
+     bottoming expressions, or integer or string literals.  That in
+     turn makes it easier to inline, with less duplication.
+     This only matters if the floated sub-expression is inside a
+     value-lambda, which in turn may be easier to inline.
+
+ (ii) (Minor) Doing so may turn a dynamic allocation (done by machine
+      instructions) into a static one. Minor because we are assuming
+      we are not escaping a value lambda.
+
+But only do so if (saves_alloc):
+     (a) the context is lazy (so we get allocation), or
+     (b) the expression is a HNF (so we get allocation), or
+     (c) the expression is bottoming and (i) applies
+         (NB: if the expression is a lambda, (b) will apply;
+              so this case only catches bottoming thunks)
 
 Examples:
 
-* Bottoming
-      f x = case x of
-              0 -> error <big thing>
-              _ -> x+1
-  Here we want to float (error <big thing>) to top level, abstracting
-  over 'x', so as to make f's RHS smaller.
-
-* HNF
-      f = case y of
-            True  -> p:q
-            False -> blah
-  We may as well float the (p:q) so it becomes a static data structure.
-
-* Case scrutinee
+* (a) Strict.  Case scrutinee
       f = case g True of ....
   Don't float (g True) to top level; then we have the admin of a
   top-level thunk to worry about, with zero gain.
 
-* Case alternative
+* (a) Strict.  Case alternative
       h = case y of
              True  -> g True
              False -> False
   Don't float (g True) to the top level
 
+* (b) HNF
+      f = case y of
+            True  -> p:q
+            False -> blah
+  We may as well float the (p:q) so it becomes a static data structure.
+
+* (c) Bottoming expressions; see also Note [Bottoming floats]
+      f x = case x of
+              0 -> error <big thing>
+              _ -> x+1
+  Here we want to float (error <big thing>) to top level, abstracting
+  over 'x', so as to make f's RHS smaller.
+
+  But (#22494) if it's more like
+       foo = case error <thing> of { ... }
+  then there is no point in floating; we are never going to inline
+  'foo' anyway.  So float bottoming things only if they escape
+  a lambda.
+
 * Arguments
      t = f (g True)
   Prior to Apr 22 we didn't float (g True) to the top if f was strict.
@@ -912,7 +923,7 @@ But, as ever, we need to be careful:
 (1) We want to float a bottoming
     expression even if it has free variables:
         f = \x. g (let v = h x in error ("urk" ++ v))
-    Then we'd like to abstract over 'x' can float the whole arg of g:
+    Then we'd like to abstract over 'x', and float the whole arg of g:
         lvl = \x. let v = h x in error ("urk" ++ v)
         f = \x. g (lvl x)
     To achieve this we pass is_bot to destLevel
@@ -921,6 +932,12 @@ But, as ever, we need to be careful:
     bottom.  Instead we treat the /body/ of such a function specially,
     via point (1).  For example:
         f = \x. ....(\y z. if x then error y else error z)....
+    If we float the whole lambda thus
+        lvl = \x. \y z. if x then error y else error z
+        f = \x. ...(lvl x)...
+    we may well end up eta-expanding that PAP to
+        f = \x. ...(\y z. lvl x y z)...
+
     ===>
         lvl = \x z y. if b then error y else error z
         f = \x. ...(\y z. lvl x z y)...
@@ -1402,7 +1419,7 @@ destLevel :: LevelEnv
           -> TyCoVarSet -- Free in the /type/ of the term
                         -- (a subset of the previous argument)
           -> Bool   -- True <=> is function
-          -> Bool   -- True <=> is bottom
+          -> Bool   -- True <=> looks like \x1..xn.bottom (n>=0)
           -> Bool   -- True <=> is a join point
           -> Level
 -- INVARIANT: if is_join=True then result >= join_ceiling
@@ -1419,7 +1436,7 @@ destLevel env fvs fvs_ty is_function is_bot is_join
 
   | is_bot              -- Send bottoming bindings to the top
   = as_far_as_poss      -- regardless; see Note [Bottoming floats]
-                        -- Esp Bottoming floats (1)
+                        -- Esp Bottoming floats (1) and (3)
 
   | Just n_args <- floatLams env
   , n_args > 0  -- n=0 case handled uniformly by the 'otherwise' case
diff --git a/testsuite/tests/simplCore/should_compile/T22494.hs b/testsuite/tests/simplCore/should_compile/T22494.hs
new file mode 100644
index 0000000000..09f7ad0d45
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T22494.hs
@@ -0,0 +1,8 @@
+module T22494 where
+
+-- After simplification we should get foo more or less as-is
+-- and not
+--     lvl = error "wombat"
+--     foo = case lvl of { ... }
+
+foo = case error "wombat" of { True -> "fred"; False -> "bill" }
diff --git a/testsuite/tests/simplCore/should_compile/T22494.stderr b/testsuite/tests/simplCore/should_compile/T22494.stderr
new file mode 100644
index 0000000000..58a07f3093
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T22494.stderr
@@ -0,0 +1,126 @@
+
+==================== Tidy Core ====================
+Result size of Tidy Core
+  = {terms: 63, types: 27, coercions: 4, joins: 0/0}
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+lvl :: GHC.Prim.Addr#
+[GblId, Unf=OtherCon []]
+lvl = "error"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+lvl1 :: [Char]
+[GblId]
+lvl1 = GHC.CString.unpackCString# lvl
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+T22494.$trModule4 :: GHC.Prim.Addr#
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 20 0}]
+T22494.$trModule4 = "main"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+lvl2 :: [Char]
+[GblId]
+lvl2 = GHC.CString.unpackCString# T22494.$trModule4
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+T22494.$trModule2 :: GHC.Prim.Addr#
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 30 0}]
+T22494.$trModule2 = "T22494"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+lvl3 :: [Char]
+[GblId]
+lvl3 = GHC.CString.unpackCString# T22494.$trModule2
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+lvl4 :: GHC.Prim.Addr#
+[GblId, Unf=OtherCon []]
+lvl4 = "T22494.hs"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+lvl5 :: [Char]
+[GblId]
+lvl5 = GHC.CString.unpackCString# lvl4
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+lvl6 :: Int
+[GblId, Unf=OtherCon []]
+lvl6 = GHC.Types.I# 8#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+lvl7 :: Int
+[GblId, Unf=OtherCon []]
+lvl7 = GHC.Types.I# 12#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+lvl8 :: Int
+[GblId, Unf=OtherCon []]
+lvl8 = GHC.Types.I# 17#
+
+-- RHS size: {terms: 8, types: 0, coercions: 0, joins: 0/0}
+lvl9 :: GHC.Stack.Types.SrcLoc
+[GblId, Unf=OtherCon []]
+lvl9 = GHC.Stack.Types.SrcLoc lvl2 lvl3 lvl5 lvl6 lvl7 lvl6 lvl8
+
+-- RHS size: {terms: 4, types: 0, coercions: 0, joins: 0/0}
+lvl10 :: GHC.Stack.Types.CallStack
+[GblId, Unf=OtherCon []]
+lvl10
+  = GHC.Stack.Types.PushCallStack
+      lvl1 lvl9 GHC.Stack.Types.EmptyCallStack
+
+-- RHS size: {terms: 1, types: 0, coercions: 0, joins: 0/0}
+lvl11 :: GHC.Prim.Addr#
+[GblId, Unf=OtherCon []]
+lvl11 = "wombat"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+lvl12 :: [Char]
+[GblId]
+lvl12 = GHC.CString.unpackCString# lvl11
+
+-- RHS size: {terms: 4, types: 3, coercions: 4, joins: 0/0}
+foo :: String
+[GblId, Str=b, Cpr=b]
+foo
+  = case error
+           @GHC.Types.LiftedRep
+           @Bool
+           (lvl10
+            `cast` (Sym (GHC.Classes.N:IP[0]
+                             <"callStack">_N <GHC.Stack.Types.CallStack>_N)
+                    :: GHC.Stack.Types.CallStack
+                       ~R# (?callStack::GHC.Stack.Types.CallStack)))
+           lvl12
+    of wild {
+    }
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+T22494.$trModule3 :: GHC.Types.TrName
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 10 10}]
+T22494.$trModule3 = GHC.Types.TrNameS T22494.$trModule4
+
+-- RHS size: {terms: 2, types: 0, coercions: 0, joins: 0/0}
+T22494.$trModule1 :: GHC.Types.TrName
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 10 10}]
+T22494.$trModule1 = GHC.Types.TrNameS T22494.$trModule2
+
+-- RHS size: {terms: 3, types: 0, coercions: 0, joins: 0/0}
+T22494.$trModule :: GHC.Types.Module
+[GblId,
+ Unf=Unf{Src=<vanilla>, TopLvl=True, Value=True, ConLike=True,
+         WorkFree=True, Expandable=True, Guidance=IF_ARGS [] 10 10}]
+T22494.$trModule
+  = GHC.Types.Module T22494.$trModule3 T22494.$trModule1
+
+
+
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index 1c169d7d6e..1622c97766 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -452,3 +452,5 @@ test('T22375', normal, compile, ['-O -ddump-simpl -dsuppress-uniques -dno-typeab
 test('T21851_2', [grep_errmsg(r'wwombat') ], multimod_compile, ['T21851_2', '-O -dno-typeable-binds -dsuppress-uniques'])
 # Should not inline m, so there shouldn't be a single YES
 test('T22317', [grep_errmsg(r'ANSWER = YES') ], compile, ['-O -dinline-check m -ddebug-output'])
+
+test('T22494', [grep_errmsg(r'case') ], compile, ['-O -ddump-simpl -dsuppress-uniques'])