Add comments explaining ProbOneShot

4 files changed, 79 insertions, 44 deletions

diff --git a/compiler/basicTypes/BasicTypes.lhs b/compiler/basicTypes/BasicTypes.lhs
index 2f86db7796..4fbfb6007a 100644
--- a/compiler/basicTypes/BasicTypes.lhs
+++ b/compiler/basicTypes/BasicTypes.lhs
@@ -155,9 +155,11 @@ type Alignment = Int -- align to next N-byte boundary (N must be a power of 2).
 -- This information may be useful in optimisation, as computations may
 -- safely be floated inside such a lambda without risk of duplicating
 -- work.
-data OneShotInfo = NoOneShotInfo -- ^ No information
-                 | ProbOneShot   -- ^ The lambda is probably applied at most once
-                 | OneShotLam    -- ^ The lambda is applied at most once.
+data OneShotInfo
+  = NoOneShotInfo -- ^ No information
+  | ProbOneShot   -- ^ The lambda is probably applied at most once
+                  -- See Note [Computing one-shot info, and ProbOneShot] in OccurAnl
+  | OneShotLam    -- ^ The lambda is applied at most once.
 
 -- | It is always safe to assume that an 'Id' has no lambda-bound variable information
 noOneShotInfo :: OneShotInfo
diff --git a/compiler/basicTypes/Demand.lhs b/compiler/basicTypes/Demand.lhs
index 2aa25ced53..f553fc2ae5 100644
--- a/compiler/basicTypes/Demand.lhs
+++ b/compiler/basicTypes/Demand.lhs
@@ -1493,6 +1493,11 @@ newtype StrictSig = StrictSig DmdType
 instance Outputable StrictSig where
    ppr (StrictSig ty) = ppr ty
 
+-- Used for printing top-level strictness pragmas in interface files
+pprIfaceStrictSig :: StrictSig -> SDoc
+pprIfaceStrictSig (StrictSig (DmdType _ dmds res))
+  = hcat (map ppr dmds) <> ppr res
+
 mkStrictSig :: DmdType -> StrictSig
 mkStrictSig dmd_ty = StrictSig dmd_ty
 
@@ -1520,29 +1525,8 @@ botSig = StrictSig botDmdType
 cprProdSig :: Arity -> StrictSig
 cprProdSig arity = StrictSig (cprProdDmdType arity)
 
-argsOneShots :: StrictSig -> Arity -> [[OneShotInfo]]
-argsOneShots (StrictSig (DmdType _ arg_ds _)) n_val_args
-  = go arg_ds
-  where
-    good_one_shot
-      | arg_ds `lengthExceeds` n_val_args = ProbOneShot
-      | otherwise                         = OneShotLam
-
-    go []               = []
-    go (arg_d : arg_ds) = argOneShots good_one_shot arg_d `cons` go arg_ds
-
-    cons [] [] = []
-    cons a  as = a:as
-
-argOneShots :: OneShotInfo -> JointDmd -> [OneShotInfo]
-argOneShots one_shot_info (JD { absd = usg })
-  = case usg of
-      Use _ arg_usg -> go arg_usg
-      _             -> []
-  where
-    go (UCall One  u) = one_shot_info : go u
-    go (UCall Many u) = NoOneShotInfo : go u
-    go _              = []
+seqStrictSig :: StrictSig -> ()
+seqStrictSig (StrictSig ty) = seqDmdType ty
 
 dmdTransformSig :: StrictSig -> CleanDemand -> DmdType
 -- (dmdTransformSig fun_sig dmd) considers a call to a function whose
@@ -1617,31 +1601,79 @@ you might do strictness analysis, but there is no inlining for the class op.
 This is weird, so I'm not worried about whether this optimises brilliantly; but
 it should not fall over.
 
