When floating, don't box an expression that's okay for speculation (#13338)

Summary:
Commit 432f952e (Float unboxed expressions by boxing) lets the float-out pass
turn, for example,

    ... (-# (remInt# x# 100000#) i#) ...

into

    let lvl :: Int
        lvl = case remInt# x# 100000# of v { __DEFAULT__ -> I# v }
    in ... (-# (case lvl of { I# v -> v }) i#) ...

But when, as in the example above, the expression that was floated out was
the argument of an application, the resulting application may no longer
satisfy the let/app invariant, because exprOkForSpeculation doesn't look
far enough inside the definition of lvl.

Solution: When the expression we floated out was okay for speculation, don't
bother boxing it. It will be evaluated earlier, and that's okay by assumption.
Fixes the let/app invariant and is cheaper too.

Test Plan: make slowtest TEST=T13338

Reviewers: simonpj, austin, bgamari

Subscribers: thomie

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

3 files changed, 21 insertions, 1 deletions

diff --git a/compiler/simplCore/SetLevels.hs b/compiler/simplCore/SetLevels.hs
index 22d4048767..0e067cc002 100644
--- a/compiler/simplCore/SetLevels.hs
+++ b/compiler/simplCore/SetLevels.hs
@@ -561,9 +561,15 @@ lvlMFE env strict_ctxt ann_expr
   =     -- Don't float it out
     lvlExpr env ann_expr
 
-  | float_is_new_lam || need_join || exprIsTopLevelBindable expr expr_ty
+  |  float_is_new_lam || need_join || exprIsTopLevelBindable expr expr_ty
+  || exprOkForSpeculation expr && not (isTopLvl dest_lvl)
          -- 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
+         -- or is okay for speculation (we'll now evaluate it earlier).
+         -- In the last case we _must not_ wrap, because it could violate
+         -- the let/app invariant (Trac #13338).
+         -- But we can't float an unboxed thing to top level; so don't float
+         -- it all, as in lvlBind. (See "Don't break let/app" below.)
   = do { expr1 <- lvlFloatRhs abs_vars dest_lvl rhs_env NonRecursive join_arity_maybe ann_expr
                   -- Treat the expr just like a right-hand side
        ; var <- newLvlVar expr1 join_arity_maybe
@@ -578,6 +584,7 @@ lvlMFE env strict_ctxt ann_expr
   | escapes_value_lam
   , not (exprIsCheap expr)  -- Boxing/unboxing isn't worth
                             -- it for cheap expressions
+  , not (exprOkForSpeculation expr && isTopLvl dest_lvl) -- Don't break let/app
   , Just (tc, _) <- splitTyConApp_maybe expr_ty
   , Just dc <- boxingDataCon_maybe tc
   , let dc_res_ty = dataConOrigResTy dc  -- No free type variables
diff --git a/testsuite/tests/simplCore/should_compile/T13338.hs b/testsuite/tests/simplCore/should_compile/T13338.hs
new file mode 100644
index 0000000000..347a9d7dae
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/T13338.hs
@@ -0,0 +1,12 @@
+{-# LANGUAGE MagicHash #-}
+
+module T13338 where
+
+import GHC.Exts
+
+magic# :: Int# -> Bool
+magic# x# = True
+{-# NOINLINE magic# #-}
+
+f :: Int# -> Int -> Int
+f x# n = length [ i | i@(I# i#) <- [0..n], magic# (remInt# x# 100000# -# i#) ]
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index 53f5ade353..23cd77cf53 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -244,3 +244,4 @@ test('T13317',
      normal,
      run_command,
      ['$MAKE -s --no-print-directory T13317'])
+test('T13338', only_ways(['optasm']), compile, ['-dcore-lint'])