Let the simplifier know that seq# forces

Add a special case in `simplAlt` to record that the result of
`seq#` is in WHNF.

Reviewers: simonmar, bgamari, simonpj

Reviewed By: simonpj

Subscribers: simonpj, rwbarton, thomie, carter

GHC Trac Issues: #15226

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

4 files changed, 123 insertions, 37 deletions

diff --git a/compiler/coreSyn/CoreSyn.hs b/compiler/coreSyn/CoreSyn.hs
index c2aeabefe6..4dd70b0c99 100644
--- a/compiler/coreSyn/CoreSyn.hs
+++ b/compiler/coreSyn/CoreSyn.hs
@@ -5,6 +5,7 @@
 
 {-# LANGUAGE CPP, DeriveDataTypeable, FlexibleContexts #-}
 {-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE BangPatterns #-}
 
 -- | CoreSyn holds all the main data types for use by for the Glasgow Haskell Compiler midsection
 module CoreSyn (
@@ -40,7 +41,7 @@ module CoreSyn (
         bindersOf, bindersOfBinds, rhssOfBind, rhssOfAlts,
         collectBinders, collectTyBinders, collectTyAndValBinders,
         collectNBinders,
-        collectArgs, collectArgsTicks, flattenBinds,
+        collectArgs, stripNArgs, collectArgsTicks, flattenBinds,
 
         exprToType, exprToCoercion_maybe,
         applyTypeToArg,
@@ -2044,6 +2045,15 @@ collectArgs expr
     go (App f a) as = go f (a:as)
     go e         as = (e, as)
 
+-- | Attempt to remove the last N arguments of a function call.
+-- Strip off any ticks encountered along the way and any ticks
+-- at the end.
+stripNArgs :: Word -> Expr a -> Maybe (Expr a)
+stripNArgs !n (Tick _ e) = stripNArgs n e
+stripNArgs 0 e = Just e
+stripNArgs n (App f _) = stripNArgs (n - 1) f
+stripNArgs _ _ = Nothing
+
 -- | Like @collectArgs@, but also collects looks through floatable
 -- ticks if it means that we can find more arguments.
 collectArgsTicks :: (Tickish Id -> Bool) -> Expr b
diff --git a/compiler/simplCore/Simplify.hs b/compiler/simplCore/Simplify.hs
index 6d1b434b8f..89e7df2495 100644
--- a/compiler/simplCore/Simplify.hs
+++ b/compiler/simplCore/Simplify.hs
@@ -28,7 +28,9 @@ import Name             ( mkSystemVarName, isExternalName, getOccFS )
 import Coercion hiding  ( substCo, substCoVar )
 import OptCoercion      ( optCoercion )
 import FamInstEnv       ( topNormaliseType_maybe )
-import DataCon          ( DataCon, dataConWorkId, dataConRepStrictness, dataConRepArgTys )
+import DataCon          ( DataCon, dataConWorkId, dataConRepStrictness
+                        , dataConRepArgTys, isUnboxedTupleCon
+                        , StrictnessMark (..) )
 import CoreMonad        ( Tick(..), SimplMode(..) )
 import CoreSyn
 import Demand           ( StrictSig(..), dmdTypeDepth, isStrictDmd )
@@ -50,6 +52,7 @@ import Pair
 import Util
 import ErrUtils
 import Module          ( moduleName, pprModuleName )
+import PrimOp          ( PrimOp (SeqOp) )
 
 
 {-
@@ -2599,11 +2602,8 @@ simplAlt env scrut' _ case_bndr' cont' (LitAlt lit, bndrs, rhs)
         ; return (LitAlt lit, [], rhs') }
 
 simplAlt env scrut' _ case_bndr' cont' (DataAlt con, vs, rhs)
-  = do  {       -- Deal with the pattern-bound variables
-                -- Mark the ones that are in ! positions in the
-                -- data constructor as certainly-evaluated.
-                -- NB: simplLamBinders preserves this eval info
-        ; let vs_with_evals = add_evals (dataConRepStrictness con)
+  = do  { -- See Note [Adding evaluatedness info to pattern-bound variables]
+          let vs_with_evals = addEvals scrut' con vs
         ; (env', vs') <- simplLamBndrs env vs_with_evals
 
                 -- Bind the case-binder to (con args)
@@ -2614,37 +2614,73 @@ simplAlt env scrut' _ case_bndr' cont' (DataAlt con, vs, rhs)
         ; env'' <- addAltUnfoldings env' scrut' case_bndr' con_app
         ; rhs' <- simplExprC env'' rhs cont'
         ; return (DataAlt con, vs', rhs') }
-  where
-        -- add_evals records the evaluated-ness of the bound variables of
-        -- a case pattern.  This is *important*.  Consider
-        --      data T = T !Int !Int
-        --
-        --      case x of { T a b -> T (a+1) b }
-        --
-        -- We really must record that b is already evaluated so that we don't
-        -- go and re-evaluate it when constructing the result.
-        -- See Note [Data-con worker strictness] in MkId.hs
-    add_evals the_strs
-        = go vs the_strs
+
+-- Note [Adding evaluatedness info to pattern-bound variables]
+-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+--
+-- add_evals records the evaluated-ness of the bound variables of
+-- a case pattern.  This is *important*.  Consider
+--
+--      data T = T !Int !Int
+--
+--      case x of { T a b -> T (a+1) b }
+--
+-- We really must record that b is already evaluated so that we don't
+-- go and re-evaluate it when constructing the result.
+-- See Note [Data-con worker strictness] in MkId.hs
+--
+-- NB: simplLamBinders preserves this eval info
+--
+-- In addition to handling data constructor fields with !s, add_evals
+-- also records the fact that the result of seq# is always in WHNF.
+-- in
+--
+--   case seq# v s of
+--     (# s', v' #) -> E
+--
+-- we want the compiler to be aware that v' is in WHNF in E. See #15226.
+-- We don't record that v itself is in WHNF (and we can't do it here).
+-- I don't know if we should attempt to do so.
+
+addEvals :: Maybe OutExpr -> DataCon -> [Id] -> [Id]
+-- See Note [Adding evaluatedness info to pattern-bound variables]
+addEvals scrut con vs
+  -- Deal with seq# applications
+  | Just scr <- scrut
+  , isUnboxedTupleCon con
+  , [s,x] <- vs
+    -- Use stripNArgs rather than collectArgsTicks to avoid building
+    -- a list of arguments only to throw it away immediately.
+  , Just (Var f) <- stripNArgs 4 scr
+  , Just SeqOp <- isPrimOpId_maybe f
+  , let x' = zapIdOccInfoAndSetEvald MarkedStrict x
+  = [s, x']
+
+  -- Deal with banged datacon fields
+addEvals _scrut con vs = go vs the_strs
+    where
+      the_strs = dataConRepStrictness con
+
+      go [] [] = []
+      go (v:vs') strs | isTyVar v = v : go vs' strs
+      go (v:vs') (str:strs) = zapIdOccInfoAndSetEvald str v : go vs' strs
+      go _ _ = pprPanic "Simplify.addEvals"
+                (ppr con $$
+                 ppr vs  $$
+                 ppr_with_length (map strdisp the_strs) $$
+                 ppr_with_length (dataConRepArgTys con) $$
+                 ppr_with_length (dataConRepStrictness con))
         where
-          go [] [] = []
-          go (v:vs') strs | isTyVar v = v : go vs' strs
-          go (v:vs') (str:strs) = zap str v : go vs' strs
-          go _ _ = pprPanic "cat_evals"
-                    (ppr con $$
-                     ppr vs  $$
-                     ppr_with_length the_strs $$
-                     ppr_with_length (dataConRepArgTys con) $$
-                     ppr_with_length (dataConRepStrictness con))
-            where
-              ppr_with_length list
-                = ppr list <+> parens (text "length =" <+> ppr (length list))
-                                    -- NB: If this panic triggers, note that
-                                    -- NoStrictnessMark doesn't print!
-
-          zap str v = setCaseBndrEvald str $ -- Add eval'dness info
-                      zapIdOccInfo v         -- And kill occ info;
-                                             -- see Note [Case alternative occ info]
+          ppr_with_length list
+            = ppr list <+> parens (text "length =" <+> ppr (length list))
+          strdisp MarkedStrict = "MarkedStrict"
+          strdisp NotMarkedStrict = "NotMarkedStrict"
+
+zapIdOccInfoAndSetEvald :: StrictnessMark -> Id -> Id
+zapIdOccInfoAndSetEvald str v =
+  setCaseBndrEvald str $ -- Add eval'dness info
+  zapIdOccInfo v         -- And kill occ info;
+                         -- see Note [Case alternative occ info]
 
 addAltUnfoldings :: SimplEnv -> Maybe OutExpr -> OutId -> OutExpr -> SimplM SimplEnv
 addAltUnfoldings env scrut case_bndr con_app
diff --git a/testsuite/tests/perf/should_run/T15226.hs b/testsuite/tests/perf/should_run/T15226.hs
new file mode 100644
index 0000000000..4c09114b89
--- /dev/null
+++ b/testsuite/tests/perf/should_run/T15226.hs
@@ -0,0 +1,30 @@
+-- T15226
+import Control.Exception (evaluate)
+
+-- Just in case Prelude.repeat changes for some reason.
+import Prelude hiding (repeat)
+
+-- We want to be sure that the compiler *doesn't* know that
+-- all the elements of the list are in WHNF, because if it
+-- does, PrelRules may erase the seq#'s altogether.
+repeat :: a -> [a]
+repeat a = res
+  where res = a : res
+{-# NOINLINE repeat #-}  -- Belt *and* suspenders
+
+silly :: [Int] -> IO ()
+silly = foldr go (pure ())
+  where
+    go x r = do
+      x' <- evaluate x
+      evaluate (x' + 3)  -- GHC should know that x' has been evaluated,
+                         -- so this calculation will be erased entirely.
+                         -- Otherwise, we'll create a thunk to pass to
+                         -- evaluate.
+      r
+
+main :: IO ()
+-- 10,000,000 repetitions take only a twentieth of a second,
+-- but allocations go up dramatically if the result is not
+-- known evaluated.
+main = silly $ take 10000000 $ repeat 1
diff --git a/testsuite/tests/perf/should_run/all.T b/testsuite/tests/perf/should_run/all.T
index 9fd997f633..b248dd56f7 100644
--- a/testsuite/tests/perf/should_run/all.T
+++ b/testsuite/tests/perf/should_run/all.T
@@ -574,3 +574,13 @@ test('T14936',
                        (wordsize(64), 51792, 5) ])],
      compile_and_run,
      ['-O2'])
+
+test('T15226',
+    [stats_num_field('bytes allocated',
+                    [ (wordsize(64), 41040, 5) ]),
+                    # 2018-06-06   41040  Let the simplifier know the result
+                    #                     of seq# is in WHNF
+                    # initial  400041040
+     only_ways(['normal'])],
+    compile_and_run,
+    ['-O'])