-Note [Non-full application] 
-~~~~~~~~~~~~~~~~~~~~~~~~~~~ 
-If a function having bottom as its demand result is applied to a less
-number of arguments than its syntactic arity, we cannot say for sure
-that it is going to diverge. This is the reason why we use the
-function appIsBottom, which, given a strictness signature and a number
-of arguments, says conservatively if the function is going to diverge
-or not.
+\begin{code}
+argsOneShots :: StrictSig -> Arity -> [[OneShotInfo]]
+-- See Note [Computing one-shot info, and ProbOneShot]
+argsOneShots (StrictSig (DmdType _ arg_ds _)) n_val_args
+  = go arg_ds
+  where
+    unsaturated_call = arg_ds `lengthExceeds` n_val_args
+    good_one_shot
+      | unsaturated_call = ProbOneShot
+      | otherwise        = OneShotLam
+
+    go []               = []
+    go (arg_d : arg_ds) = argOneShots good_one_shot arg_d `cons` go arg_ds
+
+    -- Avoid list tail like [ [], [], [] ]
+    cons [] [] = []
+    cons a  as = a:as
+
+argOneShots :: OneShotInfo -> JointDmd -> [OneShotInfo]
+argOneShots one_shot_info (JD { absd = usg })
+  = case usg of
+      Use _ arg_usg -> go arg_usg
+      _             -> []
+  where
+    go (UCall One  u) = one_shot_info : go u
+    go (UCall Many u) = NoOneShotInfo : go u
+    go _              = []
+\end{code}
+
+Note [Computing one-shot info, and ProbOneShot]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Consider a call
+    f (\pqr. e1) (\xyz. e2) e3
+where f has usage signature
+    C1(C(C1(U))) C1(U) U
+Then argsOneShots returns a [[OneShotInfo]] of
+    [[OneShot,NoOneShotInfo,OneShot],  [OneShot]]
+The occurrence analyser propagates this one-shot infor to the
+binders \pqr and \xyz; see Note [Use one-shot information] in OccurAnal.
+
+But suppose f was not saturated, so the call looks like
+    f (\pqr. e1) (\xyz. e2)
+The in principle this partial application might be shared, and
+the (\prq.e1) abstraction might be called more than once.  So
+we can't mark them OneShot. But instead we return
+    [[ProbOneShot,NoOneShotInfo,ProbOneShot],  [ProbOneShot]]
+The occurrence analyser propagates this to the \pqr and \xyz
+binders.
+
+How is it used?  Well, it's quite likely that the partial application
+of f is not shared, so the float-out pass (in SetLevels.lvlLamBndrs)
+does not float MFEs out of a ProbOneShot lambda.  That currently is
+the only way that ProbOneShot is used.
+
 
 \begin{code}
 -- appIsBottom returns true if an application to n args would diverge
+-- See Note [Unsaturated applications]
 appIsBottom :: StrictSig -> Int -> Bool
 appIsBottom (StrictSig (DmdType _ ds res)) n
             | isBotRes res                      = not $ lengthExceeds ds n 
 appIsBottom _                                 _ = False
-
-seqStrictSig :: StrictSig -> ()
-seqStrictSig (StrictSig ty) = seqDmdType ty
-
--- Used for printing top-level strictness pragmas in interface files
-pprIfaceStrictSig :: StrictSig -> SDoc
-pprIfaceStrictSig (StrictSig (DmdType _ dmds res))
-  = hcat (map ppr dmds) <> ppr res
 \end{code}
 
+Note [Unsaturated applications]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+If a function having bottom as its demand result is applied to a less
+number of arguments than its syntactic arity, we cannot say for sure
+that it is going to diverge. This is the reason why we use the
+function appIsBottom, which, given a strictness signature and a number
+of arguments, says conservatively if the function is going to diverge
+or not.
+
 Zap absence or one-shot information, under control of flags
 
 \begin{code}
diff --git a/compiler/simplCore/OccurAnal.lhs b/compiler/simplCore/OccurAnal.lhs
index 0fbd2cac99..ef212bca85 100644
--- a/compiler/simplCore/OccurAnal.lhs
+++ b/compiler/simplCore/OccurAnal.lhs
@@ -1380,13 +1380,13 @@ The occurrrence analyser propagates one-shot-lambda information in two situation
     Propagate one-shot info from the strictness signature of 'build' to
     the \cn
 
-  * Let-bindings:  eg   let f = \c. let ... in \n -> blah 
+  * 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 
+much earlier, taking advantage of the strictness signature of
 imported functions.
 
 Note [Binders in case alternatives]
diff --git a/compiler/simplCore/SetLevels.lhs b/compiler/simplCore/SetLevels.lhs
index e5cd42ec30..b8726d93a4 100644
--- a/compiler/simplCore/SetLevels.lhs
+++ b/compiler/simplCore/SetLevels.lhs
@@ -827,6 +827,7 @@ lvlLamBndrs env lvl bndrs
     is_major bndr = isId bndr && not (isProbablyOneShotLambda bndr)
        -- The "probably" part says "don't float things out of a
        -- probable one-shot lambda"
+       -- See Note [Computing one-shot info] in Demand.lhs
 
 
 lvlBndrs :: LevelEnv -> Level -> [CoreBndr] -> (LevelEnv, [LevelledBndr])