CallArity.hs 27.9 KB
Newer Older
1 2 3 4 5 6
--
-- Copyright (c) 2014 Joachim Breitner
--

module CallArity
    ( callArityAnalProgram
7
    , callArityRHS -- for testing
8 9 10 11 12 13 14 15 16
    ) where

import VarSet
import VarEnv
import DynFlags ( DynFlags )

import BasicTypes
import CoreSyn
import Id
Joachim Breitner's avatar
Joachim Breitner committed
17
import CoreArity ( typeArity )
18
import CoreUtils ( exprIsHNF, exprIsTrivial )
Joachim Breitner's avatar
Joachim Breitner committed
19 20
--import Outputable
import UnVarGraph
Joachim Breitner's avatar
Joachim Breitner committed
21
import Demand
22

23
import Control.Arrow ( first, second )
24 25 26 27


{-
%************************************************************************
28
%*                                                                      *
29
              Call Arity Analyis
30
%*                                                                      *
31 32 33 34 35 36
%************************************************************************

Note [Call Arity: The goal]
~~~~~~~~~~~~~~~~~~~~~~~~~~~

The goal of this analysis is to find out if we can eta-expand a local function,
Jan Stolarek's avatar
Jan Stolarek committed
37
based on how it is being called. The motivating example is this code,
38 39 40 41 42 43 44 45 46 47 48 49
which comes up when we implement foldl using foldr, and do list fusion:

    let go = \x -> let d = case ... of
                              False -> go (x+1)
                              True  -> id
                   in \z -> d (x + z)
    in go 1 0

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.
Jan Stolarek's avatar
Jan Stolarek committed
50
Furthermore, an analysis that only looks at the RHS of go cannot be sufficient
51 52 53 54 55 56 57 58 59 60 61 62
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.

So `callArityAnalProgram` looks at the whole let expression to figure out if
all calls are nice, i.e. have a high enough arity. It then stores the result in
the `calledArity` field of the `IdInfo` of `go`, which the next simplifier
phase will eta-expand.

The specification of the `calledArity` field is:

    No work will be lost if you eta-expand me to the arity in `calledArity`.

Joachim Breitner's avatar
Joachim Breitner committed
63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151
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₂ ∪...)
152 153 154 155
 * App e₁ e₂ (and analogously Case scrut alts), with non-trivial e₂:
   We get the results from both sides, with the argument evaluted at most once.
   Additionally, anything called by e₁ can possibly be called with anything
   from e₂.
Joachim Breitner's avatar
Joachim Breitner committed
156
   Return: C(e₁) ∪ C(e₂) ∪ (fv e₁) × (fv e₂)
157 158 159 160
 * App e₁ x:
   As this is already in A-normal form, CorePrep will not separately lambda
   bind (and hence share) x. So we conservatively assume multiple calls to x here
   Return: C(e₁) ∪ (fv e₁) × {x} ∪ {(x,x)}
Joachim Breitner's avatar
Joachim Breitner committed
161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176
 * 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
177
   with (call this set Sᵢ). There are two cases:
Joachim Breitner's avatar
Joachim Breitner committed
178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203
    * 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:
204 205 206 207 208 209 210 211 212 213 214 215 216 217 218

        let n = case .. of .. -- A thunk!
        in n 0 + n 1

    vs.

        let n = case .. of ..
        in case .. of T -> n 0
                      F -> n 1

    We are only allowed to eta-expand `n` if it is going to be called at most
    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.


Joachim Breitner's avatar
Joachim Breitner committed
219 220
Why the co-call graph?
----------------------
221

Joachim Breitner's avatar
Joachim Breitner committed
222 223
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
224 225 226 227 228 229 230 231 232

        let n = case .. of ..
        in case .. of
            True -> let go = \y -> case .. of
                                     True -> go (y + n 1)
                                     False > n
                    in go 1
            False -> n

