Do not inline or apply rules on LHS of rules
[ghc.git] / compiler / simplCore / CallArity.hs
index b43d1fe..c2a5ad0 100644 (file)
@@ -4,6 +4,7 @@
 
 module CallArity
     ( callArityAnalProgram
+    , callArityRHS -- for testing
     ) where
 
 import VarSet
@@ -13,24 +14,27 @@ import DynFlags ( DynFlags )
 import BasicTypes
 import CoreSyn
 import Id
-import CoreArity
+import CoreArity ( typeArity )
+import CoreUtils ( exprIsHNF )
+--import Outputable
+import UnVarGraph
+import Demand
 
-import Control.Arrow ( second )
-import Data.Maybe ( isJust )
+import Control.Arrow ( first, second )
 
 
 {-
 %************************************************************************
-%*                                                                     *
+%*                                                                      *
               Call Arity Analyis
-%*                                                                     *
+%*                                                                      *
 %************************************************************************
 
 Note [Call Arity: The goal]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 The goal of this analysis is to find out if we can eta-expand a local function,
-based on how it is being called. The motivating example is code this this,
+based on how it is being called. The motivating example is this code,
 which comes up when we implement foldl using foldr, and do list fusion:
 
     let go = \x -> let d = case ... of
@@ -43,7 +47,7 @@ If we do not eta-expand `go` to have arity 2, we are going to allocate a lot of
 partial function applications, which would be bad.
 
 The function `go` has a type of arity two, but only one lambda is manifest.
-Further more, an analysis that only looks at the RHS of go cannot be sufficient
+Furthermore, an analysis that only looks at the RHS of go cannot be sufficient
 to eta-expand go: If `go` is ever called with one argument (and the result used
 multiple times), we would be doing the work in `...` multiple times.
 
@@ -56,46 +60,142 @@ The specification of the `calledArity` field is:
 
     No work will be lost if you eta-expand me to the arity in `calledArity`.
 
-The specification of the analysis
----------------------------------
-
-The analysis only does a conservative approximation, there are plenty of
-situations where eta-expansion would be ok, but we do not catch it. We are
-content if all the code that foldl-via-foldr generates is being optimized
-sufficiently.
-
-The work-hourse of the analysis is the function `callArityAnal`, with the
-following type:
-
-    type CallArityEnv = VarEnv (Maybe Arity)
-    callArityAnal ::
-        Arity ->  -- The arity this expression is called with
-        VarSet -> -- The set of interesting variables
-        CoreExpr ->  -- The expression to analyse
-        (CallArityEnv, CoreExpr)
-
-and the following specification:
-
-  (callArityEnv, expr') = callArityEnv arity interestingIds expr
-
-                            <=>
-
-  Assume the expression `expr` is being passed `arity` arguments. Then it calls
-  the functions mentioned in `interestingIds` according to `callArityEnv`:
-    * The domain of `callArityEnv` is a subset of `interestingIds`.
-    * Any variable from interestingIds that is not mentioned in the `callArityEnv`
-      is absent, i.e. not called at all.
-    * Of all the variables that are mapped to a non-Nothing value by `callArityEnv`,
-      at most one is being called, with at least that many arguments.
-    * Nothing can be said about variables mapped to Noting.
-  Furthermore, expr' is expr with the callArity field of the `IdInfo` updated.
-
-The (pointwise) top of the domain is `Nothing`; the least upper bound coincides
-with the mininum on `Maybe Int` with the usual `Ord` instance for `Maybe`.
-
-The at-most-once is important for various reasons:
-
- 1. Consider:
+What we want to know for a variable
+-----------------------------------
+
+For every let-bound variable we'd like to know:
+  1. A lower bound on the arity of all calls to the variable, and
+  2. whether the variable is being called at most once or possible multiple
+     times.
+
+It is always ok to lower the arity, or pretend that there are multiple calls.
+In particular, "Minimum arity 0 and possible called multiple times" is always
+correct.
+
+
+What we want to know from an expression
+---------------------------------------
+
+In order to obtain that information for variables, we analyize expression and
+obtain bits of information:
+
+ I.  The arity analysis:
+     For every variable, whether it is absent, or called,
+     and if called, which what arity.
+
+ II. The Co-Called analysis:
+     For every two variables, whether there is a possibility that both are being
+     called.
+     We obtain as a special case: For every variables, whether there is a
+     possibility that it is being called twice.
+
+For efficiency reasons, we gather this information only for a set of
+*interesting variables*, to avoid spending time on, e.g., variables from pattern matches.
+
+The two analysis are not completely independent, as a higher arity can improve
+the information about what variables are being called once or multiple times.
+
+Note [Analysis I: The arity analyis]
+------------------------------------
+
+The arity analysis is quite straight forward: The information about an
+expression is an
+    VarEnv Arity
+where absent variables are bound to Nothing and otherwise to a lower bound to
+their arity.
+
+When we analyize an expression, we analyize it with a given context arity.
+Lambdas decrease and applications increase the incoming arity. Analysizing a
+variable will put that arity in the environment. In lets or cases all the
+results from the various subexpressions are lubed, which takes the point-wise
+minimum (considering Nothing an infinity).
+
+
+Note [Analysis II: The Co-Called analysis]
+------------------------------------------
+
+The second part is more sophisticated. For reasons explained below, it is not
+sufficient to simply know how often an expression evalutes a variable. Instead
+we need to know which variables are possibly called together.
+
+The data structure here is an undirected graph of variables, which is provided
+by the abstract
+    UnVarGraph
+
+It is safe to return a larger graph, i.e. one with more edges. The worst case
+(i.e. the least useful and always correct result) is the complete graph on all
+free variables, which means that anything can be called together with anything
+(including itself).
+
+Notation for the following:
+C(e)  is the co-called result for e.
+G₁∪G₂ is the union of two graphs
+fv    is the set of free variables (conveniently the domain of the arity analysis result)
+S₁×S₂ is the complete bipartite graph { {a,b} | a ∈ S₁, b ∈ S₂ }
+S²    is the complete graph on the set of variables S, S² = S×S
+C'(e) is a variant for bound expression:
+      If e is called at most once, or it is and stays a thunk (after the analysis),
+      it is simply C(e). Otherwise, the expression can be called multiple times
+      and we return (fv e)²
+
+The interesting cases of the analysis:
+ * Var v:
+   No other variables are being called.
+   Return {} (the empty graph)
+ * Lambda v e, under arity 0:
+   This means that e can be evaluated many times and we cannot get
+   any useful co-call information.
+   Return (fv e)²
+ * Case alternatives alt₁,alt₂,...:
+   Only one can be execuded, so
+   Return (alt₁ ∪ alt₂ ∪...)
+ * App e₁ e₂ (and analogously Case scrut alts):
+   We get the results from both sides. Additionally, anything called by e₁ can
+   possibly be called with anything from e₂.
+   Return: C(e₁) ∪ C(e₂) ∪ (fv e₁) × (fv e₂)
+ * Let v = rhs in body:
+   In addition to the results from the subexpressions, add all co-calls from
+   everything that the body calls together with v to everthing that is called
+   by v.
+   Return: C'(rhs) ∪ C(body) ∪ (fv rhs) × {v'| {v,v'} ∈ C(body)}
+ * Letrec v₁ = rhs₁ ... vₙ = rhsₙ in body
+   Tricky.
+   We assume that it is really mutually recursive, i.e. that every variable
+   calls one of the others, and that this is strongly connected (otherwise we
+   return an over-approximation, so that's ok), see note [Recursion and fixpointing].
+
+   Let V = {v₁,...vₙ}.
+   Assume that the vs have been analysed with an incoming demand and
+   cardinality consistent with the final result (this is the fixed-pointing).
+   Again we can use the results from all subexpressions.
+   In addition, for every variable vᵢ, we need to find out what it is called
+   with (call this set Sᵢ). There are two cases:
+    * If vᵢ is a function, we need to go through all right-hand-sides and bodies,
+      and collect every variable that is called together with any variable from V:
+      Sᵢ = {v' | j ∈ {1,...,n},      {v',vⱼ} ∈ C'(rhs₁) ∪ ... ∪ C'(rhsₙ) ∪ C(body) }
+    * If vᵢ is a thunk, then its rhs is evaluated only once, so we need to
+      exclude it from this set:
+      Sᵢ = {v' | j ∈ {1,...,n}, j≠i, {v',vⱼ} ∈ C'(rhs₁) ∪ ... ∪ C'(rhsₙ) ∪ C(body) }
+   Finally, combine all this:
+   Return: C(body) ∪
+           C'(rhs₁) ∪ ... ∪ C'(rhsₙ) ∪
+           (fv rhs₁) × S₁) ∪ ... ∪ (fv rhsₙ) × Sₙ)
+
+Using the result: Eta-Expansion
+-------------------------------
+
+We use the result of these two analyses to decide whether we can eta-expand the
+rhs of a let-bound variable.
+
+If the variable is already a function (exprIsHNF), and all calls to the
+variables have a higher arity than the current manifest arity (i.e. the number
+of lambdas), expand.
+
+If the variable is a thunk we must be careful: Eta-Expansion will prevent
+sharing of work, so this is only safe if there is at most one call to the
+function. Therefore, we check whether {v,v} ∈ G.
+
+    Example:
 
         let n = case .. of .. -- A thunk!
         in n 0 + n 1
