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

3 files changed, 68 insertions, 121 deletions

diff --git a/compiler/stranal/DmdAnal.hs b/compiler/stranal/DmdAnal.hs
index 6ef911f6c0..20f65d5904 100644
--- a/compiler/stranal/DmdAnal.hs
+++ b/compiler/stranal/DmdAnal.hs
@@ -200,13 +200,9 @@ 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 one_shot var
+        (lam_ty, var')   = annotateLamIdBndr env notArgOfDfun body_ty var
     in
     (postProcessUnsat defer_and_use lam_ty, Lam var' body')
 
@@ -260,17 +256,13 @@ 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 annotated_rhs) body')
+  = (body_ty2, Let (NonRec id2 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.
@@ -307,25 +299,6 @@ 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
@@ -432,23 +405,6 @@ 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.
-
 
 ************************************************************************
 *                                                                      *
@@ -749,23 +705,22 @@ 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 Many bndr
+      | isId bndr = annotateLamIdBndr env args_of_dfun dmd_ty 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 one_shot id
+annotateLamIdBndr env arg_of_dfun dmd_ty 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, setOneShotness one_shot (setIdDemandInfo id dmd))
+    (final_ty, setIdDemandInfo id dmd)
   where
       -- Watch out!  See note [Lambda-bound unfoldings]
     final_ty = case maybeUnfoldingTemplate (idUnfolding id) of
diff --git a/compiler/stranal/WorkWrap.hs b/compiler/stranal/WorkWrap.hs
index 8a5ed67513..7fde65f7a0 100644
--- a/compiler/stranal/WorkWrap.hs
+++ b/compiler/stranal/WorkWrap.hs
@@ -10,7 +10,6 @@ module WorkWrap ( wwTopBinds ) where
 import CoreSyn
 import CoreUnfold       ( certainlyWillInline, mkWwInlineRule, mkWorkerUnfolding )
 import CoreUtils        ( exprType, exprIsHNF )
-import CoreArity        ( exprArity )
 import Var
 import Id
 import IdInfo
@@ -330,7 +329,7 @@ splitFun :: DynFlags -> FamInstEnvs -> Id -> IdInfo -> [Demand] -> DmdResult ->
 splitFun dflags fam_envs fn_id fn_info wrap_dmds res_info rhs
   = WARN( not (wrap_dmds `lengthIs` arity), ppr fn_id <+> (ppr arity $$ ppr wrap_dmds $$ ppr res_info) ) do
     -- The arity should match the signature
-    stuff <- mkWwBodies dflags fam_envs fun_ty wrap_dmds res_info one_shots
+    stuff <- mkWwBodies dflags fam_envs fun_ty wrap_dmds res_info
     case stuff of
       Just (work_demands, wrap_fn, work_fn) -> do
         work_uniq <- getUniqueM
@@ -360,8 +359,18 @@ splitFun dflags fam_envs fn_id fn_info wrap_dmds res_info rhs
                                 -- Even though we may not be at top level,
                                 -- it's ok to give it an empty DmdEnv
 
-                        `setIdArity` exprArity work_rhs
+                        `setIdDemandInfo` worker_demand
+
+                        `setIdArity` work_arity
                                 -- Set the arity so that the Core Lint check that the
+
+            work_arity = length work_demands
+
+            -- See Note [Demand on the Worker]
+            single_call = saturatedByOneShots arity (demandInfo fn_info)
+            worker_demand | single_call = mkWorkerDemand work_arity
+                          | otherwise   = topDmd
+
                                 -- arity is consistent with the demand type goes through
 
             wrap_act  = ActiveAfter "0" 0
@@ -380,6 +389,8 @@ splitFun dflags fam_envs fn_id fn_info wrap_dmds res_info rhs
                                 -- Zap any loop-breaker-ness, to avoid bleating from Lint
                                 -- about a loop breaker with an INLINE rule
 
+
+
         return $ [(work_id, work_rhs), (wrap_id, wrap_rhs)]
             -- Worker first, because wrapper mentions it
 
@@ -396,20 +407,39 @@ splitFun dflags fam_envs fn_id fn_info wrap_dmds res_info rhs
                        Just _  -> topRes    -- Cpr stuff done by wrapper; kill it here
                        Nothing -> res_info  -- Preserve exception/divergence
 
-    one_shots = get_one_shots rhs
-
--- If the original function has one-shot arguments, it is important to
--- make the wrapper and worker have corresponding one-shot arguments too.
--- Otherwise we spuriously float stuff out of case-expression join points,
--- which is very annoying.
-get_one_shots :: Expr Var -> [OneShotInfo]
-get_one_shots (Lam b e)
-  | isId b    = idOneShotInfo b : get_one_shots e
-  | otherwise = get_one_shots e
-get_one_shots (Tick _ e) = get_one_shots e
-get_one_shots _          = []
 
 {-
+Note [Demand on the worker]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+If the original function is called once, according to its demand info, then
+so is the worker. This is important so that the occurrence analyser can
+attach OneShot annotations to the worker’s lambda binders.
+
+
+Example:
+
+  -- Original function
+  f [Demand=<L,1*C1(U)>] :: (a,a) -> a
+  f = \p -> ...
+
+  -- Wrapper
+  f [Demand=<L,1*C1(U)>] :: a -> a -> a
+  f = \p -> case p of (a,b) -> $wf a b
+
+  -- Worker
+  $wf [Demand=<L,1*C1(C1(U))>] :: Int -> Int
+  $wf = \a b -> ...
+
+We need to check whether the original function is called once, with
+sufficiently many arguments. This is done using saturatedByOneShots, which
+takes the arity of the original function (resp. the wrapper) and the demand on
+the original function.
+
+The demand on the worker is then calculated using mkWorkerDemand, and always of
+the form [Demand=<L,1*(C1(...(C1(U))))>]
+
+
 Note [Do not split void functions]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 Consider this rather common form of binding:
diff --git a/compiler/stranal/WwLib.hs b/compiler/stranal/WwLib.hs
index 3d9ab8365a..7c85036c1f 100644
--- a/compiler/stranal/WwLib.hs
+++ b/compiler/stranal/WwLib.hs
@@ -25,7 +25,7 @@ import TysWiredIn       ( tupleDataCon )
 import Type
 import Coercion
 import FamInstEnv
-import BasicTypes       ( Boxity(..), OneShotInfo(..), worstOneShot )
+import BasicTypes       ( Boxity(..) )
 import Literal          ( absentLiteralOf )
 import TyCon
 import UniqSupply
@@ -111,7 +111,6 @@ mkWwBodies :: DynFlags
            -> Type                                  -- Type of original function
            -> [Demand]                              -- Strictness of original function
            -> DmdResult                             -- Info about function result
-           -> [OneShotInfo]                         -- One-shot-ness of the function, value args only
            -> UniqSM (Maybe ([Demand],              -- Demands for worker (value) args
                              Id -> CoreExpr,        -- Wrapper body, lacking only the worker Id
                              CoreExpr -> CoreExpr)) -- Worker body, lacking the original function rhs
@@ -127,22 +126,20 @@ mkWwBodies :: DynFlags
 --                        let x = (a,b) in
 --                        E
 
-mkWwBodies dflags fam_envs fun_ty demands res_info one_shots
-  = do  { let arg_info = demands `zip` (one_shots ++ repeat NoOneShotInfo)
-              all_one_shots = foldr (worstOneShot . snd) OneShotLam arg_info
-        ; (wrap_args, wrap_fn_args, work_fn_args, res_ty) <- mkWWargs emptyTCvSubst fun_ty arg_info
+mkWwBodies dflags fam_envs fun_ty demands res_info
+  = do  { (wrap_args, wrap_fn_args, work_fn_args, res_ty) <- mkWWargs emptyTCvSubst fun_ty demands
         ; (useful1, work_args, wrap_fn_str, work_fn_str) <- mkWWstr dflags fam_envs wrap_args
 
         -- Do CPR w/w.  See Note [Always do CPR w/w]
         ; (useful2, wrap_fn_cpr, work_fn_cpr, cpr_res_ty)
               <- mkWWcpr (gopt Opt_CprAnal dflags) fam_envs res_ty res_info
 
-        ; let (work_lam_args, work_call_args) = mkWorkerArgs dflags work_args all_one_shots cpr_res_ty
+        ; let (work_lam_args, work_call_args) = mkWorkerArgs dflags work_args cpr_res_ty
               worker_args_dmds = [idDemandInfo v | v <- work_call_args, isId v]
               wrapper_body = wrap_fn_args . wrap_fn_cpr . wrap_fn_str . applyToVars work_call_args . Var
               worker_body = mkLams work_lam_args. work_fn_str . work_fn_cpr . work_fn_args
 
-        ; if useful1 && not (only_one_void_argument) || useful2
+        ; if useful1 && not only_one_void_argument || useful2
           then return (Just (worker_args_dmds, wrapper_body, worker_body))
           else return Nothing
         }
@@ -196,24 +193,20 @@ We use the state-token type which generates no code.
 -}
 
 mkWorkerArgs :: DynFlags -> [Var]
-             -> OneShotInfo  -- Whether all arguments are one-shot
              -> Type    -- Type of body
              -> ([Var], -- Lambda bound args
                  [Var]) -- Args at call site
-mkWorkerArgs dflags args all_one_shot res_ty
+mkWorkerArgs dflags args res_ty
     | any isId args || not needsAValueLambda
     = (args, args)
     | otherwise
-    = (args ++ [newArg], args ++ [voidPrimId])
+    = (args ++ [voidArgId], args ++ [voidPrimId])
     where
       needsAValueLambda =
         isUnliftedType res_ty
         || not (gopt Opt_FunToThunk dflags)
            -- see Note [Protecting the last value argument]
 
-      -- see Note [All One-Shot Arguments of a Worker]
-      newArg = setIdOneShotInfo voidArgId all_one_shot
-
 {-
 Note [Protecting the last value argument]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -228,29 +221,6 @@ create a space leak. 2) It can prevent inlining *under a lambda*. If w/w
 removes the last argument from a function f, then f now looks like a thunk, and
 so f can't be inlined *under a lambda*.
 
-Note [All One-Shot Arguments of a Worker]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-Sometimes, derived join-points are just lambda-lifted thunks, whose
-only argument is of the unit type and is never used. This might
-interfere with the absence analysis, basing on which results these
-never-used arguments are eliminated in the worker. The additional
-argument `all_one_shot` of `mkWorkerArgs` is to prevent this.
-
-Example.  Suppose we have
-   foo = \p(one-shot) q(one-shot). y + 3
-Then we drop the unused args to give
-   foo   = \pq. $wfoo void#
-   $wfoo = \void(one-shot). y + 3
-
-But suppse foo didn't have all one-shot args:
-   foo = \p(not-one-shot) q(one-shot). expensive y + 3
-Then we drop the unused args to give
-   foo   = \pq. $wfoo void#
-   $wfoo = \void(not-one-shot). y + 3
-
-If we made the void-arg one-shot we might inline an expensive
-computation for y, which would be terrible!
-
 
 ************************************************************************
 *                                                                      *
@@ -292,23 +262,23 @@ the \x to get what we want.
 mkWWargs :: TCvSubst            -- Freshening substitution to apply to the type
                                 --   See Note [Freshen type variables]
          -> Type                -- The type of the function
-         -> [(Demand,OneShotInfo)]     -- Demands and one-shot info for value arguments
+         -> [Demand]     -- Demands and one-shot info for value arguments
          -> UniqSM  ([Var],            -- Wrapper args
                      CoreExpr -> CoreExpr,      -- Wrapper fn
                      CoreExpr -> CoreExpr,      -- Worker fn
                      Type)                      -- Type of wrapper body
 
-mkWWargs subst fun_ty arg_info
-  | null arg_info
+mkWWargs subst fun_ty demands
+  | null demands
   = return ([], id, id, substTy subst fun_ty)
 
-  | ((dmd,one_shot):arg_info') <- arg_info
+  | (dmd:demands') <- demands
   , Just (arg_ty, fun_ty') <- splitFunTy_maybe fun_ty
   = do  { uniq <- getUniqueM
         ; let arg_ty' = substTy subst arg_ty
-              id = mk_wrap_arg uniq arg_ty' dmd one_shot
+              id = mk_wrap_arg uniq arg_ty' dmd
         ; (wrap_args, wrap_fn_args, work_fn_args, res_ty)
-              <- mkWWargs subst fun_ty' arg_info'
+              <- mkWWargs subst fun_ty' demands'
         ; return (id : wrap_args,
                   Lam id . wrap_fn_args,
                   work_fn_args . (`App` varToCoreExpr id),
@@ -319,7 +289,7 @@ mkWWargs subst fun_ty arg_info
                 -- This substTyVarBndr clones the type variable when necy
                 -- See Note [Freshen type variables]
         ; (wrap_args, wrap_fn_args, work_fn_args, res_ty)
-             <- mkWWargs subst' fun_ty' arg_info
+             <- mkWWargs subst' fun_ty' demands
         ; return (tv' : wrap_args,
                   Lam tv' . wrap_fn_args,
                   work_fn_args . (`mkTyApps` [mkTyVarTy tv']),
@@ -335,7 +305,7 @@ mkWWargs subst fun_ty arg_info
         -- simply coerces.
 
   = do { (wrap_args, wrap_fn_args, work_fn_args, res_ty)
-            <-  mkWWargs subst rep_ty arg_info
+            <-  mkWWargs subst rep_ty demands
         ; return (wrap_args,
                   \e -> Cast (wrap_fn_args e) (mkSymCo co),
                   \e -> work_fn_args (Cast e co),
@@ -348,11 +318,10 @@ mkWWargs subst fun_ty arg_info
 applyToVars :: [Var] -> CoreExpr -> CoreExpr
 applyToVars vars fn = mkVarApps fn vars
 
-mk_wrap_arg :: Unique -> Type -> Demand -> OneShotInfo -> Id
-mk_wrap_arg uniq ty dmd one_shot
+mk_wrap_arg :: Unique -> Type -> Demand -> Id
+mk_wrap_arg uniq ty dmd
   = mkSysLocalOrCoVar (fsLit "w") uniq ty
        `setIdDemandInfo` dmd
-       `setIdOneShotInfo` one_shot
 
 {-
 Note [Freshen type variables]
@@ -472,7 +441,7 @@ mkWWstr_one dflags fam_envs arg
       -- See Note [mkWWstr and unsafeCoerce]
   =  do { (uniq1:uniqs) <- getUniquesM
         ; let   unpk_args      = zipWith mk_ww_local uniqs inst_con_arg_tys
-                unpk_args_w_ds = zipWithEqual "mkWWstr" set_worker_arg_info unpk_args cs
+                unpk_args_w_ds = zipWithEqual "mkWWstr" setIdDemandInfo unpk_args cs
                 unbox_fn       = mkUnpackCase (Var arg) co uniq1
                                               data_con unpk_args
                 rebox_fn       = Let (NonRec arg con_app)
@@ -486,13 +455,6 @@ mkWWstr_one dflags fam_envs arg
 
   where
     dmd = idDemandInfo arg
-    one_shot = idOneShotInfo arg
-        -- If the wrapper argument is a one-shot lambda, then
-        -- so should (all) the corresponding worker arguments be
-        -- This bites when we do w/w on a case join point
-    set_worker_arg_info worker_arg demand
-      = worker_arg `setIdDemandInfo`  demand
-                   `setIdOneShotInfo` one_shot
 
 ----------------------
 nop_fn :: CoreExpr -> CoreExpr