Fix binary-trees regression from unnecessary floating in CorePrep.wip/D2471

Summary:
In 0d3bf62092de83375025edca6f7242812338542d, I handled lazy @(Int -> Int) f x
correctly, but failed to handle lazy @Int (f x) (we need
to collect arguments in f x).  Additionally, if we have
lazy @Int (case ...) (or anything that's not an application,
we have to bounce back to cpeRhsE.

Signed-off-by: Edward Z. Yang <ezyang@cs.stanford.edu>

Test Plan: validate

Reviewers: simonpj, nomeata, austin, bgamari

Subscribers: thomie

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

4 files changed, 76 insertions, 12 deletions

diff --git a/compiler/coreSyn/CorePrep.hs b/compiler/coreSyn/CorePrep.hs
index 0d82be5abc..b3d78172dc 100644
--- a/compiler/coreSyn/CorePrep.hs
+++ b/compiler/coreSyn/CorePrep.hs
@@ -668,13 +668,8 @@ cpeApp :: CorePrepEnv -> CoreExpr -> UniqSM (Floats, CpeRhs)
 -- May return a CpeRhs because of saturating primops
 cpeApp top_env expr
   = do { let (terminal, args, depth) = collect_args expr
-       ; (head, app, floats) <- cpe_app top_env terminal args depth
-
-        -- Now deal with the function
-       ; case head of
-           Just fn_id -> do { sat_app <- maybeSaturate fn_id app depth
-                            ; return (floats, sat_app) }
-           _other              -> return (floats, app) }
+       ; cpe_app top_env terminal args depth
+       }
 
   where
     -- We have a nested data structure of the form
@@ -702,11 +697,25 @@ cpeApp top_env expr
             -> CoreExpr
             -> [CpeArg]
             -> Int
-            -> UniqSM (Maybe Id, CpeApp, Floats)
+            -> UniqSM (Floats, CpeRhs)
     cpe_app env (Var f) (CpeArg Type{} : CpeArg arg : args) depth
         | f `hasKey` lazyIdKey          -- Replace (lazy a) with a, and
        || f `hasKey` noinlineIdKey      -- Replace (noinline a) with a
-        = cpe_app env arg args (depth - 1)
+        -- Consider the code:
+        --
+        --      lazy (f x) y
+        --
+        -- We need to make sure that we need to recursively collect arguments on
+        -- "f x", otherwise we'll float "f x" out (it's not a variable) and
+        -- end up with this awful -ddump-prep:
+        --
+        --      case f x of f_x {
+        --        __DEFAULT -> f_x y
+        --      }
+        --
+        -- rather than the far superior "f x y".  Test case is par01.
+        = let (terminal, args', depth') = collect_args arg
+          in cpe_app env terminal (args' ++ args) (depth + depth' - 1)
     cpe_app env (Var f) [CpeArg _runtimeRep@Type{}, CpeArg _type@Type{}, CpeArg arg] 1
         | f `hasKey` runRWKey
         -- Replace (runRW# f) by (f realWorld#), beta reducing if possible (this
@@ -724,7 +733,7 @@ cpeApp top_env expr
            -- cpe_ExprIsTrivial).  But note that we need the type of the
            -- expression, not the id.
            ; (app, floats) <- rebuild_app args e2 (exprType e2) emptyFloats stricts
-           ; return (hd, app, floats) }
+           ; mb_saturate hd app floats depth }
         where
           stricts = case idStrictness v of
                             StrictSig (DmdType _ demands _)
@@ -737,16 +746,27 @@ cpeApp top_env expr
                 -- Here, we can't evaluate the arg strictly, because this
                 -- partial application might be seq'd
 
+        -- We inlined into something that's not a var and has no args.
+        -- Bounce it back up to cpeRhsE.
+    cpe_app env fun [] _ = cpeRhsE env fun
+
         -- N-variable fun, better let-bind it
-    cpe_app env fun args _
+    cpe_app env fun args depth
       = do { (fun_floats, fun') <- cpeArg env evalDmd fun ty
                           -- The evalDmd says that it's sure to be evaluated,
                           -- so we'll end up case-binding it
            ; (app, floats) <- rebuild_app args fun' ty fun_floats []
-           ; return (Nothing, app, floats) }
+           ; mb_saturate Nothing app floats depth }
         where
           ty = exprType fun
 
+    -- Saturate if necessary
+    mb_saturate head app floats depth =
+       case head of
+         Just fn_id -> do { sat_app <- maybeSaturate fn_id app depth
+                          ; return (floats, sat_app) }
+         _other              -> return (floats, app)
+
     -- Deconstruct and rebuild the application, floating any non-atomic
     -- arguments to the outside.  We collect the type of the expression,
     -- the head of the application, and the number of actual value arguments,
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index 92f9af4797..e2e0bb6c31 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -243,3 +243,4 @@ test('T12076sat', normal, compile, ['-O'])
 test('T12212', normal, compile, ['-O'])
 
 test('noinline01', only_ways(['optasm']), compile, ['-ddump-stg -dsuppress-uniques -O'])
+test('par01', only_ways(['optasm']), compile, ['-ddump-prep -dsuppress-uniques -O2'])
diff --git a/testsuite/tests/simplCore/should_compile/par01.hs b/testsuite/tests/simplCore/should_compile/par01.hs
new file mode 100644
index 0000000000..e67fb13a4b
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/par01.hs
@@ -0,0 +1,10 @@
+module Par01 where
+
+import GHC.Conc
+
+-- The smoking gun in -ddump-prep is:
+--  case Par01.depth d of sat { __DEFAULT -> sat }
+-- this should never happen!
+
+depth :: Int -> Int
+depth d = d `par` depth d
diff --git a/testsuite/tests/simplCore/should_compile/par01.stderr b/testsuite/tests/simplCore/should_compile/par01.stderr
new file mode 100644
index 0000000000..90d467f71c
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/par01.stderr
@@ -0,0 +1,33 @@
+
+==================== CorePrep ====================
+Result size of CorePrep = {terms: 18, types: 8, coercions: 0}
+
+Rec {
+-- RHS size: {terms: 7, types: 3, coercions: 0}
+Par01.depth [Occ=LoopBreaker] :: GHC.Types.Int -> GHC.Types.Int
+[GblId, Arity=1, Caf=NoCafRefs, Str=<L,U>, Unf=OtherCon []]
+Par01.depth =
+  \ (d :: GHC.Types.Int) ->
+    case GHC.Prim.par# @ GHC.Types.Int d of { __DEFAULT ->
+    Par01.depth d
+    }
+end Rec }
+
+-- RHS size: {terms: 2, types: 0, coercions: 0}
+Par01.$trModule2 :: GHC.Types.TrName
+[GblId, Caf=NoCafRefs, Str=m1, Unf=OtherCon []]
+Par01.$trModule2 = GHC.Types.TrNameS "main"#
+
+-- RHS size: {terms: 2, types: 0, coercions: 0}
+Par01.$trModule1 :: GHC.Types.TrName
+[GblId, Caf=NoCafRefs, Str=m1, Unf=OtherCon []]
+Par01.$trModule1 = GHC.Types.TrNameS "Par01"#
+
+-- RHS size: {terms: 3, types: 0, coercions: 0}
+Par01.$trModule :: GHC.Types.Module
+[GblId, Caf=NoCafRefs, Str=m, Unf=OtherCon []]
+Par01.$trModule =
+  GHC.Types.Module Par01.$trModule2 Par01.$trModule1
+
+
+