@@ -110,24 +210,12 @@ The at-most-once is important for various reasons:
     once in the body of the outer let. So we need to know, for each variable
     individually, that it is going to be called at most once.
 
- 2. We need to know it for non-thunks as well, because they might call a thunk:
 
-        let n = case .. of ..
-            f x = n (x+1)
-        in f 1 + f 2
+Why the co-call graph?
+----------------------
 
-    vs.
-
-        let n = case .. of ..
-            f x = n (x+1)
-        in case .. of T -> f 0
-                      F -> f 1
-
-    Here, the body of f calls n exactly once, but f itself is being called
-    multiple times, so eta-expansion is not allowed.
-
- 3. We need to know that at most one of the interesting functions is being
-    called, because of recursion. Consider:
+Why is it not sufficient to simply remember which variables are called once and
+which are called multiple times? It would be in the previous example, but consider
 
         let n = case .. of ..
         in case .. of
@@ -137,7 +225,7 @@ The at-most-once is important for various reasons:
                     in go 1
             False -> n
 
-    vs.
+vs.
 
         let n = case .. of ..
         in case .. of
@@ -147,134 +235,204 @@ The at-most-once is important for various reasons:
                     in go 1
             False -> n
 
-    In both cases, the body and the rhs of the inner let call n at most once.
-    But only in the second case that holds for the whole expression! The
-    crucial difference is that in the first case, the rhs of `go` can call
-    *both* `go` and `n`, and hence can call `n` multiple times as it recurses,
-    while in the second case it calls `go` or `n`, but not both.
+In both cases, the body and the rhs of the inner let call n at most once.
+But only in the second case that holds for the whole expression! The
+crucial difference is that in the first case, the rhs of `go` can call
+*both* `go` and `n`, and hence can call `n` multiple times as it recurses,
+while in the second case find out that `go` and `n` are not called together.
 
