--
-- Copyright (c) 2014 Joachim Breitner
--

{-# LANGUAGE BangPatterns #-}

module GHC.Core.Opt.CallArity
    ( callArityAnalProgram
    , callArityRHS -- for testing
    ) where

import GHC.Prelude

import GHC.Types.Var.Set
import GHC.Types.Var.Env

import GHC.Types.Basic
import GHC.Core
import GHC.Types.Id
import GHC.Core.Opt.Arity ( typeArity )
import GHC.Core.Utils ( exprIsCheap, exprIsTrivial )
import GHC.Data.Graph.UnVar
import GHC.Types.Demand
import GHC.Utils.Misc

import Control.Arrow ( first, second )


{-
%************************************************************************
%*                                                                      *
              Call Arity Analysis
%*                                                                      *
%************************************************************************

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 this code,
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.
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.

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`.

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 possibly multiple
     times.

It is always okay to lower the arity, or pretend that there are multiple calls.
In particular, "Minimum arity 0 and possibly called multiple times" is always
correct.


What we want to know from an expression
---------------------------------------

In order to obtain that information for variables, we analyze expression and
obtain bits of information:

 I.  The arity analysis:
     For every variable, whether it is absent, or called,
     and if called, with 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 variable, 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 analysis]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

The arity analysis is quite straightforward: 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 analyze an expression, we analyze it with a given context arity.
Lambdas decrease and applications increase the incoming arity. Analysing a
variable will put that arity in the environment. In `let`s or `case`s all the
results from the various subexpressions are lub'd, 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 evaluates 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 executed, so
   Return (alt₁ ∪ alt₂ ∪...)
 * App e₁ e₂ (and analogously Case scrut alts), with non-trivial e₂:
   We get the results from both sides, with the argument evaluated at most once.
   Additionally, anything called by e₁ can possibly be called with anything
   from e₂.
   Return: C(e₁) ∪ C(e₂) ∪ (fv e₁) × (fv e₂)
 * 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)}
 * 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 everything 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 (exprIsCheap), 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

    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.


Why the co-call graph?
----------------------

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
            True -> let go = \y -> case .. of
                                     True -> go (y + n 1)
                                     False > n
                    in go 1
            False -> n

vs.

        let n = case .. of ..
        in case .. of
            True -> let go = \y -> case .. of
                                     True -> go (y+1)
                                     False > n
                    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 find out that `go` and `n` are not called together.


Why co-call information for functions?
--------------------------------------

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.

    Example:

        let n = case .. of ..
            f x = n (x+1)
        in f 1 + f 2

    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.


Note [Analysis type signature]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

The workhorse 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

                            <=>

  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
every time 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 arguments.
  * 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 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 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 reanalyzed, we are done.
    Otherwise, repeat from step 3.


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 multiple times.

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

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

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.


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
    [typeArity invariants] in GHC.Core.Opt.Arity).

    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.

Note [What is a thunk]
~~~~~~~~~~~~~~~~~~~~~~

Originally, everything that is not in WHNF (`exprIsWHNF`) is considered a
thunk, not eta-expanded, to avoid losing any sharing. This is also how the
published papers on Call Arity describe it.

In practice, there are thunks that do a just little work, such as
pattern-matching on a variable, and the benefits of eta-expansion likely
outweigh the cost of doing that repeatedly. Therefore, this implementation of
Call Arity considers everything that is not cheap (`exprIsCheap`) as a thunk.

Note [Call Arity and Join Points]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

The Call Arity analysis does not care about join points, and treats them just
like normal functions. This is ok.

The analysis *could* make use of the fact that join points are always evaluated
in the same context as the join-binding they are defined in and are always
one-shot, and handle join points separately, as suggested in
https://gitlab.haskell.org/ghc/ghc/issues/13479#note_134870.
This *might* be more efficient (for example, join points would not have to be
considered interesting variables), but it would also add redundant code. So for
now we do not do that.

The simplifier never eta-expands join points (it instead pushes extra arguments from
an eta-expanded context into the join point’s RHS), so the call arity
annotation on join points is not actually used. As it would be equally valid
(though less efficient) to eta-expand join points, this is the simplifier's
choice, and hence Call Arity sets the call arity for join points as well.
-}

-- Main entry point

callArityAnalProgram :: CoreProgram -> CoreProgram
callArityAnalProgram :: CoreProgram -> CoreProgram
callArityAnalProgram CoreProgram
binds = CoreProgram
binds'
  where
    (CallArityRes
_, CoreProgram
binds') = [Id] -> VarSet -> CoreProgram -> (CallArityRes, CoreProgram)
callArityTopLvl [] VarSet
emptyVarSet CoreProgram
binds

-- See Note [Analysing top-level binds]
callArityTopLvl :: [Var] -> VarSet -> [CoreBind] -> (CallArityRes, [CoreBind])
callArityTopLvl :: [Id] -> VarSet -> CoreProgram -> (CallArityRes, CoreProgram)
callArityTopLvl [Id]
exported VarSet
_ []
    = ( CallArityRes -> CallArityRes
calledMultipleTimes (CallArityRes -> CallArityRes) -> CallArityRes -> CallArityRes
forall a b. (a -> b) -> a -> b
$ (UnVarGraph
emptyUnVarGraph, [(Id, Arity)] -> VarEnv Arity
forall a. [(Id, a)] -> VarEnv a
mkVarEnv ([(Id, Arity)] -> VarEnv Arity) -> [(Id, Arity)] -> VarEnv Arity
forall a b. (a -> b) -> a -> b
$ [(Id
v, Arity
0) | Id
v <- [Id]
exported])
      , [] )
callArityTopLvl [Id]
exported VarSet
int1 (CoreBind
b:CoreProgram
bs)
    = (CallArityRes
ae2, CoreBind
b'CoreBind -> CoreProgram -> CoreProgram
forall a. a -> [a] -> [a]
:CoreProgram
bs')
  where
    int2 :: [Id]
int2 = CoreBind -> [Id]
forall b. Bind b -> [b]
bindersOf CoreBind
b
    exported' :: [Id]
exported' = (Id -> Bool) -> [Id] -> [Id]
forall a. (a -> Bool) -> [a] -> [a]
filter Id -> Bool
isExportedId [Id]
int2 [Id] -> [Id] -> [Id]
forall a. [a] -> [a] -> [a]
++ [Id]
exported
    int' :: VarSet
int' = VarSet
int1 VarSet -> CoreBind -> VarSet
`addInterestingBinds` CoreBind
b
    (CallArityRes
ae1, CoreProgram
bs') = [Id] -> VarSet -> CoreProgram -> (CallArityRes, CoreProgram)
callArityTopLvl [Id]
exported' VarSet
int' CoreProgram
bs
    (CallArityRes
ae2, CoreBind
b')  = VarSet
-> CallArityRes -> VarSet -> CoreBind -> (CallArityRes, CoreBind)
callArityBind (CoreBind -> VarSet
boringBinds CoreBind
b) CallArityRes
ae1 VarSet
int1 CoreBind
b


callArityRHS :: CoreExpr -> CoreExpr
callArityRHS :: CoreExpr -> CoreExpr
callArityRHS = (CallArityRes, CoreExpr) -> CoreExpr
forall a b. (a, b) -> b
snd ((CallArityRes, CoreExpr) -> CoreExpr)
-> (CoreExpr -> (CallArityRes, CoreExpr)) -> CoreExpr -> CoreExpr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Arity -> VarSet -> CoreExpr -> (CallArityRes, CoreExpr)
callArityAnal Arity
0 VarSet
emptyVarSet

-- 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
    (CallArityRes, CoreExpr)
        -- How this expression uses its interesting variables
        -- and the expression with IdInfo updated

-- The trivial base cases
callArityAnal :: Arity -> VarSet -> CoreExpr -> (CallArityRes, CoreExpr)
callArityAnal Arity
_     VarSet
_   e :: CoreExpr
e@(Lit Literal
_)
    = (CallArityRes
emptyArityRes, CoreExpr
e)
callArityAnal Arity
_     VarSet
_   e :: CoreExpr
e@(Type Type
_)
    = (CallArityRes
emptyArityRes, CoreExpr
e)
callArityAnal Arity
_     VarSet
_   e :: CoreExpr
e@(Coercion Coercion
_)
    = (CallArityRes
emptyArityRes, CoreExpr
e)
-- The transparent cases
callArityAnal Arity
arity VarSet
int (Tick CoreTickish
t CoreExpr
e)
    = (CoreExpr -> CoreExpr)
-> (CallArityRes, CoreExpr) -> (CallArityRes, CoreExpr)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (CoreTickish -> CoreExpr -> CoreExpr
forall b. CoreTickish -> Expr b -> Expr b
Tick CoreTickish
t) ((CallArityRes, CoreExpr) -> (CallArityRes, CoreExpr))
-> (CallArityRes, CoreExpr) -> (CallArityRes, CoreExpr)
forall a b. (a -> b) -> a -> b
$ Arity -> VarSet -> CoreExpr -> (CallArityRes, CoreExpr)
callArityAnal Arity
arity VarSet
int CoreExpr
e
callArityAnal Arity
arity VarSet
int (Cast CoreExpr
e Coercion
co)
    = (CoreExpr -> CoreExpr)
-> (CallArityRes, CoreExpr) -> (CallArityRes, CoreExpr)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (\CoreExpr
e -> CoreExpr -> Coercion -> CoreExpr
forall b. Expr b -> Coercion -> Expr b
Cast CoreExpr
e Coercion
co) ((CallArityRes, CoreExpr) -> (CallArityRes, CoreExpr))
-> (CallArityRes, CoreExpr) -> (CallArityRes, CoreExpr)
forall a b. (a -> b) -> a -> b
$ Arity -> VarSet -> CoreExpr -> (CallArityRes, CoreExpr)
callArityAnal Arity
arity VarSet
int CoreExpr
e

-- The interesting case: Variables, Lambdas, Lets, Applications, Cases
callArityAnal Arity
arity VarSet
int e :: CoreExpr
e@(Var Id
v)
    | Id
v Id -> VarSet -> Bool
`elemVarSet` VarSet
int
    = (Id -> Arity -> CallArityRes
unitArityRes Id
v Arity
arity, CoreExpr
e)
    | Bool
otherwise
    = (CallArityRes
emptyArityRes, CoreExpr
e)

-- Non-value lambdas are ignored
callArityAnal Arity
arity VarSet
int (Lam Id
v CoreExpr
e) | Bool -> Bool
not (Id -> Bool
isId Id
v)
    = (CoreExpr -> CoreExpr)
-> (CallArityRes, CoreExpr) -> (CallArityRes, CoreExpr)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (Id -> CoreExpr -> CoreExpr
forall b. b -> Expr b -> Expr b
Lam Id
v) ((CallArityRes, CoreExpr) -> (CallArityRes, CoreExpr))
-> (CallArityRes, CoreExpr) -> (CallArityRes, CoreExpr)
forall a b. (a -> b) -> a -> b
$ Arity -> VarSet -> CoreExpr -> (CallArityRes, CoreExpr)
callArityAnal Arity
arity (VarSet
int VarSet -> Id -> VarSet
`delVarSet` Id
v) CoreExpr
e

-- We have a lambda that may be called multiple times, so its free variables
-- can all be co-called.
callArityAnal Arity
0     VarSet
int (Lam Id
v CoreExpr
e)
    = (CallArityRes
ae', Id -> CoreExpr -> CoreExpr
forall b. b -> Expr b -> Expr b
Lam Id
v CoreExpr
e')
  where
    (CallArityRes
ae, CoreExpr
e') = Arity -> VarSet -> CoreExpr -> (CallArityRes, CoreExpr)
callArityAnal Arity
0 (VarSet
int VarSet -> Id -> VarSet
`delVarSet` Id
v) CoreExpr
e
    ae' :: CallArityRes
ae' = CallArityRes -> CallArityRes
calledMultipleTimes CallArityRes
ae
-- We have a lambda that we are calling. decrease arity.
callArityAnal Arity
arity VarSet
int (Lam Id
v CoreExpr
e)
    = (CallArityRes
ae, Id -> CoreExpr -> CoreExpr
forall b. b -> Expr b -> Expr b
Lam Id
v CoreExpr
e')
  where
    (CallArityRes
ae, CoreExpr
e') = Arity -> VarSet -> CoreExpr -> (CallArityRes, CoreExpr)
callArityAnal (Arity
arity Arity -> Arity -> Arity
forall a. Num a => a -> a -> a
- Arity
1) (VarSet
int VarSet -> Id -> VarSet
`delVarSet` Id
v) CoreExpr
e

-- Application. Increase arity for the called expression, nothing to know about
-- the second
callArityAnal Arity
arity VarSet
int (App CoreExpr
e (Type Type
t))
    = (CoreExpr -> CoreExpr)
-> (CallArityRes, CoreExpr) -> (CallArityRes, CoreExpr)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (\CoreExpr
e -> CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
e (Type -> CoreExpr
forall b. Type -> Expr b
Type Type
t)) ((CallArityRes, CoreExpr) -> (CallArityRes, CoreExpr))
-> (CallArityRes, CoreExpr) -> (CallArityRes, CoreExpr)
forall a b. (a -> b) -> a -> b
$ Arity -> VarSet -> CoreExpr -> (CallArityRes, CoreExpr)
callArityAnal Arity
arity VarSet
int CoreExpr
e
callArityAnal Arity
arity VarSet
int (App CoreExpr
e1 CoreExpr
e2)
    = (CallArityRes
final_ae, CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
e1' CoreExpr
e2')
  where
    (CallArityRes
ae1, CoreExpr
e1') = Arity -> VarSet -> CoreExpr -> (CallArityRes, CoreExpr)
callArityAnal (Arity
arity Arity -> Arity -> Arity
forall a. Num a => a -> a -> a
+ Arity
1) VarSet
int CoreExpr
e1
    (CallArityRes
ae2, CoreExpr
e2') = Arity -> VarSet -> CoreExpr -> (CallArityRes, CoreExpr)
callArityAnal Arity
0           VarSet
int CoreExpr
e2
    -- 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' :: CallArityRes
ae2' | CoreExpr -> Bool
exprIsTrivial CoreExpr
e2 = CallArityRes -> CallArityRes
calledMultipleTimes CallArityRes
ae2
         | Bool
otherwise        = CallArityRes
ae2
    final_ae :: CallArityRes
final_ae = CallArityRes
ae1 CallArityRes -> CallArityRes -> CallArityRes
`both` CallArityRes
ae2'

-- Case expression.
callArityAnal Arity
arity VarSet
int (Case CoreExpr
scrut Id
bndr Type
ty [Alt Id]
alts)
    = -- pprTrace "callArityAnal:Case"
      --          (vcat [ppr scrut, ppr final_ae])
      (CallArityRes
final_ae, CoreExpr -> Id -> Type -> [Alt Id] -> CoreExpr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case CoreExpr
scrut' Id
bndr Type
ty [Alt Id]
alts')
  where
    ([CallArityRes]
alt_aes, [Alt Id]
alts') = [(CallArityRes, Alt Id)] -> ([CallArityRes], [Alt Id])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(CallArityRes, Alt Id)] -> ([CallArityRes], [Alt Id]))
-> [(CallArityRes, Alt Id)] -> ([CallArityRes], [Alt Id])
forall a b. (a -> b) -> a -> b
$ (Alt Id -> (CallArityRes, Alt Id))
-> [Alt Id] -> [(CallArityRes, Alt Id)]
forall a b. (a -> b) -> [a] -> [b]
map Alt Id -> (CallArityRes, Alt Id)
go [Alt Id]
alts
    go :: Alt Id -> (CallArityRes, Alt Id)
go (Alt AltCon
dc [Id]
bndrs CoreExpr
e) = let (CallArityRes
ae, CoreExpr
e') = Arity -> VarSet -> CoreExpr -> (CallArityRes, CoreExpr)
callArityAnal Arity
arity (VarSet
int VarSet -> [Id] -> VarSet
`delVarSetList` (Id
bndrId -> [Id] -> [Id]
forall a. a -> [a] -> [a]
:[Id]
bndrs)) CoreExpr
e
                          in  (CallArityRes
ae, AltCon -> [Id] -> CoreExpr -> Alt Id
forall b. AltCon -> [b] -> Expr b -> Alt b
Alt AltCon
dc [Id]
bndrs CoreExpr
e')
    alt_ae :: CallArityRes
alt_ae = [CallArityRes] -> CallArityRes
lubRess [CallArityRes]
alt_aes
    (CallArityRes
scrut_ae, CoreExpr
scrut') = Arity -> VarSet -> CoreExpr -> (CallArityRes, CoreExpr)
callArityAnal Arity
0 VarSet
int CoreExpr
scrut
    final_ae :: CallArityRes
final_ae = CallArityRes
scrut_ae CallArityRes -> CallArityRes -> CallArityRes
`both` CallArityRes
alt_ae

-- For lets, use callArityBind
callArityAnal Arity
arity VarSet
int (Let CoreBind
bind CoreExpr
e)
  = -- pprTrace "callArityAnal:Let"
    --          (vcat [ppr v, ppr arity, ppr n, ppr final_ae ])
    (CallArityRes
final_ae, CoreBind -> CoreExpr -> CoreExpr
forall b. Bind b -> Expr b -> Expr b
Let CoreBind
bind' CoreExpr
e')
  where
    int_body :: VarSet
int_body = VarSet
int VarSet -> CoreBind -> VarSet
`addInterestingBinds` CoreBind
bind
    (CallArityRes
ae_body, CoreExpr
e') = Arity -> VarSet -> CoreExpr -> (CallArityRes, CoreExpr)
callArityAnal Arity
arity VarSet
int_body CoreExpr
e
    (CallArityRes
final_ae, CoreBind
bind') = VarSet
-> CallArityRes -> VarSet -> CoreBind -> (CallArityRes, CoreBind)
callArityBind (CoreBind -> VarSet
boringBinds CoreBind
bind) CallArityRes
ae_body VarSet
int CoreBind
bind

-- Which bindings should we look at?
-- See Note [Which variables are interesting]
isInteresting :: Var -> Bool
isInteresting :: Id -> Bool
isInteresting Id
v = Type -> Arity
typeArity (Id -> Type
idType Id
v) Arity -> Arity -> Bool
forall a. Ord a => a -> a -> Bool
> Arity
0

interestingBinds :: CoreBind -> [Var]
interestingBinds :: CoreBind -> [Id]
interestingBinds = (Id -> Bool) -> [Id] -> [Id]
forall a. (a -> Bool) -> [a] -> [a]
filter Id -> Bool
isInteresting ([Id] -> [Id]) -> (CoreBind -> [Id]) -> CoreBind -> [Id]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoreBind -> [Id]
forall b. Bind b -> [b]
bindersOf

boringBinds :: CoreBind -> VarSet
boringBinds :: CoreBind -> VarSet
boringBinds = [Id] -> VarSet
mkVarSet ([Id] -> VarSet) -> (CoreBind -> [Id]) -> CoreBind -> VarSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Id -> Bool) -> [Id] -> [Id]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Id -> Bool) -> Id -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Id -> Bool
isInteresting) ([Id] -> [Id]) -> (CoreBind -> [Id]) -> CoreBind -> [Id]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoreBind -> [Id]
forall b. Bind b -> [b]
bindersOf

addInterestingBinds :: VarSet -> CoreBind -> VarSet
addInterestingBinds :: VarSet -> CoreBind -> VarSet
addInterestingBinds VarSet
int CoreBind
bind
    = VarSet
int VarSet -> [Id] -> VarSet
`delVarSetList`    CoreBind -> [Id]
forall b. Bind b -> [b]
bindersOf CoreBind
bind -- Possible shadowing
          VarSet -> [Id] -> VarSet
`extendVarSetList` CoreBind -> [Id]
interestingBinds CoreBind
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 :: VarSet
-> CallArityRes -> VarSet -> CoreBind -> (CallArityRes, CoreBind)
callArityBind VarSet
boring_vars CallArityRes
ae_body VarSet
int (NonRec Id
v CoreExpr
rhs)
  | Bool
otherwise
  = -- pprTrace "callArityBind:NonRec"
    --          (vcat [ppr v, ppr ae_body, ppr int, ppr ae_rhs, ppr safe_arity])
    (CallArityRes
final_ae, Id -> CoreExpr -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec Id
v' CoreExpr
rhs')
  where
    is_thunk :: Bool
is_thunk = Bool -> Bool
not (CoreExpr -> Bool
exprIsCheap CoreExpr
rhs) -- see Note [What is a thunk]
    -- If v is boring, we will not find it in ae_body, but always assume (0, False)
    boring :: Bool
boring = Id
v Id -> VarSet -> Bool
`elemVarSet` VarSet
boring_vars

    (Arity
arity, Bool
called_once)
        | Bool
boring    = (Arity
0, Bool
False) -- See Note [Taking boring variables into account]
        | Bool
otherwise = CallArityRes -> Id -> (Arity, Bool)
lookupCallArityRes CallArityRes
ae_body Id
v
    safe_arity :: Arity
safe_arity | Bool
called_once = Arity
arity
               | Bool
is_thunk    = Arity
0      -- A thunk! Do not eta-expand
               | Bool
otherwise   = Arity
arity

    -- See Note [Trimming arity]
    trimmed_arity :: Arity
trimmed_arity = Id -> Arity -> Arity
trimArity Id
v Arity
safe_arity

    (CallArityRes
ae_rhs, CoreExpr
rhs') = Arity -> VarSet -> CoreExpr -> (CallArityRes, CoreExpr)
callArityAnal Arity
trimmed_arity VarSet
int CoreExpr
rhs


    ae_rhs' :: CallArityRes
ae_rhs'| Bool
called_once     = CallArityRes
ae_rhs
           | Arity
safe_arity Arity -> Arity -> Bool
forall a. Eq a => a -> a -> Bool
== Arity
0 = CallArityRes
ae_rhs -- If it is not a function, its body is evaluated only once
           | Bool
otherwise       = CallArityRes -> CallArityRes
calledMultipleTimes CallArityRes
ae_rhs

    called_by_v :: UnVarSet
called_by_v = CallArityRes -> UnVarSet
domRes CallArityRes
ae_rhs'
    called_with_v :: UnVarSet
called_with_v
        | Bool
boring    = CallArityRes -> UnVarSet
domRes CallArityRes
ae_body
        | Bool
otherwise = CallArityRes -> Id -> UnVarSet
calledWith CallArityRes
ae_body Id
v UnVarSet -> Id -> UnVarSet
`delUnVarSet` Id
v
    final_ae :: CallArityRes
final_ae = UnVarSet -> UnVarSet -> CallArityRes -> CallArityRes
addCrossCoCalls UnVarSet
called_by_v UnVarSet
called_with_v (CallArityRes -> CallArityRes) -> CallArityRes -> CallArityRes
forall a b. (a -> b) -> a -> b
$ CallArityRes
ae_rhs' CallArityRes -> CallArityRes -> CallArityRes
`lubRes` Id -> CallArityRes -> CallArityRes
resDel Id
v CallArityRes
ae_body

    v' :: Id
v' = Id
v Id -> Arity -> Id
`setIdCallArity` Arity
trimmed_arity


-- Recursive let. See Note [Recursion and fixpointing]
callArityBind VarSet
boring_vars CallArityRes
ae_body VarSet
int b :: CoreBind
b@(Rec [(Id, CoreExpr)]
binds)
  = -- (if length binds > 300 then
    -- pprTrace "callArityBind:Rec"
    --           (vcat [ppr (Rec binds'), ppr ae_body, ppr int, ppr ae_rhs]) else id) $
    (CallArityRes
final_ae, [(Id, CoreExpr)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec [(Id, CoreExpr)]
binds')
  where
    -- See Note [Taking boring variables into account]
    any_boring :: Bool
any_boring = (Id -> Bool) -> [Id] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Id -> VarSet -> Bool
`elemVarSet` VarSet
boring_vars) [ Id
i | (Id
i, CoreExpr
_) <- [(Id, CoreExpr)]
binds]

    int_body :: VarSet
int_body = VarSet
int VarSet -> CoreBind -> VarSet
`addInterestingBinds` CoreBind
b
    (CallArityRes
ae_rhs, [(Id, CoreExpr)]
binds') = [(Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)]
-> (CallArityRes, [(Id, CoreExpr)])
fix [(Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)]
initial_binds
    final_ae :: CallArityRes
final_ae = CoreBind -> [Id]
forall b. Bind b -> [b]
bindersOf CoreBind
b [Id] -> CallArityRes -> CallArityRes
`resDelList` CallArityRes
ae_rhs

    initial_binds :: [(Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)]
initial_binds = [(Id
i,Maybe (Bool, Arity, CallArityRes)
forall a. Maybe a
Nothing,CoreExpr
e) | (Id
i,CoreExpr
e) <- [(Id, CoreExpr)]
binds]

    fix :: [(Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)] -> (CallArityRes, [(Id, CoreExpr)])
    fix :: [(Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)]
-> (CallArityRes, [(Id, CoreExpr)])
fix [(Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)]
ann_binds
        | -- pprTrace "callArityBind:fix" (vcat [ppr ann_binds, ppr any_change, ppr ae]) $
          Bool
any_change
        = [(Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)]
-> (CallArityRes, [(Id, CoreExpr)])
fix [(Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)]
ann_binds'
        | Bool
otherwise
        = (CallArityRes
ae, ((Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)
 -> (Id, CoreExpr))
-> [(Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)]
-> [(Id, CoreExpr)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Id
i, Maybe (Bool, Arity, CallArityRes)
_, CoreExpr
e) -> (Id
i, CoreExpr
e)) [(Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)]
ann_binds')
      where
        aes_old :: [(Id, CallArityRes)]
aes_old = [ (Id
i,CallArityRes
ae) | (Id
i, Just (Bool
_,Arity
_,CallArityRes
ae), CoreExpr
_) <- [(Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)]
ann_binds ]
        ae :: CallArityRes
ae = Bool -> [(Id, CallArityRes)] -> CallArityRes -> CallArityRes
callArityRecEnv Bool
any_boring [(Id, CallArityRes)]
aes_old CallArityRes
ae_body

        rerun :: (Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)
-> (Bool, (Id, Maybe (Bool, Arity, CallArityRes), CoreExpr))
rerun (Id
i, Maybe (Bool, Arity, CallArityRes)
mbLastRun, CoreExpr
rhs)
            | Id
i Id -> VarSet -> Bool
`elemVarSet` VarSet
int_body Bool -> Bool -> Bool
&& Bool -> Bool
not (Id
i Id -> UnVarSet -> Bool
`elemUnVarSet` CallArityRes -> UnVarSet
domRes CallArityRes
ae)
            -- No call to this yet, so do nothing
            = (Bool
False, (Id
i, Maybe (Bool, Arity, CallArityRes)
forall a. Maybe a
Nothing, CoreExpr
rhs))

            | Just (Bool
old_called_once, Arity
old_arity, CallArityRes
_) <- Maybe (Bool, Arity, CallArityRes)
mbLastRun
            , Bool
called_once Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
== Bool
old_called_once
            , Arity
new_arity Arity -> Arity -> Bool
forall a. Eq a => a -> a -> Bool
== Arity
old_arity
            -- No change, no need to re-analyze
            = (Bool
False, (Id
i, Maybe (Bool, Arity, CallArityRes)
mbLastRun, CoreExpr
rhs))

            | Bool
otherwise
            -- We previously analyzed this with a different arity (or not at all)
            = let is_thunk :: Bool
is_thunk = Bool -> Bool
not (CoreExpr -> Bool
exprIsCheap CoreExpr
rhs) -- see Note [What is a thunk]

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

                  -- See Note [Trimming arity]
                  trimmed_arity :: Arity
trimmed_arity = Id -> Arity -> Arity
trimArity Id
i Arity
safe_arity

                  (CallArityRes
ae_rhs, CoreExpr
rhs') = Arity -> VarSet -> CoreExpr -> (CallArityRes, CoreExpr)
callArityAnal Arity
trimmed_arity VarSet
int_body CoreExpr
rhs

                  ae_rhs' :: CallArityRes
ae_rhs' | Bool
called_once     = CallArityRes
ae_rhs
                          | Arity
safe_arity Arity -> Arity -> Bool
forall a. Eq a => a -> a -> Bool
== Arity
0 = CallArityRes
ae_rhs -- If it is not a function, its body is evaluated only once
                          | Bool
otherwise       = CallArityRes -> CallArityRes
calledMultipleTimes CallArityRes
ae_rhs

                  i' :: Id
i' = Id
i Id -> Arity -> Id
`setIdCallArity` Arity
trimmed_arity

              in (Bool
True, (Id
i', (Bool, Arity, CallArityRes) -> Maybe (Bool, Arity, CallArityRes)
forall a. a -> Maybe a
Just (Bool
called_once, Arity
new_arity, CallArityRes
ae_rhs'), CoreExpr
rhs'))
          where
            -- See Note [Taking boring variables into account]
            (Arity
new_arity, Bool
called_once) | Id
i Id -> VarSet -> Bool
`elemVarSet` VarSet
boring_vars = (Arity
0, Bool
False)
                                     | Bool
otherwise                  = CallArityRes -> Id -> (Arity, Bool)
lookupCallArityRes CallArityRes
ae Id
i

        ([Bool]
changes, [(Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)]
ann_binds') = [(Bool, (Id, Maybe (Bool, Arity, CallArityRes), CoreExpr))]
-> ([Bool], [(Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Bool, (Id, Maybe (Bool, Arity, CallArityRes), CoreExpr))]
 -> ([Bool], [(Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)]))
-> [(Bool, (Id, Maybe (Bool, Arity, CallArityRes), CoreExpr))]
-> ([Bool], [(Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)])
forall a b. (a -> b) -> a -> b
$ ((Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)
 -> (Bool, (Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)))
-> [(Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)]
-> [(Bool, (Id, Maybe (Bool, Arity, CallArityRes), CoreExpr))]
forall a b. (a -> b) -> [a] -> [b]
map (Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)
-> (Bool, (Id, Maybe (Bool, Arity, CallArityRes), CoreExpr))
rerun [(Id, Maybe (Bool, Arity, CallArityRes), CoreExpr)]
ann_binds
        any_change :: Bool
any_change = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or [Bool]
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 :: Bool -> [(Id, CallArityRes)] -> CallArityRes -> CallArityRes
callArityRecEnv Bool
any_boring [(Id, CallArityRes)]
ae_rhss CallArityRes
ae_body
    = -- (if length ae_rhss > 300 then pprTrace "callArityRecEnv" (vcat [ppr ae_rhss, ppr ae_body, ppr ae_new]) else id) $
      CallArityRes
ae_new
  where
    vars :: [Id]
vars = ((Id, CallArityRes) -> Id) -> [(Id, CallArityRes)] -> [Id]
forall a b. (a -> b) -> [a] -> [b]
map (Id, CallArityRes) -> Id
forall a b. (a, b) -> a
fst [(Id, CallArityRes)]
ae_rhss

    ae_combined :: CallArityRes
ae_combined = [CallArityRes] -> CallArityRes
lubRess (((Id, CallArityRes) -> CallArityRes)
-> [(Id, CallArityRes)] -> [CallArityRes]
forall a b. (a -> b) -> [a] -> [b]
map (Id, CallArityRes) -> CallArityRes
forall a b. (a, b) -> b
snd [(Id, CallArityRes)]
ae_rhss) CallArityRes -> CallArityRes -> CallArityRes
`lubRes` CallArityRes
ae_body

    cross_calls :: UnVarGraph
cross_calls
        -- See Note [Taking boring variables into account]
        | Bool
any_boring               = UnVarSet -> UnVarGraph
completeGraph (CallArityRes -> UnVarSet
domRes CallArityRes
ae_combined)
        -- Also, calculating cross_calls is expensive. Simply be conservative
        -- if the mutually recursive group becomes too large.
        | [(Id, CallArityRes)] -> Arity -> Bool
forall a. [a] -> Arity -> Bool
lengthExceeds [(Id, CallArityRes)]
ae_rhss Arity
25 = UnVarSet -> UnVarGraph
completeGraph (CallArityRes -> UnVarSet
domRes CallArityRes
ae_combined)
        | Bool
otherwise                = [UnVarGraph] -> UnVarGraph
unionUnVarGraphs ([UnVarGraph] -> UnVarGraph) -> [UnVarGraph] -> UnVarGraph
forall a b. (a -> b) -> a -> b
$ ((Id, CallArityRes) -> UnVarGraph)
-> [(Id, CallArityRes)] -> [UnVarGraph]
forall a b. (a -> b) -> [a] -> [b]
map (Id, CallArityRes) -> UnVarGraph
cross_call [(Id, CallArityRes)]
ae_rhss
    cross_call :: (Id, CallArityRes) -> UnVarGraph
cross_call (Id
v, CallArityRes
ae_rhs) = UnVarSet -> UnVarSet -> UnVarGraph
completeBipartiteGraph UnVarSet
called_by_v UnVarSet
called_with_v
      where
        is_thunk :: Bool
is_thunk = Id -> Arity
idCallArity Id
v Arity -> Arity -> Bool
forall a. Eq a => a -> a -> Bool
== Arity
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 :: CallArityRes
ae_before_v | Bool
is_thunk  = [CallArityRes] -> CallArityRes
lubRess (((Id, CallArityRes) -> CallArityRes)
-> [(Id, CallArityRes)] -> [CallArityRes]
forall a b. (a -> b) -> [a] -> [b]
map (Id, CallArityRes) -> CallArityRes
forall a b. (a, b) -> b
snd ([(Id, CallArityRes)] -> [CallArityRes])
-> [(Id, CallArityRes)] -> [CallArityRes]
forall a b. (a -> b) -> a -> b
$ ((Id, CallArityRes) -> Bool)
-> [(Id, CallArityRes)] -> [(Id, CallArityRes)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
/= Id
v) (Id -> Bool)
-> ((Id, CallArityRes) -> Id) -> (Id, CallArityRes) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Id, CallArityRes) -> Id
forall a b. (a, b) -> a
fst) [(Id, CallArityRes)]
ae_rhss) CallArityRes -> CallArityRes -> CallArityRes
`lubRes` CallArityRes
ae_body
                    | Bool
otherwise = CallArityRes
ae_combined
        -- What do we want to know from these?
        -- Which calls can happen next to any recursive call.
        called_with_v :: UnVarSet
called_with_v
            = [UnVarSet] -> UnVarSet
unionUnVarSets ([UnVarSet] -> UnVarSet) -> [UnVarSet] -> UnVarSet
forall a b. (a -> b) -> a -> b
$ (Id -> UnVarSet) -> [Id] -> [UnVarSet]
forall a b. (a -> b) -> [a] -> [b]
map (CallArityRes -> Id -> UnVarSet
calledWith CallArityRes
ae_before_v) [Id]
vars
        called_by_v :: UnVarSet
called_by_v = CallArityRes -> UnVarSet
domRes CallArityRes
ae_rhs

    ae_new :: CallArityRes
ae_new = (UnVarGraph -> UnVarGraph) -> CallArityRes -> CallArityRes
forall b c d. (b -> c) -> (b, d) -> (c, d)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (UnVarGraph
cross_calls UnVarGraph -> UnVarGraph -> UnVarGraph
`unionUnVarGraph`) CallArityRes
ae_combined

-- See Note [Trimming arity]
trimArity :: Id -> Arity -> Arity
trimArity :: Id -> Arity -> Arity
trimArity Id
v Arity
a = [Arity] -> Arity
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [Arity
a, Arity
max_arity_by_type, Arity
max_arity_by_strsig]
  where
    max_arity_by_type :: Arity
max_arity_by_type = Type -> Arity
typeArity (Id -> Type
idType Id
v)
    max_arity_by_strsig :: Arity
max_arity_by_strsig
        | Divergence -> Bool
isDeadEndDiv Divergence
result_info = [Demand] -> Arity
forall a. [a] -> Arity
forall (t :: * -> *) a. Foldable t => t a -> Arity
length [Demand]
demands
        | Bool
otherwise = Arity
a

    ([Demand]
demands, Divergence
result_info) = DmdSig -> ([Demand], Divergence)
splitDmdSig (Id -> DmdSig
idDmdSig Id
v)

---------------------------------------
-- Functions related to CallArityRes --
---------------------------------------

-- Result type for the two analyses.
-- See Note [Analysis I: The arity analysis]
-- and Note [Analysis II: The Co-Called analysis]
type CallArityRes = (UnVarGraph, VarEnv Arity)

emptyArityRes :: CallArityRes
emptyArityRes :: CallArityRes
emptyArityRes = (UnVarGraph
emptyUnVarGraph, VarEnv Arity
forall a. VarEnv a
emptyVarEnv)

unitArityRes :: Var -> Arity -> CallArityRes
unitArityRes :: Id -> Arity -> CallArityRes
unitArityRes Id
v Arity
arity = (UnVarGraph
emptyUnVarGraph, Id -> Arity -> VarEnv Arity
forall a. Id -> a -> VarEnv a
unitVarEnv Id
v Arity
arity)

resDelList :: [Var] -> CallArityRes -> CallArityRes
resDelList :: [Id] -> CallArityRes -> CallArityRes
resDelList [Id]
vs CallArityRes
ae = (CallArityRes -> Id -> CallArityRes)
-> CallArityRes -> [Id] -> CallArityRes
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((Id -> CallArityRes -> CallArityRes)
-> CallArityRes -> Id -> CallArityRes
forall a b c. (a -> b -> c) -> b -> a -> c
flip Id -> CallArityRes -> CallArityRes
resDel) CallArityRes
ae [Id]
vs

resDel :: Var -> CallArityRes -> CallArityRes
resDel :: Id -> CallArityRes -> CallArityRes
resDel Id
v (!UnVarGraph
g, !VarEnv Arity
ae) = (UnVarGraph
g UnVarGraph -> Id -> UnVarGraph
`delNode` Id
v, VarEnv Arity
ae VarEnv Arity -> Id -> VarEnv Arity
forall a. VarEnv a -> Id -> VarEnv a
`delVarEnv` Id
v)

domRes :: CallArityRes -> UnVarSet
domRes :: CallArityRes -> UnVarSet
domRes (UnVarGraph
_, VarEnv Arity
ae) = VarEnv Arity -> UnVarSet
forall elt. VarEnv elt -> UnVarSet
varEnvDomain VarEnv Arity
ae

-- In the result, find out the minimum arity and whether the variable is called
-- at most once.
lookupCallArityRes :: CallArityRes -> Var -> (Arity, Bool)
lookupCallArityRes :: CallArityRes -> Id -> (Arity, Bool)
lookupCallArityRes (UnVarGraph
g, VarEnv Arity
ae) Id
v
    = case VarEnv Arity -> Id -> Maybe Arity
forall a. VarEnv a -> Id -> Maybe a
lookupVarEnv VarEnv Arity
ae Id
v of
        Just Arity
a -> (Arity
a, Bool -> Bool
not (UnVarGraph
g UnVarGraph -> Id -> Bool
`hasLoopAt` Id
v))
        Maybe Arity
Nothing -> (Arity
0, Bool
False)

calledWith :: CallArityRes -> Var -> UnVarSet
calledWith :: CallArityRes -> Id -> UnVarSet
calledWith (UnVarGraph
g, VarEnv Arity
_) Id
v = UnVarGraph -> Id -> UnVarSet
neighbors UnVarGraph
g Id
v

addCrossCoCalls :: UnVarSet -> UnVarSet -> CallArityRes -> CallArityRes
addCrossCoCalls :: UnVarSet -> UnVarSet -> CallArityRes -> CallArityRes
addCrossCoCalls UnVarSet
set1 UnVarSet
set2 = (UnVarGraph -> UnVarGraph) -> CallArityRes -> CallArityRes
forall b c d. (b -> c) -> (b, d) -> (c, d)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (UnVarSet -> UnVarSet -> UnVarGraph
completeBipartiteGraph UnVarSet
set1 UnVarSet
set2 UnVarGraph -> UnVarGraph -> UnVarGraph
`unionUnVarGraph`)

-- Replaces the co-call graph by a complete graph (i.e. no information)
calledMultipleTimes :: CallArityRes -> CallArityRes
calledMultipleTimes :: CallArityRes -> CallArityRes
calledMultipleTimes CallArityRes
res = (UnVarGraph -> UnVarGraph) -> CallArityRes -> CallArityRes
forall b c d. (b -> c) -> (b, d) -> (c, d)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (UnVarGraph -> UnVarGraph -> UnVarGraph
forall a b. a -> b -> a
const (UnVarSet -> UnVarGraph
completeGraph (CallArityRes -> UnVarSet
domRes CallArityRes
res))) CallArityRes
res

-- Used for application and cases
both :: CallArityRes -> CallArityRes -> CallArityRes
both :: CallArityRes -> CallArityRes -> CallArityRes
both CallArityRes
r1 CallArityRes
r2 = UnVarSet -> UnVarSet -> CallArityRes -> CallArityRes
addCrossCoCalls (CallArityRes -> UnVarSet
domRes CallArityRes
r1) (CallArityRes -> UnVarSet
domRes CallArityRes
r2) (CallArityRes -> CallArityRes) -> CallArityRes -> CallArityRes
forall a b. (a -> b) -> a -> b
$ CallArityRes
r1 CallArityRes -> CallArityRes -> CallArityRes
`lubRes` CallArityRes
r2

-- Used when combining results from alternative cases; take the minimum
lubRes :: CallArityRes -> CallArityRes -> CallArityRes
lubRes :: CallArityRes -> CallArityRes -> CallArityRes
lubRes (UnVarGraph
g1, VarEnv Arity
ae1) (UnVarGraph
g2, VarEnv Arity
ae2) = (UnVarGraph
g1 UnVarGraph -> UnVarGraph -> UnVarGraph
`unionUnVarGraph` UnVarGraph
g2, VarEnv Arity
ae1 VarEnv Arity -> VarEnv Arity -> VarEnv Arity
`lubArityEnv` VarEnv Arity
ae2)

lubArityEnv :: VarEnv Arity -> VarEnv Arity -> VarEnv Arity
lubArityEnv :: VarEnv Arity -> VarEnv Arity -> VarEnv Arity
lubArityEnv = (Arity -> Arity -> Arity)
-> VarEnv Arity -> VarEnv Arity -> VarEnv Arity
forall a. (a -> a -> a) -> VarEnv a -> VarEnv a -> VarEnv a
plusVarEnv_C Arity -> Arity -> Arity
forall a. Ord a => a -> a -> a
min

lubRess :: [CallArityRes] -> CallArityRes
lubRess :: [CallArityRes] -> CallArityRes
lubRess = (CallArityRes -> CallArityRes -> CallArityRes)
-> CallArityRes -> [CallArityRes] -> CallArityRes
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' CallArityRes -> CallArityRes -> CallArityRes
lubRes CallArityRes
emptyArityRes