{-
(c) The University of Glasgow 2006
(c) The GRASP/AQUA Project, Glasgow University, 1992-1998


        Arity and eta expansion
-}

{-# LANGUAGE CPP #-}

{-# OPTIONS_GHC -Wno-incomplete-record-updates #-}

-- | Arity and eta expansion
module GHC.Core.Opt.Arity
   ( manifestArity, joinRhsArity, exprArity, typeArity
   , exprEtaExpandArity, findRhsArity
   , etaExpand, etaExpandAT
   , etaExpandToJoinPoint, etaExpandToJoinPointRule
   , exprBotStrictness_maybe
   , ArityType(..), expandableArityType, arityTypeArity
   , maxWithArity, isBotArityType, idArityType
   )
where

#include "HsVersions.h"

import GHC.Prelude

import GHC.Core
import GHC.Core.FVs
import GHC.Core.Utils
import GHC.Core.Subst
import GHC.Types.Demand
import GHC.Types.Var
import GHC.Types.Var.Env
import GHC.Types.Id
import GHC.Core.Type as Type
import GHC.Core.TyCon     ( initRecTc, checkRecTc )
import GHC.Core.Predicate ( isDictTy )
import GHC.Core.Coercion as Coercion
import GHC.Core.Multiplicity
import GHC.Types.Var.Set
import GHC.Types.Basic
import GHC.Types.Unique
import GHC.Driver.Session ( DynFlags, GeneralFlag(..), gopt )
import GHC.Utils.Outputable
import GHC.Data.FastString
import GHC.Utils.Misc     ( lengthAtLeast )

{-
************************************************************************
*                                                                      *
              manifestArity and exprArity
*                                                                      *
************************************************************************

exprArity is a cheap-and-cheerful version of exprEtaExpandArity.
It tells how many things the expression can be applied to before doing
any work.  It doesn't look inside cases, lets, etc.  The idea is that
exprEtaExpandArity will do the hard work, leaving something that's easy
for exprArity to grapple with.  In particular, Simplify uses exprArity to
compute the ArityInfo for the Id.

Originally I thought that it was enough just to look for top-level lambdas, but
it isn't.  I've seen this

        foo = PrelBase.timesInt

We want foo to get arity 2 even though the eta-expander will leave it
unchanged, in the expectation that it'll be inlined.  But occasionally it
isn't, because foo is blacklisted (used in a rule).

Similarly, see the ok_note check in exprEtaExpandArity.  So
        f = __inline_me (\x -> e)
won't be eta-expanded.

And in any case it seems more robust to have exprArity be a bit more intelligent.
But note that   (\x y z -> f x y z)
should have arity 3, regardless of f's arity.
-}

manifestArity :: CoreExpr -> Arity
-- ^ manifestArity sees how many leading value lambdas there are,
--   after looking through casts
manifestArity :: CoreExpr -> JoinArity
manifestArity (Lam Id
v CoreExpr
e) | Id -> Bool
isId Id
v        = JoinArity
1 JoinArity -> JoinArity -> JoinArity
forall a. Num a => a -> a -> a
+ CoreExpr -> JoinArity
manifestArity CoreExpr
e
                        | Bool
otherwise     = CoreExpr -> JoinArity
manifestArity CoreExpr
e
manifestArity (Tick Tickish Id
t CoreExpr
e) | Bool -> Bool
not (Tickish Id -> Bool
forall id. Tickish id -> Bool
tickishIsCode Tickish Id
t) =  CoreExpr -> JoinArity
manifestArity CoreExpr
e
manifestArity (Cast CoreExpr
e CoercionR
_)                = CoreExpr -> JoinArity
manifestArity CoreExpr
e
manifestArity CoreExpr
_                         = JoinArity
0

joinRhsArity :: CoreExpr -> JoinArity
-- Join points are supposed to have manifestly-visible
-- lambdas at the top: no ticks, no casts, nothing
-- Moreover, type lambdas count in JoinArity
joinRhsArity :: CoreExpr -> JoinArity
joinRhsArity (Lam Id
_ CoreExpr
e) = JoinArity
1 JoinArity -> JoinArity -> JoinArity
forall a. Num a => a -> a -> a
+ CoreExpr -> JoinArity
joinRhsArity CoreExpr
e
joinRhsArity CoreExpr
_         = JoinArity
0


---------------
exprArity :: CoreExpr -> Arity
-- ^ An approximate, fast, version of 'exprEtaExpandArity'
exprArity :: CoreExpr -> JoinArity
exprArity CoreExpr
e = CoreExpr -> JoinArity
go CoreExpr
e
  where
    go :: CoreExpr -> JoinArity
go (Var Id
v)                     = Id -> JoinArity
idArity Id
v
    go (Lam Id
x CoreExpr
e) | Id -> Bool
isId Id
x          = CoreExpr -> JoinArity
go CoreExpr
e JoinArity -> JoinArity -> JoinArity
forall a. Num a => a -> a -> a
+ JoinArity
1
                 | Bool
otherwise       = CoreExpr -> JoinArity
go CoreExpr
e
    go (Tick Tickish Id
t CoreExpr
e) | Bool -> Bool
not (Tickish Id -> Bool
forall id. Tickish id -> Bool
tickishIsCode Tickish Id
t) = CoreExpr -> JoinArity
go CoreExpr
e
    go (Cast CoreExpr
e CoercionR
co)                 = JoinArity -> Type -> JoinArity
trim_arity (CoreExpr -> JoinArity
go CoreExpr
e) (CoercionR -> Type
coercionRKind CoercionR
co)
                                        -- Note [exprArity invariant]
    go (App CoreExpr
e (Type Type
_))            = CoreExpr -> JoinArity
go CoreExpr
e
    go (App CoreExpr
f CoreExpr
a) | CoreExpr -> Bool
exprIsTrivial CoreExpr
a = (CoreExpr -> JoinArity
go CoreExpr
f JoinArity -> JoinArity -> JoinArity
forall a. Num a => a -> a -> a
- JoinArity
1) JoinArity -> JoinArity -> JoinArity
forall a. Ord a => a -> a -> a
`max` JoinArity
0
        -- See Note [exprArity for applications]
        -- NB: coercions count as a value argument

    go CoreExpr
_                           = JoinArity
0

    trim_arity :: Arity -> Type -> Arity
    trim_arity :: JoinArity -> Type -> JoinArity
trim_arity JoinArity
arity Type
ty = JoinArity
arity JoinArity -> JoinArity -> JoinArity
forall a. Ord a => a -> a -> a
`min` [OneShotInfo] -> JoinArity
forall (t :: * -> *) a. Foldable t => t a -> JoinArity
length (Type -> [OneShotInfo]
typeArity Type
ty)

---------------
typeArity :: Type -> [OneShotInfo]
-- How many value arrows are visible in the type?
-- We look through foralls, and newtypes
-- See Note [exprArity invariant]
typeArity :: Type -> [OneShotInfo]
typeArity Type
ty
  = RecTcChecker -> Type -> [OneShotInfo]
go RecTcChecker
initRecTc Type
ty
  where
    go :: RecTcChecker -> Type -> [OneShotInfo]
go RecTcChecker
rec_nts Type
ty
      | Just (Id
_, Type
ty')  <- Type -> Maybe (Id, Type)
splitForAllTy_maybe Type
ty
      = RecTcChecker -> Type -> [OneShotInfo]
go RecTcChecker
rec_nts Type
ty'

      | Just (Type
_,Type
arg,Type
res) <- Type -> Maybe (Type, Type, Type)
splitFunTy_maybe Type
ty
      = Type -> OneShotInfo
typeOneShot Type
arg OneShotInfo -> [OneShotInfo] -> [OneShotInfo]
forall a. a -> [a] -> [a]
: RecTcChecker -> Type -> [OneShotInfo]
go RecTcChecker
rec_nts Type
res

      | Just (TyCon
tc,[Type]
tys) <- HasDebugCallStack => Type -> Maybe (TyCon, [Type])
Type -> Maybe (TyCon, [Type])
splitTyConApp_maybe Type
ty
      , Just (Type
ty', CoercionR
_) <- TyCon -> [Type] -> Maybe (Type, CoercionR)
instNewTyCon_maybe TyCon
tc [Type]
tys
      , Just RecTcChecker
rec_nts' <- RecTcChecker -> TyCon -> Maybe RecTcChecker
checkRecTc RecTcChecker
rec_nts TyCon
tc  -- See Note [Expanding newtypes]
                                                -- in GHC.Core.TyCon
--   , not (isClassTyCon tc)    -- Do not eta-expand through newtype classes
--                              -- See Note [Newtype classes and eta expansion]
--                              (no longer required)
      = RecTcChecker -> Type -> [OneShotInfo]
go RecTcChecker
rec_nts' Type
ty'
        -- Important to look through non-recursive newtypes, so that, eg
        --      (f x)   where f has arity 2, f :: Int -> IO ()
        -- Here we want to get arity 1 for the result!
        --
        -- AND through a layer of recursive newtypes
        -- e.g. newtype Stream m a b = Stream (m (Either b (a, Stream m a b)))

      | Bool
otherwise
      = []

---------------
exprBotStrictness_maybe :: CoreExpr -> Maybe (Arity, StrictSig)
-- A cheap and cheerful function that identifies bottoming functions
-- and gives them a suitable strictness signatures.  It's used during
-- float-out
exprBotStrictness_maybe :: CoreExpr -> Maybe (JoinArity, StrictSig)
exprBotStrictness_maybe CoreExpr
e
  = case ArityType -> Maybe JoinArity
getBotArity (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
e) of
        Maybe JoinArity
Nothing -> Maybe (JoinArity, StrictSig)
forall a. Maybe a
Nothing
        Just JoinArity
ar -> (JoinArity, StrictSig) -> Maybe (JoinArity, StrictSig)
forall a. a -> Maybe a
Just (JoinArity
ar, JoinArity -> StrictSig
sig JoinArity
ar)
  where
    env :: ArityEnv
env    = AE :: CheapFun -> Bool -> IdSet -> ArityEnv
AE { ae_ped_bot :: Bool
ae_ped_bot = Bool
True
                , ae_cheap_fn :: CheapFun
ae_cheap_fn = \ CoreExpr
_ Maybe Type
_ -> Bool
False
                , ae_joins :: IdSet
ae_joins = IdSet
emptyVarSet }
    sig :: JoinArity -> StrictSig
sig JoinArity
ar = [Demand] -> Divergence -> StrictSig
mkClosedStrictSig (JoinArity -> Demand -> [Demand]
forall a. JoinArity -> a -> [a]
replicate JoinArity
ar Demand
topDmd) Divergence
botDiv

{-
Note [exprArity invariant]
~~~~~~~~~~~~~~~~~~~~~~~~~~
exprArity has the following invariants:

  (1) If typeArity (exprType e) = n,
      then manifestArity (etaExpand e n) = n

      That is, etaExpand can always expand as much as typeArity says
      So the case analysis in etaExpand and in typeArity must match

  (2) exprArity e <= typeArity (exprType e)

  (3) Hence if (exprArity e) = n, then manifestArity (etaExpand e n) = n

      That is, if exprArity says "the arity is n" then etaExpand really
      can get "n" manifest lambdas to the top.

Why is this important?  Because
  - In GHC.Iface.Tidy we use exprArity to fix the *final arity* of
    each top-level Id, and in
  - In CorePrep we use etaExpand on each rhs, so that the visible lambdas
    actually match that arity, which in turn means
    that the StgRhs has the right number of lambdas

An alternative would be to do the eta-expansion in GHC.Iface.Tidy, at least
for top-level bindings, in which case we would not need the trim_arity
in exprArity.  That is a less local change, so I'm going to leave it for today!

Note [Newtype classes and eta expansion]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    NB: this nasty special case is no longer required, because
    for newtype classes we don't use the class-op rule mechanism
    at all.  See Note [Single-method classes] in GHC.Tc.TyCl.Instance. SLPJ May 2013

-------- Old out of date comments, just for interest -----------
We have to be careful when eta-expanding through newtypes.  In general
it's a good idea, but annoyingly it interacts badly with the class-op
rule mechanism.  Consider

   class C a where { op :: a -> a }
   instance C b => C [b] where
     op x = ...

These translate to

   co :: forall a. (a->a) ~ C a

   $copList :: C b -> [b] -> [b]
   $copList d x = ...

   $dfList :: C b -> C [b]
   {-# DFunUnfolding = [$copList] #-}
   $dfList d = $copList d |> co@[b]

Now suppose we have:

   dCInt :: C Int

   blah :: [Int] -> [Int]
   blah = op ($dfList dCInt)

Now we want the built-in op/$dfList rule will fire to give
   blah = $copList dCInt

But with eta-expansion 'blah' might (and in #3772, which is
slightly more complicated, does) turn into

   blah = op (\eta. ($dfList dCInt |> sym co) eta)

and now it is *much* harder for the op/$dfList rule to fire, because
exprIsConApp_maybe won't hold of the argument to op.  I considered
trying to *make* it hold, but it's tricky and I gave up.

The test simplCore/should_compile/T3722 is an excellent example.
-------- End of old out of date comments, just for interest -----------


Note [exprArity for applications]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
When we come to an application we check that the arg is trivial.
   eg  f (fac x) does not have arity 2,
                 even if f has arity 3!

* We require that is trivial rather merely cheap.  Suppose f has arity 2.
  Then    f (Just y)
  has arity 0, because if we gave it arity 1 and then inlined f we'd get
          let v = Just y in \w. <f-body>
  which has arity 0.  And we try to maintain the invariant that we don't
  have arity decreases.

*  The `max 0` is important!  (\x y -> f x) has arity 2, even if f is
   unknown, hence arity 0


************************************************************************
*                                                                      *
           Computing the "arity" of an expression
*                                                                      *
************************************************************************

Note [Definition of arity]
~~~~~~~~~~~~~~~~~~~~~~~~~~
The "arity" of an expression 'e' is n if
   applying 'e' to *fewer* than n *value* arguments
   converges rapidly

Or, to put it another way

   there is no work lost in duplicating the partial
   application (e x1 .. x(n-1))

In the divergent case, no work is lost by duplicating because if the thing
is evaluated once, that's the end of the program.

Or, to put it another way, in any context C

   C[ (\x1 .. xn. e x1 .. xn) ]
         is as efficient as
   C[ e ]

It's all a bit more subtle than it looks:

Note [One-shot lambdas]
~~~~~~~~~~~~~~~~~~~~~~~
Consider one-shot lambdas
                let x = expensive in \y z -> E
We want this to have arity 1 if the \y-abstraction is a 1-shot lambda.

Note [Dealing with bottom]
~~~~~~~~~~~~~~~~~~~~~~~~~~
A Big Deal with computing arities is expressions like

   f = \x -> case x of
               True  -> \s -> e1
               False -> \s -> e2

This happens all the time when f :: Bool -> IO ()
In this case we do eta-expand, in order to get that \s to the
top, and give f arity 2.

This isn't really right in the presence of seq.  Consider
        (f bot) `seq` 1

This should diverge!  But if we eta-expand, it won't.  We ignore this
"problem" (unless -fpedantic-bottoms is on), because being scrupulous
would lose an important transformation for many programs. (See
#5587 for an example.)

Consider also
        f = \x -> error "foo"
Here, arity 1 is fine.  But if it is
        f = \x -> case x of
                        True  -> error "foo"
                        False -> \y -> x+y
then we want to get arity 2.  Technically, this isn't quite right, because
        (f True) `seq` 1
should diverge, but it'll converge if we eta-expand f.  Nevertheless, we
do so; it improves some programs significantly, and increasing convergence
isn't a bad thing.  Hence the ABot/ATop in ArityType.

So these two transformations aren't always the Right Thing, and we
have several tickets reporting unexpected behaviour resulting from
this transformation.  So we try to limit it as much as possible:

 (1) Do NOT move a lambda outside a known-bottom case expression
       case undefined of { (a,b) -> \y -> e }
     This showed up in #5557

 (2) Do NOT move a lambda outside a case if all the branches of
     the case are known to return bottom.
        case x of { (a,b) -> \y -> error "urk" }
     This case is less important, but the idea is that if the fn is
     going to diverge eventually anyway then getting the best arity
     isn't an issue, so we might as well play safe

 (3) Do NOT move a lambda outside a case unless
     (a) The scrutinee is ok-for-speculation, or
     (b) more liberally: the scrutinee is cheap (e.g. a variable), and
         -fpedantic-bottoms is not enforced (see #2915 for an example)

Of course both (1) and (2) are readily defeated by disguising the bottoms.

4. Note [Newtype arity]
~~~~~~~~~~~~~~~~~~~~~~~~
Non-recursive newtypes are transparent, and should not get in the way.
We do (currently) eta-expand recursive newtypes too.  So if we have, say

        newtype T = MkT ([T] -> Int)

Suppose we have
        e = coerce T f
where f has arity 1.  Then: etaExpandArity e = 1;
that is, etaExpandArity looks through the coerce.

When we eta-expand e to arity 1: eta_expand 1 e T
we want to get:                  coerce T (\x::[T] -> (coerce ([T]->Int) e) x)

  HOWEVER, note that if you use coerce bogusly you can ge
        coerce Int negate
  And since negate has arity 2, you might try to eta expand.  But you can't
  decompose Int to a function type.   Hence the final case in eta_expand.

Note [The state-transformer hack]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Suppose we have
        f = e
where e has arity n.  Then, if we know from the context that f has
a usage type like
        t1 -> ... -> tn -1-> t(n+1) -1-> ... -1-> tm -> ...
then we can expand the arity to m.  This usage type says that
any application (x e1 .. en) will be applied to uniquely to (m-n) more args
Consider f = \x. let y = <expensive>
                 in case x of
                      True  -> foo
                      False -> \(s:RealWorld) -> e
where foo has arity 1.  Then we want the state hack to
apply to foo too, so we can eta expand the case.

Then we expect that if f is applied to one arg, it'll be applied to two
(that's the hack -- we don't really know, and sometimes it's false)
See also Id.isOneShotBndr.

Note [State hack and bottoming functions]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
It's a terrible idea to use the state hack on a bottoming function.
Here's what happens (#2861):

  f :: String -> IO T
  f = \p. error "..."

Eta-expand, using the state hack:

  f = \p. (\s. ((error "...") |> g1) s) |> g2
  g1 :: IO T ~ (S -> (S,T))
  g2 :: (S -> (S,T)) ~ IO T

Extrude the g2

  f' = \p. \s. ((error "...") |> g1) s
  f = f' |> (String -> g2)

Discard args for bottomming function

  f' = \p. \s. ((error "...") |> g1 |> g3
  g3 :: (S -> (S,T)) ~ (S,T)

Extrude g1.g3

  f'' = \p. \s. (error "...")
  f' = f'' |> (String -> S -> g1.g3)

And now we can repeat the whole loop.  Aargh!  The bug is in applying the
state hack to a function which then swallows the argument.

This arose in another guise in #3959.  Here we had

     catch# (throw exn >> return ())

Note that (throw :: forall a e. Exn e => e -> a) is called with [a = IO ()].
After inlining (>>) we get

     catch# (\_. throw {IO ()} exn)

We must *not* eta-expand to

     catch# (\_ _. throw {...} exn)

because 'catch#' expects to get a (# _,_ #) after applying its argument to
a State#, not another function!

In short, we use the state hack to allow us to push let inside a lambda,
but not to introduce a new lambda.


Note [ArityType]
~~~~~~~~~~~~~~~~
ArityType is the result of a compositional analysis on expressions,
from which we can decide the real arity of the expression (extracted
with function exprEtaExpandArity).

Here is what the fields mean. If an arbitrary expression 'f' has
ArityType 'at', then

 * If at = ABot n, then (f x1..xn) definitely diverges. Partial
   applications to fewer than n args may *or may not* diverge.

   We allow ourselves to eta-expand bottoming functions, even
   if doing so may lose some `seq` sharing,
       let x = <expensive> in \y. error (g x y)
       ==> \y. let x = <expensive> in error (g x y)

 * If at = ATop as, and n=length as,
   then expanding 'f' to (\x1..xn. f x1 .. xn) loses no sharing,
   assuming the calls of f respect the one-shot-ness of
   its definition.

   NB 'f' is an arbitrary expression, eg (f = g e1 e2).  This 'f'
   can have ArityType as ATop, with length as > 0, only if e1 e2 are
   themselves.

 * In both cases, f, (f x1), ... (f x1 ... f(n-1)) are definitely
   really functions, or bottom, but *not* casts from a data type, in
   at least one case branch.  (If it's a function in one case branch but
   an unsafe cast from a data type in another, the program is bogus.)
   So eta expansion is dynamically ok; see Note [State hack and
   bottoming functions], the part about catch#

Example:
      f = \x\y. let v = <expensive> in
          \s(one-shot) \t(one-shot). blah
      'f' has ArityType [ManyShot,ManyShot,OneShot,OneShot]
      The one-shot-ness means we can, in effect, push that
      'let' inside the \st.


Suppose f = \xy. x+y
Then  f             :: AT [False,False] ATop
      f v           :: AT [False]       ATop
      f <expensive> :: AT []            ATop

-------------------- Main arity code ----------------------------
-}


data ArityType   -- See Note [ArityType]
  = ATop [OneShotInfo]
  | ABot Arity
  deriving( ArityType -> ArityType -> Bool
(ArityType -> ArityType -> Bool)
-> (ArityType -> ArityType -> Bool) -> Eq ArityType
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ArityType -> ArityType -> Bool
$c/= :: ArityType -> ArityType -> Bool
== :: ArityType -> ArityType -> Bool
$c== :: ArityType -> ArityType -> Bool
Eq )
     -- There is always an explicit lambda
     -- to justify the [OneShot], or the Arity

instance Outputable ArityType where
  ppr :: ArityType -> SDoc
ppr (ATop [OneShotInfo]
os) = String -> SDoc
text String
"ATop" SDoc -> SDoc -> SDoc
<> SDoc -> SDoc
parens (JoinArity -> SDoc
forall a. Outputable a => a -> SDoc
ppr ([OneShotInfo] -> JoinArity
forall (t :: * -> *) a. Foldable t => t a -> JoinArity
length [OneShotInfo]
os))
  ppr (ABot JoinArity
n)  = String -> SDoc
text String
"ABot" SDoc -> SDoc -> SDoc
<> SDoc -> SDoc
parens (JoinArity -> SDoc
forall a. Outputable a => a -> SDoc
ppr JoinArity
n)

arityTypeArity :: ArityType -> Arity
-- The number of value args for the arity type
arityTypeArity :: ArityType -> JoinArity
arityTypeArity (ATop [OneShotInfo]
oss) = [OneShotInfo] -> JoinArity
forall (t :: * -> *) a. Foldable t => t a -> JoinArity
length [OneShotInfo]
oss
arityTypeArity (ABot JoinArity
ar)  = JoinArity
ar

expandableArityType :: ArityType -> Bool
-- True <=> eta-expansion will add at least one lambda
expandableArityType :: ArityType -> Bool
expandableArityType (ATop [OneShotInfo]
oss) = Bool -> Bool
not ([OneShotInfo] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [OneShotInfo]
oss)
expandableArityType (ABot JoinArity
ar)  = JoinArity
ar JoinArity -> JoinArity -> Bool
forall a. Eq a => a -> a -> Bool
/= JoinArity
0

isBotArityType :: ArityType -> Bool
isBotArityType :: ArityType -> Bool
isBotArityType (ABot {}) = Bool
True
isBotArityType (ATop {}) = Bool
False

arityTypeOneShots :: ArityType -> [OneShotInfo]
arityTypeOneShots :: ArityType -> [OneShotInfo]
arityTypeOneShots (ATop [OneShotInfo]
oss) = [OneShotInfo]
oss
arityTypeOneShots (ABot JoinArity
ar)  = JoinArity -> OneShotInfo -> [OneShotInfo]
forall a. JoinArity -> a -> [a]
replicate JoinArity
ar OneShotInfo
OneShotLam
   -- If we are diveging or throwing an exception anyway
   -- it's fine to push redexes inside the lambdas

botArityType :: ArityType
botArityType :: ArityType
botArityType = JoinArity -> ArityType
ABot JoinArity
0   -- Unit for andArityType

maxWithArity :: ArityType -> Arity -> ArityType
maxWithArity :: ArityType -> JoinArity -> ArityType
maxWithArity at :: ArityType
at@(ABot {}) JoinArity
_   = ArityType
at
maxWithArity at :: ArityType
at@(ATop [OneShotInfo]
oss) JoinArity
ar
     | [OneShotInfo]
oss [OneShotInfo] -> JoinArity -> Bool
forall a. [a] -> JoinArity -> Bool
`lengthAtLeast` JoinArity
ar = ArityType
at
     | Bool
otherwise              = [OneShotInfo] -> ArityType
ATop (JoinArity -> [OneShotInfo] -> [OneShotInfo]
forall a. JoinArity -> [a] -> [a]
take JoinArity
ar ([OneShotInfo]
oss [OneShotInfo] -> [OneShotInfo] -> [OneShotInfo]
forall a. [a] -> [a] -> [a]
++ OneShotInfo -> [OneShotInfo]
forall a. a -> [a]
repeat OneShotInfo
NoOneShotInfo))

vanillaArityType :: ArityType
vanillaArityType :: ArityType
vanillaArityType = [OneShotInfo] -> ArityType
ATop []      -- Totally uninformative

-- ^ The Arity returned is the number of value args the
-- expression can be applied to without doing much work
exprEtaExpandArity :: DynFlags -> CoreExpr -> ArityType
-- exprEtaExpandArity is used when eta expanding
--      e  ==>  \xy -> e x y
exprEtaExpandArity :: DynFlags -> CoreExpr -> ArityType
exprEtaExpandArity DynFlags
dflags CoreExpr
e
  = ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
e
  where
    env :: ArityEnv
env = AE :: CheapFun -> Bool -> IdSet -> ArityEnv
AE { ae_cheap_fn :: CheapFun
ae_cheap_fn = DynFlags -> CheapAppFun -> CheapFun
mk_cheap_fn DynFlags
dflags CheapAppFun
isCheapApp
             , ae_ped_bot :: Bool
ae_ped_bot  = GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_PedanticBottoms DynFlags
dflags
             , ae_joins :: IdSet
ae_joins    = IdSet
emptyVarSet }

getBotArity :: ArityType -> Maybe Arity
-- Arity of a divergent function
getBotArity :: ArityType -> Maybe JoinArity
getBotArity (ABot JoinArity
n) = JoinArity -> Maybe JoinArity
forall a. a -> Maybe a
Just JoinArity
n
getBotArity ArityType
_        = Maybe JoinArity
forall a. Maybe a
Nothing

mk_cheap_fn :: DynFlags -> CheapAppFun -> CheapFun
mk_cheap_fn :: DynFlags -> CheapAppFun -> CheapFun
mk_cheap_fn DynFlags
dflags CheapAppFun
cheap_app
  | Bool -> Bool
not (GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_DictsCheap DynFlags
dflags)
  = \CoreExpr
e Maybe Type
_     -> CheapAppFun -> CoreExpr -> Bool
exprIsCheapX CheapAppFun
cheap_app CoreExpr
e
  | Bool
otherwise
  = \CoreExpr
e Maybe Type
mb_ty -> CheapAppFun -> CoreExpr -> Bool
exprIsCheapX CheapAppFun
cheap_app CoreExpr
e
             Bool -> Bool -> Bool
|| case Maybe Type
mb_ty of
                  Maybe Type
Nothing -> Bool
False
                  Just Type
ty -> Type -> Bool
isDictTy Type
ty


----------------------
findRhsArity :: DynFlags -> Id -> CoreExpr -> Arity -> ArityType
-- This implements the fixpoint loop for arity analysis
-- See Note [Arity analysis]
-- If findRhsArity e = (n, is_bot) then
--  (a) any application of e to <n arguments will not do much work,
--      so it is safe to expand e  ==>  (\x1..xn. e x1 .. xn)
--  (b) if is_bot=True, then e applied to n args is guaranteed bottom
findRhsArity :: DynFlags -> Id -> CoreExpr -> JoinArity -> ArityType
findRhsArity DynFlags
dflags Id
bndr CoreExpr
rhs JoinArity
old_arity
  = ArityType -> ArityType
go (CheapAppFun -> ArityType
get_arity CheapAppFun
init_cheap_app)
       -- We always call exprEtaExpandArity once, but usually
       -- that produces a result equal to old_arity, and then
       -- we stop right away (since arities should not decrease)
       -- Result: the common case is that there is just one iteration
  where
    init_cheap_app :: CheapAppFun
    init_cheap_app :: CheapAppFun
init_cheap_app Id
fn JoinArity
n_val_args
      | Id
fn Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
bndr = Bool
True   -- On the first pass, this binder gets infinite arity
      | Bool
otherwise  = CheapAppFun
isCheapApp Id
fn JoinArity
n_val_args

    go :: ArityType -> ArityType
    go :: ArityType -> ArityType
go ArityType
cur_atype
      | JoinArity
cur_arity JoinArity -> JoinArity -> Bool
forall a. Ord a => a -> a -> Bool
<= JoinArity
old_arity = ArityType
cur_atype
      | ArityType
new_atype ArityType -> ArityType -> Bool
forall a. Eq a => a -> a -> Bool
== ArityType
cur_atype = ArityType
cur_atype
      | Bool
otherwise =
#if defined(DEBUG)
                    pprTrace "Exciting arity"
                       (vcat [ ppr bndr <+> ppr cur_atype <+> ppr new_atype
                             , ppr rhs])
#endif
                    ArityType -> ArityType
go ArityType
new_atype
      where
        new_atype :: ArityType
new_atype = CheapAppFun -> ArityType
get_arity CheapAppFun
cheap_app

        cur_arity :: JoinArity
cur_arity = ArityType -> JoinArity
arityTypeArity ArityType
cur_atype
        cheap_app :: CheapAppFun
        cheap_app :: CheapAppFun
cheap_app Id
fn JoinArity
n_val_args
          | Id
fn Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
bndr = JoinArity
n_val_args JoinArity -> JoinArity -> Bool
forall a. Ord a => a -> a -> Bool
< JoinArity
cur_arity
          | Bool
otherwise  = CheapAppFun
isCheapApp Id
fn JoinArity
n_val_args

    get_arity :: CheapAppFun -> ArityType
    get_arity :: CheapAppFun -> ArityType
get_arity CheapAppFun
cheap_app = ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
rhs
      where
         env :: ArityEnv
env = AE :: CheapFun -> Bool -> IdSet -> ArityEnv
AE { ae_cheap_fn :: CheapFun
ae_cheap_fn = DynFlags -> CheapAppFun -> CheapFun
mk_cheap_fn DynFlags
dflags CheapAppFun
cheap_app
                  , ae_ped_bot :: Bool
ae_ped_bot  = GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_PedanticBottoms DynFlags
dflags
                  , ae_joins :: IdSet
ae_joins    = IdSet
emptyVarSet }

{-
Note [Arity analysis]
~~~~~~~~~~~~~~~~~~~~~
The motivating example for arity analysis is this:

  f = \x. let g = f (x+1)
          in \y. ...g...

What arity does f have?  Really it should have arity 2, but a naive
look at the RHS won't see that.  You need a fixpoint analysis which
says it has arity "infinity" the first time round.

This example happens a lot; it first showed up in Andy Gill's thesis,
fifteen years ago!  It also shows up in the code for 'rnf' on lists
in #4138.

The analysis is easy to achieve because exprEtaExpandArity takes an
argument
     type CheapFun = CoreExpr -> Maybe Type -> Bool
used to decide if an expression is cheap enough to push inside a
lambda.  And exprIsCheapX in turn takes an argument
     type CheapAppFun = Id -> Int -> Bool
which tells when an application is cheap. This makes it easy to
write the analysis loop.

The analysis is cheap-and-cheerful because it doesn't deal with
mutual recursion.  But the self-recursive case is the important one.

Note [Eta expanding through dictionaries]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
If the experimental -fdicts-cheap flag is on, we eta-expand through
dictionary bindings.  This improves arities. Thereby, it also
means that full laziness is less prone to floating out the
application of a function to its dictionary arguments, which
can thereby lose opportunities for fusion.  Example:
        foo :: Ord a => a -> ...
     foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....
        -- So foo has arity 1

     f = \x. foo dInt $ bar x

The (foo DInt) is floated out, and makes ineffective a RULE
     foo (bar x) = ...

One could go further and make exprIsCheap reply True to any
dictionary-typed expression, but that's more work.
-}

arityLam :: Id -> ArityType -> ArityType
arityLam :: Id -> ArityType -> ArityType
arityLam Id
id (ATop [OneShotInfo]
as) = [OneShotInfo] -> ArityType
ATop (Id -> OneShotInfo
idStateHackOneShotInfo Id
id OneShotInfo -> [OneShotInfo] -> [OneShotInfo]
forall a. a -> [a] -> [a]
: [OneShotInfo]
as)
arityLam Id
_  (ABot JoinArity
n)  = JoinArity -> ArityType
ABot (JoinArity
nJoinArity -> JoinArity -> JoinArity
forall a. Num a => a -> a -> a
+JoinArity
1)

floatIn :: Bool -> ArityType -> ArityType
-- We have something like (let x = E in b),
-- where b has the given arity type.
floatIn :: Bool -> ArityType -> ArityType
floatIn Bool
_     (ABot JoinArity
n)  = JoinArity -> ArityType
ABot JoinArity
n
floatIn Bool
True  (ATop [OneShotInfo]
as) = [OneShotInfo] -> ArityType
ATop [OneShotInfo]
as
floatIn Bool
False (ATop [OneShotInfo]
as) = [OneShotInfo] -> ArityType
ATop ((OneShotInfo -> Bool) -> [OneShotInfo] -> [OneShotInfo]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile OneShotInfo -> Bool
isOneShotInfo [OneShotInfo]
as)
   -- If E is not cheap, keep arity only for one-shots

arityApp :: ArityType -> Bool -> ArityType
-- Processing (fun arg) where at is the ArityType of fun,
-- Knock off an argument and behave like 'let'
arityApp :: ArityType -> Bool -> ArityType
arityApp (ABot JoinArity
0)      Bool
_     = JoinArity -> ArityType
ABot JoinArity
0
arityApp (ABot JoinArity
n)      Bool
_     = JoinArity -> ArityType
ABot (JoinArity
nJoinArity -> JoinArity -> JoinArity
forall a. Num a => a -> a -> a
-JoinArity
1)
arityApp (ATop [])     Bool
_     = [OneShotInfo] -> ArityType
ATop []
arityApp (ATop (OneShotInfo
_:[OneShotInfo]
as)) Bool
cheap = Bool -> ArityType -> ArityType
floatIn Bool
cheap ([OneShotInfo] -> ArityType
ATop [OneShotInfo]
as)

andArityType :: ArityType -> ArityType -> ArityType   -- Used for branches of a 'case'
-- This is least upper bound in the ArityType lattice
andArityType :: ArityType -> ArityType -> ArityType
andArityType (ABot JoinArity
n1) (ABot JoinArity
n2)  = JoinArity -> ArityType
ABot (JoinArity
n1 JoinArity -> JoinArity -> JoinArity
forall a. Ord a => a -> a -> a
`max` JoinArity
n2) -- Note [ABot branches: use max]
andArityType (ATop [OneShotInfo]
as)  (ABot JoinArity
_)  = [OneShotInfo] -> ArityType
ATop [OneShotInfo]
as
andArityType (ABot JoinArity
_)   (ATop [OneShotInfo]
bs) = [OneShotInfo] -> ArityType
ATop [OneShotInfo]
bs
andArityType (ATop [OneShotInfo]
as)  (ATop [OneShotInfo]
bs) = [OneShotInfo] -> ArityType
ATop ([OneShotInfo]
as [OneShotInfo] -> [OneShotInfo] -> [OneShotInfo]
`combine` [OneShotInfo]
bs)
  where      -- See Note [Combining case branches]
    combine :: [OneShotInfo] -> [OneShotInfo] -> [OneShotInfo]
combine (OneShotInfo
a:[OneShotInfo]
as) (OneShotInfo
b:[OneShotInfo]
bs) = (OneShotInfo
a OneShotInfo -> OneShotInfo -> OneShotInfo
`bestOneShot` OneShotInfo
b) OneShotInfo -> [OneShotInfo] -> [OneShotInfo]
forall a. a -> [a] -> [a]
: [OneShotInfo] -> [OneShotInfo] -> [OneShotInfo]
combine [OneShotInfo]
as [OneShotInfo]
bs
    combine []     [OneShotInfo]
bs     = (OneShotInfo -> Bool) -> [OneShotInfo] -> [OneShotInfo]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile OneShotInfo -> Bool
isOneShotInfo [OneShotInfo]
bs
    combine [OneShotInfo]
as     []     = (OneShotInfo -> Bool) -> [OneShotInfo] -> [OneShotInfo]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile OneShotInfo -> Bool
isOneShotInfo [OneShotInfo]
as

{- Note [ABot branches: use max]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider   case x of
             True  -> \x.  error "urk"
             False -> \xy. error "urk2"

Remember: ABot n means "if you apply to n args, it'll definitely diverge".
So we need (ABot 2) for the whole thing, the /max/ of the ABot arities.

Note [Combining case branches]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider
  go = \x. let z = go e0
               go2 = \x. case x of
                           True  -> z
                           False -> \s(one-shot). e1
           in go2 x
We *really* want to eta-expand go and go2.
When combining the branches of the case we have
     ATop [] `andAT` ATop [OneShotLam]
and we want to get ATop [OneShotLam].  But if the inner
lambda wasn't one-shot we don't want to do this.
(We need a proper arity analysis to justify that.)

So we combine the best of the two branches, on the (slightly dodgy)
basis that if we know one branch is one-shot, then they all must be.

Note [Arity trimming]
~~~~~~~~~~~~~~~~~~~~~
Consider ((\x y. blah) |> co), where co :: (Int->Int->Int) ~ (Int -> F a) , and
F is some type family.

Because of Note [exprArity invariant], item (2), we must return with arity at
most 1, because typeArity (Int -> F a) = 1.  So we have to trim the result of
calling arityType on (\x y. blah).  Failing to do so, and hence breaking the
exprArity invariant, led to #5441.

How to trim?  For ATop, it's easy.  But we must take great care with ABot.
Suppose the expression was (\x y. error "urk"), we'll get (ABot 2).  We
absolutely must not trim that to (ABot 1), because that claims that
((\x y. error "urk") |> co) diverges when given one argument, which it
absolutely does not. And Bad Things happen if we think something returns bottom
when it doesn't (#16066).

So, do not reduce the 'n' in (ABot n); rather, switch (conservatively) to ATop.

Historical note: long ago, we unconditionally switched to ATop when we
encountered a cast, but that is far too conservative: see #5475
-}

---------------------------
type CheapFun = CoreExpr -> Maybe Type -> Bool
        -- How to decide if an expression is cheap
        -- If the Maybe is Just, the type is the type
        -- of the expression; Nothing means "don't know"

data ArityEnv
  = AE { ArityEnv -> CheapFun
ae_cheap_fn :: CheapFun
       , ArityEnv -> Bool
ae_ped_bot  :: Bool       -- True <=> be pedantic about bottoms
       , ArityEnv -> IdSet
ae_joins    :: IdSet      -- In-scope join points
                                   -- See Note [Eta-expansion and join points]
  }

extendJoinEnv :: ArityEnv -> [JoinId] -> ArityEnv
extendJoinEnv :: ArityEnv -> [Id] -> ArityEnv
extendJoinEnv env :: ArityEnv
env@(AE { ae_joins :: ArityEnv -> IdSet
ae_joins = IdSet
joins }) [Id]
join_ids
  = ArityEnv
env { ae_joins :: IdSet
ae_joins = IdSet
joins IdSet -> [Id] -> IdSet
`extendVarSetList` [Id]
join_ids }

----------------
arityType :: ArityEnv -> CoreExpr -> ArityType

arityType :: ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env (Cast CoreExpr
e CoercionR
co)
  = case ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
e of
      ATop [OneShotInfo]
os -> [OneShotInfo] -> ArityType
ATop (JoinArity -> [OneShotInfo] -> [OneShotInfo]
forall a. JoinArity -> [a] -> [a]
take JoinArity
co_arity [OneShotInfo]
os)  -- See Note [Arity trimming]
      ABot JoinArity
n | JoinArity
co_arity JoinArity -> JoinArity -> Bool
forall a. Ord a => a -> a -> Bool
< JoinArity
n -> [OneShotInfo] -> ArityType
ATop (JoinArity -> OneShotInfo -> [OneShotInfo]
forall a. JoinArity -> a -> [a]
replicate JoinArity
co_arity OneShotInfo
noOneShotInfo)
             | Bool
otherwise    -> JoinArity -> ArityType
ABot JoinArity
n
  where
    co_arity :: JoinArity
co_arity = [OneShotInfo] -> JoinArity
forall (t :: * -> *) a. Foldable t => t a -> JoinArity
length (Type -> [OneShotInfo]
typeArity (CoercionR -> Type
coercionRKind CoercionR
co))
    -- See Note [exprArity invariant] (2); must be true of
    -- arityType too, since that is how we compute the arity
    -- of variables, and they in turn affect result of exprArity
    -- #5441 is a nice demo
    -- However, do make sure that ATop -> ATop and ABot -> ABot!
    --   Casts don't affect that part. Getting this wrong provoked #5475

arityType ArityEnv
env (Var Id
v)
  | Id
v Id -> IdSet -> Bool
`elemVarSet` ArityEnv -> IdSet
ae_joins ArityEnv
env
  = ArityType
botArityType  -- See Note [Eta-expansion and join points]
  | Bool
otherwise
  = Id -> ArityType
idArityType Id
v

        -- Lambdas; increase arity
arityType ArityEnv
env (Lam Id
x CoreExpr
e)
  | Id -> Bool
isId Id
x    = Id -> ArityType -> ArityType
arityLam Id
x (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
e)
  | Bool
otherwise = ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
e

        -- Applications; decrease arity, except for types
arityType ArityEnv
env (App CoreExpr
fun (Type Type
_))
   = ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
fun
arityType ArityEnv
env (App CoreExpr
fun CoreExpr
arg )
   = ArityType -> Bool -> ArityType
arityApp (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
fun) (ArityEnv -> CheapFun
ae_cheap_fn ArityEnv
env CoreExpr
arg Maybe Type
forall a. Maybe a
Nothing)

        -- Case/Let; keep arity if either the expression is cheap
        -- or it's a 1-shot lambda
        -- The former is not really right for Haskell
        --      f x = case x of { (a,b) -> \y. e }
        --  ===>
        --      f x y = case x of { (a,b) -> e }
        -- The difference is observable using 'seq'
        --
arityType ArityEnv
env (Case CoreExpr
scrut Id
_ Type
_ [Alt Id]
alts)
  | CoreExpr -> Bool
exprIsDeadEnd CoreExpr
scrut Bool -> Bool -> Bool
|| [Alt Id] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Alt Id]
alts
  = ArityType
botArityType    -- Do not eta expand
                    -- See Note [Dealing with bottom (1)]
  | Bool
otherwise
  = case ArityType
alts_type of
     ABot JoinArity
n  | JoinArity
nJoinArity -> JoinArity -> Bool
forall a. Ord a => a -> a -> Bool
>JoinArity
0       -> [OneShotInfo] -> ArityType
ATop []       -- Don't eta expand
             | Bool
otherwise -> ArityType
botArityType  -- if RHS is bottomming
                                          -- See Note [Dealing with bottom (2)]

     ATop [OneShotInfo]
as | Bool -> Bool
not (ArityEnv -> Bool
ae_ped_bot ArityEnv
env)    -- See Note [Dealing with bottom (3)]
             , ArityEnv -> CheapFun
ae_cheap_fn ArityEnv
env CoreExpr
scrut Maybe Type
forall a. Maybe a
Nothing -> [OneShotInfo] -> ArityType
ATop [OneShotInfo]
as
             | CoreExpr -> Bool
exprOkForSpeculation CoreExpr
scrut    -> [OneShotInfo] -> ArityType
ATop [OneShotInfo]
as
             | Bool
otherwise                     -> [OneShotInfo] -> ArityType
ATop ((OneShotInfo -> Bool) -> [OneShotInfo] -> [OneShotInfo]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile OneShotInfo -> Bool
isOneShotInfo [OneShotInfo]
as)
  where
    alts_type :: ArityType
alts_type = (ArityType -> ArityType -> ArityType) -> [ArityType] -> ArityType
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 ArityType -> ArityType -> ArityType
andArityType [ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
rhs | (AltCon
_,[Id]
_,CoreExpr
rhs) <- [Alt Id]
alts]

arityType ArityEnv
env (Let (NonRec Id
j CoreExpr
rhs) CoreExpr
body)
  | Just JoinArity
join_arity <- Id -> Maybe JoinArity
isJoinId_maybe Id
j
  , ([Id]
_, CoreExpr
rhs_body)   <- JoinArity -> CoreExpr -> ([Id], CoreExpr)
forall b. JoinArity -> Expr b -> ([b], Expr b)
collectNBinders JoinArity
join_arity CoreExpr
rhs
  = -- See Note [Eta-expansion and join points]
    ArityType -> ArityType -> ArityType
andArityType (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
rhs_body)
                 (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env' CoreExpr
body)
  where
     env' :: ArityEnv
env' = ArityEnv -> [Id] -> ArityEnv
extendJoinEnv ArityEnv
env [Id
j]

arityType ArityEnv
env (Let (Rec [(Id, CoreExpr)]
pairs) CoreExpr
body)
  | ((Id
j,CoreExpr
_):[(Id, CoreExpr)]
_) <- [(Id, CoreExpr)]
pairs
  , Id -> Bool
isJoinId Id
j
  = -- See Note [Eta-expansion and join points]
    ((Id, CoreExpr) -> ArityType -> ArityType)
-> ArityType -> [(Id, CoreExpr)] -> ArityType
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (ArityType -> ArityType -> ArityType
andArityType (ArityType -> ArityType -> ArityType)
-> ((Id, CoreExpr) -> ArityType)
-> (Id, CoreExpr)
-> ArityType
-> ArityType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Id, CoreExpr) -> ArityType
do_one) (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env' CoreExpr
body) [(Id, CoreExpr)]
pairs
  where
    env' :: ArityEnv
env' = ArityEnv -> [Id] -> ArityEnv
extendJoinEnv ArityEnv
env (((Id, CoreExpr) -> Id) -> [(Id, CoreExpr)] -> [Id]
forall a b. (a -> b) -> [a] -> [b]
map (Id, CoreExpr) -> Id
forall a b. (a, b) -> a
fst [(Id, CoreExpr)]
pairs)
    do_one :: (Id, CoreExpr) -> ArityType
do_one (Id
j,CoreExpr
rhs)
      | Just JoinArity
arity <- Id -> Maybe JoinArity
isJoinId_maybe Id
j
      = ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env' (CoreExpr -> ArityType) -> CoreExpr -> ArityType
forall a b. (a -> b) -> a -> b
$ ([Id], CoreExpr) -> CoreExpr
forall a b. (a, b) -> b
snd (([Id], CoreExpr) -> CoreExpr) -> ([Id], CoreExpr) -> CoreExpr
forall a b. (a -> b) -> a -> b
$ JoinArity -> CoreExpr -> ([Id], CoreExpr)
forall b. JoinArity -> Expr b -> ([b], Expr b)
collectNBinders JoinArity
arity CoreExpr
rhs
      | Bool
otherwise
      = String -> SDoc -> ArityType
forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"arityType:joinrec" ([(Id, CoreExpr)] -> SDoc
forall a. Outputable a => a -> SDoc
ppr [(Id, CoreExpr)]
pairs)

arityType ArityEnv
env (Let Bind Id
b CoreExpr
e)
  = Bool -> ArityType -> ArityType
floatIn (Bind Id -> Bool
cheap_bind Bind Id
b) (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
e)
  where
    cheap_bind :: Bind Id -> Bool
cheap_bind (NonRec Id
b CoreExpr
e) = (Id, CoreExpr) -> Bool
is_cheap (Id
b,CoreExpr
e)
    cheap_bind (Rec [(Id, CoreExpr)]
prs)    = ((Id, CoreExpr) -> Bool) -> [(Id, CoreExpr)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Id, CoreExpr) -> Bool
is_cheap [(Id, CoreExpr)]
prs
    is_cheap :: (Id, CoreExpr) -> Bool
is_cheap (Id
b,CoreExpr
e) = ArityEnv -> CheapFun
ae_cheap_fn ArityEnv
env CoreExpr
e (Type -> Maybe Type
forall a. a -> Maybe a
Just (Id -> Type
idType Id
b))

arityType ArityEnv
env (Tick Tickish Id
t CoreExpr
e)
  | Bool -> Bool
not (Tickish Id -> Bool
forall id. Tickish id -> Bool
tickishIsCode Tickish Id
t)     = ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
e

arityType ArityEnv
_ CoreExpr
_ = ArityType
vanillaArityType

{- Note [Eta-expansion and join points]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider this (#18328)

  f x = join j y = case y of
                      True -> \a. blah
                      False -> \b. blah
        in case x of
              A -> j True
              B -> \c. blah
              C -> j False

and suppose the join point is too big to inline.  Now, what is the
arity of f?  If we inlined the join point, we'd definitely say "arity
2" because we are prepared to push case-scrutinisation inside a
lambda.  But currently the join point totally messes all that up,
because (thought of as a vanilla let-binding) the arity pinned on 'j'
is just 1.

Why don't we eta-expand j?  Because of
Note [Do not eta-expand join points] in GHC.Core.Opt.Simplify.Utils

Even if we don't eta-expand j, why is its arity only 1?
See invariant 2b in Note [Invariants on join points] in GHC.Core.

So we do this:

* Treat the RHS of a join-point binding, /after/ stripping off
  join-arity lambda-binders, as very like the body of the let.
  More precisely, do andArityType with the arityType from the
  body of the let.

* Dually, when we come to a /call/ of a join point, just no-op
  by returning botArityType, the bottom element of ArityType,
  which so that: bot `andArityType` x = x

* This works if the join point is bound in the expression we are
  taking the arityType of.  But if it's bound further out, it makes
  no sense to say that (say) the arityType of (j False) is ABot 0.
  Bad things happen.  So we keep track of the in-scope join-point Ids
  in ae_join.

This will make f, above, have arity 2. Then, we'll eta-expand it thus:

  f x eta = (join j y = ... in case x of ...) eta

and the Simplify will automatically push that application of eta into
the join points.

An alternative (roughly equivalent) idea would be to carry an
environment mapping let-bound Ids to their ArityType.
-}

idArityType :: Id -> ArityType
idArityType :: Id -> ArityType
idArityType Id
v
  | StrictSig
strict_sig <- Id -> StrictSig
idStrictness Id
v
  , Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ StrictSig -> Bool
isTopSig StrictSig
strict_sig
  , ([Demand]
ds, Divergence
res) <- StrictSig -> ([Demand], Divergence)
splitStrictSig StrictSig
strict_sig
  , let arity :: JoinArity
arity = [Demand] -> JoinArity
forall (t :: * -> *) a. Foldable t => t a -> JoinArity
length [Demand]
ds
  = if Divergence -> Bool
isDeadEndDiv Divergence
res then JoinArity -> ArityType
ABot JoinArity
arity
                        else [OneShotInfo] -> ArityType
ATop (JoinArity -> [OneShotInfo] -> [OneShotInfo]
forall a. JoinArity -> [a] -> [a]
take JoinArity
arity [OneShotInfo]
one_shots)
  | Bool
otherwise
  = [OneShotInfo] -> ArityType
ATop (JoinArity -> [OneShotInfo] -> [OneShotInfo]
forall a. JoinArity -> [a] -> [a]
take (Id -> JoinArity
idArity Id
v) [OneShotInfo]
one_shots)
  where
    one_shots :: [OneShotInfo]  -- One-shot-ness derived from the type
    one_shots :: [OneShotInfo]
one_shots = Type -> [OneShotInfo]
typeArity (Id -> Type
idType Id
v)

{-
%************************************************************************
%*                                                                      *
              The main eta-expander
%*                                                                      *
%************************************************************************

We go for:
   f = \x1..xn -> N  ==>   f = \x1..xn y1..ym -> N y1..ym
                                 (n >= 0)

where (in both cases)

        * The xi can include type variables

        * The yi are all value variables

        * N is a NORMAL FORM (i.e. no redexes anywhere)
          wanting a suitable number of extra args.

The biggest reason for doing this is for cases like

        f = \x -> case x of
                    True  -> \y -> e1
                    False -> \y -> e2

Here we want to get the lambdas together.  A good example is the nofib
program fibheaps, which gets 25% more allocation if you don't do this
eta-expansion.

We may have to sandwich some coerces between the lambdas
to make the types work.   exprEtaExpandArity looks through coerces
when computing arity; and etaExpand adds the coerces as necessary when
actually computing the expansion.

Note [No crap in eta-expanded code]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The eta expander is careful not to introduce "crap".  In particular,
given a CoreExpr satisfying the 'CpeRhs' invariant (in CorePrep), it
returns a CoreExpr satisfying the same invariant. See Note [Eta
expansion and the CorePrep invariants] in CorePrep.

This means the eta-expander has to do a bit of on-the-fly
simplification but it's not too hard.  The alternative, of relying on
a subsequent clean-up phase of the Simplifier to de-crapify the result,
means you can't really use it in CorePrep, which is painful.

Note [Eta expansion for join points]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The no-crap rule is very tiresome to guarantee when
we have join points. Consider eta-expanding
   let j :: Int -> Int -> Bool
       j x = e
   in b

The simple way is
  \(y::Int). (let j x = e in b) y

The no-crap way is
  \(y::Int). let j' :: Int -> Bool
                 j' x = e y
             in b[j'/j] y
where I have written to stress that j's type has
changed.  Note that (of course!) we have to push the application
inside the RHS of the join as well as into the body.  AND if j
has an unfolding we have to push it into there too.  AND j might
be recursive...

So for now I'm abandoning the no-crap rule in this case. I think
that for the use in CorePrep it really doesn't matter; and if
it does, then CoreToStg.myCollectArgs will fall over.

(Moreover, I think that casts can make the no-crap rule fail too.)

Note [Eta expansion and SCCs]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Note that SCCs are not treated specially by etaExpand.  If we have
        etaExpand 2 (\x -> scc "foo" e)
        = (\xy -> (scc "foo" e) y)
So the costs of evaluating 'e' (not 'e y') are attributed to "foo"

Note [Eta expansion and source notes]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
CorePrep puts floatable ticks outside of value applications, but not
type applications. As a result we might be trying to eta-expand an
expression like

  (src<...> v) @a

which we want to lead to code like

  \x -> src<...> v @a x

This means that we need to look through type applications and be ready
to re-add floats on the top.

Note [Eta expansion with ArityType]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The etaExpandAT function takes an ArityType (not just an Arity) to
guide eta-expansion.  Why? Because we want to preserve one-shot info.
Consider
  foo = \x. case x of
              True  -> (\s{os}. blah) |> co
              False -> wubble
We'll get an ArityType for foo of (ATop [NoOneShot,OneShot]).

Then we want to eta-expand to
  foo = \x. (\eta{os}. (case x of ...as before...) eta) |> some_co

That 'eta' binder is fresh, and we really want it to have the
one-shot flag from the inner \s{osf}.  By expanding with the
ArityType gotten from analysing the RHS, we achieve this neatly.

This makes a big difference to the one-shot monad trick;
see Note [The one-shot state monad trick] in GHC.Core.Unify.
-}

-- | @etaExpand n e@ returns an expression with
-- the same meaning as @e@, but with arity @n@.
--
-- Given:
--
-- > e' = etaExpand n e
--
-- We should have that:
--
-- > ty = exprType e = exprType e'
etaExpand   :: Arity     -> CoreExpr -> CoreExpr
etaExpandAT :: ArityType -> CoreExpr -> CoreExpr

etaExpand :: JoinArity -> CoreExpr -> CoreExpr
etaExpand   JoinArity
n  CoreExpr
orig_expr = [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand (JoinArity -> OneShotInfo -> [OneShotInfo]
forall a. JoinArity -> a -> [a]
replicate JoinArity
n OneShotInfo
NoOneShotInfo) CoreExpr
orig_expr
etaExpandAT :: ArityType -> CoreExpr -> CoreExpr
etaExpandAT ArityType
at CoreExpr
orig_expr = [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand (ArityType -> [OneShotInfo]
arityTypeOneShots ArityType
at)      CoreExpr
orig_expr
                           -- See Note [Eta expansion with ArityType]

-- etaExpand arity e = res
-- Then 'res' has at least 'arity' lambdas at the top
-- See Note [Eta expansion with ArityType]
--
-- etaExpand deals with for-alls. For example:
--              etaExpand 1 E
-- where  E :: forall a. a -> a
-- would return
--      (/\b. \y::a -> E b y)
--
-- It deals with coerces too, though they are now rare
-- so perhaps the extra code isn't worth it

eta_expand :: [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand :: [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand [OneShotInfo]
one_shots CoreExpr
orig_expr
  = [OneShotInfo] -> CoreExpr -> CoreExpr
go [OneShotInfo]
one_shots CoreExpr
orig_expr
  where
      -- Strip off existing lambdas and casts before handing off to mkEtaWW
      -- Note [Eta expansion and SCCs]
    go :: [OneShotInfo] -> CoreExpr -> CoreExpr
go [] CoreExpr
expr = CoreExpr
expr
    go oss :: [OneShotInfo]
oss@(OneShotInfo
_:[OneShotInfo]
oss1) (Lam Id
v CoreExpr
body) | Id -> Bool
isTyVar Id
v = Id -> CoreExpr -> CoreExpr
forall b. b -> Expr b -> Expr b
Lam Id
v ([OneShotInfo] -> CoreExpr -> CoreExpr
go [OneShotInfo]
oss  CoreExpr
body)
                                 | Bool
otherwise = Id -> CoreExpr -> CoreExpr
forall b. b -> Expr b -> Expr b
Lam Id
v ([OneShotInfo] -> CoreExpr -> CoreExpr
go [OneShotInfo]
oss1 CoreExpr
body)
    go [OneShotInfo]
oss (Cast CoreExpr
expr CoercionR
co) = CoreExpr -> CoercionR -> CoreExpr
forall b. Expr b -> CoercionR -> Expr b
Cast ([OneShotInfo] -> CoreExpr -> CoreExpr
go [OneShotInfo]
oss CoreExpr
expr) CoercionR
co

    go [OneShotInfo]
oss CoreExpr
expr
      = -- pprTrace "ee" (vcat [ppr orig_expr, ppr expr, ppr etas]) $
        CoreExpr -> CoreExpr
retick (CoreExpr -> CoreExpr) -> CoreExpr -> CoreExpr
forall a b. (a -> b) -> a -> b
$ [EtaInfo] -> CoreExpr -> CoreExpr
etaInfoAbs [EtaInfo]
etas (Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
etaInfoApp Subst
subst' CoreExpr
sexpr [EtaInfo]
etas)
      where
          in_scope :: InScopeSet
in_scope = IdSet -> InScopeSet
mkInScopeSet (CoreExpr -> IdSet
exprFreeVars CoreExpr
expr)
          (InScopeSet
in_scope', [EtaInfo]
etas) = [OneShotInfo]
-> SDoc -> InScopeSet -> Type -> (InScopeSet, [EtaInfo])
mkEtaWW [OneShotInfo]
oss (CoreExpr -> SDoc
forall a. Outputable a => a -> SDoc
ppr CoreExpr
orig_expr) InScopeSet
in_scope (CoreExpr -> Type
exprType CoreExpr
expr)
          subst' :: Subst
subst' = InScopeSet -> Subst
mkEmptySubst InScopeSet
in_scope'

          -- Find ticks behind type apps.
          -- See Note [Eta expansion and source notes]
          (CoreExpr
expr', [CoreExpr]
args) = CoreExpr -> (CoreExpr, [CoreExpr])
forall b. Expr b -> (Expr b, [Expr b])
collectArgs CoreExpr
expr
          ([Tickish Id]
ticks, CoreExpr
expr'') = (Tickish Id -> Bool) -> CoreExpr -> ([Tickish Id], CoreExpr)
forall b. (Tickish Id -> Bool) -> Expr b -> ([Tickish Id], Expr b)
stripTicksTop Tickish Id -> Bool
forall id. Tickish id -> Bool
tickishFloatable CoreExpr
expr'
          sexpr :: CoreExpr
sexpr = (CoreExpr -> CoreExpr -> CoreExpr)
-> CoreExpr -> [CoreExpr] -> CoreExpr
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App CoreExpr
expr'' [CoreExpr]
args
          retick :: CoreExpr -> CoreExpr
retick CoreExpr
expr = (Tickish Id -> CoreExpr -> CoreExpr)
-> CoreExpr -> [Tickish Id] -> CoreExpr
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Tickish Id -> CoreExpr -> CoreExpr
mkTick CoreExpr
expr [Tickish Id]
ticks

                                -- Abstraction    Application
--------------
data EtaInfo = EtaVar Var       -- /\a. []        [] a
                                -- \x.  []        [] x
             | EtaCo Coercion   -- [] |> sym co   [] |> co

instance Outputable EtaInfo where
   ppr :: EtaInfo -> SDoc
ppr (EtaVar Id
v) = String -> SDoc
text String
"EtaVar" SDoc -> SDoc -> SDoc
<+> Id -> SDoc
forall a. Outputable a => a -> SDoc
ppr Id
v
   ppr (EtaCo CoercionR
co) = String -> SDoc
text String
"EtaCo"  SDoc -> SDoc -> SDoc
<+> CoercionR -> SDoc
forall a. Outputable a => a -> SDoc
ppr CoercionR
co

pushCoercion :: Coercion -> [EtaInfo] -> [EtaInfo]
pushCoercion :: CoercionR -> [EtaInfo] -> [EtaInfo]
pushCoercion CoercionR
co1 (EtaCo CoercionR
co2 : [EtaInfo]
eis)
  | CoercionR -> Bool
isReflCo CoercionR
co = [EtaInfo]
eis
  | Bool
otherwise   = CoercionR -> EtaInfo
EtaCo CoercionR
co EtaInfo -> [EtaInfo] -> [EtaInfo]
forall a. a -> [a] -> [a]
: [EtaInfo]
eis
  where
    co :: CoercionR
co = CoercionR
co1 CoercionR -> CoercionR -> CoercionR
`mkTransCo` CoercionR
co2

pushCoercion CoercionR
co [EtaInfo]
eis = CoercionR -> EtaInfo
EtaCo CoercionR
co EtaInfo -> [EtaInfo] -> [EtaInfo]
forall a. a -> [a] -> [a]
: [EtaInfo]
eis

--------------
etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr
etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr
etaInfoAbs []               CoreExpr
expr = CoreExpr
expr
etaInfoAbs (EtaVar Id
v : [EtaInfo]
eis) CoreExpr
expr = Id -> CoreExpr -> CoreExpr
forall b. b -> Expr b -> Expr b
Lam Id
v ([EtaInfo] -> CoreExpr -> CoreExpr
etaInfoAbs [EtaInfo]
eis CoreExpr
expr)
etaInfoAbs (EtaCo CoercionR
co : [EtaInfo]
eis) CoreExpr
expr = CoreExpr -> CoercionR -> CoreExpr
forall b. Expr b -> CoercionR -> Expr b
Cast ([EtaInfo] -> CoreExpr -> CoreExpr
etaInfoAbs [EtaInfo]
eis CoreExpr
expr) (CoercionR -> CoercionR
mkSymCo CoercionR
co)

--------------
etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
-- (etaInfoApp s e eis) returns something equivalent to
--             ((substExpr s e) `appliedto` eis)

etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
etaInfoApp Subst
subst (Lam Id
v1 CoreExpr
e) (EtaVar Id
v2 : [EtaInfo]
eis)
  = Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
etaInfoApp (Subst -> Id -> Id -> Subst
GHC.Core.Subst.extendSubstWithVar Subst
subst Id
v1 Id
v2) CoreExpr
e [EtaInfo]
eis

etaInfoApp Subst
subst (Cast CoreExpr
e CoercionR
co1) [EtaInfo]
eis
  = Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
etaInfoApp Subst
subst CoreExpr
e (CoercionR -> [EtaInfo] -> [EtaInfo]
pushCoercion CoercionR
co' [EtaInfo]
eis)
  where
    co' :: CoercionR
co' = HasCallStack => Subst -> CoercionR -> CoercionR
Subst -> CoercionR -> CoercionR
GHC.Core.Subst.substCo Subst
subst CoercionR
co1

etaInfoApp Subst
subst (Case CoreExpr
e Id
b Type
ty [Alt Id]
alts) [EtaInfo]
eis
  = CoreExpr -> Id -> Type -> [Alt Id] -> CoreExpr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case (Subst -> CoreExpr -> CoreExpr
subst_expr Subst
subst CoreExpr
e) Id
b1 Type
ty' [Alt Id]
alts'
  where
    (Subst
subst1, Id
b1) = Subst -> Id -> (Subst, Id)
substBndr Subst
subst Id
b
    alts' :: [Alt Id]
alts' = (Alt Id -> Alt Id) -> [Alt Id] -> [Alt Id]
forall a b. (a -> b) -> [a] -> [b]
map Alt Id -> Alt Id
forall {a}. (a, [Id], CoreExpr) -> (a, [Id], CoreExpr)
subst_alt [Alt Id]
alts
    ty' :: Type
ty'   = Type -> [EtaInfo] -> Type
etaInfoAppTy (Subst -> Type -> Type
GHC.Core.Subst.substTy Subst
subst Type
ty) [EtaInfo]
eis
    subst_alt :: (a, [Id], CoreExpr) -> (a, [Id], CoreExpr)
subst_alt (a
con, [Id]
bs, CoreExpr
rhs) = (a
con, [Id]
bs', Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
etaInfoApp Subst
subst2 CoreExpr
rhs [EtaInfo]
eis)
              where
                 (Subst
subst2,[Id]
bs') = Subst -> [Id] -> (Subst, [Id])
substBndrs Subst
subst1 [Id]
bs

etaInfoApp Subst
subst (Let Bind Id
b CoreExpr
e) [EtaInfo]
eis
  | Bool -> Bool
not (Bind Id -> Bool
isJoinBind Bind Id
b)
    -- See Note [Eta expansion for join points]
  = Bind Id -> CoreExpr -> CoreExpr
forall b. Bind b -> Expr b -> Expr b
Let Bind Id
b' (Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
etaInfoApp Subst
subst' CoreExpr
e [EtaInfo]
eis)
  where
    (Subst
subst', Bind Id
b') = HasDebugCallStack => Subst -> Bind Id -> (Subst, Bind Id)
Subst -> Bind Id -> (Subst, Bind Id)
substBindSC Subst
subst Bind Id
b

etaInfoApp Subst
subst (Tick Tickish Id
t CoreExpr
e) [EtaInfo]
eis
  = Tickish Id -> CoreExpr -> CoreExpr
forall b. Tickish Id -> Expr b -> Expr b
Tick (Subst -> Tickish Id -> Tickish Id
substTickish Subst
subst Tickish Id
t) (Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
etaInfoApp Subst
subst CoreExpr
e [EtaInfo]
eis)

etaInfoApp Subst
subst CoreExpr
expr [EtaInfo]
_
  | (Var Id
fun, [CoreExpr]
_) <- CoreExpr -> (CoreExpr, [CoreExpr])
forall b. Expr b -> (Expr b, [Expr b])
collectArgs CoreExpr
expr
  , Var Id
fun' <- HasDebugCallStack => Subst -> Id -> CoreExpr
Subst -> Id -> CoreExpr
lookupIdSubst Subst
subst Id
fun
  , Id -> Bool
isJoinId Id
fun'
  = Subst -> CoreExpr -> CoreExpr
subst_expr Subst
subst CoreExpr
expr

etaInfoApp Subst
subst CoreExpr
e [EtaInfo]
eis
  = CoreExpr -> [EtaInfo] -> CoreExpr
forall {b}. Expr b -> [EtaInfo] -> Expr b
go (Subst -> CoreExpr -> CoreExpr
subst_expr Subst
subst CoreExpr
e) [EtaInfo]
eis
  where
    go :: Expr b -> [EtaInfo] -> Expr b
go Expr b
e []                  = Expr b
e
    go Expr b
e (EtaVar Id
v    : [EtaInfo]
eis) = Expr b -> [EtaInfo] -> Expr b
go (Expr b -> Expr b -> Expr b
forall b. Expr b -> Expr b -> Expr b
App Expr b
e (Id -> Expr b
forall b. Id -> Expr b
varToCoreExpr Id
v)) [EtaInfo]
eis
    go Expr b
e (EtaCo CoercionR
co    : [EtaInfo]
eis) = Expr b -> [EtaInfo] -> Expr b
go (Expr b -> CoercionR -> Expr b
forall b. Expr b -> CoercionR -> Expr b
Cast Expr b
e CoercionR
co) [EtaInfo]
eis


--------------
etaInfoAppTy :: Type -> [EtaInfo] -> Type
-- If                    e :: ty
-- then   etaInfoApp e eis :: etaInfoApp ty eis
etaInfoAppTy :: Type -> [EtaInfo] -> Type
etaInfoAppTy Type
ty []               = Type
ty
etaInfoAppTy Type
ty (EtaVar Id
v : [EtaInfo]
eis) = Type -> [EtaInfo] -> Type
etaInfoAppTy (Type -> CoreExpr -> Type
applyTypeToArg Type
ty (Id -> CoreExpr
forall b. Id -> Expr b
varToCoreExpr Id
v)) [EtaInfo]
eis
etaInfoAppTy Type
_  (EtaCo CoercionR
co : [EtaInfo]
eis) = Type -> [EtaInfo] -> Type
etaInfoAppTy (CoercionR -> Type
coercionRKind CoercionR
co) [EtaInfo]
eis

--------------
-- | @mkEtaWW n _ fvs ty@ will compute the 'EtaInfo' necessary for eta-expanding
-- an expression @e :: ty@ to take @n@ value arguments, where @fvs@ are the
-- free variables of @e@.
--
-- Note that this function is entirely unconcerned about cost centres and other
-- semantically-irrelevant source annotations, so call sites must take care to
-- preserve that info. See Note [Eta expansion and SCCs].
mkEtaWW
  :: [OneShotInfo]
  -- ^ How many value arguments to eta-expand
  -> SDoc
  -- ^ The pretty-printed original expression, for warnings.
  -> InScopeSet
  -- ^ A super-set of the free vars of the expression to eta-expand.
  -> Type
  -> (InScopeSet, [EtaInfo])
  -- ^ The variables in 'EtaInfo' are fresh wrt. to the incoming 'InScopeSet'.
  -- The outgoing 'InScopeSet' extends the incoming 'InScopeSet' with the
  -- fresh variables in 'EtaInfo'.

mkEtaWW :: [OneShotInfo]
-> SDoc -> InScopeSet -> Type -> (InScopeSet, [EtaInfo])
mkEtaWW [OneShotInfo]
orig_oss SDoc
ppr_orig_expr InScopeSet
in_scope Type
orig_ty
  = JoinArity
-> [OneShotInfo]
-> TCvSubst
-> Type
-> [EtaInfo]
-> (InScopeSet, [EtaInfo])
go JoinArity
0 [OneShotInfo]
orig_oss TCvSubst
empty_subst Type
orig_ty []
  where
    empty_subst :: TCvSubst
empty_subst = InScopeSet -> TCvSubst
mkEmptyTCvSubst InScopeSet
in_scope

    go :: Int                -- For fresh names
       -> [OneShotInfo]      -- Number of value args to expand to
       -> TCvSubst -> Type   -- We are really looking at subst(ty)
       -> [EtaInfo]          -- Accumulating parameter
       -> (InScopeSet, [EtaInfo])
    go :: JoinArity
-> [OneShotInfo]
-> TCvSubst
-> Type
-> [EtaInfo]
-> (InScopeSet, [EtaInfo])
go JoinArity
_ [] TCvSubst
subst Type
_ [EtaInfo]
eis       -- See Note [exprArity invariant]
       ----------- Done!  No more expansion needed
       = (TCvSubst -> InScopeSet
getTCvInScope TCvSubst
subst, [EtaInfo] -> [EtaInfo]
forall a. [a] -> [a]
reverse [EtaInfo]
eis)

    go JoinArity
n oss :: [OneShotInfo]
oss@(OneShotInfo
one_shot:[OneShotInfo]
oss1) TCvSubst
subst Type
ty [EtaInfo]
eis       -- See Note [exprArity invariant]
       ----------- Forall types  (forall a. ty)
       | Just (Id
tcv,Type
ty') <- Type -> Maybe (Id, Type)
splitForAllTy_maybe Type
ty
       , (TCvSubst
subst', Id
tcv') <- HasCallStack => TCvSubst -> Id -> (TCvSubst, Id)
TCvSubst -> Id -> (TCvSubst, Id)
Type.substVarBndr TCvSubst
subst Id
tcv
       , let oss' :: [OneShotInfo]
oss' | Id -> Bool
isTyVar Id
tcv = [OneShotInfo]
oss
                  | Bool
otherwise   = [OneShotInfo]
oss1
         -- A forall can bind a CoVar, in which case
         -- we consume one of the [OneShotInfo]
       = JoinArity
-> [OneShotInfo]
-> TCvSubst
-> Type
-> [EtaInfo]
-> (InScopeSet, [EtaInfo])
go JoinArity
n [OneShotInfo]
oss' TCvSubst
subst' Type
ty' (Id -> EtaInfo
EtaVar Id
tcv' EtaInfo -> [EtaInfo] -> [EtaInfo]
forall a. a -> [a] -> [a]
: [EtaInfo]
eis)

       ----------- Function types  (t1 -> t2)
       | Just (Type
mult, Type
arg_ty, Type
res_ty) <- Type -> Maybe (Type, Type, Type)
splitFunTy_maybe Type
ty
       , Bool -> Bool
not (Type -> Bool
isTypeLevPoly Type
arg_ty)
          -- See Note [Levity polymorphism invariants] in GHC.Core
          -- See also test case typecheck/should_run/EtaExpandLevPoly

       , (TCvSubst
subst', Id
eta_id) <- JoinArity -> TCvSubst -> Scaled Type -> (TCvSubst, Id)
freshEtaId JoinArity
n TCvSubst
subst (Type -> Type -> Scaled Type
forall a. Type -> a -> Scaled a
Scaled Type
mult Type
arg_ty)
          -- Avoid free vars of the original expression

       , let eta_id' :: Id
eta_id' = Id
eta_id Id -> OneShotInfo -> Id
`setIdOneShotInfo` OneShotInfo
one_shot
       = JoinArity
-> [OneShotInfo]
-> TCvSubst
-> Type
-> [EtaInfo]
-> (InScopeSet, [EtaInfo])
go (JoinArity
nJoinArity -> JoinArity -> JoinArity
forall a. Num a => a -> a -> a
+JoinArity
1) [OneShotInfo]
oss1 TCvSubst
subst' Type
res_ty (Id -> EtaInfo
EtaVar Id
eta_id' EtaInfo -> [EtaInfo] -> [EtaInfo]
forall a. a -> [a] -> [a]
: [EtaInfo]
eis)

       ----------- Newtypes
       -- Given this:
       --      newtype T = MkT ([T] -> Int)
       -- Consider eta-expanding this
       --      eta_expand 1 e T
       -- We want to get
       --      coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
       | Just (CoercionR
co, Type
ty') <- Type -> Maybe (CoercionR, Type)
topNormaliseNewType_maybe Type
ty
       , let co' :: CoercionR
co' = HasCallStack => TCvSubst -> CoercionR -> CoercionR
TCvSubst -> CoercionR -> CoercionR
Coercion.substCo TCvSubst
subst CoercionR
co
             -- Remember to apply the substitution to co (#16979)
             -- (or we could have applied to ty, but then
             --  we'd have had to zap it for the recursive call)
       = JoinArity
-> [OneShotInfo]
-> TCvSubst
-> Type
-> [EtaInfo]
-> (InScopeSet, [EtaInfo])
go JoinArity
n [OneShotInfo]
oss TCvSubst
subst Type
ty' (CoercionR -> [EtaInfo] -> [EtaInfo]
pushCoercion CoercionR
co' [EtaInfo]
eis)

       | Bool
otherwise       -- We have an expression of arity > 0,
                         -- but its type isn't a function, or a binder
                         -- is levity-polymorphic
       = WARN( True, (ppr orig_oss <+> ppr orig_ty) $$ ppr_orig_expr )
         (TCvSubst -> InScopeSet
getTCvInScope TCvSubst
subst, [EtaInfo] -> [EtaInfo]
forall a. [a] -> [a]
reverse [EtaInfo]
eis)
        -- This *can* legitimately happen:
        -- e.g.  coerce Int (\x. x) Essentially the programmer is
        -- playing fast and loose with types (Happy does this a lot).
        -- So we simply decline to eta-expand.  Otherwise we'd end up
        -- with an explicit lambda having a non-function type



------------
subst_expr :: Subst -> CoreExpr -> CoreExpr
-- Apply a substitution to an expression.  We use substExpr
-- not substExprSC (short-cutting substitution) because
-- we may be changing the types of join points, so applying
-- the in-scope set is necessary.
--
-- ToDo: we could instead check if we actually *are*
-- changing any join points' types, and if not use substExprSC.
subst_expr :: Subst -> CoreExpr -> CoreExpr
subst_expr = HasDebugCallStack => Subst -> CoreExpr -> CoreExpr
Subst -> CoreExpr -> CoreExpr
substExpr


--------------

-- | Split an expression into the given number of binders and a body,
-- eta-expanding if necessary. Counts value *and* type binders.
etaExpandToJoinPoint :: JoinArity -> CoreExpr -> ([CoreBndr], CoreExpr)
etaExpandToJoinPoint :: JoinArity -> CoreExpr -> ([Id], CoreExpr)
etaExpandToJoinPoint JoinArity
join_arity CoreExpr
expr
  = JoinArity -> [Id] -> CoreExpr -> ([Id], CoreExpr)
go JoinArity
join_arity [] CoreExpr
expr
  where
    go :: JoinArity -> [Id] -> CoreExpr -> ([Id], CoreExpr)
go JoinArity
0 [Id]
rev_bs CoreExpr
e         = ([Id] -> [Id]
forall a. [a] -> [a]
reverse [Id]
rev_bs, CoreExpr
e)
    go JoinArity
n [Id]
rev_bs (Lam Id
b CoreExpr
e) = JoinArity -> [Id] -> CoreExpr -> ([Id], CoreExpr)
go (JoinArity
nJoinArity -> JoinArity -> JoinArity
forall a. Num a => a -> a -> a
-JoinArity
1) (Id
b Id -> [Id] -> [Id]
forall a. a -> [a] -> [a]
: [Id]
rev_bs) CoreExpr
e
    go JoinArity
n [Id]
rev_bs CoreExpr
e         = case JoinArity -> CoreExpr -> ([Id], CoreExpr)
etaBodyForJoinPoint JoinArity
n CoreExpr
e of
                              ([Id]
bs, CoreExpr
e') -> ([Id] -> [Id]
forall a. [a] -> [a]
reverse [Id]
rev_bs [Id] -> [Id] -> [Id]
forall a. [a] -> [a] -> [a]
++ [Id]
bs, CoreExpr
e')

etaExpandToJoinPointRule :: JoinArity -> CoreRule -> CoreRule
etaExpandToJoinPointRule :: JoinArity -> CoreRule -> CoreRule
etaExpandToJoinPointRule JoinArity
_ rule :: CoreRule
rule@(BuiltinRule {})
  = WARN(True, (sep [text "Can't eta-expand built-in rule:", ppr rule]))
      -- How did a local binding get a built-in rule anyway? Probably a plugin.
    CoreRule
rule
etaExpandToJoinPointRule JoinArity
join_arity rule :: CoreRule
rule@(Rule { ru_bndrs :: CoreRule -> [Id]
ru_bndrs = [Id]
bndrs, ru_rhs :: CoreRule -> CoreExpr
ru_rhs = CoreExpr
rhs
                                               , ru_args :: CoreRule -> [CoreExpr]
ru_args  = [CoreExpr]
args })
  | JoinArity
need_args JoinArity -> JoinArity -> Bool
forall a. Eq a => a -> a -> Bool
== JoinArity
0
  = CoreRule
rule
  | JoinArity
need_args JoinArity -> JoinArity -> Bool
forall a. Ord a => a -> a -> Bool
< JoinArity
0
  = String -> SDoc -> CoreRule
forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"etaExpandToJoinPointRule" (JoinArity -> SDoc
forall a. Outputable a => a -> SDoc
ppr JoinArity
join_arity SDoc -> SDoc -> SDoc
$$ CoreRule -> SDoc
forall a. Outputable a => a -> SDoc
ppr CoreRule
rule)
  | Bool
otherwise
  = CoreRule
rule { ru_bndrs :: [Id]
ru_bndrs = [Id]
bndrs [Id] -> [Id] -> [Id]
forall a. [a] -> [a] -> [a]
++ [Id]
new_bndrs, ru_args :: [CoreExpr]
ru_args = [CoreExpr]
args [CoreExpr] -> [CoreExpr] -> [CoreExpr]
forall a. [a] -> [a] -> [a]
++ [CoreExpr]
forall {b}. [Expr b]
new_args
         , ru_rhs :: CoreExpr
ru_rhs = CoreExpr
new_rhs }
  where
    need_args :: JoinArity
need_args = JoinArity
join_arity JoinArity -> JoinArity -> JoinArity
forall a. Num a => a -> a -> a
- [CoreExpr] -> JoinArity
forall (t :: * -> *) a. Foldable t => t a -> JoinArity
length [CoreExpr]
args
    ([Id]
new_bndrs, CoreExpr
new_rhs) = JoinArity -> CoreExpr -> ([Id], CoreExpr)
etaBodyForJoinPoint JoinArity
need_args CoreExpr
rhs
    new_args :: [Expr b]
new_args = [Id] -> [Expr b]
forall b. [Id] -> [Expr b]
varsToCoreExprs [Id]
new_bndrs

-- Adds as many binders as asked for; assumes expr is not a lambda
etaBodyForJoinPoint :: Int -> CoreExpr -> ([CoreBndr], CoreExpr)
etaBodyForJoinPoint :: JoinArity -> CoreExpr -> ([Id], CoreExpr)
etaBodyForJoinPoint JoinArity
need_args CoreExpr
body
  = JoinArity
-> Type -> TCvSubst -> [Id] -> CoreExpr -> ([Id], CoreExpr)
forall {b}.
JoinArity -> Type -> TCvSubst -> [Id] -> Expr b -> ([Id], Expr b)
go JoinArity
need_args (CoreExpr -> Type
exprType CoreExpr
body) (CoreExpr -> TCvSubst
init_subst CoreExpr
body) [] CoreExpr
body
  where
    go :: JoinArity -> Type -> TCvSubst -> [Id] -> Expr b -> ([Id], Expr b)
go JoinArity
0 Type
_  TCvSubst
_     [Id]
rev_bs Expr b
e
      = ([Id] -> [Id]
forall a. [a] -> [a]
reverse [Id]
rev_bs, Expr b
e)
    go JoinArity
n Type
ty TCvSubst
subst [Id]
rev_bs Expr b
e
      | Just (Id
tv, Type
res_ty) <- Type -> Maybe (Id, Type)
splitForAllTy_maybe Type
ty
      , let (TCvSubst
subst', Id
tv') = HasCallStack => TCvSubst -> Id -> (TCvSubst, Id)
TCvSubst -> Id -> (TCvSubst, Id)
Type.substVarBndr TCvSubst
subst Id
tv
      = JoinArity -> Type -> TCvSubst -> [Id] -> Expr b -> ([Id], Expr b)
go (JoinArity
nJoinArity -> JoinArity -> JoinArity
forall a. Num a => a -> a -> a
-JoinArity
1) Type
res_ty TCvSubst
subst' (Id
tv' Id -> [Id] -> [Id]
forall a. a -> [a] -> [a]
: [Id]
rev_bs) (Expr b
e Expr b -> Expr b -> Expr b
forall b. Expr b -> Expr b -> Expr b
`App` Id -> Expr b
forall b. Id -> Expr b
varToCoreExpr Id
tv')
      | Just (Type
mult, Type
arg_ty, Type
res_ty) <- Type -> Maybe (Type, Type, Type)
splitFunTy_maybe Type
ty
      , let (TCvSubst
subst', Id
b) = JoinArity -> TCvSubst -> Scaled Type -> (TCvSubst, Id)
freshEtaId JoinArity
n TCvSubst
subst (Type -> Type -> Scaled Type
forall a. Type -> a -> Scaled a
Scaled Type
mult Type
arg_ty)
      = JoinArity -> Type -> TCvSubst -> [Id] -> Expr b -> ([Id], Expr b)
go (JoinArity
nJoinArity -> JoinArity -> JoinArity
forall a. Num a => a -> a -> a
-JoinArity
1) Type
res_ty TCvSubst
subst' (Id
b Id -> [Id] -> [Id]
forall a. a -> [a] -> [a]
: [Id]
rev_bs) (Expr b
e Expr b -> Expr b -> Expr b
forall b. Expr b -> Expr b -> Expr b
`App` Id -> Expr b
forall b. Id -> Expr b
Var Id
b)
      | Bool
otherwise
      = String -> SDoc -> ([Id], Expr b)
forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"etaBodyForJoinPoint" (SDoc -> ([Id], Expr b)) -> SDoc -> ([Id], Expr b)
forall a b. (a -> b) -> a -> b
$ JoinArity -> SDoc
int JoinArity
need_args SDoc -> SDoc -> SDoc
$$
                                         CoreExpr -> SDoc
forall a. Outputable a => a -> SDoc
ppr CoreExpr
body SDoc -> SDoc -> SDoc
$$ Type -> SDoc
forall a. Outputable a => a -> SDoc
ppr (CoreExpr -> Type
exprType CoreExpr
body)

    init_subst :: CoreExpr -> TCvSubst
init_subst CoreExpr
e = InScopeSet -> TCvSubst
mkEmptyTCvSubst (IdSet -> InScopeSet
mkInScopeSet (CoreExpr -> IdSet
exprFreeVars CoreExpr
e))

--------------
freshEtaId :: Int -> TCvSubst -> Scaled Type -> (TCvSubst, Id)
-- Make a fresh Id, with specified type (after applying substitution)
-- It should be "fresh" in the sense that it's not in the in-scope set
-- of the TvSubstEnv; and it should itself then be added to the in-scope
-- set of the TvSubstEnv
--
-- The Int is just a reasonable starting point for generating a unique;
-- it does not necessarily have to be unique itself.
freshEtaId :: JoinArity -> TCvSubst -> Scaled Type -> (TCvSubst, Id)
freshEtaId JoinArity
n TCvSubst
subst Scaled Type
ty
      = (TCvSubst
subst', Id
eta_id')
      where
        Scaled Type
mult' Type
ty' = HasCallStack => TCvSubst -> Scaled Type -> Scaled Type
TCvSubst -> Scaled Type -> Scaled Type
Type.substScaledTyUnchecked TCvSubst
subst Scaled Type
ty
        eta_id' :: Id
eta_id' = InScopeSet -> Id -> Id
uniqAway (TCvSubst -> InScopeSet
getTCvInScope TCvSubst
subst) (Id -> Id) -> Id -> Id
forall a b. (a -> b) -> a -> b
$
                  FastString -> Unique -> Type -> Type -> Id
mkSysLocalOrCoVar (String -> FastString
fsLit String
"eta") (JoinArity -> Unique
mkBuiltinUnique JoinArity
n) Type
mult' Type
ty'
                  -- "OrCoVar" since this can be used to eta-expand
                  -- coercion abstractions
        subst' :: TCvSubst
subst'  = TCvSubst -> Id -> TCvSubst
extendTCvInScope TCvSubst
subst Id
eta_id'