-Note [Which variables are interesting]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-Unfortunately, the set of interesting variables is not irrelevant for the
-precision of the analysis. Consider this example
+Why co-call information for functions?
+--------------------------------------
 
-    let n = ... :: Int
-    in  let go = \x -> let d = case ... of
-                                  False -> go (x+1)
-                                  True  -> id
-                       in \z -> d (x + z)
-        in go n 0
+Although for eta-expansion we need the information only for thunks, we still
+need to know whether functions are being called once or multiple times, and
+together with what other functions.
 
-Of course, `go` should be interesting. If we consider `n` as interesting as
-well, then the body of the second let will return
-    { go |-> Nothing , n |-> Just 0 }
-or
-    { go |-> 2, n |-> Nothing}.
-Only the latter is useful, but it is hard to decide that locally.
+    Example:
 
-So the heuristics is:
+        let n = case .. of ..
+            f x = n (x+1)
+        in f 1 + f 2
 
-    Variables are interesting if their RHS has a lower exprArity than
-    typeArity.
+    vs.
 
-(which is precisely the those variables where this analysis can actually cause
-some eta-expansion.)
+        let n = case .. of ..
+            f x = n (x+1)
+        in case .. of T -> f 0
+                      F -> f 1
 
-But this is not uniformly a win. Consider:
+    Here, the body of f calls n exactly once, but f itself is being called
+    multiple times, so eta-expansion is not allowed.
 
-    let go = \x -> let d = case ... of
-                              False -> go (x+1)
-                              True  -> id
-                       n x = d (x+1)
-                   in \z -> n (x + z)
-    in go n 0
 
-Now `n` is not going to be considered interesting (its type is `Int -> Int`).
-But this will prevent us from detecting how the body of the let calls `d`, and
-we will not find out anything.
+Note [Analysis type signature]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The work-hourse of the analysis is the function `callArityAnal`, with the
+following type:
+
+    type CallArityRes = (UnVarGraph, VarEnv Arity)
+    callArityAnal ::
+        Arity ->  -- The arity this expression is called with
+        VarSet -> -- The set of interesting variables
+        CoreExpr ->  -- The expression to analyse
+        (CallArityRes, CoreExpr)
 
-It might be possible to be smarter here; this needs find-tuning as we find more
-examples.
+and the following specification:
 
