Revert "Demand Analyzer: Do not set OneShot information"

This reverts commit 28fe0eea4d161b707f67aae26fddaa2e60d8a901 due to
various regressions. I’m not sure why my local
./validate --slow
run did not catch this, though.

3 files changed, 60 insertions, 30 deletions

diff --git a/compiler/simplCore/OccurAnal.hs b/compiler/simplCore/OccurAnal.hs
index 91332c4416..6628ee70ee 100644
--- a/compiler/simplCore/OccurAnal.hs
+++ b/compiler/simplCore/OccurAnal.hs
@@ -1395,29 +1395,19 @@ markManyIf False uds = uds
 {-
 Note [Use one-shot information]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-The occurrrence analyser propagates one-shot-lambda information in two
-situations:
-
-  * Applications:  eg   build (\c n -> blah)
-
+The occurrrence analyser propagates one-shot-lambda information in two situation
+  * Applications:  eg   build (\cn -> blah)
     Propagate one-shot info from the strictness signature of 'build' to
-    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.
+    the \cn
 
   * 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)
 
-    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.
+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.
 
 Note [Binders in case alternatives]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -1544,7 +1534,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 (\c n -> e)).  Here the occurrence analyser
+        -- This happens in (build (\cn -> 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
 
@@ -1565,13 +1555,8 @@ 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
-               -- 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]
-
+                            where
+                               bndr' = updOneShotInfo bndr one_shot
        | otherwise
       = go ctxt bndrs (bndr:rev_bndrs)
 
diff --git a/compiler/stranal/DmdAnal.hs b/compiler/stranal/DmdAnal.hs
index 20f65d5904..6ef911f6c0 100644
--- a/compiler/stranal/DmdAnal.hs
+++ b/compiler/stranal/DmdAnal.hs
@@ -200,9 +200,13 @@ dmdAnal' env dmd (Lam var body)
   = let (body_dmd, defer_and_use) = peelCallDmd dmd
           -- body_dmd: a demand to analyze the body
 
+        one_shot         = useCount (getUseDmd defer_and_use)
+          -- one_shot: one-shotness of the lambda
+          --           hence, cardinality of its free vars
+
         env'             = extendSigsWithLam env var
         (body_ty, body') = dmdAnal env' body_dmd body
-        (lam_ty, var')   = annotateLamIdBndr env notArgOfDfun body_ty var
+        (lam_ty, var')   = annotateLamIdBndr env notArgOfDfun body_ty one_shot var
     in
     (postProcessUnsat defer_and_use lam_ty, Lam var' body')
 
@@ -256,13 +260,17 @@ dmdAnal' env dmd (Case scrut case_bndr ty alts)
     (res_ty, Case scrut' case_bndr' ty alts')
 
 dmdAnal' env dmd (Let (NonRec id rhs) body)
-  = (body_ty2, Let (NonRec id2 rhs') body')
+  = (body_ty2, Let (NonRec id2 annotated_rhs) body')
   where
     (sig, lazy_fv, id1, rhs') = dmdAnalRhs NotTopLevel Nothing env id rhs
     (body_ty, body')          = dmdAnal (extendAnalEnv NotTopLevel env id sig) dmd body
     (body_ty1, id2)           = annotateBndr env body_ty id1
     body_ty2                  = addLazyFVs body_ty1 lazy_fv
 
+    -- Annotate top-level lambdas at RHS basing on the aggregated demand info
+    -- See Note [Annotating lambdas at right-hand side]
+    annotated_rhs = annLamWithShotness (idDemandInfo id2) rhs'
+
         -- If the actual demand is better than the vanilla call
         -- demand, you might think that we might do better to re-analyse
         -- the RHS with the stronger demand.
@@ -299,6 +307,25 @@ io_hack_reqd scrut con bndrs
   | otherwise
   = False
 
+annLamWithShotness :: Demand -> CoreExpr -> CoreExpr
+annLamWithShotness d e
+  | Just u <- cleanUseDmd_maybe d
+  = go u e
+  | otherwise = e
+  where
+    go u e
+      | Just (c, u') <- peelUseCall u
+      , Lam bndr body <- e
+      = if isTyVar bndr
+        then Lam bndr                    (go u  body)
+        else Lam (setOneShotness c bndr) (go u' body)
+      | otherwise
+      = e
+
+setOneShotness :: Count -> Id -> Id
+setOneShotness One  bndr = setOneShotLambda bndr
+setOneShotness Many bndr = bndr
+
 dmdAnalAlt :: AnalEnv -> CleanDemand -> Id -> Alt Var -> (DmdType, Alt Var)
 dmdAnalAlt env dmd case_bndr (con,bndrs,rhs)
   | null bndrs    -- Literals, DEFAULT, and nullary constructors
@@ -405,6 +432,23 @@ free variable |y|. Conversely, if the demand on |h| is unleashed right
 on the spot, we will get the desired result, namely, that |f| is
 strict in |y|.
 
+Note [Annotating lambdas at right-hand side]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Let us take a look at the following example:
+
+g f = let x = 100
+          h = \y -> f x y
+       in h 5
+
+One can see that |h| is called just once, therefore the RHS of h can
+be annotated as a one-shot lambda. This is done by the function
+annLamWithShotness *a posteriori*, i.e., basing on the aggregated
+usage demand on |h| from the body of |let|-expression, which is C1(U)
+in this case.
+
+In other words, for locally-bound lambdas we can infer
+one-shotness.
+
 
 ************************************************************************
 *                                                                      *
@@ -705,22 +749,23 @@ annotateLamBndrs :: AnalEnv -> DFunFlag -> DmdType -> [Var] -> (DmdType, [Var])
 annotateLamBndrs env args_of_dfun ty bndrs = mapAccumR annotate ty bndrs
   where
     annotate dmd_ty bndr
-      | isId bndr = annotateLamIdBndr env args_of_dfun dmd_ty bndr
+      | isId bndr = annotateLamIdBndr env args_of_dfun dmd_ty Many bndr
       | otherwise = (dmd_ty, bndr)
 
 annotateLamIdBndr :: AnalEnv
                   -> DFunFlag   -- is this lambda at the top of the RHS of a dfun?
                   -> DmdType    -- Demand type of body
+                  -> Count      -- One-shot-ness of the lambda
                   -> Id         -- Lambda binder
                   -> (DmdType,  -- Demand type of lambda
                       Id)       -- and binder annotated with demand
 
-annotateLamIdBndr env arg_of_dfun dmd_ty id
+annotateLamIdBndr env arg_of_dfun dmd_ty one_shot id
 -- For lambdas we add the demand to the argument demands
 -- Only called for Ids
   = ASSERT( isId id )
     -- pprTrace "annLamBndr" (vcat [ppr id, ppr _dmd_ty]) $
-    (final_ty, setIdDemandInfo id dmd)
+    (final_ty, setOneShotness one_shot (setIdDemandInfo id dmd))
   where
       -- Watch out!  See note [Lambda-bound unfoldings]
     final_ty = case maybeUnfoldingTemplate (idUnfolding id) of
diff --git a/testsuite/tests/stranal/should_compile/all.T b/testsuite/tests/stranal/should_compile/all.T
index dabc9fcd84..3ac075b716 100644
--- a/testsuite/tests/stranal/should_compile/all.T
+++ b/testsuite/tests/stranal/should_compile/all.T
@@ -45,6 +45,6 @@ test('T9208', when(compiler_debugged(), expect_broken(9208)), compile, [''])
 # Hence the above expect_broken.  See comments in the Trac ticket
 
 test('T10694', [ grepCoreString(r'Str=') ],   compile, ['-dppr-cols=200 -ddump-simpl'])
-test('T11770', [ checkCoreString('OneShot') ], compile, ['-ddump-simpl'])
+test('T11770', [ expect_broken(117700), checkCoreString("OneShot") ], compile, ['-ddump-simpl'])