Joachim Breitner's avatar
Joachim Breitner committed
233
vs.
234 235 236 237 238 239 240 241 242

        let n = case .. of ..
        in case .. of
            True -> let go = \y -> case .. of
                                     True -> go (y+1)
                                     False > n
                    in go 1
            False -> n

Joachim Breitner's avatar
Joachim Breitner committed
243 244 245 246 247
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.
248 249


Joachim Breitner's avatar
Joachim Breitner committed
250 251
Why co-call information for functions?
--------------------------------------
252

Joachim Breitner's avatar
Joachim Breitner committed
253 254 255
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.
256

Joachim Breitner's avatar
Joachim Breitner committed
257
    Example:
258

Joachim Breitner's avatar
Joachim Breitner committed
259 260 261
        let n = case .. of ..
            f x = n (x+1)
        in f 1 + f 2
262

Joachim Breitner's avatar
Joachim Breitner committed
263
    vs.
264

Joachim Breitner's avatar
Joachim Breitner committed
265 266 267 268
        let n = case .. of ..
            f x = n (x+1)
        in case .. of T -> f 0
                      F -> f 1
269

Joachim Breitner's avatar
Joachim Breitner committed
270 271
    Here, the body of f calls n exactly once, but f itself is being called
    multiple times, so eta-expansion is not allowed.
272 273


Joachim Breitner's avatar
Joachim Breitner committed
274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289
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)

and the following specification:

  ((coCalls, callArityEnv), expr') = callArityEnv arity interestingIds expr
290

Joachim Breitner's avatar
Joachim Breitner committed
291
                            <=>
292

Joachim Breitner's avatar
Joachim Breitner committed
293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311
  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)
312

