Demand Analyzer: Do not set OneShot information (second try)

as suggested in ticket:11770#comment:1. This code was buggy
(#11770), and the occurrence analyzer does the same job anyways.

This also elaborates the notes in the occurrence analyzer accordingly.

Previously, the worker/wrapper code would go through lengths to transfer
the oneShot annotations from the original function to both the worker
and the wrapper. We now simply transfer the demand on the worker, and
let the subsequent occurrence analyzer push this onto the lambda
binders.

This also requires the occurrence analyzer to do this more reliably.
Previously, it would not hand out OneShot annotatoins to things that
would not `certainly_inline` (and it might not have mattered, as the
Demand Analysis might have handed out the annotations). Now we hand out
one-shot annotations unconditionally.

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

1 files changed, 29 insertions, 13 deletions

diff --git a/compiler/simplCore/OccurAnal.hs b/compiler/simplCore/OccurAnal.hs
index 6628ee70ee..41a6f7fa71 100644
--- a/compiler/simplCore/OccurAnal.hs
+++ b/compiler/simplCore/OccurAnal.hs
@@ -1104,12 +1104,13 @@ occAnalNonRecRhs :: OccEnv
 occAnalNonRecRhs env bndr rhs
   = occAnal rhs_env rhs
   where
+    -- See Note [Cascading inlines]
+    env1 | certainly_inline = env
+         | otherwise        = rhsCtxt env
+
     -- See Note [Use one-shot info]
-    env1 = env { occ_one_shots = argOneShots OneShotLam dmd }
+    rhs_env = env1 { occ_one_shots = argOneShots OneShotLam dmd }
 
-    -- See Note [Cascading inlines]
-    rhs_env | certainly_inline = env1
-            | otherwise        = rhsCtxt env1
 
     certainly_inline -- See Note [Cascading inlines]
       = case idOccInfo bndr of
@@ -1395,19 +1396,29 @@ markManyIf False uds = uds
 {-
 Note [Use one-shot information]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-The occurrrence analyser propagates one-shot-lambda information in two situation
-  * Applications:  eg   build (\cn -> blah)
+The occurrrence analyser propagates one-shot-lambda information in two
+situations:
+
+  * Applications:  eg   build (\c n -> blah)
+
     Propagate one-shot info from the strictness signature of 'build' to
-    the \cn
+    the \c n.
+
+    This strictness signature can come from a module interface, in the case of
+    an imported function, or from a previous run of the demand analyser.
 
   * Let-bindings:  eg   let f = \c. let ... in \n -> blah
                         in (build f, build f)
+
     Propagate one-shot info from the demanand-info on 'f' to the
     lambdas in its RHS (which may not be syntactically at the top)
 
-Some of this is done by the demand analyser, but this way it happens
-much earlier, taking advantage of the strictness signature of
-imported functions.
+    This information must have come from a previous run of the demanand
+    analyser.
+
+Previously, the demand analyser would *also* set the one-shot information, but
+that code was buggy (see #11770), so doing it only in on place, namely here, is
+saner.
 
 Note [Binders in case alternatives]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -1534,7 +1545,7 @@ oneShotGroup :: OccEnv -> [CoreBndr]
              -> ( OccEnv
                 , [CoreBndr] )
         -- The result binders have one-shot-ness set that they might not have had originally.
-        -- This happens in (build (\cn -> e)).  Here the occurrence analyser
+        -- This happens in (build (\c n -> e)).  Here the occurrence analyser
         -- linearity context knows that c,n are one-shot, and it records that fact in
         -- the binder. This is useful to guide subsequent float-in/float-out tranformations
 
@@ -1555,8 +1566,13 @@ oneShotGroup env@(OccEnv { occ_one_shots = ctxt }) bndrs
       = case ctxt of
           []                -> go []   bndrs (bndr : rev_bndrs)
           (one_shot : ctxt) -> go ctxt bndrs (bndr': rev_bndrs)
-                            where
-                               bndr' = updOneShotInfo bndr one_shot
+            where
+               bndr' = updOneShotInfo bndr one_shot
+               -- Use updOneShotInfo, not setOneShotInfo, as pre-existing
+               -- one-shot info might be better than what we can infer, e.g.
+               -- due to explicit use of the magic 'oneShot' function.
+               -- See Note [The oneShot function]
+
        | otherwise
       = go ctxt bndrs (bndr:rev_bndrs)