Improve the handling of used-once stuffbetter-ho-cardinality

Joachim and I are committing this onto a branch so that we can share it,
but we expect to do a bit more work before merging it onto head.

Nofib staus:
  - Most programs, no change
  - A few improve
  - A couple get worse (cacheprof, tak, rfib)
Investigating the "get worse" set is what's holding up putting this
on head.

The major issue is this.  Consider

    map (f g) ys

where f's demand signature looks like

   f :: <L,C1(C1(U))> -> <L,U> -> .

So 'f' is not saturated.  What demand do we place on g?
Answer
        C(C1(U))
That is, the inner C1 should stay, even though f is not saturated.

I found that this made a significant difference in the demand signatures
inferred in GHC.IO, which uses lots of higher-order exception handlers.

I also had to add used-once demand signatures for some of the
'catch' primops, so that we know their handlers are only called once.

3 files changed, 50 insertions, 38 deletions

diff --git a/compiler/simplCore/OccurAnal.lhs b/compiler/simplCore/OccurAnal.lhs
index 6106388fa4..11391a3553 100644
--- a/compiler/simplCore/OccurAnal.lhs
+++ b/compiler/simplCore/OccurAnal.lhs
@@ -1084,7 +1084,7 @@ occAnalNonRecRhs env bndr rhs
   = occAnal rhs_env rhs
   where
     -- See Note [Use one-shot info]
-    env1 = env { occ_one_shots = argOneShots dmd }
+    env1 = env { occ_one_shots = argOneShots OneShotLam dmd }
 
     -- See Note [Cascading inlines]
     rhs_env | certainly_inline = env1
@@ -1234,13 +1234,14 @@ occAnal env expr@(Lam _ _)
         (final_usage, tagged_binders) = tagLamBinders body_usage binders'
                       -- Use binders' to put one-shot info on the lambdas
 
