Fix arityType: -fpedantic-bottoms, join points, etcwip/T21694a

This MR fixes #21694, #21755.  It also makes sure that #21948 and
fix to #21694.

* For #21694 the underlying problem was that we were calling arityType
  on an expression that had free join points.  This is a Bad Bad Idea.
  See Note [No free join points in arityType].

* To make "no free join points in arityType" work out I had to avoid
  trying to use eta-expansion for runRW#. This entailed a few changes
  in the Simplifier's treatment of runRW#.  See
  GHC.Core.Opt.Simplify.Iteration Note [No eta-expansion in runRW#]

* I also made andArityType work correctly with -fpedantic-bottoms;
  see Note [Combining case branches: andWithTail].

* Rewrote Note [Combining case branches: optimistic one-shot-ness]

* arityType previously treated join points differently to other
  let-bindings. This patch makes them unform; arityType analyses
  the RHS of all bindings to get its ArityType, and extends am_sigs.

  I realised that, now we have am_sigs giving the ArityType for
  let-bound Ids, we don't need the (pre-dating) special code in
  arityType for join points. But instead we need to extend the env for
  Rec bindings, which weren't doing before.  More uniform now.  See
  Note [arityType for let-bindings].

  This meant we could get rid of ae_joins, and in fact get rid of
  EtaExpandArity altogether.  Simpler.

* And finally, it was the strange treatment of join-point Ids in
  arityType (involving a fake ABot type) that led to a serious bug:
  #21755.  Fixed by this refactoring, which treats them uniformly;
  but without breaking #18328.

  In fact, the arity for recursive join bindings is pretty tricky;
  see the long Note [Arity for recursive join bindings]
  in GHC.Core.Opt.Simplify.Utils.  That led to more refactoring,
  including deciding that an Id could have an Arity that is bigger
  than its JoinArity; see Note [Invariants on join points], item
  2(b) in GHC.Core

* Make sure that the "demand threshold" for join points in DmdAnal
  is no bigger than the join-arity.  In GHC.Core.Opt.DmdAnal see
  Note [Demand signatures are computed for a threshold arity based on idArity]

* I moved GHC.Core.Utils.exprIsDeadEnd into GHC.Core.Opt.Arity,
  where it more properly belongs.

* Remove an old, redundant hack in FloatOut.  The old Note was
  Note [Bottoming floats: eta expansion] in GHC.Core.Opt.SetLevels.

Compile time improves very slightly on average:

Metrics: compile_time/bytes allocated
---------------------------------------------------------------------------------------
  T18223(normal) ghc/alloc    725,808,720    747,839,216  +3.0%  BAD
  T6048(optasm)  ghc/alloc    105,006,104    101,599,472  -3.2% GOOD
  geo. mean                                          -0.2%
  minimum                                            -3.2%
  maximum                                            +3.0%

For some reason Windows was better

   T10421(normal) ghc/alloc    125,888,360    124,129,168  -1.4% GOOD
   T18140(normal) ghc/alloc     85,974,520     83,884,224  -2.4% GOOD
  T18698b(normal) ghc/alloc    236,764,568    234,077,288  -1.1% GOOD
   T18923(normal) ghc/alloc     75,660,528     73,994,512  -2.2% GOOD
    T6048(optasm) ghc/alloc    112,232,512    108,182,520  -3.6% GOOD
  geo. mean                                          -0.6%

I had a quick look at T18223 but it is knee deep in coercions and
the size of everything looks similar before and after.  I decided
to accept that 3% increase in exchange for goodness elsewhere.

Metric Decrease:
    T10421
    T18140
    T18698b
    T18923
    T6048

Metric Increase:
    T18223

1 files changed, 20 insertions, 29 deletions

diff --git a/compiler/GHC/Core/Opt/SetLevels.hs b/compiler/GHC/Core/Opt/SetLevels.hs
index 85ac7e2e86..9645a10340 100644
--- a/compiler/GHC/Core/Opt/SetLevels.hs
+++ b/compiler/GHC/Core/Opt/SetLevels.hs
@@ -104,7 +104,7 @@ import GHC.Types.Unique.DSet  ( getUniqDSet )
 import GHC.Types.Var.Env
 import GHC.Types.Literal      ( litIsTrivial )
 import GHC.Types.Demand       ( DmdSig, prependArgsDmdSig )
-import GHC.Types.Cpr          ( mkCprSig, botCpr )
+import GHC.Types.Cpr          ( CprSig, prependArgsCprSig )
 import GHC.Types.Name         ( getOccName, mkSystemVarName )
 import GHC.Types.Name.Occurrence ( occNameString )
 import GHC.Types.Unique       ( hasKey )
@@ -659,9 +659,7 @@ lvlMFE env strict_ctxt ann_expr
          -- No wrapping needed if the type is lifted, or is a literal string
          -- or if we are wrapping it in one or more value lambdas
   = do { expr1 <- lvlFloatRhs abs_vars dest_lvl rhs_env NonRecursive
-                              (isJust mb_bot_str)
-                              join_arity_maybe
-                              ann_expr
+                              is_bot_lam join_arity_maybe ann_expr
                   -- Treat the expr just like a right-hand side
        ; var <- newLvlVar expr1 join_arity_maybe is_mk_static
        ; let var2 = annotateBotStr var float_n_lams mb_bot_str
@@ -702,6 +700,7 @@ lvlMFE env strict_ctxt ann_expr
     fvs          = freeVarsOf ann_expr
     fvs_ty       = tyCoVarsOfType expr_ty
     is_bot       = isBottomThunk mb_bot_str
+    is_bot_lam   = isJust mb_bot_str
     is_function  = isFunction ann_expr
     mb_bot_str   = exprBotStrictness_maybe expr
                            -- See Note [Bottoming floats]
@@ -750,10 +749,10 @@ hasFreeJoin :: LevelEnv -> DVarSet -> Bool
 hasFreeJoin env fvs
   = not (maxFvLevel isJoinId env fvs == tOP_LEVEL)
 
-isBottomThunk :: Maybe (Arity, s) -> Bool
+isBottomThunk :: Maybe (Arity, DmdSig, CprSig) -> Bool
 -- See Note [Bottoming floats] (2)
-isBottomThunk (Just (0, _)) = True   -- Zero arity
-isBottomThunk _             = False
+isBottomThunk (Just (0, _, _)) = True   -- Zero arity
+isBottomThunk _                = False
 
 {- Note [Floating to the top]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -977,16 +976,6 @@ Id, *immediately*, for three reasons:
     thing is based on the cheap-and-cheerful exprIsDeadEnd, I'm not sure
     that it'll nail all such cases.
 
-Note [Bottoming floats: eta expansion] c.f Note [Bottoming floats]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-Tiresomely, though, the simplifier has an invariant that the manifest
-arity of the RHS should be the same as the arity; but we can't call
-etaExpand during GHC.Core.Opt.SetLevels because it works over a decorated form of
-CoreExpr.  So we do the eta expansion later, in GHC.Core.Opt.FloatOut.
-But we should only eta-expand if the RHS doesn't already have the right
-exprArity, otherwise we get unnecessary top-level bindings if the RHS was
-trivial after the next run of the Simplifier.
-
 Note [Case MFEs]
 ~~~~~~~~~~~~~~~~
 We don't float a case expression as an MFE from a strict context.  Why not?
@@ -1008,17 +997,18 @@ answer.
 
 -}
 
-annotateBotStr :: Id -> Arity -> Maybe (Arity, DmdSig) -> Id
+annotateBotStr :: Id -> Arity -> Maybe (Arity, DmdSig, CprSig) -> Id
 -- See Note [Bottoming floats] for why we want to add
 -- bottoming information right now
 --
 -- n_extra are the number of extra value arguments added during floating
-annotateBotStr id n_extra mb_str
-  = case mb_str of
-      Nothing           -> id
-      Just (arity, sig) -> id `setIdArity`  (arity + n_extra)
-                              `setIdDmdSig` prependArgsDmdSig n_extra sig
-                              `setIdCprSig` mkCprSig (arity + n_extra) botCpr
+annotateBotStr id n_extra mb_bot_str
+  | Just (arity, str_sig, cpr_sig) <- mb_bot_str
+  = id `setIdArity`  (arity + n_extra)
+       `setIdDmdSig` prependArgsDmdSig n_extra str_sig
+       `setIdCprSig` prependArgsCprSig n_extra cpr_sig
+  | otherwise
+  = id
 
 notWorthFloating :: CoreExpr -> [Var] -> Bool
 -- Returns True if the expression would be replaced by
@@ -1127,7 +1117,7 @@ lvlBind env (AnnNonRec bndr rhs)
           -- bit brutal, but unlifted bindings aren't expensive either
 
   = -- No float
-    do { rhs' <- lvlRhs env NonRecursive is_bot mb_join_arity rhs
+    do { rhs' <- lvlRhs env NonRecursive is_bot_lam mb_join_arity rhs
        ; let  bind_lvl        = incMinorLvl (le_ctxt_lvl env)
               (env', [bndr']) = substAndLvlBndrs NonRecursive env bind_lvl [bndr]
        ; return (NonRec bndr' rhs', env') }
@@ -1136,7 +1126,7 @@ lvlBind env (AnnNonRec bndr rhs)
   | null abs_vars
   = do {  -- No type abstraction; clone existing binder
          rhs' <- lvlFloatRhs [] dest_lvl env NonRecursive
-                             is_bot mb_join_arity rhs
+                             is_bot_lam mb_join_arity rhs
        ; (env', [bndr']) <- cloneLetVars NonRecursive env dest_lvl [bndr]
        ; let bndr2 = annotateBotStr bndr' 0 mb_bot_str
        ; return (NonRec (TB bndr2 (FloatMe dest_lvl)) rhs', env') }
@@ -1144,7 +1134,7 @@ lvlBind env (AnnNonRec bndr rhs)
   | otherwise
   = do {  -- Yes, type abstraction; create a new binder, extend substitution, etc
          rhs' <- lvlFloatRhs abs_vars dest_lvl env NonRecursive
-                             is_bot mb_join_arity rhs
+                             is_bot_lam mb_join_arity rhs
        ; (env', [bndr']) <- newPolyBndrs dest_lvl env abs_vars [bndr]
        ; let bndr2 = annotateBotStr bndr' n_extra mb_bot_str
        ; return (NonRec (TB bndr2 (FloatMe dest_lvl)) rhs', env') }
@@ -1155,11 +1145,12 @@ lvlBind env (AnnNonRec bndr rhs)
     rhs_fvs    = freeVarsOf rhs
     bind_fvs   = rhs_fvs `unionDVarSet` dIdFreeVars bndr
     abs_vars   = abstractVars dest_lvl env bind_fvs
-    dest_lvl   = destLevel env bind_fvs ty_fvs (isFunction rhs) is_bot is_join
+    dest_lvl   = destLevel env bind_fvs ty_fvs (isFunction rhs) is_bot_lam is_join
 
     deann_rhs  = deAnnotate rhs
     mb_bot_str = exprBotStrictness_maybe deann_rhs
-    is_bot     = isJust mb_bot_str
+    is_bot_lam = isJust mb_bot_str
+        -- is_bot_lam: looks like (\xy. bot), maybe zero lams
         -- NB: not isBottomThunk!  See Note [Bottoming floats] point (3)
 
     n_extra    = count isId abs_vars