+  ((coCalls, callArityEnv), expr') = callArityEnv arity interestingIds expr
+
+                            <=>
+
+  Assume the expression `expr` is being passed `arity` arguments. Then it holds that
+    * The domain of `callArityEnv` is a subset of `interestingIds`.
+    * Any variable from `interestingIds` that is not mentioned in the `callArityEnv`
+      is absent, i.e. not called at all.
+    * Every call from `expr` to a variable bound to n in `callArityEnv` has at
+      least n value arguments.
+    * For two interesting variables `v1` and `v2`, they are not adjacent in `coCalls`,
+      then in no execution of `expr` both are being called.
+  Furthermore, expr' is expr with the callArity field of the `IdInfo` updated.
+
+
+Note [Which variables are interesting]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The analysis would quickly become prohibitive expensive if we would analyse all
+variables; for most variables we simply do not care about how often they are
+called, i.e. variables bound in a pattern match. So interesting are variables that are
+ * top-level or let bound
+ * and possibly functions (typeArity > 0)
+
+Note [Taking boring variables into account]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+If we decide that the variable bound in `let x = e1 in e2` is not interesting,
+the analysis of `e2` will not report anything about `x`. To ensure that
+`callArityBind` does still do the right thing we have to take that into account
+everytime we would be lookup up `x` in the analysis result of `e2`.
+  * Instead of calling lookupCallArityRes, we return (0, True), indicating
+    that this variable might be called many times with no variables.
+  * Instead of checking `calledWith x`, we assume that everything can be called
+    with it.
+  * In the recursive case, when calclulating the `cross_calls`, if there is
+    any boring variable in the recursive group, we ignore all co-call-results
+    and directly go to a very conservative assumption.
+
+The last point has the nice side effect that the relatively expensive
+integration of co-call results in a recursive groups is often skipped. This
+helped to avoid the compile time blowup in some real-world code with large
+recursive groups (#10293).
 
 Note [Recursion and fixpointing]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-For a recursive let, we begin by analysing the body, using the same incoming
-arity as for the whole expression.
- * If we do not get useful information about how we are calling the rhs, we
-   analyse the rhs using an incoming demand of 0 (which is always ok), and use
-   `forgetGoodCalls` to ignore any information coming from the rhs.
- * If we do get useful information from the body, we use that as the incoming
-   demand on the rhs. Then we check if the rhs calls itself with the same arity.
-   - If so, we are done.
-   - If not, we re-analise the rhs with the reduced arity. We do that until
-     we are down to the exprArity, which then is certainly correct.
-   We can `lubEnv` the results from the body and the rhs: The body calls *either*
-   the rhs *or* one of the other mentioned variables. Similarly, the rhs calls
-   *either* itself again *or* one of the other mentioned variables. This precision
-   is required!
+For a mutually recursive let, we begin by
+ 1. analysing the body, using the same incoming arity as for the whole expression.
+ 2. Then we iterate, memoizing for each of the bound variables the last
+    analysis call, i.e. incoming arity, whether it is called once, and the CallArityRes.
+ 3. We combine the analysis result from the body and the memoized results for
+    the arguments (if already present).
+ 4. For each variable, we find out the incoming arity and whether it is called
+    once, based on the the current analysis result. If this differs from the
+    memoized results, we re-analyse the rhs and update the memoized table.
+ 5. If nothing had to be reanalized, we are done.
+    Otherwise, repeat from step 3.
 
-We do not analyse mutually recursive functions. This can be done once we see it
-in the wild.
 
-Note [Case and App: Which side to take?]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Note [Thunks in recursive groups]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-Combining the case branches is easy, just `lubEnv` them – at most one branch is
-taken.
+We never eta-expand a thunk in a recursive group, on the grounds that if it is
+part of a recursive group, then it will be called multipe times.
 
-But how to combine that with the information coming from the scrunitee? Very
-similarly, how to combine the information from the callee and argument of an
-`App`?
+This is not necessarily true, e.g.  it would be safe to eta-expand t2 (but not
+t1) in the follwing code:
 
-It would not be correct to just `lubEnv` then: `f n` obviously calls *both* `f`
-and `n`. We need to forget about the calls from one side using `forgetGoodCalls`. But
-which one?
+  let go x = t1
+      t1 = if ... then t2 else ...
+      t2 = if ... then go 1 else ...
+  in go 0
 
-Both are correct, and sometimes one and sometimes the other is more precise
-(also see example in [Which variables are interesting]).
+Detecting this would require finding out what variables are only ever called
+from thunks. While this is certainly possible, we yet have to see this to be
+relevant in the wild.
 
-So currently, we first check the scrunitee (resp. the callee) if the returned
-value has any usesful information, and if so, we use that; otherwise we use the
-information from the alternatives (resp. the argument).
 
-It might be smarter to look for “more important” variables first, i.e. the
-innermost recursive variable.
+Note [Analysing top-level binds]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+We can eta-expand top-level-binds if they are not exported, as we see all calls
+to them. The plan is as follows: Treat the top-level binds as nested lets around
+a body representing “all external calls”, which returns a pessimistic
+CallArityRes (the co-call graph is the complete graph, all arityies 0).
+
+Note [Trimming arity]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+In the Call Arity papers, we are working on an untyped lambda calculus with no
+other id annotations, where eta-expansion is always possible. But this is not
+the case for Core!
+ 1. We need to ensure the invariant
+      callArity e <= typeArity (exprType e)
+    for the same reasons that exprArity needs this invariant (see Note
+    [exprArity invariant] in CoreArity).
+
+    If we are not doing that, a too-high arity annotation will be stored with
+    the id, confusing the simplifier later on.
+
+ 2. Eta-expanding a right hand side might invalidate existing annotations. In
+    particular, if an id has a strictness annotation of <...><...>b, then
+    passing two arguments to it will definitely bottom out, so the simplifier
+    will throw away additional parameters. This conflicts with Call Arity! So
+    we ensure that we never eta-expand such a value beyond the number of
+    arguments mentioned in the strictness signature.
+    See #10176 for a real-world-example.
 
 -}
 
+-- Main entry point
+
 callArityAnalProgram :: DynFlags -> CoreProgram -> CoreProgram
-callArityAnalProgram _dflags = map callArityBind
+callArityAnalProgram _dflags binds = binds'
+  where
+    (_, binds') = callArityTopLvl [] emptyVarSet binds
+
+-- See Note [Analysing top-level-binds]
+callArityTopLvl :: [Var] -> VarSet -> [CoreBind] -> (CallArityRes, [CoreBind])
+callArityTopLvl exported _ []
+    = ( calledMultipleTimes $ (emptyUnVarGraph, mkVarEnv $ [(v, 0) | v <- exported])
+      , [] )
+callArityTopLvl exported int1 (b:bs)
+    = (ae2, b':bs')
+  where
+    int2 = bindersOf b
+    exported' = filter isExportedId int2 ++ exported
+    int' = int1 `addInterestingBinds` b
+    (ae1, bs') = callArityTopLvl exported' int' bs
+    (ae2, b')  = callArityBind (boringBinds b) ae1 int1 b
 
-callArityBind :: CoreBind -> CoreBind
-callArityBind (NonRec id rhs) = NonRec id (callArityRHS rhs) 
-callArityBind (Rec binds) = Rec $ map (\(id,rhs) -> (id, callArityRHS rhs)) binds
 
 callArityRHS :: CoreExpr -> CoreExpr
 callArityRHS = snd . callArityAnal 0 emptyVarSet
 
-
-type CallArityEnv = VarEnv (Maybe Arity)
-
+-- The main analysis function. See Note [Analysis type signature]
 callArityAnal ::
     Arity ->  -- The arity this expression is called with
     VarSet -> -- The set of interesting variables
     CoreExpr ->  -- The expression to analyse
-    (CallArityEnv, CoreExpr)
-        -- How this expression uses its interesting variables:
-        --   Just n  => a tail call with that arity
-        --   Nothing => other uses
+    (CallArityRes, CoreExpr)
+        -- How this expression uses its interesting variables
         -- and the expression with IdInfo updated
 
 -- The trivial base cases
 callArityAnal _     _   e@(Lit _)
-    = (emptyVarEnv, e)
+    = (emptyArityRes, e)
 callArityAnal _     _   e@(Type _)
-    = (emptyVarEnv, e)
+    = (emptyArityRes, e)
 callArityAnal _     _   e@(Coercion _)
-    = (emptyVarEnv, e)
+    = (emptyArityRes, e)
 -- The transparent cases
 callArityAnal arity int (Tick t e)
     = second (Tick t) $ callArityAnal arity int e
@@ -284,123 +442,39 @@ callArityAnal arity int (Cast e co)
 -- The interesting case: Variables, Lambdas, Lets, Applications, Cases
 callArityAnal arity int e@(Var v)
     | v `elemVarSet` int
-    = (unitVarEnv v (Just arity), e)
+    = (unitArityRes v arity, e)
     | otherwise
-    = (emptyVarEnv, e)
+    = (emptyArityRes, e)
+
+-- Non-value lambdas are ignored
+callArityAnal arity int (Lam v e) | not (isId v)
+    = second (Lam v) $ callArityAnal arity (int `delVarSet` v) e
 
--- We have a lambda that we are not sure to call. Tail calls therein
--- are no longer tail calls
+-- We have a lambda that may be called multiple times, so its free variables
+-- can all be co-called.
 callArityAnal 0     int (Lam v e)
     = (ae', Lam v e')
   where
-    (ae, e') = callArityAnal 0 int e
-    ae' = forgetGoodCalls ae
+    (ae, e') = callArityAnal 0 (int `delVarSet` v) e
+    ae' = calledMultipleTimes ae
 -- We have a lambda that we are calling. decrease arity.
 callArityAnal arity int (Lam v e)
     = (ae, Lam v e')
   where
-    (ae, e') = callArityAnal (arity - 1) int e
-
--- Boring non-recursive let, i.e. no eta expansion possible. do not be smart about this
--- See Note [Which variables are interesting]
-callArityAnal arity int (Let (NonRec v rhs) e)
-    | exprArity rhs >= length (typeArity (idType v))
-    = (ae_final, Let (NonRec v rhs') e')
-  where
-    (ae_rhs, rhs') = callArityAnal 0 int rhs
-    (ae_body, e')  = callArityAnal arity int e
-    ae_body' = ae_body `delVarEnv` v
-    ae_final = forgetGoodCalls ae_rhs `lubEnv` ae_body'
-
--- Non-recursive let. Find out how the body calls the rhs, analise that,
--- and combine the results, convervatively using both
-callArityAnal arity int (Let (NonRec v rhs) e)
-
-    -- We are tail-calling into the rhs. So a tail-call in the RHS is a
-    -- tail-call for everything
-    | Just n <- rhs_arity
-    = let (ae_rhs, rhs') = callArityAnal n int rhs
-          final_ae       = ae_rhs `lubEnv` ae_body'
-          v'             = v `setIdCallArity` n
-      in -- pprTrace "callArityAnal:LetNonRecTailCall"
-         --          (vcat [ppr v, ppr arity, ppr n, ppr final_ae ])
-         (final_ae, Let (NonRec v' rhs') e')
-
-    -- We are calling the rhs in any other way (or not at all), so kill the
-    -- tail-call information from there
-    | otherwise
-    = let (ae_rhs, rhs') = callArityAnal 0 int rhs
-          final_ae = forgetGoodCalls ae_rhs `lubEnv` ae_body'
-          v'             = v `setIdCallArity` 0
-      in -- pprTrace "callArityAnal:LetNonRecNonTailCall"
-         --          (vcat [ppr v, ppr arity, ppr final_ae ])
-         (final_ae, Let (NonRec v' rhs') e')
-  where
-    int_body = int `extendVarSet` v
-    (ae_body, e') = callArityAnal arity int_body e
-    ae_body' = ae_body `delVarEnv` v
-    rhs_arity = lookupWithDefaultVarEnv ae_body Nothing v
-
--- Boring recursive let, i.e. no eta expansion possible. do not be smart about this
-callArityAnal arity int (Let (Rec [(v,rhs)]) e)
-    | exprArity rhs >= length (typeArity (idType v))
-    = (ae_final, Let (Rec [(v,rhs')]) e')
-  where
-    (ae_rhs, rhs') = callArityAnal 0 int rhs
-    (ae_body, e')  = callArityAnal arity int e
-    ae_final = (forgetGoodCalls ae_rhs `lubEnv` ae_body) `delVarEnv` v
-
--- Recursive let.
--- See Note [Recursion and fixpointing]
-callArityAnal arity int (Let (Rec [(v,rhs)]) e)
-    -- We are tail-calling into the rhs. So a tail-call in the RHS is a
-    -- tail-call for everything
-    | Just n <- rhs_arity
-    = let (ae_rhs, rhs_arity', rhs') = callArityFix n int_body v rhs
-          final_ae = (ae_rhs `lubEnv` ae_body) `delVarEnv` v
-          v'             = v `setIdCallArity` rhs_arity'
-      in -- pprTrace "callArityAnal:LetRecTailCall"
-         --          (vcat [ppr v, ppr arity, ppr n, ppr rhs_arity', ppr final_ae ])
-         (final_ae, Let (Rec [(v',rhs')]) e')
-    -- We are calling the body in any other way (or not at all), so kill the
-    -- tail-call information from there. No need to iterate there.
-    | otherwise
-    = let (ae_rhs, rhs') = callArityAnal 0 int_body rhs
-          final_ae = (forgetGoodCalls ae_rhs `lubEnv` ae_body) `delVarEnv` v
-          v'             = v `setIdCallArity` 0
-      in -- pprTrace "callArityAnal:LetRecNonTailCall"
-         --          (vcat [ppr v, ppr arity, ppr final_ae ])
-         (final_ae, Let (Rec [(v',rhs')]) e')
-  where
-    int_body = int `extendVarSet` v
-    (ae_body, e') = callArityAnal arity int_body e
-    rhs_arity = lookupWithDefaultVarEnv ae_body Nothing v
-
--- Mutual recursion. Do nothing serious here, for now
-callArityAnal arity int (Let (Rec binds) e)
-    = (final_ae, Let (Rec binds') e')
-  where
-    (aes, binds') = unzip $ map go binds
-    go (i,e) = let (ae,e') = callArityAnal 0 int e
-               in (forgetGoodCalls ae, (i,e'))
-    (ae, e') = callArityAnal arity int e
-    final_ae = foldl lubEnv ae aes `delVarEnvList` map fst binds
+    (ae, e') = callArityAnal (arity - 1) (int `delVarSet` v) e
 
 -- Application. Increase arity for the called expresion, nothing to know about
 -- the second
+callArityAnal arity int (App e (Type t))
+    = second (\e -> App e (Type t)) $ callArityAnal arity int e
 callArityAnal arity int (App e1 e2)
     = (final_ae, App e1' e2')
   where
     (ae1, e1') = callArityAnal (arity + 1) int e1
     (ae2, e2') = callArityAnal 0           int e2
-    -- See Note [Case and App: Which side to take?]
-    final_ae = ae1 `useBetterOf` ae2
-
--- Case expression. Here we decide whether
--- we want to look at calls from the scrunitee or the alternatives;
--- one of them we set to Nothing.
--- Naive idea: If there are interesting calls in the scrunitee,
--- zap the alternatives
+    final_ae = ae1 `both` ae2
+
+-- Case expression.
 callArityAnal arity int (Case scrut bndr ty alts)
     = -- pprTrace "callArityAnal:Case"
       --          (vcat [ppr scrut, ppr final_ae])
@@ -409,51 +483,236 @@ callArityAnal arity int (Case scrut bndr ty alts)
     (alt_aes, alts') = unzip $ map go alts
     go (dc, bndrs, e) = let (ae, e') = callArityAnal arity int e
                         in  (ae, (dc, bndrs, e'))
-    alt_ae = foldl lubEnv emptyVarEnv alt_aes
+    alt_ae = lubRess alt_aes
     (scrut_ae, scrut') = callArityAnal 0 int scrut
-    -- See Note [Case and App: Which side to take?]
-    final_ae = scrut_ae `useBetterOf` alt_ae
+    final_ae = scrut_ae `both` alt_ae
 
-callArityFix :: Arity -> VarSet -> Id -> CoreExpr -> (CallArityEnv, Arity, CoreExpr)
-callArityFix arity int v e
+-- For lets, use callArityBind
+callArityAnal arity int (Let bind e)
+  = -- pprTrace "callArityAnal:Let"
+    --          (vcat [ppr v, ppr arity, ppr n, ppr final_ae ])
+    (final_ae, Let bind' e')
+  where
+    int_body = int `addInterestingBinds` bind
+    (ae_body, e') = callArityAnal arity int_body e
+    (final_ae, bind') = callArityBind (boringBinds bind) ae_body int bind
 
-    | arity <= min_arity
-    -- The incoming arity is already lower than the exprArity, so we can
-    -- ignore the arity coming from the RHS
-    = (ae `delVarEnv` v, 0, e')
+-- Which bindings should we look at?
+-- See Note [Which variables are interesting]
+isInteresting :: Var -> Bool
+isInteresting v = 0 < length (typeArity (idType v))
+
+interestingBinds :: CoreBind -> [Var]
+interestingBinds = filter isInteresting . bindersOf
+
+boringBinds :: CoreBind -> VarSet
+boringBinds = mkVarSet . filter (not . isInteresting) . bindersOf
+
+addInterestingBinds :: VarSet -> CoreBind -> VarSet
+addInterestingBinds int bind
+    = int `delVarSetList`    bindersOf bind -- Possible shadowing
+          `extendVarSetList` interestingBinds bind
+
+-- Used for both local and top-level binds
+-- Second argument is the demand from the body
+callArityBind :: VarSet -> CallArityRes -> VarSet -> CoreBind -> (CallArityRes, CoreBind)
+-- Non-recursive let
+callArityBind boring_vars ae_body int (NonRec v rhs)
+  | otherwise
+  = -- pprTrace "callArityBind:NonRec"
+    --          (vcat [ppr v, ppr ae_body, ppr int, ppr ae_rhs, ppr safe_arity])
+    (final_ae, NonRec v' rhs')
+  where
+    is_thunk = not (exprIsHNF rhs)
+    -- If v is boring, we will not find it in ae_body, but always assume (0, False)
+    boring = v `elemVarSet` boring_vars
 
-    | otherwise
-    = case new_arity of
-        -- Not nicely recursive, rerun with arity 0
-        -- (which will do at most one iteration, see above)
-        -- (Or not recursive at all, but that was hopefully handled by the simplifier before)
-        Nothing -> callArityFix 0 int v e
-
-        Just n -> if n < arity
-            -- RHS puts a lower arity on itself, but still a nice call, so try with that
-            then callArityFix n int v e
-
-            -- RHS calls itself with at least as many arguments as the body of
-            -- the let: Great!
-            else (ae `delVarEnv` v, n, e')
+    (arity, called_once)
+        | boring    = (0, False) -- See Note [Taking boring variables into account]
+        | otherwise = lookupCallArityRes ae_body v
+    safe_arity | called_once = arity
+               | is_thunk    = 0      -- A thunk! Do not eta-expand
+               | otherwise   = arity
+
+    -- See Note [Trimming arity]
+    trimmed_arity = trimArity v safe_arity
+
+    (ae_rhs, rhs') = callArityAnal trimmed_arity int rhs
+
+
+    ae_rhs'| called_once     = ae_rhs
+           | safe_arity == 0 = ae_rhs -- If it is not a function, its body is evaluated only once
+           | otherwise       = calledMultipleTimes ae_rhs
+
+    called_by_v = domRes ae_rhs'
+    called_with_v
+        | boring    = domRes ae_body
+        | otherwise = calledWith ae_body v `delUnVarSet` v
+    final_ae = addCrossCoCalls called_by_v called_with_v $ ae_rhs' `lubRes` resDel v ae_body
+
+    v' = v `setIdCallArity` trimmed_arity
+
+
+-- Recursive let. See Note [Recursion and fixpointing]
+callArityBind boring_vars ae_body int b@(Rec binds)
+  = -- (if length binds > 300 then
+    -- pprTrace "callArityBind:Rec"
+    --           (vcat [ppr (Rec binds'), ppr ae_body, ppr int, ppr ae_rhs]) else id) $
+    (final_ae, Rec binds')
+  where
+    -- See Note [Taking boring variables into account]
+    any_boring = any (`elemVarSet` boring_vars) [ i | (i, _) <- binds]
+
+    int_body = int `addInterestingBinds` b
+    (ae_rhs, binds') = fix initial_binds
+    final_ae = bindersOf b `resDelList` ae_rhs
+
+    initial_binds = [(i,Nothing,e) | (i,e) <- binds]
+
+    fix :: [(Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)] -> (CallArityRes, [(Id, CoreExpr)])
+    fix ann_binds
+        | -- pprTrace "callArityBind:fix" (vcat [ppr ann_binds, ppr any_change, ppr ae]) $
+          any_change
+        = fix ann_binds'
+        | otherwise
+        = (ae, map (\(i, _, e) -> (i, e)) ann_binds')
+      where
+        aes_old = [ (i,ae) | (i, Just (_,_,ae), _) <- ann_binds ]
+        ae = callArityRecEnv any_boring aes_old ae_body
+
+        rerun (i, mbLastRun, rhs)
+            | i `elemVarSet` int_body && not (i `elemUnVarSet` domRes ae)
+            -- No call to this yet, so do nothing
+            = (False, (i, Nothing, rhs))
+
+            | Just (old_called_once, old_arity, _) <- mbLastRun
+            , called_once == old_called_once
+            , new_arity == old_arity
+            -- No change, no need to re-analize
+            = (False, (i, mbLastRun, rhs))
+
+            | otherwise
+            -- We previously analized this with a different arity (or not at all)
+            = let is_thunk = not (exprIsHNF rhs)
+
+                  safe_arity | is_thunk    = 0  -- See Note [Thunks in recursive groups]
+                             | otherwise   = new_arity
+
+                  -- See Note [Trimming arity]
+                  trimmed_arity = trimArity i safe_arity
+
+                  (ae_rhs, rhs') = callArityAnal trimmed_arity int_body rhs
+
+                  ae_rhs' | called_once     = ae_rhs
+                          | safe_arity == 0 = ae_rhs -- If it is not a function, its body is evaluated only once
+                          | otherwise       = calledMultipleTimes ae_rhs
+
+              in (True, (i `setIdCallArity` trimmed_arity, Just (called_once, new_arity, ae_rhs'), rhs'))
+          where
+            -- See Note [Taking boring variables into account]
+            (new_arity, called_once) | i `elemVarSet` boring_vars = (0, False)
+                                     | otherwise                  = lookupCallArityRes ae i
+
+        (changes, ann_binds') = unzip $ map rerun ann_binds
+        any_change = or changes
+
+-- Combining the results from body and rhs, (mutually) recursive case
+-- See Note [Analysis II: The Co-Called analysis]
+callArityRecEnv :: Bool -> [(Var, CallArityRes)] -> CallArityRes -> CallArityRes
+callArityRecEnv any_boring ae_rhss ae_body
+    = -- (if length ae_rhss > 300 then pprTrace "callArityRecEnv" (vcat [ppr ae_rhss, ppr ae_body, ppr ae_new]) else id) $
+      ae_new
+  where
+    vars = map fst ae_rhss
+
+    ae_combined = lubRess (map snd ae_rhss) `lubRes` ae_body
+
+    cross_calls
+        -- See Note [Taking boring variables into account]
+        | any_boring          = completeGraph (domRes ae_combined)
+        -- Also, calculating cross_calls is expensive. Simply be conservative
+        -- if the mutually recursive group becomes too large.
+        | length ae_rhss > 25 = completeGraph (domRes ae_combined)
+        | otherwise           = unionUnVarGraphs $ map cross_call ae_rhss
+    cross_call (v, ae_rhs) = completeBipartiteGraph called_by_v called_with_v
+      where
+        is_thunk = idCallArity v == 0
+        -- What rhs are relevant as happening before (or after) calling v?
+        --    If v is a thunk, everything from all the _other_ variables
+        --    If v is not a thunk, everything can happen.
+        ae_before_v | is_thunk  = lubRess (map snd $ filter ((/= v) . fst) ae_rhss) `lubRes` ae_body
+                    | otherwise = ae_combined
+        -- What do we want to know from these?
+        -- Which calls can happen next to any recursive call.
+        called_with_v
+            = unionUnVarSets $ map (calledWith ae_before_v) vars
+        called_by_v = domRes ae_rhs
+
+    ae_new = first (cross_calls `unionUnVarGraph`) ae_combined
+
+-- See Note [Trimming arity]
+trimArity :: Id -> Arity -> Arity
+trimArity v a = minimum [a, max_arity_by_type, max_arity_by_strsig]
   where
-    (ae, e') = callArityAnal arity int e
-    new_arity = lookupWithDefaultVarEnv ae Nothing v
-    min_arity = exprArity e
+    max_arity_by_type = length (typeArity (idType v))
+    max_arity_by_strsig
+        | isBotRes result_info = length demands
+        | otherwise = a
+
+    (demands, result_info) = splitStrictSig (idStrictness v)
+
+---------------------------------------
+-- Functions related to CallArityRes --
+---------------------------------------
+
+-- Result type for the two analyses.
+-- See Note [Analysis I: The arity analyis]
+-- and Note [Analysis II: The Co-Called analysis]
+type CallArityRes = (UnVarGraph, VarEnv Arity)
 
+emptyArityRes :: CallArityRes
+emptyArityRes = (emptyUnVarGraph, emptyVarEnv)
 
-anyGoodCalls :: VarEnv (Maybe Arity) -> Bool
-anyGoodCalls = foldVarEnv ((||) . isJust) False
+unitArityRes :: Var -> Arity -> CallArityRes
+unitArityRes v arity = (emptyUnVarGraph, unitVarEnv v arity)
 
-forgetGoodCalls :: VarEnv (Maybe Arity) -> VarEnv (Maybe Arity)
-forgetGoodCalls = mapVarEnv (const Nothing)
+resDelList :: [Var] -> CallArityRes -> CallArityRes
+resDelList vs ae = foldr resDel ae vs
 
--- See Note [Case and App: Which side to take?]
-useBetterOf :: CallArityEnv -> CallArityEnv -> CallArityEnv
-useBetterOf ae1 ae2 | anyGoodCalls ae1 = ae1 `lubEnv` forgetGoodCalls ae2
-useBetterOf ae1 ae2 | otherwise        = forgetGoodCalls ae1 `lubEnv` ae2
+resDel :: Var -> CallArityRes -> CallArityRes
+resDel v (g, ae) = (g `delNode` v, ae `delVarEnv` v)
+
+domRes :: CallArityRes -> UnVarSet
+domRes (_, ae) = varEnvDom ae
+
+-- In the result, find out the minimum arity and whether the variable is called
+-- at most once.
+lookupCallArityRes :: CallArityRes -> Var -> (Arity, Bool)
+lookupCallArityRes (g, ae) v
+    = case lookupVarEnv ae v of
+        Just a -> (a, not (v `elemUnVarSet` (neighbors g v)))
+        Nothing -> (0, False)
+
+calledWith :: CallArityRes -> Var -> UnVarSet
+calledWith (g, _) v = neighbors g v
+
+addCrossCoCalls :: UnVarSet -> UnVarSet -> CallArityRes -> CallArityRes
+addCrossCoCalls set1 set2 = first (completeBipartiteGraph set1 set2 `unionUnVarGraph`)
+
+-- Replaces the co-call graph by a complete graph (i.e. no information)
+calledMultipleTimes :: CallArityRes -> CallArityRes
+calledMultipleTimes res = first (const (completeGraph (domRes res))) res
+
+-- Used for application and cases
+both :: CallArityRes -> CallArityRes -> CallArityRes
+both r1 r2 = addCrossCoCalls (domRes r1) (domRes r2) $ r1 `lubRes` r2
 
 -- Used when combining results from alternative cases; take the minimum
-lubEnv :: CallArityEnv -> CallArityEnv -> CallArityEnv
-lubEnv = plusVarEnv_C min
+lubRes :: CallArityRes -> CallArityRes -> CallArityRes
+lubRes (g1, ae1) (g2, ae2) = (g1 `unionUnVarGraph` g2, ae1 `lubArityEnv` ae2)
+
+lubArityEnv :: VarEnv Arity -> VarEnv Arity -> VarEnv Arity
+lubArityEnv = plusVarEnv_C min
 
+lubRess :: [CallArityRes] -> CallArityRes
+lubRess = foldl lubRes emptyArityRes