Careful tweaking to exprOkForSpeculation

This patch does several things:

* Make exprOkForSpeculation ignore evaluatedness of variables
  See the Note [exprOkForSpeculation and evaluated variables]
  This means that the binder-swap transformation no longer
  invaliates the let/app invariant.

* Make exprOkForSpeculation return False for
     DataToTagOp and SeqOp.
  See Note [exprOkForSpeculation and SeqOp/DataToTagOp]

* Remove the 'can_fail' property from dataToTag#; it was
  always a hack (described in the old Note [dataToTag#] in
  primops.txt.pp), and now its not necessary because of the
  fixes above.

* Make SetLevels use exprIsHNF, /not/ exprOkForSpeculation,
  when floating single-alternative cases.  See SetLevels
  Note [Floating single-alternative cases]

* Fix a buglet in FloatIn; probably never bites in practice
  See Note [Dead bindings]

Collectively, these changes finally fix Trac #15696.

6 files changed, 149 insertions, 111 deletions

diff --git a/compiler/coreSyn/CoreUtils.hs b/compiler/coreSyn/CoreUtils.hs
index aa77592ef0..348179d460 100644
--- a/compiler/coreSyn/CoreUtils.hs
+++ b/compiler/coreSyn/CoreUtils.hs
@@ -1556,25 +1556,28 @@ app_ok primop_ok fun args
               -- Often there is a literal divisor, and this
               -- can get rid of a thunk in an inner loop
 
-        | SeqOp <- op    -- See Note [seq# and expr_ok]
-        -> all (expr_ok primop_ok) args
+        | SeqOp <- op  -- See Note [exprOkForSpeculation and SeqOp/DataToTagOp]
+        -> False       --     for the special cases for SeqOp and DataToTagOp
+        | DataToTagOp <- op
+        -> False
 
         | otherwise
         -> primop_ok op  -- Check the primop itself
-        && and (zipWith arg_ok arg_tys args)  -- Check the arguments
+        && and (zipWith primop_arg_ok arg_tys args)  -- Check the arguments
 
       _other -> isUnliftedType (idType fun)          -- c.f. the Var case of exprIsHNF
              || idArity fun > n_val_args             -- Partial apps
-             || (n_val_args == 0 &&
-                 isEvaldUnfolding (idUnfolding fun)) -- Let-bound values
+             -- NB: even in the nullary case, do /not/ check
+             --     for evaluated-ness of the fun;
+             --     see Note [exprOkForSpeculation and evaluated variables]
              where
                n_val_args = valArgCount args
   where
     (arg_tys, _) = splitPiTys (idType fun)
 
-    arg_ok :: TyBinder -> CoreExpr -> Bool
-    arg_ok (Named _) _ = True   -- A type argument
-    arg_ok (Anon ty) arg        -- A term argument
+    primop_arg_ok :: TyBinder -> CoreExpr -> Bool
+    primop_arg_ok (Named _) _ = True   -- A type argument
+    primop_arg_ok (Anon ty) arg        -- A term argument
        | isUnliftedType ty = expr_ok primop_ok arg
        | otherwise         = True  -- See Note [Primops with lifted arguments]
 
@@ -1638,14 +1641,34 @@ But we restrict it sharply:
   arise.
 
 * We restrict it to exhaustive alternatives. A non-exhaustive
-  case manifestly isn't ok-for-speculation. Consider
+  case manifestly isn't ok-for-speculation. for example,
+  this is a valid program (albeit a slightly dodgy one)
+    let v = case x of { B -> ...; C -> ... }
+    in case x of
+         A -> ...
+         _ ->  ...v...v....
+  Should v be considered ok-for-speculation?  Its scrutinee may be
+  evaluated, but the alternatives are incomplete so we should not
+  evaluate it strictly.
+
+  Now, all this is for lifted types, but it'd be the same for any
+  finite unlifted type. We don't have many of them, but we might
+  add unlifted algebraic types in due course.
+
+
+----- Historical note: Trac #15696: --------
+  Previously SetLevels used exprOkForSpeculation to guide
+  floating of single-alternative cases; it now uses exprIsHNF
+  Note [Floating single-alternative cases].
+
+  But in those days, consider
     case e of x { DEAFULT ->
       ...(case x of y
             A -> ...
             _ -> ...(case (case x of { B -> p; C -> p }) of
                        I# r -> blah)...
-  If SetLevesls considers the inner nested case as ok-for-speculation
-  it can do case-floating (see Note [Floating cases] in SetLevels).
+  If SetLevels considers the inner nested case as
+  ok-for-speculation it can do case-floating (in SetLevels).
   So we'd float to:
     case e of x { DEAFULT ->
     case (case x of { B -> p; C -> p }) of I# r ->
@@ -1654,19 +1677,6 @@ But we restrict it sharply:
             _ -> ...blah...)...
   which is utterly bogus (seg fault); see Trac #5453.
 
-  Similarly, this is a valid program (albeit a slightly dodgy one)
-    let v = case x of { B -> ...; C -> ... }
-    in case x of
-         A -> ...
-         _ ->  ...v...v....
-  Should v be considered ok-for-speculation?  Its scrutinee may be
-  evaluated, but the alternatives are incomplete so we should not
-  evaluate it strictly.
-
-  Now, all this is for lifted types, but it'd be the same for any
-  finite unlifted type. We don't have many of them, but we might
-  add unlifted algebraic types in due course.
-
 ----- Historical note: Trac #3717: --------
     foo :: Int -> Int
     foo 0 = 0
@@ -1699,27 +1709,54 @@ evaluate them.  Indeed, in general primops are, well, primitive
 and do not perform evaluation.
 
 Bottom line:
-  * in exprOkForSpeculation we simply ignore all lifted arguments.
-  * except see Note [seq# and expr_ok] for an exception
+  * In exprOkForSpeculation we simply ignore all lifted arguments.
+  * In the rare case of primops that /do/ evaluate their arguments,
+    (namely DataToTagOp and SeqOp) return False; see
+    Note [exprOkForSpeculation and evaluated variables]
+
+Note [exprOkForSpeculation and SeqOp/DataToTagOp]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Most primops with lifted arguments don't evaluate them
+(see Note [Primops with lifted arguments]), so we can ignore
+that argument entirely when doing exprOkForSpeculation.
+
+But DataToTagOp and SeqOp are exceptions to that rule.
+For reasons described in Note [exprOkForSpeculation and
+evaluated variables], we simply return False for them.
+
+Not doing this made #5129 go bad.
+Lots of discussion in #15696.
+
+Note [exprOkForSpeculation and evaluated variables]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Recall that
+  seq#       :: forall a s. a -> State# s -> (# State# s, a #)
+  dataToTag# :: forall a.   a -> Int#
+must always evaluate their first argument.
 
+Now consider these examples:
+ * case x of y { DEFAULT -> ....y.... }
+   Should 'y' (alone) be considered ok-for-speculation?
 
-Note [seq# and expr_ok]
-~~~~~~~~~~~~~~~~~~~~~~~
-Recall that
-   seq# :: forall a s . a -> State# s -> (# State# s, a #)
-must always evaluate its first argument.  So it's really a
-counter-example to Note [Primops with lifted arguments]. In
-the case of seq# we must check the argument to seq#.  Remember
-item (d) of the specification of exprOkForSpeculation:
+ * case x of y { DEFAULT -> ....f (dataToTag# y)... }
+   Should (dataToTag# y) be considered ok-for-spec?
+
+You could argue 'yes', because in the case alternative we know that
+'y' is evaluated.  But the binder-swap transformation, which is
+extremely useful for float-out, changes these expressions to
+   case x of y { DEFAULT -> ....x.... }
+   case x of y { DEFAULT -> ....f (dataToTag# x)... }
 
-  -- Precisely, it returns @True@ iff:
-  --  a) The expression guarantees to terminate,
-         ...
-  --  d) without throwing a Haskell exception
+And now the expression does not obey the let/app invariant!  Yikes!
+Moreover we really might float (f (dataToTag# x)) outside the case,
+and then it really, really doesn't obey the let/app invariant.
 
-The lack of this special case caused Trac #5129 to go bad again.
-See comment:24 and following
+The solution is simple: exprOkForSpeculation does not try to take
+advantage of the evaluated-ness of (lifted) varaibles.  And it returns
+False (always) for DataToTagOp and SeqOp.
 
+Note that exprIsHNF /can/ and does take advantage of evaluated-ness;
+it doesn't have the trickiness of the let/app invariant to worry about.
 
 ************************************************************************
 *                                                                      *
@@ -1781,6 +1818,8 @@ exprIsHNFlike is_con is_con_unf = is_hnf_like
                              -- and (e.g.) primops that don't have unfoldings
       || is_con_unf (idUnfolding v)
         -- Check the thing's unfolding; it might be bound to a value
+        --   or to a guaranteed-evaluated variable (isEvaldUnfolding)
+        --   Contrast with Note [exprOkForSpeculation and evaluated variables]
         -- We don't look through loop breakers here, which is a bit conservative
         -- but otherwise I worry that if an Id's unfolding is just itself,
         -- we could get an infinite loop
@@ -1800,8 +1839,8 @@ exprIsHNFlike is_con is_con_unf = is_hnf_like
     is_hnf_like (Let _ e)        = is_hnf_like e  -- Lazy let(rec)s don't affect us
     is_hnf_like _                = False
 
-    -- There is at least one value argument
-    -- 'n' is number of value args to which the expression is applied
+    -- 'n' is the number of value args to which the expression is applied
+    -- And n>0: there is at least one value argument
     app_is_value :: CoreExpr -> Int -> Bool
     app_is_value (Var f)    nva = id_app_is_value f nva
     app_is_value (Tick _ f) nva = app_is_value f nva
@@ -1809,7 +1848,7 @@ exprIsHNFlike is_con is_con_unf = is_hnf_like
     app_is_value (App f a)  nva
       | isValArg a              = app_is_value f (nva + 1)
       | otherwise               = app_is_value f nva
-    app_is_value _ _ = False
+    app_is_value _          _   = False
 
     id_app_is_value id n_val_args
        = is_con id
diff --git a/compiler/coreSyn/MkCore.hs b/compiler/coreSyn/MkCore.hs
index 73c2e7c134..4fa824a7dc 100644
--- a/compiler/coreSyn/MkCore.hs
+++ b/compiler/coreSyn/MkCore.hs
@@ -549,7 +549,7 @@ data FloatBind
   = FloatLet  CoreBind
   | FloatCase CoreExpr Id AltCon [Var]
       -- case e of y { C ys -> ... }
-      -- See Note [Floating cases] in SetLevels
+      -- See Note [Floating single-alternative cases] in SetLevels
 
 instance Outputable FloatBind where
   ppr (FloatLet b) = text "LET" <+> ppr b
diff --git a/compiler/prelude/PrelRules.hs b/compiler/prelude/PrelRules.hs
index ce269e36f6..cd04074545 100644
--- a/compiler/prelude/PrelRules.hs
+++ b/compiler/prelude/PrelRules.hs
@@ -1042,7 +1042,36 @@ dataToTagRule = a `mplus` b
       ASSERT( not (isNewTyCon (dataConTyCon dc)) ) return ()
       return $ mkIntVal dflags (toInteger (dataConTagZ dc))
 
-{-
+{- Note [dataToTag# magic]
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+The primop dataToTag# is unusual because it evaluates its argument.
+Only `SeqOp` shares that property.  (Other primops do not do anything
+as fancy as argument evaluation.)  The special handling for dataToTag#
+is:
+
+* CoreUtils.exprOkForSpeculation has a special case for DataToTagOp,
+  (actually in app_ok).  Most primops with lifted arguments do not
+  evaluate those arguments, but DataToTagOp and SeqOp are two
+  exceptions.  We say that they are /never/ ok-for-speculation,
+  regardless of the evaluated-ness of their argument.
+  See CoreUtils Note [exprOkForSpeculation and SeqOp/DataToTagOp]
+
+* There is a special case for DataToTagOp in StgCmmExpr.cgExpr,
+  that evaluates its argument and then extracts the tag from
+  the returned value.
+
+* An application like (dataToTag# (Just x)) is optimised by
+  dataToTagRule in PrelRules.
+
+* A case expression like
+     case (dataToTag# e) of <alts>
+  gets transformed t
+     case e of <transformed alts>
+  by PrelRules.caseRules; see Note [caseRules for dataToTag]
+
+See Trac #15696 for a long saga.
+
+
 ************************************************************************
 *                                                                      *
 \subsection{Rules for seq# and spark#}
@@ -1097,7 +1126,7 @@ Implementing seq#.  The compiler has magic for SeqOp in
 - StgCmmExpr.cgExpr, and cgCase: special case for seq#
 
 - CoreUtils.exprOkForSpeculation;
-  see Note [seq# and expr_ok] in CoreUtils
+  see Note [exprOkForSpeculation and SeqOp/DataToTagOp] in CoreUtils
 
 - Simplify.addEvals records evaluated-ness for the result; see
   Note [Adding evaluatedness info to pattern-bound variables]
diff --git a/compiler/prelude/primops.txt.pp b/compiler/prelude/primops.txt.pp
index bf69776166..eb635fb215 100644
--- a/compiler/prelude/primops.txt.pp
+++ b/compiler/prelude/primops.txt.pp
@@ -3150,65 +3150,12 @@ section "Tag to enum stuff"
 primop  DataToTagOp "dataToTag#" GenPrimOp
    a -> Int#  -- Zero-indexed; the first constructor has tag zero
    with
-   can_fail   = True -- See Note [dataToTag#]
    strictness = { \ _arity -> mkClosedStrictSig [evalDmd] topRes }
+   -- See Note [dataToTag# magic] in PrelRules
 
 primop  TagToEnumOp "tagToEnum#" GenPrimOp
    Int# -> a
 
--- Note [dataToTag#]
--- ~~~~~~~~~~~~~~~~~
--- The primop dataToTag# is unusual because it evaluates its argument.
--- Only `SeqOp` shares that property.  (Other primops do not do anything
--- as fancy as argument evaluation.)  The special handling for dataToTag#
--- is:
---
--- * CoreUtils.exprOkForSpeculation has a special case for DataToTagOp,
---   (actually in app_ok).  Most primops with lifted arguments do not
---   evaluate those arguments, but DataToTagOp and SeqOp are two
---   exceptions.  We say that they are /never/ ok-for-speculation,
---   regardless of the evaluated-ness of their argument.
---   See CoreUtils Note [PrimOps that evaluate their arguments]
---
--- * There is a special case for DataToTagOp in StgCmmExpr.cgExpr,
---   that evaluates its argument and then extracts the tag from
---   the returned value.
---
--- * An application like (dataToTag# (Just x)) is optimised by
---   dataToTagRule in PrelRules.
---
--- * A case expression like
---      case (dataToTag# e) of <alts>
---   gets transformed t
---      case e of <transformed alts>
---   by PrelRules.caseRules; see Note [caseRules for dataToTag]
---
--- See Trac #15696 for a long saga.
---
--- Note [dataToTag# hack]
--- ~~~~~~~~~~~~~~~~~~~~~~
--- (This a temporary hack: see Trac #15696 commment:60.)
---
--- dataToTag# evaluates its argument, so we don't want to float it out.
--- Consider:
---
---     \z. case x of y -> let v = dataToTag# y in ...
---
--- To improve floating, the FloatOut pass (deliberately) does a
--- binder-swap on the case, to give
---
---     \z. case x of y -> let v = dataToTag# x in ...
---
--- Now FloatOut might float that v-binding outside the \z
---
---     let v = dataToTag# x in \z. case x of y -> ...
---
--- But that is bad because that might mean x gets evaluated much too early!
---
--- Solution: make dataToTag# into a can_fail primop.  That will stop it floating
--- (see Note [PrimOp can_fail and has_side_effects] in PrimOp).  It's a bit of
--- a hack but never mind.
-
 ------------------------------------------------------------------------
 section "Bytecode operations"
         {Support for manipulating bytecode objects used by the interpreter and
diff --git a/compiler/simplCore/FloatIn.hs b/compiler/simplCore/FloatIn.hs
index 2593b1d7a1..e8c7ef2460 100644
--- a/compiler/simplCore/FloatIn.hs
+++ b/compiler/simplCore/FloatIn.hs
@@ -142,7 +142,8 @@ fiExpr :: DynFlags
        -> CoreExprWithFVs   -- Input expr
        -> CoreExpr          -- Result
 
-fiExpr _ to_drop (_, AnnLit lit)     = ASSERT( null to_drop ) Lit lit
+fiExpr _ to_drop (_, AnnLit lit)     = wrapFloats to_drop (Lit lit)
+                                       -- See Note [Dead bindings]
 fiExpr _ to_drop (_, AnnType ty)     = ASSERT( null to_drop ) Type ty
 fiExpr _ to_drop (_, AnnVar v)       = wrapFloats to_drop (Var v)
 fiExpr _ to_drop (_, AnnCoercion co) = wrapFloats to_drop (Coercion co)
@@ -202,7 +203,15 @@ fiExpr dflags to_drop ann_expr@(_,AnnApp {})
       where
        (arg_ty, res_ty) = splitFunTy fun_ty
 
-{-
+{- Note [Dead bindings]
+~~~~~~~~~~~~~~~~~~~~~~~
+At a literal we won't usually have any floated bindings; the
+only way that can happen is if the binding wrapped the literal
+/in the original input program/.  e.g.
+   case x of { DEFAULT -> 1# }
+But, while this may be unusual it is not actually wrong, and it did
+once happen (Trac #15696).
+
 Note [Do not destroy the let/app invariant]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 Watch out for
diff --git a/compiler/simplCore/SetLevels.hs b/compiler/simplCore/SetLevels.hs
index a63ed27407..c3802bec71 100644
--- a/compiler/simplCore/SetLevels.hs
+++ b/compiler/simplCore/SetLevels.hs
@@ -460,13 +460,12 @@ lvlCase :: LevelEnv             -- Level of in-scope names/tyvars
         -> [CoreAltWithFVs]     -- Input alternatives
         -> LvlM LevelledExpr    -- Result expression
 lvlCase env scrut_fvs scrut' case_bndr ty alts
+  -- See Note [Floating single-alternative cases]
   | [(con@(DataAlt {}), bs, body)] <- alts
-  , exprOkForSpeculation (deTagExpr scrut')
-                                  -- See Note [Check the output scrutinee for okForSpec]
+  , exprIsHNF (deTagExpr scrut')  -- See Note [Check the output scrutinee for okForSpec]
   , not (isTopLvl dest_lvl)       -- Can't have top-level cases
   , not (floatTopLvlOnly env)     -- Can float anywhere
-  =     -- See Note [Floating cases]
-        -- Always float the case if possible
+  =     -- Always float the case if possible
         -- Unlike lets we don't insist that it escapes a value lambda
     do { (env1, (case_bndr' : bs')) <- cloneCaseBndrs env dest_lvl (case_bndr : bs)
        ; let rhs_env = extendCaseBndrEnv env1 case_bndr scrut'
@@ -492,15 +491,15 @@ lvlCase env scrut_fvs scrut' case_bndr ty alts
       where
         (new_env, bs') = substAndLvlBndrs NonRecursive alts_env incd_lvl bs
 
-{-
-Note [Floating cases]
-~~~~~~~~~~~~~~~~~~~~~
+{- Note [Floating single-alternative cases]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 Consider this:
   data T a = MkT !a
   f :: T Int -> blah
   f x vs = case x of { MkT y ->
              let f vs = ...(case y of I# w -> e)...f..
              in f vs
+
 Here we can float the (case y ...) out, because y is sure
 to be evaluated, to give
   f x vs = case x of { MkT y ->
@@ -512,13 +511,28 @@ That saves unboxing it every time round the loop.  It's important in
 some DPH stuff where we really want to avoid that repeated unboxing in
 the inner loop.
 
-Things to note
+Things to note:
+
+ * The test we perform is exprIsHNF, and /not/ exprOkForSpeculation.
+
+     - exrpIsHNF catches the key case of an evaluated variable
+
+     - exprOkForSpeculaion is /false/ of an evaluated varaible;
+       See Note [exprOkForSpeculation and evaluated variables] in CoreUtils
+       So we'd actually miss the key case!
+
+     - Nothing is gained from the extra generality of exprOkForSpeculation
+       since we only consider floating a case whose single alternative
+       is a DataAlt   K a b -> rhs
+
  * We can't float a case to top level
+
  * It's worth doing this float even if we don't float
    the case outside a value lambda.  Example
      case x of {
        MkT y -> (case y of I# w2 -> ..., case y of I# w2 -> ...)
    If we floated the cases out we could eliminate one of them.
+
  * We only do this with a single-alternative case
 
 Note [Check the output scrutinee for okForSpec]
@@ -535,7 +549,7 @@ speculation here, but the former is not -- and indeed we can't float
 the inner case out, at least not unless x is also evaluated at its
 binding site.  See Trac #5453.
 
-That's why we apply exprOkForSpeculation to scrut' and not to scrut.
+That's why we apply exprIsHNF to scrut' and not to scrut.
 -}
 
 lvlNonTailMFE :: LevelEnv             -- Level of in-scope names/tyvars