-        really_final_usage | linear    = final_usage
-                           | otherwise = mapVarEnv markInsideLam final_usage
+        really_final_usage
+          | all isOneShotBndr binders' = final_usage
+          | otherwise = mapVarEnv markInsideLam final_usage
     in
     (really_final_usage, mkLams tagged_binders body') }
   where
-    (binders, body)              = collectBinders expr
-    (env_body, binders', linear) = oneShotGroup env binders
+    (binders, body)      = collectBinders expr
+    (env_body, binders') = oneShotGroup env binders
 
 occAnal env (Case scrut bndr ty alts)
   = case occ_anal_scrut scrut alts     of { (scrut_usage, scrut') ->
@@ -1332,15 +1333,16 @@ occAnalApp env (Var fun, args)
     in
     (fun_uds +++ final_args_uds, mkApps (Var fun) args') }
   where
-    fun_uds  = mkOneOcc env fun (valArgCount args > 0)
-    is_exp = isExpandableApp fun (valArgCount args)
+    n_val_args = valArgCount args
+    fun_uds    = mkOneOcc env fun (n_val_args > 0)
+    is_exp     = isExpandableApp fun n_val_args
            -- See Note [CONLIKE pragma] in BasicTypes
            -- The definition of is_exp should match that in
            -- Simplify.prepareRhs
 
-    one_shots  = argsOneShots (idStrictness fun) (valArgCount args)
+    one_shots  = argsOneShots (idStrictness fun) n_val_args
                  -- See Note [Use one-shot info]
-        
+
     args_stuff = occAnalArgs env args one_shots
 
                         -- (foldr k z xs) may call k many times, but it never
@@ -1466,15 +1468,11 @@ instance Outputable OccEncl where
   ppr OccRhs     = ptext (sLit "occRhs")
   ppr OccVanilla = ptext (sLit "occVanilla")
 
-type OneShots = [Bool]
+type OneShots = [OneShotInfo]
         -- []           No info
         --
-        -- True:ctxt    Analysing a function-valued expression that will be
-        --                      applied just once
-        --
-        -- False:ctxt   Analysing a function-valued expression that may
-        --                      be applied many times; but when it is,
-        --                      the OneShots inside applies
+        -- one_shot_info:ctxt    Analysing a function-valued expression that
+        --                       will be applied as described by one_shot_info
 
 initOccEnv :: (Activation -> Bool) -> OccEnv
 initOccEnv active_rule
@@ -1502,38 +1500,37 @@ isRhsEnv (OccEnv { occ_encl = OccVanilla }) = False
 
 oneShotGroup :: OccEnv -> [CoreBndr] 
              -> ( OccEnv
-                , [CoreBndr]
-                , Bool )   -- True <=> all binders are one-shot
+                , [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
         -- 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
 
 oneShotGroup env@(OccEnv { occ_one_shots = ctxt }) bndrs
-  = go ctxt bndrs [] True
+  = go ctxt bndrs []
   where
-    go ctxt [] rev_bndrs linear 
+    go ctxt [] rev_bndrs
       = ( env { occ_one_shots = ctxt, occ_encl = OccVanilla }
-        , reverse rev_bndrs
-        , linear )
+        , reverse rev_bndrs )
 
-    go ctxt (bndr:bndrs) rev_bndrs lin_acc
+    go [] bndrs rev_bndrs
+      = ( env { occ_one_shots = [], occ_encl = OccVanilla }
+        , reverse rev_bndrs ++ bndrs )
+
+    go ctxt (bndr:bndrs) rev_bndrs
       | isId bndr
+      
       = case ctxt of
-          []            -> go [] bndrs (bndr:rev_bndrs) (lin_acc && one_shot)
-          (linear : ctxt) 
-             | one_shot  -> go ctxt bndrs (bndr : rev_bndrs) lin_acc
-             | linear    -> go ctxt bndrs (bndr': rev_bndrs) lin_acc
-             | otherwise -> go ctxt bndrs (bndr : rev_bndrs) False
-      | otherwise
-      = go ctxt bndrs (bndr:rev_bndrs) lin_acc
-      where
-        one_shot = isOneShotBndr bndr
-        bndr'    = setOneShotLambda bndr
+          []                -> go []   bndrs (bndr : rev_bndrs)
+          (one_shot : ctxt) -> go ctxt bndrs (bndr': rev_bndrs)
+                            where
+                               bndr' = updOneShotInfo bndr one_shot
+       | otherwise
+      = go ctxt bndrs (bndr:rev_bndrs)
 
 addAppCtxt :: OccEnv -> [Arg CoreBndr] -> OccEnv
 addAppCtxt env@(OccEnv { occ_one_shots = ctxt }) args
-  = env { occ_one_shots = replicate (valArgCount args) True ++ ctxt }
+  = env { occ_one_shots = replicate (valArgCount args) OneShotLam ++ ctxt }
 \end{code}
 
 
diff --git a/compiler/simplCore/SetLevels.lhs b/compiler/simplCore/SetLevels.lhs
index 2fca56cf17..7bcc53f6de 100644
--- a/compiler/simplCore/SetLevels.lhs
+++ b/compiler/simplCore/SetLevels.lhs
@@ -815,7 +815,9 @@ lvlLamBndrs lvl bndrs
     new_lvl | any is_major bndrs = incMajorLvl lvl
             | otherwise          = incMinorLvl lvl
 
-    is_major bndr = isId bndr && not (isOneShotLambda bndr)
+    is_major bndr = isId bndr && not (isProbablyOneShotLambda bndr)
+       -- The "probably" part says "don't float things out of a
+       -- probable one-shot lambda"
 \end{code}
 
 \begin{code}
diff --git a/compiler/simplCore/SimplUtils.lhs b/compiler/simplCore/SimplUtils.lhs
index 6c7dcc2042..36f292deb3 100644
--- a/compiler/simplCore/SimplUtils.lhs
+++ b/compiler/simplCore/SimplUtils.lhs
@@ -1195,8 +1195,11 @@ tryEtaExpandRhs env bndr rhs
   = do { dflags <- getDynFlags
        ; (new_arity, new_rhs) <- try_expand dflags
 
-       ; WARN( new_arity < old_arity || new_arity < _dmd_arity,
-               (ptext (sLit "Arity decrease:") <+> (ppr bndr <+> ppr old_arity
+       ; WARN( new_arity < old_arity,
+               (ptext (sLit "Arity decrease:") <+> (ppr bndr
+                <+> ppr old_arity <+> ppr new_arity) $$ ppr new_rhs) )
+         WARN( new_arity < _dmd_arity,
+               (ptext (sLit "Arity less than dmd sig arity:") <+> (ppr bndr
                 <+> ppr new_arity <+> ppr _dmd_arity) $$ ppr new_rhs) )
                         -- Note [Arity decrease]
          return (new_arity, new_rhs) }
@@ -1211,13 +1214,23 @@ tryEtaExpandRhs env bndr rhs
       = do { tick (EtaExpansion bndr)
            ; return (new_arity, etaExpand new_arity rhs) }
       | otherwise
-      = return (manifest_arity, rhs)
+      = return (exprArity rhs, rhs)   -- See Note [Return exprArity, not manifestArity]
 
     manifest_arity = manifestArity rhs
     old_arity  = idArity bndr
     _dmd_arity = length $ fst $ splitStrictSig $ idStrictness bndr
 \end{code}
 
+Note [Return exprArity, not manifestArity]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Consider
+  f = \xy. blah
+  g = f 2
+The f will get arity 2, and we want g to get arity 1, even though
+exprEtaExpandArity (and hence findArity) may not eta-expand it.
+Hence tryEtaExpand should return (exprArity (f 2)), not its
+manifest arity (which is zero).
+
 Note [Eta-expanding at let bindings]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 We now eta expand at let-bindings, which is where the payoff comes.