dmdAnal: Move handling of datacon strictness to mkWWstr_one

Previously datacon strictness was accounted for when we demand analysed a case
analysis. However, this results in pessimistic demands in some cases. For
instance, consider the program (from T10482)

    data family Bar a
    data instance Bar (a, b) = BarPair !(Bar a) !(Bar b)
    newtype instance Bar Int = Bar Int

    foo :: Bar ((Int, Int), Int) -> Int -> Int
    foo f k =
      case f of
        BarPair x y -> case burble of
                          True -> case x of
                                    BarPair p q -> ...
                          False -> ...

We really should be able to assume that `p` is already evaluated since it came
from a strict field of BarPair.

However, as written the demand analyser can not conclude this since we may end
up in the False branch of the case on `burble` (which places no demand on `x`).
By accounting for the data con strictness later, applied to the demand of the
RHS, we get the strict demand signature we want.

See Note [Add demands for strict constructors] for a more comprehensive
discussion.

Test Plan: Validate

Reviewers: simonpj, osa1, goldfire

Subscribers: rwbarton, carter

GHC Trac Issues: #15696

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

2 files changed, 81 insertions, 59 deletions

diff --git a/compiler/stranal/DmdAnal.hs b/compiler/stranal/DmdAnal.hs
index 995911939f..0b8133d98f 100644
--- a/compiler/stranal/DmdAnal.hs
+++ b/compiler/stranal/DmdAnal.hs
@@ -250,7 +250,7 @@ dmdAnal' env dmd (Case scrut case_bndr ty [(DataAlt dc, bndrs, rhs)])
 
         -- Compute demand on the scrutinee
         -- See Note [Demand on scrutinee of a product case]