313 314
Note [Taking boring variables into account]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
315 316 317

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
318 319 320 321 322 323 324 325 326 327 328 329 330 331
`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).
332

333 334 335
Note [Recursion and fixpointing]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Joachim Breitner's avatar
Joachim Breitner committed
336 337 338 339 340 341 342 343 344 345 346
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.
347

348 349 350 351 352 353 354 355 356 357 358 359 360 361 362

Note [Thunks in recursive groups]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

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.

This is not necessarily true, e.g.  it would be safe to eta-expand t2 (but not
t1) in the follwing code:

  let go x = t1
      t1 = if ... then t2 else ...
      t2 = if ... then go 1 else ...
  in go 0

Joachim Breitner's avatar
Joachim Breitner committed
363
Detecting this would require finding out what variables are only ever called
364 365 366 367
from thunks. While this is certainly possible, we yet have to see this to be
relevant in the wild.


368 369 370 371 372
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
Joachim Breitner's avatar
Joachim Breitner committed
373 374
a body representing “all external calls”, which returns a pessimistic
CallArityRes (the co-call graph is the complete graph, all arityies 0).
375

Joachim Breitner's avatar
Joachim Breitner committed
376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391
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
392
    passing two arguments to it will definitely bottom out, so the simplifier
Joachim Breitner's avatar
Joachim Breitner committed
393 394 395 396 397
    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.

398 399
-}

Joachim Breitner's avatar
Joachim Breitner committed
400 401
-- Main entry point

402
callArityAnalProgram :: DynFlags -> CoreProgram -> CoreProgram
403 404 405 406 407
callArityAnalProgram _dflags binds = binds'
  where
    (_, binds') = callArityTopLvl [] emptyVarSet binds

-- See Note [Analysing top-level-binds]
Joachim Breitner's avatar
Joachim Breitner committed
408
callArityTopLvl :: [Var] -> VarSet -> [CoreBind] -> (CallArityRes, [CoreBind])
409
callArityTopLvl exported _ []
Joachim Breitner's avatar
Joachim Breitner committed
410 411
    = ( calledMultipleTimes $ (emptyUnVarGraph, mkVarEnv $ [(v, 0) | v <- exported])
      , [] )
412 413 414
callArityTopLvl exported int1 (b:bs)
    = (ae2, b':bs')
  where
Joachim Breitner's avatar
Joachim Breitner committed
415
    int2 = bindersOf b
416
    exported' = filter isExportedId int2 ++ exported
Joachim Breitner's avatar
Joachim Breitner committed
417
    int' = int1 `addInterestingBinds` b
418
    (ae1, bs') = callArityTopLvl exported' int' bs
419
    (ae2, b')  = callArityBind (boringBinds b) ae1 int1 b
420 421 422 423 424


callArityRHS :: CoreExpr -> CoreExpr
callArityRHS = snd . callArityAnal 0 emptyVarSet

Joachim Breitner's avatar
Joachim Breitner committed
425
-- The main analysis function. See Note [Analysis type signature]
426 427 428 429
callArityAnal ::
    Arity ->  -- The arity this expression is called with
    VarSet -> -- The set of interesting variables
    CoreExpr ->  -- The expression to analyse
Joachim Breitner's avatar
Joachim Breitner committed
430
    (CallArityRes, CoreExpr)
431
        -- How this expression uses its interesting variables
432 433 434 435
        -- and the expression with IdInfo updated

-- The trivial base cases
callArityAnal _     _   e@(Lit _)
Joachim Breitner's avatar
Joachim Breitner committed
436
    = (emptyArityRes, e)
437
callArityAnal _     _   e@(Type _)
Joachim Breitner's avatar
Joachim Breitner committed
438
    = (emptyArityRes, e)
439
callArityAnal _     _   e@(Coercion _)
Joachim Breitner's avatar
Joachim Breitner committed
440
    = (emptyArityRes, e)
441 442 443 444 445 446 447 448 449
-- The transparent cases
callArityAnal arity int (Tick t e)
    = second (Tick t) $ callArityAnal arity int e
callArityAnal arity int (Cast e co)
    = second (\e -> Cast e co) $ callArityAnal arity int e

-- The interesting case: Variables, Lambdas, Lets, Applications, Cases
callArityAnal arity int e@(Var v)
    | v `elemVarSet` int
Joachim Breitner's avatar
Joachim Breitner committed
450
    = (unitArityRes v arity, e)
451
    | otherwise
Joachim Breitner's avatar
Joachim Breitner committed
452
    = (emptyArityRes, e)
453

454 455
-- Non-value lambdas are ignored
callArityAnal arity int (Lam v e) | not (isId v)
Joachim Breitner's avatar
Joachim Breitner committed
456
    = second (Lam v) $ callArityAnal arity (int `delVarSet` v) e
457

Joachim Breitner's avatar
Joachim Breitner committed
458 459
-- We have a lambda that may be called multiple times, so its free variables
-- can all be co-called.
460 461 462
callArityAnal 0     int (Lam v e)
    = (ae', Lam v e')
  where
Joachim Breitner's avatar
Joachim Breitner committed
463
    (ae, e') = callArityAnal 0 (int `delVarSet` v) e
Joachim Breitner's avatar
Joachim Breitner committed
464
    ae' = calledMultipleTimes ae
465 466 467 468
-- We have a lambda that we are calling. decrease arity.
callArityAnal arity int (Lam v e)
    = (ae, Lam v e')
  where
Joachim Breitner's avatar
Joachim Breitner committed
469
    (ae, e') = callArityAnal (arity - 1) (int `delVarSet` v) e
470 471 472

-- Application. Increase arity for the called expresion, nothing to know about
-- the second
473 474
callArityAnal arity int (App e (Type t))
    = second (\e -> App e (Type t)) $ callArityAnal arity int e
475 476 477 478 479
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
480 481 482 483 484 485
    -- If the argument is trivial (e.g. a variable), then it will _not_ be
    -- let-bound in the Core to STG transformation (CorePrep actually),
    -- so no sharing will happen here, and we have to assume many calls.
    ae2' | exprIsTrivial e2 = calledMultipleTimes ae2
         | otherwise        = ae2
    final_ae = ae1 `both` ae2'
486

Joachim Breitner's avatar
Joachim Breitner committed
487
-- Case expression.
488 489 490 491 492 493 494 495
callArityAnal arity int (Case scrut bndr ty alts)
    = -- pprTrace "callArityAnal:Case"
      --          (vcat [ppr scrut, ppr final_ae])
      (final_ae, Case scrut' bndr ty alts')
  where
    (alt_aes, alts') = unzip $ map go alts
    go (dc, bndrs, e) = let (ae, e') = callArityAnal arity int e
                        in  (ae, (dc, bndrs, e'))
Joachim Breitner's avatar
Joachim Breitner committed
496
    alt_ae = lubRess alt_aes
497
    (scrut_ae, scrut') = callArityAnal 0 int scrut
Joachim Breitner's avatar
Joachim Breitner committed
498 499 500 501 502 503 504 505 506 507
    final_ae = scrut_ae `both` alt_ae

-- 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
508
    (final_ae, bind') = callArityBind (boringBinds bind) ae_body int bind
Joachim Breitner's avatar
Joachim Breitner committed
509

510 511
-- Which bindings should we look at?
-- See Note [Which variables are interesting]
512 513 514
isInteresting :: Var -> Bool
isInteresting v = 0 < length (typeArity (idType v))

515
interestingBinds :: CoreBind -> [Var]
516 517 518 519
interestingBinds = filter isInteresting . bindersOf

boringBinds :: CoreBind -> VarSet
boringBinds = mkVarSet . filter (not . isInteresting) . bindersOf
520

Joachim Breitner's avatar
Joachim Breitner committed
521 522 523 524 525
addInterestingBinds :: VarSet -> CoreBind -> VarSet
addInterestingBinds int bind
    = int `delVarSetList`    bindersOf bind -- Possible shadowing
          `extendVarSetList` interestingBinds bind

526
-- Used for both local and top-level binds
527 528
-- Second argument is the demand from the body
callArityBind :: VarSet -> CallArityRes -> VarSet -> CoreBind -> (CallArityRes, CoreBind)
529
-- Non-recursive let
530
callArityBind boring_vars ae_body int (NonRec v rhs)
Joachim Breitner's avatar
Joachim Breitner committed
531
  | otherwise
532 533 534 535
  = -- pprTrace "callArityBind:NonRec"
    --          (vcat [ppr v, ppr ae_body, ppr int, ppr ae_rhs, ppr safe_arity])
    (final_ae, NonRec v' rhs')
  where
536
    is_thunk = not (exprIsHNF rhs)
537 538
    -- If v is boring, we will not find it in ae_body, but always assume (0, False)
    boring = v `elemVarSet` boring_vars
539

540 541 542
    (arity, called_once)
        | boring    = (0, False) -- See Note [Taking boring variables into account]
        | otherwise = lookupCallArityRes ae_body v
543 544 545
    safe_arity | called_once = arity
               | is_thunk    = 0      -- A thunk! Do not eta-expand
               | otherwise   = arity
Joachim Breitner's avatar
Joachim Breitner committed
546 547 548 549 550 551

    -- See Note [Trimming arity]
    trimmed_arity = trimArity v safe_arity

    (ae_rhs, rhs') = callArityAnal trimmed_arity int rhs

552 553 554 555 556

    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

557 558 559 560 561 562
    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

Joachim Breitner's avatar
Joachim Breitner committed
563
    v' = v `setIdCallArity` trimmed_arity
564 565


566
-- Recursive let. See Note [Recursion and fixpointing]
567 568 569 570
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) $
Joachim Breitner's avatar
Joachim Breitner committed
571
    (final_ae, Rec binds')
572
  where
573 574 575
    -- See Note [Taking boring variables into account]
    any_boring = any (`elemVarSet` boring_vars) [ i | (i, _) <- binds]

Joachim Breitner's avatar
Joachim Breitner committed
576
    int_body = int `addInterestingBinds` b
Joachim Breitner's avatar
Joachim Breitner committed
577 578
    (ae_rhs, binds') = fix initial_binds
    final_ae = bindersOf b `resDelList` ae_rhs
579

Joachim Breitner's avatar
Joachim Breitner committed
580
    initial_binds = [(i,Nothing,e) | (i,e) <- binds]
581

Joachim Breitner's avatar
Joachim Breitner committed
582 583 584 585 586
    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'
Joachim Breitner's avatar
Joachim Breitner committed
587
        | otherwise
Joachim Breitner's avatar
Joachim Breitner committed
588
        = (ae, map (\(i, _, e) -> (i, e)) ann_binds')
Joachim Breitner's avatar
Joachim Breitner committed
589
      where
Joachim Breitner's avatar
Joachim Breitner committed
590
        aes_old = [ (i,ae) | (i, Just (_,_,ae), _) <- ann_binds ]
591
        ae = callArityRecEnv any_boring aes_old ae_body
Joachim Breitner's avatar
Joachim Breitner committed
592 593 594 595 596 597 598 599 600 601 602 603 604 605

        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)
606 607 608 609 610
            = let is_thunk = not (exprIsHNF rhs)

                  safe_arity | is_thunk    = 0  -- See Note [Thunks in recursive groups]
                             | otherwise   = new_arity

Joachim Breitner's avatar
Joachim Breitner committed
611 612 613 614
                  -- See Note [Trimming arity]
                  trimmed_arity = trimArity i safe_arity

                  (ae_rhs, rhs') = callArityAnal trimmed_arity int_body rhs
615 616 617 618 619

                  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

Joachim Breitner's avatar
Joachim Breitner committed
620
              in (True, (i `setIdCallArity` trimmed_arity, Just (called_once, new_arity, ae_rhs'), rhs'))
Joachim Breitner's avatar
Joachim Breitner committed
621
          where
622 623 624
            -- See Note [Taking boring variables into account]
            (new_arity, called_once) | i `elemVarSet` boring_vars = (0, False)
                                     | otherwise                  = lookupCallArityRes ae i
Joachim Breitner's avatar
Joachim Breitner committed
625 626 627 628 629 630

        (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]
631 632 633
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) $
Joachim Breitner's avatar
Joachim Breitner committed
634 635 636
      ae_new
  where
    vars = map fst ae_rhss
637

Joachim Breitner's avatar
Joachim Breitner committed
638
    ae_combined = lubRess (map snd ae_rhss) `lubRes` ae_body
639

640 641 642
    cross_calls
        -- See Note [Taking boring variables into account]
        | any_boring          = completeGraph (domRes ae_combined)
643 644 645
        -- 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)
646
        | otherwise           = unionUnVarGraphs $ map cross_call ae_rhss
Joachim Breitner's avatar
Joachim Breitner committed
647 648 649 650 651 652 653 654 655 656 657 658 659 660 661 662
    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

663 664 665 666 667 668 669 670 671 672 673
-- See Note [Trimming arity]
trimArity :: Id -> Arity -> Arity
trimArity v a = minimum [a, max_arity_by_type, max_arity_by_strsig]
  where
    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)

Joachim Breitner's avatar
Joachim Breitner committed
674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 704 705 706 707 708 709 710 711 712 713 714 715 716 717 718
---------------------------------------
-- 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)

unitArityRes :: Var -> Arity -> CallArityRes
unitArityRes v arity = (emptyUnVarGraph, unitVarEnv v arity)

resDelList :: [Var] -> CallArityRes -> CallArityRes
resDelList vs ae = foldr resDel ae vs

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
719 720

-- Used when combining results from alternative cases; take the minimum
Joachim Breitner's avatar
Joachim Breitner committed
721 722 723 724 725
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
726

Joachim Breitner's avatar
Joachim Breitner committed
727 728
lubRess :: [CallArityRes] -> CallArityRes
lubRess = foldl lubRes emptyArityRes