-        scrut_dmd          = mkProdDmd (addDataConStrictness dc id_dmds)
+        scrut_dmd          = mkProdDmd id_dmds
         (scrut_ty, scrut') = dmdAnal env scrut_dmd scrut
         res_ty             = alt_ty3 `bothDmdType` toBothDmdArg scrut_ty
         case_bndr'         = setIdDemandInfo case_bndr case_bndr_dmd
@@ -1214,17 +1214,6 @@ extendEnvForProdAlt env scrut case_bndr dc bndrs
     is_var (Var v)    = isLocalId v
     is_var _          = False
 
-addDataConStrictness :: DataCon -> [Demand] -> [Demand]
--- See Note [Add demands for strict constructors]
-addDataConStrictness con ds
-  = ASSERT2( equalLength strs ds, ppr con $$ ppr strs $$ ppr ds )
-    zipWith add ds strs
-  where
-    strs = dataConRepStrictness con
-    add dmd str | isMarkedStrict str
-                , not (isAbsDmd dmd) = strictifyDmd dmd
-                | otherwise          = dmd
-
 findBndrsDmds :: AnalEnv -> DmdType -> [Var] -> (DmdType, [Demand])
 -- Return the demands on the Ids in the [Var]
 findBndrsDmds env dmd_ty bndrs
@@ -1308,8 +1297,8 @@ binders the CPR property.  Specifically
                    | otherwise = x
 
    For $wf2 we are going to unbox the MkT *and*, since it is strict, the
-   first argument of the MkT; see Note [Add demands for strict constructors].
-   But then we don't want box it up again when returning it!  We want
+   first argument of the MkT; see Note [Add demands for strict constructors]
+   in WwLib. But then we don't want box it up again when returning it! We want
    'f2' to have the CPR property, so we give 'x' the CPR property.
 
  * It's a bit delicate because if this case is scrutinising something other
@@ -1325,50 +1314,6 @@ binders the CPR property.  Specifically
    sub-component thereof.  But it's simple, and nothing terrible
    happens if we get it wrong.  e.g. Trac #10694.
 
-Note [Add demands for strict constructors]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-Consider this program (due to Roman):
-
-    data X a = X !a
-
-    foo :: X Int -> Int -> Int
-    foo (X a) n = go 0
-     where
-       go i | i < n     = a + go (i+1)
-            | otherwise = 0
-
-We want the worker for 'foo' too look like this:
-
-    $wfoo :: Int# -> Int# -> Int#
-
-with the first argument unboxed, so that it is not eval'd each time
-around the 'go' loop (which would otherwise happen, since 'foo' is not
-strict in 'a').  It is sound for the wrapper to pass an unboxed arg
-because X is strict, so its argument must be evaluated.  And if we
-*don't* pass an unboxed argument, we can't even repair it by adding a
-`seq` thus:
-
-    foo (X a) n = a `seq` go 0
-
-because the seq is discarded (very early) since X is strict!
-
-We achieve the effect using addDataConStrictness.  It is called at a
-case expression, such as the pattern match on (X a) in the example
-above.  After computing how 'a' is used in the alternatives, we add an
-extra 'seqDmd' to it.  The case alternative isn't itself strict in the
-sub-components, but simply evaluating the scrutinee to HNF does force
-those sub-components.
-
-If the argument is not used at all in the alternative (i.e. it is
-Absent), then *don't* add a 'seqDmd'.  If we do, it makes it look used
-and hence it'll be passed to the worker when it doesn't need to be.
-Hence the isAbsDmd test in addDataConStrictness.
-
-There is the usual danger of reboxing, which as usual we ignore. But
-if X is monomorphic, and has an UNPACK pragma, then this optimisation
-is even more important.  We don't want the wrapper to rebox an unboxed
-argument, and pass an Int to $wfoo!
-
 
 Note [Initial CPR for strict binders]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
diff --git a/compiler/stranal/WwLib.hs b/compiler/stranal/WwLib.hs
index f01dc6c385..ce036c8c26 100644
--- a/compiler/stranal/WwLib.hs
+++ b/compiler/stranal/WwLib.hs
@@ -614,7 +614,9 @@ mkWWstr_one dflags fam_envs has_inlineable_prag arg
   , cs `equalLength` inst_con_arg_tys
       -- See Note [mkWWstr and unsafeCoerce]
   = do { (uniq1:uniqs) <- getUniquesM
-        ; let   unpk_args = zipWith3 mk_ww_arg uniqs inst_con_arg_tys cs
+        ; let   -- See Note [Add demands for strict constructors]
+                cs'       = addDataConStrictness data_con cs
+                unpk_args = zipWith3 mk_ww_arg uniqs inst_con_arg_tys cs'
                 unbox_fn  = mkUnpackCase (Var arg) co uniq1
                                          data_con unpk_args
                 arg_no_unf = zapStableUnfolding arg
@@ -638,7 +640,82 @@ mkWWstr_one dflags fam_envs has_inlineable_prag arg
 nop_fn :: CoreExpr -> CoreExpr
 nop_fn body = body
 
+addDataConStrictness :: DataCon -> [Demand] -> [Demand]
+-- See Note [Add demands for strict constructors]
+addDataConStrictness con ds
+  = ASSERT2( equalLength strs ds, ppr con $$ ppr strs $$ ppr ds )
+    zipWith add ds strs
+  where
+    strs = dataConRepStrictness con
+    add dmd str | isMarkedStrict str
+                , not (isAbsDmd dmd) = dmd `bothDmd` seqDmd
+                | otherwise          = dmd
+
 {-
+Note [Add demands for strict constructors]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Consider this program (due to Roman):
+
+    data X a = X !a
+
+    foo :: X Int -> Int -> Int
+    foo (X a) n = go 0
+     where
+       go i | i < n     = a + go (i+1)
+            | otherwise = 0
+
+We want the worker for 'foo' too look like this:
+
+    $wfoo :: Int# -> Int# -> Int#
+
+with the first argument unboxed, so that it is not eval'd each time
+around the 'go' loop (which would otherwise happen, since 'foo' is not
+strict in 'a').  It is sound for the wrapper to pass an unboxed arg
+because X is strict, so its argument must be evaluated.  And if we
+*don't* pass an unboxed argument, we can't even repair it by adding a
+`seq` thus:
+
+    foo (X a) n = a `seq` go 0
+
+So here's what we do
+
+* We leave the demand-analysis alone. The demand on 'a' in the definition of
+  'foo' is <L, U(U)>; the strictness info is Lazy because foo's body may or may
+  not evaluate 'a'; but the usage info says that 'a' is unpacked and its content
+  is used.
+
+* During worker/wrapper, if we unpack a strict constructor (as we do for 'foo'),
+  we use 'strictifyDemand' to bump up the strictness on the strict arguments of
+  the data constructor. That in turn means that, if the usage info supports
+  doing so (i.e. splitProdDmd_maybe returns Just), we will unpack that argument
+  -- even though the original demand (e.g. on 'a') was lazy.
+
+The net effect is that the w/w transformation is more aggressive about unpacking
+the strict arguments of a data constructor, when that eagerness is supported by
+the usage info.
+
+This works in nested situations like
+
+    data family Bar a
+    data instance Bar (a, b) = BarPair !(Bar a) !(Bar b)
+    newtype instance Bar Int = Bar Int
+
+    foo :: Bar ((Int, Int), Int) -> Int -> Int
+    foo f k =
+      case f of
+        BarPair x y -> case burble of
+                         True -> case x of
+                                   BarPair p q -> ...
+                         False -> ...
+
+The extra eagerness lets us produce a worker of type:
+
+    $wfoo :: Int# -> Int# -> Int# -> Int -> Int
+    $wfoo p# q# y# = ...
+
+even though the `case x` is only lazily evaluated
+
+
 Note [mkWWstr and unsafeCoerce]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 By using unsafeCoerce, it is possible to make the number of demands fail to