{-
(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
   , exprBotStrictness_maybe

   -- ** ArityType
   , ArityType(..), mkBotArityType, mkManifestArityType, expandableArityType
   , arityTypeArity, maxWithArity, idArityType

   -- ** Join points
   , etaExpandToJoinPoint, etaExpandToJoinPointRule

   -- ** Coercions and casts
   , pushCoArg, pushCoArgs, pushCoValArg, pushCoTyArg
   , pushCoercionIntoLambda, pushCoDataCon, collectBindersPushingCo
   )
where

#include "HsVersions.h"

import GHC.Prelude

import GHC.Driver.Ppr

import GHC.Core
import GHC.Core.FVs
import GHC.Core.Utils
import GHC.Types.Demand
import GHC.Types.Var
import GHC.Types.Var.Env
import GHC.Types.Id

-- We have two sorts of substitution:
--   GHC.Core.Subst.Subst, and GHC.Core.TyCo.TCvSubst
-- Both have substTy, substCo  Hence need for qualification
import GHC.Core.Subst    as Core
import GHC.Core.Type     as Type
import GHC.Core.Coercion as Type

import GHC.Core.DataCon
import GHC.Core.TyCon     ( tyConArity )
import GHC.Core.TyCon.RecWalk     ( initRecTc, checkRecTc )
import GHC.Core.Predicate ( isDictTy )
import GHC.Core.Multiplicity
import GHC.Types.Basic
import GHC.Types.Tickish
import GHC.Builtin.Uniques
import GHC.Driver.Session ( DynFlags, GeneralFlag(..), gopt )
import GHC.Utils.Outputable
import GHC.Utils.Panic
import GHC.Data.FastString
import GHC.Data.Pair
import GHC.Utils.Misc

{-
************************************************************************
*                                                                      *
              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 -> Int
manifestArity (Lam TyVar
v CoreExpr
e) | TyVar -> Bool
isId TyVar
v        = Int
1 forall a. Num a => a -> a -> a
+ CoreExpr -> Int
manifestArity CoreExpr
e
                        | Bool
otherwise     = CoreExpr -> Int
manifestArity CoreExpr
e
manifestArity (Tick CoreTickish
t CoreExpr
e) | Bool -> Bool
not (forall (pass :: TickishPass). GenTickish pass -> Bool
tickishIsCode CoreTickish
t) =  CoreExpr -> Int
manifestArity CoreExpr
e
manifestArity (Cast CoreExpr
e Coercion
_)                = CoreExpr -> Int
manifestArity CoreExpr
e
manifestArity CoreExpr
_                         = Int
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 -> Int
joinRhsArity (Lam TyVar
_ CoreExpr
e) = Int
1 forall a. Num a => a -> a -> a
+ CoreExpr -> Int
joinRhsArity CoreExpr
e
joinRhsArity CoreExpr
_         = Int
0


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

    go CoreExpr
_                           = Int
0

    trim_arity :: Arity -> Type -> Arity
    trim_arity :: Int -> Type -> Int
trim_arity Int
arity Type
ty = Int
arity forall a. Ord a => a -> a -> a
`min` forall (t :: * -> *) a. Foldable t => t a -> Int
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 (TyVar
_, Type
ty')  <- Type -> Maybe (TyVar, Type)
splitForAllTyCoVar_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 forall a. a -> [a] -> [a]
: RecTcChecker -> Type -> [OneShotInfo]
go RecTcChecker
rec_nts Type
res

      | Just (TyCon
tc,[Type]
tys) <- HasDebugCallStack => Type -> Maybe (TyCon, [Type])
splitTyConApp_maybe Type
ty
      , Just (Type
ty', Coercion
_) <- TyCon -> [Type] -> Maybe (Type, Coercion)
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 (Int, StrictSig)
exprBotStrictness_maybe CoreExpr
e
  = case ArityType -> Maybe Int
getBotArity (HasDebugCallStack => ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
botStrictnessArityEnv CoreExpr
e) of
        Maybe Int
Nothing -> forall a. Maybe a
Nothing
        Just Int
ar -> forall a. a -> Maybe a
Just (Int
ar, Int -> StrictSig
sig Int
ar)
  where
    sig :: Int -> StrictSig
sig Int
ar = [Demand] -> Divergence -> StrictSig
mkClosedStrictSig (forall a. Int -> a -> [a]
replicate Int
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 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).

We use the following notation:
  at  ::= \o1..on.div
  div ::= T | x | ⊥
  o   ::= ? | 1
And omit the \. if n = 0. Examples:
  \?11.T stands for @AT [NoOneShotInfo,OneShotLam,OneShotLam] topDiv@
  ⊥      stands for @AT [] botDiv@
See the 'Outputable' instance for more information. It's pretty simple.

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

 * If @at = AT [o1,..,on] botDiv@ (notation: \o1..on.⊥), 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 = AT [o1,..,on] topDiv@ (notation: \o1..on.T), then expanding 'f'
   to @\x1..xn. f x1..xn@ loses no sharing, assuming the calls of f respect
   the one-shot-ness o1..on of its definition.

   NB 'f' is an arbitrary expression, eg @f = g e1 e2@.  This 'f' can have
   arity type @AT oss _@, with @length oss > 0@, only if e1 e2 are themselves
   cheap.

 * In both cases, @f@, @f x1@, ... @f x1 ... x(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 arity type \??11.T
      The one-shot-ness means we can, in effect, push that
      'let' inside the \st.


Suppose f = \xy. x+y
Then  f             :: \??.T
      f v           :: \?.T
      f <expensive> :: T
-}


-- | The analysis lattice of arity analysis. It is isomorphic to
--
-- @
--    data ArityType'
--      = AEnd Divergence
--      | ALam OneShotInfo ArityType'
-- @
--
-- Which is easier to display the Hasse diagram for:
--
-- @
--  ALam OneShotLam at
--          |
--      AEnd topDiv
--          |
--  ALam NoOneShotInfo at
--          |
--      AEnd exnDiv
--          |
--      AEnd botDiv
-- @
--
-- where the @at@ fields of @ALam@ are inductively subject to the same order.
-- That is, @ALam os at1 < ALam os at2@ iff @at1 < at2@.
--
-- Why the strange Top element?
--   See Note [Combining case branches: optimistic one-shot-ness]
--
-- We rely on this lattice structure for fixed-point iteration in
-- 'findRhsArity'. For the semantics of 'ArityType', see Note [ArityType].
data ArityType
  = AT ![OneShotInfo] !Divergence
  -- ^ @AT oss div@ means this value can safely be eta-expanded @length oss@
  -- times, provided use sites respect the 'OneShotInfo's in @oss@.
  -- A 'OneShotLam' annotation can come from two sources:
  --     * The user annotated a lambda as one-shot with 'GHC.Exts.oneShot'
  --     * It's from a lambda binder of a type affected by `-fstate-hack`.
  --       See 'idStateHackOneShotInfo'.
  -- In both cases, 'OneShotLam' should win over 'NoOneShotInfo', see
  -- Note [Combining case branches].
  --
  -- If @div@ is dead-ending ('isDeadEndDiv'), then application to
  -- @length os@ arguments will surely diverge, similar to the situation
  -- with 'DmdType'.
  deriving ArityType -> ArityType -> Bool
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

-- | This is the BNF of the generated output:
--
-- @
-- @
--
-- We format
-- @AT [o1,..,on] topDiv@ as @\o1..on.T@ and
-- @AT [o1,..,on] botDiv@ as @\o1..on.⊥@, respectively.
-- More concretely, @AT [NOI,OS,OS] topDiv@ is formatted as @\?11.T@.
-- If the one-shot info is empty, we omit the leading @\.@.
instance Outputable ArityType where
  ppr :: ArityType -> SDoc
ppr (AT [OneShotInfo]
oss Divergence
div)
    | forall (t :: * -> *) a. Foldable t => t a -> Bool
null [OneShotInfo]
oss  = Divergence -> SDoc
pp_div Divergence
div
    | Bool
otherwise = Char -> SDoc
char Char
'\\' SDoc -> SDoc -> SDoc
<> [SDoc] -> SDoc
hcat (forall a b. (a -> b) -> [a] -> [b]
map OneShotInfo -> SDoc
pp_os [OneShotInfo]
oss) SDoc -> SDoc -> SDoc
<> SDoc
dot SDoc -> SDoc -> SDoc
<> Divergence -> SDoc
pp_div Divergence
div
    where
      pp_div :: Divergence -> SDoc
pp_div Divergence
Diverges = Char -> SDoc
char Char
'⊥'
      pp_div Divergence
ExnOrDiv = Char -> SDoc
char Char
'x'
      pp_div Divergence
Dunno    = Char -> SDoc
char Char
'T'
      pp_os :: OneShotInfo -> SDoc
pp_os OneShotInfo
OneShotLam    = Char -> SDoc
char Char
'1'
      pp_os OneShotInfo
NoOneShotInfo = Char -> SDoc
char Char
'?'

mkBotArityType :: [OneShotInfo] -> ArityType
mkBotArityType :: [OneShotInfo] -> ArityType
mkBotArityType [OneShotInfo]
oss = [OneShotInfo] -> Divergence -> ArityType
AT [OneShotInfo]
oss Divergence
botDiv

botArityType :: ArityType
botArityType :: ArityType
botArityType = [OneShotInfo] -> ArityType
mkBotArityType []

mkManifestArityType :: [Var] -> CoreExpr -> ArityType
mkManifestArityType :: [TyVar] -> CoreExpr -> ArityType
mkManifestArityType [TyVar]
bndrs CoreExpr
body
  = [OneShotInfo] -> Divergence -> ArityType
AT [OneShotInfo]
oss Divergence
div
  where
    oss :: [OneShotInfo]
oss = [TyVar -> OneShotInfo
idOneShotInfo TyVar
bndr | TyVar
bndr <- [TyVar]
bndrs, TyVar -> Bool
isId TyVar
bndr]
    div :: Divergence
div | CoreExpr -> Bool
exprIsDeadEnd CoreExpr
body = Divergence
botDiv
        | Bool
otherwise          = Divergence
topDiv

topArityType :: ArityType
topArityType :: ArityType
topArityType = [OneShotInfo] -> Divergence -> ArityType
AT [] Divergence
topDiv

-- | The number of value args for the arity type
arityTypeArity :: ArityType -> Arity
arityTypeArity :: ArityType -> Int
arityTypeArity (AT [OneShotInfo]
oss Divergence
_) = forall (t :: * -> *) a. Foldable t => t a -> Int
length [OneShotInfo]
oss

-- | True <=> eta-expansion will add at least one lambda
expandableArityType :: ArityType -> Bool
expandableArityType :: ArityType -> Bool
expandableArityType ArityType
at = ArityType -> Int
arityTypeArity ArityType
at forall a. Eq a => a -> a -> Bool
/= Int
0

-- | See Note [Dead ends] in "GHC.Types.Demand".
-- Bottom implies a dead end.
isDeadEndArityType :: ArityType -> Bool
isDeadEndArityType :: ArityType -> Bool
isDeadEndArityType (AT [OneShotInfo]
_ Divergence
div) = Divergence -> Bool
isDeadEndDiv Divergence
div

-- | Expand a non-bottoming arity type so that it has at least the given arity.
maxWithArity :: ArityType -> Arity -> ArityType
maxWithArity :: ArityType -> Int -> ArityType
maxWithArity at :: ArityType
at@(AT [OneShotInfo]
oss Divergence
div) !Int
ar
  | ArityType -> Bool
isDeadEndArityType ArityType
at    = ArityType
at
  | [OneShotInfo]
oss forall a. [a] -> Int -> Bool
`lengthAtLeast` Int
ar   = ArityType
at
  | Bool
otherwise                = [OneShotInfo] -> Divergence -> ArityType
AT (forall a. Int -> [a] -> [a]
take Int
ar forall a b. (a -> b) -> a -> b
$ [OneShotInfo]
oss forall a. [a] -> [a] -> [a]
++ forall a. a -> [a]
repeat OneShotInfo
NoOneShotInfo) Divergence
div

-- | Trim an arity type so that it has at most the given arity.
-- Any excess 'OneShotInfo's are truncated to 'topDiv', even if they end in
-- 'ABot'.
minWithArity :: ArityType -> Arity -> ArityType
minWithArity :: ArityType -> Int -> ArityType
minWithArity at :: ArityType
at@(AT [OneShotInfo]
oss Divergence
_) Int
ar
  | [OneShotInfo]
oss forall a. [a] -> Int -> Bool
`lengthAtMost` Int
ar = ArityType
at
  | Bool
otherwise             = [OneShotInfo] -> Divergence -> ArityType
AT (forall a. Int -> [a] -> [a]
take Int
ar [OneShotInfo]
oss) Divergence
topDiv

takeWhileOneShot :: ArityType -> ArityType
takeWhileOneShot :: ArityType -> ArityType
takeWhileOneShot (AT [OneShotInfo]
oss Divergence
div)
  | Divergence -> Bool
isDeadEndDiv Divergence
div = [OneShotInfo] -> Divergence -> ArityType
AT (forall a. (a -> Bool) -> [a] -> [a]
takeWhile OneShotInfo -> Bool
isOneShotInfo [OneShotInfo]
oss) Divergence
topDiv
  | Bool
otherwise        = [OneShotInfo] -> Divergence -> ArityType
AT (forall a. (a -> Bool) -> [a] -> [a]
takeWhile OneShotInfo -> Bool
isOneShotInfo [OneShotInfo]
oss) Divergence
div

-- | 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 = HasDebugCallStack => ArityEnv -> CoreExpr -> ArityType
arityType (DynFlags -> ArityEnv
findRhsArityEnv DynFlags
dflags) CoreExpr
e

getBotArity :: ArityType -> Maybe Arity
-- Arity of a divergent function
getBotArity :: ArityType -> Maybe Int
getBotArity (AT [OneShotInfo]
oss Divergence
div)
  | Divergence -> Bool
isDeadEndDiv Divergence
div = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t a -> Int
length [OneShotInfo]
oss
  | Bool
otherwise        = forall a. Maybe a
Nothing

----------------------
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 -> TyVar -> CoreExpr -> Int -> ArityType
findRhsArity DynFlags
dflags TyVar
bndr CoreExpr
rhs Int
old_arity
  = Int -> ArityType -> ArityType
go Int
0 ArityType
botArityType
      -- We always do one step, but usually that produces a result equal to
      -- old_arity, and then we stop right away, because old_arity is assumed
      -- to be sound. In other words, arities should never decrease.
      -- Result: the common case is that there is just one iteration
  where
    go :: Int -> ArityType -> ArityType
    go :: Int -> ArityType -> ArityType
go !Int
n cur_at :: ArityType
cur_at@(AT [OneShotInfo]
oss Divergence
div)
      | Bool -> Bool
not (Divergence -> Bool
isDeadEndDiv Divergence
div)           -- the "stop right away" case
      , forall (t :: * -> *) a. Foldable t => t a -> Int
length [OneShotInfo]
oss forall a. Ord a => a -> a -> Bool
<= Int
old_arity = ArityType
cur_at -- from above
      | ArityType
next_at forall a. Eq a => a -> a -> Bool
== ArityType
cur_at       = ArityType
cur_at
      | Bool
otherwise               =
         -- Warn if more than 2 iterations. Why 2? See Note [Exciting arity]
         WARN( debugIsOn && n > 2, text "Exciting arity"
                                   $$ nest 2 (
                                        ppr bndr <+> ppr cur_at <+> ppr next_at
                                        $$ ppr rhs) )
         Int -> ArityType -> ArityType
go (Int
nforall a. Num a => a -> a -> a
+Int
1) ArityType
next_at
      where
        next_at :: ArityType
next_at = ArityType -> ArityType
step ArityType
cur_at

    step :: ArityType -> ArityType
    step :: ArityType -> ArityType
step ArityType
at = -- pprTrace "step" (ppr bndr <+> ppr at <+> ppr (arityType env rhs)) $
              HasDebugCallStack => ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
rhs
      where
        env :: ArityEnv
env = ArityEnv -> TyVar -> ArityType -> ArityEnv
extendSigEnv (DynFlags -> ArityEnv
findRhsArityEnv DynFlags
dflags) TyVar
bndr ArityType
at

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

We do the necessary, quite simple fixed-point iteration in 'findRhsArity',
which assumes for a single binding 'ABot' on the first run and iterates
until it finds a stable arity type. Two wrinkles

* We often have to ask (see the Case or Let case of 'arityType') whether some
  expression is cheap. In the case of an application, that depends on the arity
  of the application head! That's why we have our own version of 'exprIsCheap',
  'myExprIsCheap', that will integrate the optimistic arity types we have on
  f and g into the cheapness check.

* Consider this (#18793)

    go = \ds. case ds of
           []     -> id
           (x:ys) -> let acc = go ys in
                     case blah of
                       True  -> acc
                       False -> \ x1 -> acc (negate x1)

  We must propagate go's optimistically large arity to @acc@, so that the
  tail call to @acc@ in the True branch has sufficient arity.  This is done
  by the 'am_sigs' field in 'FindRhsArity', and 'lookupSigEnv' in the Var case
  of 'arityType'.

Note [Exciting Arity]
~~~~~~~~~~~~~~~~~~~~~
The fixed-point iteration in 'findRhsArity' stabilises very quickly in almost
all cases. To get notified of cases where we need an usual number of iterations,
we emit a warning in debug mode, so that we can investigate and make sure that
we really can't do better. It's a gross hack, but catches real bugs (#18870).

Now, which number is "unusual"? We pick n > 2. Here's a pretty common and
expected example that takes two iterations and would ruin the specificity
of the warning (from T18937):

  f :: [Int] -> Int -> Int
  f []     = id
  f (x:xs) = let y = sum [0..x]
             in \z -> f xs (y + z)

Fixed-point iteration starts with arity type ⊥ for f. After the first
iteration, we get arity type \??.T, e.g. arity 2, because we unconditionally
'floatIn' the let-binding (see its bottom case).  After the second iteration,
we get arity type \?.T, e.g. arity 1, because now we are no longer allowed
to floatIn the non-cheap let-binding.  Which is all perfectly benign, but
means we do two iterations (well, actually 3 'step's to detect we are stable)
and don't want to emit the warning.

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 :: TyVar -> ArityType -> ArityType
arityLam TyVar
id (AT [OneShotInfo]
oss Divergence
div) = [OneShotInfo] -> Divergence -> ArityType
AT (TyVar -> OneShotInfo
idStateHackOneShotInfo TyVar
id forall a. a -> [a] -> [a]
: [OneShotInfo]
oss) Divergence
div

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
cheap ArityType
at
  | ArityType -> Bool
isDeadEndArityType ArityType
at Bool -> Bool -> Bool
|| Bool
cheap = ArityType
at
  -- If E is not cheap, keep arity only for one-shots
  | Bool
otherwise                      = ArityType -> ArityType
takeWhileOneShot ArityType
at

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 (AT (OneShotInfo
_:[OneShotInfo]
oss) Divergence
div) Bool
cheap = Bool -> ArityType -> ArityType
floatIn Bool
cheap ([OneShotInfo] -> Divergence -> ArityType
AT [OneShotInfo]
oss Divergence
div)
arityApp ArityType
at               Bool
_     = ArityType
at

-- | Least upper bound in the 'ArityType' lattice.
-- See the haddocks on 'ArityType' for the lattice.
--
-- Used for branches of a @case@.
andArityType :: ArityEnv -> ArityType -> ArityType -> ArityType
andArityType :: ArityEnv -> ArityType -> ArityType -> ArityType
andArityType ArityEnv
env (AT (OneShotInfo
lam1:[OneShotInfo]
lams1) Divergence
div1) (AT (OneShotInfo
lam2:[OneShotInfo]
lams2) Divergence
div2)
  | AT [OneShotInfo]
lams' Divergence
div' <- ArityEnv -> ArityType -> ArityType -> ArityType
andArityType ArityEnv
env ([OneShotInfo] -> Divergence -> ArityType
AT [OneShotInfo]
lams1 Divergence
div1) ([OneShotInfo] -> Divergence -> ArityType
AT [OneShotInfo]
lams2 Divergence
div2)
  = [OneShotInfo] -> Divergence -> ArityType
AT ((OneShotInfo
lam1 OneShotInfo -> OneShotInfo -> OneShotInfo
`and_lam` OneShotInfo
lam2) forall a. a -> [a] -> [a]
: [OneShotInfo]
lams') Divergence
div'
  where
    (OneShotInfo
os1) and_lam :: OneShotInfo -> OneShotInfo -> OneShotInfo
`and_lam` (OneShotInfo
os2)
      = ( OneShotInfo
os1 OneShotInfo -> OneShotInfo -> OneShotInfo
`bestOneShot` OneShotInfo
os2)
        -- bestOneShot: see Note [Combining case branches: optimistic one-shot-ness]

andArityType ArityEnv
env (AT [] Divergence
div1) ArityType
at2 = ArityEnv -> Divergence -> ArityType -> ArityType
andWithTail ArityEnv
env Divergence
div1 ArityType
at2
andArityType ArityEnv
env ArityType
at1 (AT [] Divergence
div2) = ArityEnv -> Divergence -> ArityType -> ArityType
andWithTail ArityEnv
env Divergence
div2 ArityType
at1

andWithTail :: ArityEnv -> Divergence -> ArityType -> ArityType
andWithTail :: ArityEnv -> Divergence -> ArityType -> ArityType
andWithTail ArityEnv
env Divergence
div1 at2 :: ArityType
at2@(AT [OneShotInfo]
lams2 Divergence
_)
  | Divergence -> Bool
isDeadEndDiv Divergence
div1     -- case x of { T -> error; F -> \y.e }
  = ArityType
at2        -- Note [ABot branches: max arity wins]

  | ArityEnv -> Bool
pedanticBottoms ArityEnv
env  -- Note [Combining case branches: andWithTail]
  = [OneShotInfo] -> Divergence -> ArityType
AT [] Divergence
topDiv

  | Bool
otherwise  -- case x of { T -> plusInt <expensive>; F -> \y.e }
  = [OneShotInfo] -> Divergence -> ArityType
AT [OneShotInfo]
lams2 Divergence
topDiv    -- We know div1 = topDiv
    -- See Note [Combining case branches: andWithTail]


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

Remember: \o1..on.⊥ means "if you apply to n args, it'll definitely diverge".
So we need \??.⊥ for the whole thing, the /max/ of both arities.

Note [Combining case branches: optimistic one-shot-ness]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
When combining the ArityTypes for two case branches (with andArityType)
and both ArityTypes have ATLamInfo, then we just combine their
expensive-ness and one-shot info.  The tricky point is when we have
     case x of True -> \x{one-shot). blah1
               Fale -> \y.           blah2

Since one-shot-ness is about the /consumer/ not the /producer/, we
optimistically assume that if either branch is one-shot, 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.  Surprisingly,
this means that the one-shot arity type is effectively the top element
of the lattice.

Hence the call to `bestOneShot` in `andArityType`.

Note [Combining case branches: andWithTail]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
When combining the ArityTypes for two case branches (with andArityType)
and one side or the other has run out of ATLamInfo; then we get
into `andWithTail`.

* If one branch is guaranteed bottom (isDeadEndDiv), we just take
  the other; see Note [ABot branches: max arity wins]

* Otherwise, if pedantic-bottoms is on, we just have to return
  AT [] topDiv.  E.g. if we have
    f x z = case x of True  -> \y. blah
                      False -> z
  then we can't eta-expand, because that would change the behaviour
  of (f False bottom().

* But if pedantic-bottoms is not on, we allow ourselves to push
  `z` under a lambda (much as we allow ourselves to put the `case x`
  under a lambda).  However we know nothing about the expensiveness
  or one-shot-ness of `z`, so we'd better assume it looks like
  (Expensive, NoOneShotInfo) all the way. Remembering
  Note [Combining case branches: optimistic one-shot-ness],
  we just add work to ever ATLamInfo, keeping the one-shot-ness.

Here's an example:
  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 respect the one-shot annotation provided by the
user and eta-expand go and go2.
When combining the branches of the case we have
     T `andAT` \1.T
and we want to get \1.T.
But if the inner lambda wasn't one-shot (\?.T) we don't want to do this.
(We need a usage 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.
Surprisingly, this means that the one-shot arity type is effectively the top
element of the lattice.

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?  If we end in topDiv, it's easy.  But we must take great care with
dead ends (i.e. botDiv). Suppose the expression was (\x y. error "urk"),
we'll get \??.⊥.  We absolutely must not trim that to \?.⊥, 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, if we need to trim a dead-ending arity type, switch (conservatively) to
topDiv.

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

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

-- | Each of the entry-points of the analyser ('arityType') has different
-- requirements. The entry-points are
--
--   1. 'exprBotStrictness_maybe'
--   2. 'exprEtaExpandArity'
--   3. 'findRhsArity'
--
-- For each of the entry-points, there is a separate mode that governs
--
--   1. How pedantic we are wrt. ⊥, in 'pedanticBottoms'.
--   2. Whether we store arity signatures for non-recursive let-bindings,
--      accessed in 'extendSigEnv'/'lookupSigEnv'.
--      See Note [Arity analysis] why that's important.
--   3. Which expressions we consider cheap to float inside a lambda,
--      in 'myExprIsCheap'.
data AnalysisMode
  = BotStrictness
  -- ^ Used during 'exprBotStrictness_maybe'.
  | EtaExpandArity { AnalysisMode -> Bool
am_ped_bot :: !Bool
                   , AnalysisMode -> Bool
am_dicts_cheap :: !Bool }
  -- ^ Used for finding an expression's eta-expanding arity quickly, without
  -- fixed-point iteration ('exprEtaExpandArity').
  | FindRhsArity { am_ped_bot :: !Bool
                 , am_dicts_cheap :: !Bool
                 , AnalysisMode -> IdEnv ArityType
am_sigs :: !(IdEnv ArityType) }
  -- ^ Used for regular, fixed-point arity analysis ('findRhsArity').
  --   See Note [Arity analysis] for details about fixed-point iteration.
  --   INVARIANT: Disjoint with 'ae_joins'.

data ArityEnv
  = AE
  { ArityEnv -> AnalysisMode
ae_mode   :: !AnalysisMode
  -- ^ The analysis mode. See 'AnalysisMode'.
  }

-- | The @ArityEnv@ used by 'exprBotStrictness_maybe'. Pedantic about bottoms
-- and no application is ever considered cheap.
botStrictnessArityEnv :: ArityEnv
botStrictnessArityEnv :: ArityEnv
botStrictnessArityEnv = AE { ae_mode :: AnalysisMode
ae_mode = AnalysisMode
BotStrictness }

{-
-- | The @ArityEnv@ used by 'exprEtaExpandArity'.
etaExpandArityEnv :: DynFlags -> ArityEnv
etaExpandArityEnv dflags
  = AE { ae_mode  = EtaExpandArity { am_ped_bot = gopt Opt_PedanticBottoms dflags
                                   , am_dicts_cheap = gopt Opt_DictsCheap dflags }
       , ae_joins = emptyVarSet }
-}

-- | The @ArityEnv@ used by 'findRhsArity'.
findRhsArityEnv :: DynFlags -> ArityEnv
findRhsArityEnv :: DynFlags -> ArityEnv
findRhsArityEnv DynFlags
dflags
  = AE { ae_mode :: AnalysisMode
ae_mode  = FindRhsArity { am_ped_bot :: Bool
am_ped_bot = GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_PedanticBottoms DynFlags
dflags
                                 , am_dicts_cheap :: Bool
am_dicts_cheap = GeneralFlag -> DynFlags -> Bool
gopt GeneralFlag
Opt_DictsCheap DynFlags
dflags
                                 , am_sigs :: IdEnv ArityType
am_sigs = forall a. VarEnv a
emptyVarEnv }
       }

isFindRhsArity :: ArityEnv -> Bool
isFindRhsArity :: ArityEnv -> Bool
isFindRhsArity (AE { ae_mode :: ArityEnv -> AnalysisMode
ae_mode = FindRhsArity {} }) = Bool
True
isFindRhsArity ArityEnv
_                                  = Bool
False

-- First some internal functions in snake_case for deleting in certain VarEnvs
-- of the ArityType. Don't call these; call delInScope* instead!

modifySigEnv :: (IdEnv ArityType -> IdEnv ArityType) -> ArityEnv -> ArityEnv
modifySigEnv :: (IdEnv ArityType -> IdEnv ArityType) -> ArityEnv -> ArityEnv
modifySigEnv IdEnv ArityType -> IdEnv ArityType
f env :: ArityEnv
env@AE { ae_mode :: ArityEnv -> AnalysisMode
ae_mode = am :: AnalysisMode
am@FindRhsArity{am_sigs :: AnalysisMode -> IdEnv ArityType
am_sigs = IdEnv ArityType
sigs} } =
  ArityEnv
env { ae_mode :: AnalysisMode
ae_mode = AnalysisMode
am { am_sigs :: IdEnv ArityType
am_sigs = IdEnv ArityType -> IdEnv ArityType
f IdEnv ArityType
sigs } }
modifySigEnv IdEnv ArityType -> IdEnv ArityType
_ ArityEnv
env = ArityEnv
env
{-# INLINE modifySigEnv #-}

del_sig_env :: Id -> ArityEnv -> ArityEnv -- internal!
del_sig_env :: TyVar -> ArityEnv -> ArityEnv
del_sig_env TyVar
id = (IdEnv ArityType -> IdEnv ArityType) -> ArityEnv -> ArityEnv
modifySigEnv (\IdEnv ArityType
sigs -> forall a. VarEnv a -> TyVar -> VarEnv a
delVarEnv IdEnv ArityType
sigs TyVar
id)
{-# INLINE del_sig_env #-}

del_sig_env_list :: [Id] -> ArityEnv -> ArityEnv -- internal!
del_sig_env_list :: [TyVar] -> ArityEnv -> ArityEnv
del_sig_env_list [TyVar]
ids = (IdEnv ArityType -> IdEnv ArityType) -> ArityEnv -> ArityEnv
modifySigEnv (\IdEnv ArityType
sigs -> forall a. VarEnv a -> [TyVar] -> VarEnv a
delVarEnvList IdEnv ArityType
sigs [TyVar]
ids)
{-# INLINE del_sig_env_list #-}

-- end of internal deletion functions

extendSigEnv :: ArityEnv -> Id -> ArityType -> ArityEnv
extendSigEnv :: ArityEnv -> TyVar -> ArityType -> ArityEnv
extendSigEnv ArityEnv
env TyVar
id ArityType
ar_ty
  = (IdEnv ArityType -> IdEnv ArityType) -> ArityEnv -> ArityEnv
modifySigEnv (\IdEnv ArityType
sigs -> forall a. VarEnv a -> TyVar -> a -> VarEnv a
extendVarEnv IdEnv ArityType
sigs TyVar
id ArityType
ar_ty) ArityEnv
env

delInScope :: ArityEnv -> Id -> ArityEnv
delInScope :: ArityEnv -> TyVar -> ArityEnv
delInScope ArityEnv
env TyVar
id = TyVar -> ArityEnv -> ArityEnv
del_sig_env TyVar
id ArityEnv
env

delInScopeList :: ArityEnv -> [Id] -> ArityEnv
delInScopeList :: ArityEnv -> [TyVar] -> ArityEnv
delInScopeList ArityEnv
env [TyVar]
ids = [TyVar] -> ArityEnv -> ArityEnv
del_sig_env_list [TyVar]
ids ArityEnv
env

lookupSigEnv :: ArityEnv -> Id -> Maybe ArityType
lookupSigEnv :: ArityEnv -> TyVar -> Maybe ArityType
lookupSigEnv AE{ ae_mode :: ArityEnv -> AnalysisMode
ae_mode = AnalysisMode
mode } TyVar
id = case AnalysisMode
mode of
  AnalysisMode
BotStrictness                  -> forall a. Maybe a
Nothing
  EtaExpandArity{}               -> forall a. Maybe a
Nothing
  FindRhsArity{ am_sigs :: AnalysisMode -> IdEnv ArityType
am_sigs = IdEnv ArityType
sigs } -> forall a. VarEnv a -> TyVar -> Maybe a
lookupVarEnv IdEnv ArityType
sigs TyVar
id

-- | Whether the analysis should be pedantic about bottoms.
-- 'exprBotStrictness_maybe' always is.
pedanticBottoms :: ArityEnv -> Bool
pedanticBottoms :: ArityEnv -> Bool
pedanticBottoms AE{ ae_mode :: ArityEnv -> AnalysisMode
ae_mode = AnalysisMode
mode } = case AnalysisMode
mode of
  AnalysisMode
BotStrictness                          -> Bool
True
  EtaExpandArity{ am_ped_bot :: AnalysisMode -> Bool
am_ped_bot = Bool
ped_bot } -> Bool
ped_bot
  FindRhsArity{ am_ped_bot :: AnalysisMode -> Bool
am_ped_bot = Bool
ped_bot }   -> Bool
ped_bot

-- | A version of 'exprIsCheap' that considers results from arity analysis
-- and optionally the expression's type.
-- Under 'exprBotStrictness_maybe', no expressions are cheap.
myExprIsCheap :: ArityEnv -> CoreExpr -> Maybe Type -> Bool
myExprIsCheap :: ArityEnv -> CoreExpr -> Maybe Type -> Bool
myExprIsCheap AE{ae_mode :: ArityEnv -> AnalysisMode
ae_mode = AnalysisMode
mode} CoreExpr
e Maybe Type
mb_ty = case AnalysisMode
mode of
  AnalysisMode
BotStrictness -> Bool
False
  AnalysisMode
_             -> Bool
cheap_dict Bool -> Bool -> Bool
|| CoreExpr -> Bool
cheap_fun CoreExpr
e
    where
      cheap_dict :: Bool
cheap_dict = AnalysisMode -> Bool
am_dicts_cheap AnalysisMode
mode Bool -> Bool -> Bool
&& forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Type -> Bool
isDictTy Maybe Type
mb_ty forall a. Eq a => a -> a -> Bool
== forall a. a -> Maybe a
Just Bool
True
      cheap_fun :: CoreExpr -> Bool
cheap_fun CoreExpr
e = case AnalysisMode
mode of
#if __GLASGOW_HASKELL__ <= 900
        BotStrictness                -> panic "impossible"
#endif
        EtaExpandArity{}             -> CoreExpr -> Bool
exprIsCheap CoreExpr
e
        FindRhsArity{am_sigs :: AnalysisMode -> IdEnv ArityType
am_sigs = IdEnv ArityType
sigs} -> CheapAppFun -> CoreExpr -> Bool
exprIsCheapX (IdEnv ArityType -> CheapAppFun
myIsCheapApp IdEnv ArityType
sigs) CoreExpr
e

-- | A version of 'isCheapApp' that considers results from arity analysis.
-- See Note [Arity analysis] for what's in the signature environment and why
-- it's important.
myIsCheapApp :: IdEnv ArityType -> CheapAppFun
myIsCheapApp :: IdEnv ArityType -> CheapAppFun
myIsCheapApp IdEnv ArityType
sigs TyVar
fn Int
n_val_args = case forall a. VarEnv a -> TyVar -> Maybe a
lookupVarEnv IdEnv ArityType
sigs TyVar
fn of
  -- Nothing means not a local function, fall back to regular
  -- 'GHC.Core.Utils.isCheapApp'
  Maybe ArityType
Nothing -> CheapAppFun
isCheapApp TyVar
fn Int
n_val_args

  -- `Just at` means local function with `at` as current SafeArityType.
  -- NB the SafeArityType bit: that means we can ignore the cost flags
  --    in 'lams', and just consider the length
  -- Roughly approximate what 'isCheapApp' is doing.
  Just (AT [OneShotInfo]
oss Divergence
div)
    | Divergence -> Bool
isDeadEndDiv Divergence
div -> Bool
True -- See Note [isCheapApp: bottoming functions] in GHC.Core.Utils
    | Int
n_val_args forall a. Ord a => a -> a -> Bool
< forall (t :: * -> *) a. Foldable t => t a -> Int
length [OneShotInfo]
oss -> Bool
True -- Essentially isWorkFreeApp
    | Bool
otherwise -> Bool
False

----------------
arityType :: HasDebugCallStack => ArityEnv -> CoreExpr -> ArityType
-- Precondition: all the free join points of the expression
--               are bound by the ArityEnv
-- See Note [No free join points in arityType]

arityType :: HasDebugCallStack => ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env (Cast CoreExpr
e Coercion
co)
  = ArityType -> Int -> ArityType
minWithArity (HasDebugCallStack => ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
e) Int
co_arity -- See Note [Arity trimming]
  where
    co_arity :: Int
co_arity = forall (t :: * -> *) a. Foldable t => t a -> Int
length (Type -> [OneShotInfo]
typeArity (Coercion -> Type
coercionRKind Coercion
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

arityType ArityEnv
env (Var TyVar
v)
  | Just ArityType
at <- ArityEnv -> TyVar -> Maybe ArityType
lookupSigEnv ArityEnv
env TyVar
v -- Local binding
  = ArityType
at
  | Bool
otherwise
  = ASSERT2( (not (isFindRhsArity env && isJoinId v)) , (ppr v) )
    -- All join-point should be in the ae_sigs
    -- See Note [No free join points in arityType]
    TyVar -> ArityType
idArityType TyVar
v

        -- Lambdas; increase arity
arityType ArityEnv
env (Lam TyVar
x CoreExpr
e)
  | TyVar -> Bool
isId TyVar
x    = TyVar -> ArityType -> ArityType
arityLam TyVar
x (HasDebugCallStack => ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env' CoreExpr
e)
  | Bool
otherwise = HasDebugCallStack => ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env' CoreExpr
e
  where
    env' :: ArityEnv
env' = ArityEnv -> TyVar -> ArityEnv
delInScope ArityEnv
env TyVar
x

        -- Applications; decrease arity, except for types
arityType ArityEnv
env (App CoreExpr
fun (Type Type
_))
   = HasDebugCallStack => ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
fun
arityType ArityEnv
env (App CoreExpr
fun CoreExpr
arg )
   = ArityType -> Bool -> ArityType
arityApp (HasDebugCallStack => ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
fun) (ArityEnv -> CoreExpr -> Maybe Type -> Bool
myExprIsCheap ArityEnv
env CoreExpr
arg 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 TyVar
bndr Type
_ [Alt TyVar]
alts)
  | CoreExpr -> Bool
exprIsDeadEnd CoreExpr
scrut Bool -> Bool -> Bool
|| forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Alt TyVar]
alts
  = ArityType
botArityType    -- Do not eta expand. See Note [Dealing with bottom (1)]
  | Bool -> Bool
not (ArityEnv -> Bool
pedanticBottoms ArityEnv
env)  -- See Note [Dealing with bottom (2)]
  , ArityEnv -> CoreExpr -> Maybe Type -> Bool
myExprIsCheap ArityEnv
env CoreExpr
scrut (forall a. a -> Maybe a
Just (TyVar -> Type
idType TyVar
bndr))
  = ArityType
alts_type
  | CoreExpr -> Bool
exprOkForSpeculation CoreExpr
scrut
  = ArityType
alts_type

  | Bool
otherwise                  -- In the remaining cases we may not push
  = ArityType -> ArityType
takeWhileOneShot ArityType
alts_type -- evaluation of the scrutinee in
  where
    env' :: ArityEnv
env' = ArityEnv -> TyVar -> ArityEnv
delInScope ArityEnv
env TyVar
bndr
    arity_type_alt :: Alt TyVar -> ArityType
arity_type_alt (Alt AltCon
_con [TyVar]
bndrs CoreExpr
rhs) = HasDebugCallStack => ArityEnv -> CoreExpr -> ArityType
arityType (ArityEnv -> [TyVar] -> ArityEnv
delInScopeList ArityEnv
env' [TyVar]
bndrs) CoreExpr
rhs
    alts_type :: ArityType
alts_type = forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 (ArityEnv -> ArityType -> ArityType -> ArityType
andArityType ArityEnv
env) (forall a b. (a -> b) -> [a] -> [b]
map Alt TyVar -> ArityType
arity_type_alt [Alt TyVar]
alts)

arityType ArityEnv
env (Let (NonRec TyVar
b CoreExpr
r) CoreExpr
e)
  = -- See Note [arityType for let-bindings]
    Bool -> ArityType -> ArityType
floatIn Bool
cheap_rhs (HasDebugCallStack => ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env' CoreExpr
e)
  where
    cheap_rhs :: Bool
cheap_rhs = ArityEnv -> CoreExpr -> Maybe Type -> Bool
myExprIsCheap ArityEnv
env CoreExpr
r (forall a. a -> Maybe a
Just (TyVar -> Type
idType TyVar
b))
    env' :: ArityEnv
env'      = ArityEnv -> TyVar -> ArityType -> ArityEnv
extendSigEnv ArityEnv
env TyVar
b (HasDebugCallStack => ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
r)

arityType ArityEnv
env (Let (Rec [(TyVar, CoreExpr)]
prs) CoreExpr
e)
  = Bool -> ArityType -> ArityType
floatIn (forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (TyVar, CoreExpr) -> Bool
is_cheap [(TyVar, CoreExpr)]
prs) (HasDebugCallStack => ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env' CoreExpr
e)
  where
    is_cheap :: (TyVar, CoreExpr) -> Bool
is_cheap (TyVar
b,CoreExpr
e) = ArityEnv -> CoreExpr -> Maybe Type -> Bool
myExprIsCheap ArityEnv
env' CoreExpr
e (forall a. a -> Maybe a
Just (TyVar -> Type
idType TyVar
b))
    env' :: ArityEnv
env'            = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ArityEnv -> (TyVar, CoreExpr) -> ArityEnv
extend_rec ArityEnv
env [(TyVar, CoreExpr)]
prs
    extend_rec :: ArityEnv -> (Id,CoreExpr) -> ArityEnv
    extend_rec :: ArityEnv -> (TyVar, CoreExpr) -> ArityEnv
extend_rec ArityEnv
env (TyVar
b,CoreExpr
e) = ArityEnv -> TyVar -> ArityType -> ArityEnv
extendSigEnv ArityEnv
env TyVar
b  forall a b. (a -> b) -> a -> b
$
                           [TyVar] -> CoreExpr -> ArityType
mkManifestArityType [TyVar]
bndrs CoreExpr
body
                         where
                           ([TyVar]
bndrs, CoreExpr
body) = forall b. Expr b -> ([b], Expr b)
collectBinders CoreExpr
e
      -- We can't call arityType on the RHS, because it might mention
      -- join points bound in this very letrec, and we don't want to
      -- do a fixpoint calculation here.  So we make do with the
      -- manifest arity

arityType ArityEnv
env (Tick CoreTickish
t CoreExpr
e)
  | Bool -> Bool
not (forall (pass :: TickishPass). GenTickish pass -> Bool
tickishIsCode CoreTickish
t)     = HasDebugCallStack => ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
e

arityType ArityEnv
_ CoreExpr
_ = ArityType
topArityType


{- Note [No free join points in arityType]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Suppose we call arityType on this expression (EX1)
   \x . case x of True  -> \y. e
                  False -> $j 3
where $j is a join point.  It really makes no sense to talk of the arity
of this expression, because it has a free join point.  In particular, we
can't eta-expand the expression because we'd have do the same thing to the
binding of $j, and we can't see that binding.

If we had (EX2)
   \x. join $j y = blah
       case x of True  -> \y. e
                 False -> $j 3
then it would make perfect sense: we can determine $j's ArityType, and
propagate it to the usage site as usual.

But how can we get (EX1)?  It doesn't make much sense, because $j can't
be a join point under the \x anyway.  So we make it a precondition of
arityType that the argument has no free join-point Ids.  (This is checked
with an assesrt in the Var case of arityType.)

BUT the invariant risks being invalidated by one very narrow special case: runRW#
   join $j y = blah
   runRW# (\s. case x of True  -> \y. e
                         False -> $j x)

We have special magic in OccurAnal, and Simplify to allow continuations to
move into the body of a runRW# call.

So we are careful never to attempt to eta-expand the (\s.blah) in the
argument to runRW#, at least not when there is a literal lambda there,
so that OccurAnal has seen it and allowed join points bound outside.
See Note [No eta-expansion in runRW#] in GHC.Core.Opt.Simplify.Iteration.

Note [arityType for let-bindings]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
For non-recursive let-bindings, we just get the arityType of the RHS,
and extend the environment.  That works nicely for things like this
(#18793):
  go = \ ds. case ds_a2CF of {
               []     -> id
               : y ys -> case y of { GHC.Types.I# x ->
                         let acc = go ys in
                         case x ># 42# of {
                            __DEFAULT -> acc
                            1# -> \x1. acc (negate x2)

Here we want to get a good arity for `acc`, based on the ArityType
of `go`.

All this is particularly important for 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. It's important that we extend the envt with j's ArityType,
so that we can use that information in the A/C branch of the case.

For /recursive/ bindings it's more difficult, to call arityType,
because we don't have an ArityType to put in the envt for the
recursively bound Ids.  So for non-join-point bindings we satisfy
ourselves with mkManifestArityType.  Typically we'll have eta-expanded
the binding (based on an earlier fixpoint calculation in
findRhsArity), so the manifest arity is good.

But for /recursive join points/ things are not so good.
See Note [Arity type for recursive join bindings]

See Note [Arity type for recursive join bindings]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider
  f x = joinrec j 0 = \ a b c -> (a,x,b)
                j n = j (n-1)
        in j 20

Obviously `f` should get arity 4.  But the manifest arity of `j`
is 1.  Remember, we don't eta-expand join points; see
GHC.Core.Opt.Simplify.Utils Note [Do not eta-expand join points].
And the ArityInfo on `j` will be just 1 too; see GHC.Core
Note [Invariants on join points], item (2b).  So using
Note [ArityType for let-bindings] won't work well.

We could do a fixpoint iteration, but that's a heavy hammer
to use in arityType.  So we take advantage of it being a join
point:

* Extend the ArityEnv to bind each of the recursive binders
  (all join points) to `botArityType`.  This means that any
  jump to the join point will return botArityType, which is
  unit for `andArityType`:
      botAritType `andArityType` at = at
  So it's almost as if those "jump" branches didn't exist.

* In this extended env, find the ArityType of each of the RHS, after
  stripping off the join-point binders.

* Use andArityType to combine all these RHS ArityTypes.

* Find the ArityType of the body, also in this strange extended
  environment

* And combine that into the result with andArityType.

In our example, the jump (j 20) will yield Bot, as will the jump
(j (n-1)). We'll 'and' those the ArityType of (\abc. blah).  Good!

In effect we are treating the RHSs as alternative bodies (like
in a case), and ignoring all jumps.  In this way we don't need
to take a fixpoint.  Tricky!

NB: we treat /non-recursive/ join points in the same way, but
actually it works fine to treat them uniformly with normal
let-bindings, and that takes less code.
-}

idArityType :: Id -> ArityType
idArityType :: TyVar -> ArityType
idArityType TyVar
v
  | StrictSig
strict_sig <- TyVar -> StrictSig
idStrictness TyVar
v
  , Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ StrictSig -> Bool
isTopSig StrictSig
strict_sig
  , ([Demand]
ds, Divergence
div) <- StrictSig -> ([Demand], Divergence)
splitStrictSig StrictSig
strict_sig
  , let arity :: Int
arity = forall (t :: * -> *) a. Foldable t => t a -> Int
length [Demand]
ds
  -- Every strictness signature admits an arity signature!
  = [OneShotInfo] -> Divergence -> ArityType
AT (forall a. Int -> [a] -> [a]
take Int
arity [OneShotInfo]
one_shots) Divergence
div
  | Bool
otherwise
  = [OneShotInfo] -> Divergence -> ArityType
AT (forall a. Int -> [a] -> [a]
take (TyVar -> Int
idArity TyVar
v) [OneShotInfo]
one_shots) Divergence
topDiv
  where
    one_shots :: [OneShotInfo]  -- One-shot-ness derived from the type
    one_shots :: [OneShotInfo]
one_shots = Type -> [OneShotInfo]
typeArity (TyVar -> Type
idType TyVar
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 \?1.T.

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{os}.  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.Utils.Monad.
-}

-- | @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 :: Int -> CoreExpr -> CoreExpr
etaExpand   Int
n          CoreExpr
orig_expr = [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand (forall a. Int -> a -> [a]
replicate Int
n OneShotInfo
NoOneShotInfo) CoreExpr
orig_expr
etaExpandAT :: ArityType -> CoreExpr -> CoreExpr
etaExpandAT (AT [OneShotInfo]
oss Divergence
_) CoreExpr
orig_expr = [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand [OneShotInfo]
oss                         CoreExpr
orig_expr
                           -- See Note [Eta expansion with ArityType]

-- etaExpand arity e = res
-- Then 'res' has at least 'arity' lambdas at the top
--    possibly with a cast wrapped around the outside
-- 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)

eta_expand :: [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand :: [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand [OneShotInfo]
one_shots (Cast CoreExpr
expr Coercion
co)
  = CoreExpr -> Coercion -> CoreExpr
mkCast ([OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand [OneShotInfo]
one_shots CoreExpr
expr) Coercion
co

eta_expand [OneShotInfo]
one_shots CoreExpr
orig_expr
  = [OneShotInfo] -> [TyVar] -> CoreExpr -> CoreExpr
go [OneShotInfo]
one_shots [] CoreExpr
orig_expr
  where
      -- Strip off existing lambdas and casts before handing off to mkEtaWW
      -- This is mainly to avoid spending time cloning binders and substituting
      -- when there is actually nothing to do.  It's slightly awkward to deal
      -- with casts here, apart from the topmost one, and they are rare, so
      -- if we find one we just hand off to mkEtaWW anyway
      -- Note [Eta expansion and SCCs]
    go :: [OneShotInfo] -> [TyVar] -> CoreExpr -> CoreExpr
go [] [TyVar]
_ CoreExpr
_ = CoreExpr
orig_expr  -- Already has the specified arity; no-op

    go oss :: [OneShotInfo]
oss@(OneShotInfo
_:[OneShotInfo]
oss1) [TyVar]
vs (Lam TyVar
v CoreExpr
body)
      | TyVar -> Bool
isTyVar TyVar
v = [OneShotInfo] -> [TyVar] -> CoreExpr -> CoreExpr
go [OneShotInfo]
oss  (TyVar
vforall a. a -> [a] -> [a]
:[TyVar]
vs) CoreExpr
body
      | Bool
otherwise = [OneShotInfo] -> [TyVar] -> CoreExpr -> CoreExpr
go [OneShotInfo]
oss1 (TyVar
vforall a. a -> [a] -> [a]
:[TyVar]
vs) CoreExpr
body

    go [OneShotInfo]
oss [TyVar]
rev_vs CoreExpr
expr
      = -- pprTrace "ee" (vcat [ppr orig_expr, ppr expr, pprEtaInfos etas]) $
        CoreExpr -> CoreExpr
retick forall a b. (a -> b) -> a -> b
$ EtaInfo -> CoreExpr -> CoreExpr
etaInfoAbs EtaInfo
top_eis forall a b. (a -> b) -> a -> b
$
                 InScopeSet -> CoreExpr -> EtaInfo -> CoreExpr
etaInfoApp InScopeSet
in_scope' CoreExpr
sexpr EtaInfo
eis
      where
          in_scope :: InScopeSet
in_scope = VarSet -> InScopeSet
mkInScopeSet (CoreExpr -> VarSet
exprFreeVars CoreExpr
expr)
          (InScopeSet
in_scope', eis :: EtaInfo
eis@(EI [TyVar]
eta_bndrs MCoercionR
mco))
              = [OneShotInfo]
-> SDoc -> InScopeSet -> Type -> (InScopeSet, EtaInfo)
mkEtaWW [OneShotInfo]
oss (forall a. Outputable a => a -> SDoc
ppr CoreExpr
orig_expr) InScopeSet
in_scope (CoreExpr -> Type
exprType CoreExpr
expr)
          top_bndrs :: [TyVar]
top_bndrs = forall a. [a] -> [a]
reverse [TyVar]
rev_vs
          top_eis :: EtaInfo
top_eis   = [TyVar] -> MCoercionR -> EtaInfo
EI ([TyVar]
top_bndrs forall a. [a] -> [a] -> [a]
++ [TyVar]
eta_bndrs) ([TyVar] -> MCoercionR -> MCoercionR
mkPiMCos [TyVar]
top_bndrs MCoercionR
mco)

          -- Find ticks behind type apps.
          -- See Note [Eta expansion and source notes]
          -- I don't really understand this code SLPJ May 21
          (CoreExpr
expr', [CoreExpr]
args) = forall b. Expr b -> (Expr b, [Expr b])
collectArgs CoreExpr
expr
          ([CoreTickish]
ticks, CoreExpr
expr'') = forall b.
(CoreTickish -> Bool) -> Expr b -> ([CoreTickish], Expr b)
stripTicksTop forall (pass :: TickishPass). GenTickish pass -> Bool
tickishFloatable CoreExpr
expr'
          sexpr :: CoreExpr
sexpr = forall b. Expr b -> [Expr b] -> Expr b
mkApps CoreExpr
expr'' [CoreExpr]
args
          retick :: CoreExpr -> CoreExpr
retick CoreExpr
expr = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr CoreTickish -> CoreExpr -> CoreExpr
mkTick CoreExpr
expr [CoreTickish]
ticks

{- *********************************************************************
*                                                                      *
              The EtaInfo mechanism
          mkEtaWW, etaInfoAbs, etaInfoApp
*                                                                      *
********************************************************************* -}

{- Note [The EtaInfo mechanism]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Suppose we have (e :: ty) and we want to eta-expand it to arity N.
This what eta_expand does.  We do it in two steps:

1.  mkEtaWW: from 'ty' and 'N' build a EtaInfo which describes
    the shape of the expansion necessary to expand to arity N.

2.  Build the term
       \ v1..vn.  e v1 .. vn
    where those abstractions and applications are described by
    the same EtaInfo.  Specifically we build the term

       etaInfoAbs etas (etaInfoApp in_scope e etas)

   where etas :: EtaInfo
         etaInfoAbs builds the lambdas
         etaInfoApp builds the applictions

   Note that the /same/ EtaInfo drives both etaInfoAbs and etaInfoApp

To a first approximation EtaInfo is just [Var].  But
casts complicate the question.  If we have
   newtype N a = MkN (S -> a)
and
   ty = N (N Int)
then the eta-expansion must look like
        (\x (\y. ((e |> co1) x) |> co2) y)
           |> sym co2)
        |> sym co1
where
  co1 :: N (N Int) ~ S -> N Int
  co2 :: N Int     ~ S -> Int

Blimey!  Look at all those casts.  Moreover, if the type
is very deeply nested (as happens in #18223), the repetition
of types can make the overall term very large.  So there is a big
payoff in cancelling out casts aggressively wherever possible.
(See also Note [No crap in eta-expanded code].)

This matters a lot in etaEInfoApp, where we
* Do beta-reduction on the fly
* Use getArg_maybe to get a cast out of the way,
  so that we can do beta reduction
Together this makes a big difference.  Consider when e is
   case x of
      True  -> (\x -> e1) |> c1
      False -> (\p -> e2) |> c2

When we eta-expand this to arity 1, say, etaInfoAbs will wrap
a (\eta) around the outside and use etaInfoApp to apply each
alternative to 'eta'.  We want to beta-reduce all that junk
away.

#18223 was a dramatic example in which the intermediate term was
grotesquely huge, even though the next Simplifier iteration squashed
it.  Better to kill it at birth.
-}

--------------
data EtaInfo = EI [Var] MCoercionR

-- EI bs co
-- Abstraction:  (\b1 b2 .. bn. []) |> sym co
-- Application:  ([] |> co) b1 b2 .. bn
--
--    e :: T    co :: T ~ (t1 -> t2 -> .. -> tn -> tr)
--    e = (\b1 b2 ... bn. (e |> co) b1 b2 .. bn) |> sym co


etaInfoApp :: InScopeSet -> CoreExpr -> EtaInfo -> CoreExpr
-- (etaInfoApp s e eis) returns something equivalent to
--             (substExpr s e `appliedto` eis)
-- See Note [The EtaInfo mechanism]

etaInfoApp :: InScopeSet -> CoreExpr -> EtaInfo -> CoreExpr
etaInfoApp InScopeSet
in_scope CoreExpr
expr EtaInfo
eis
  = Subst -> CoreExpr -> EtaInfo -> CoreExpr
go (InScopeSet -> Subst
mkEmptySubst InScopeSet
in_scope) CoreExpr
expr EtaInfo
eis
  where
    go :: Subst -> CoreExpr -> EtaInfo -> CoreExpr
    -- 'go' pushed down the eta-infos into the branch of a case
    -- and the body of a let; and does beta-reduction if possible
    --   go subst fun co [b1,..,bn]  returns  (subst(fun) |> co) b1 .. bn
    go :: Subst -> CoreExpr -> EtaInfo -> CoreExpr
go Subst
subst (Tick CoreTickish
t CoreExpr
e) EtaInfo
eis
      = forall b. CoreTickish -> Expr b -> Expr b
Tick (Subst -> CoreTickish -> CoreTickish
substTickish Subst
subst CoreTickish
t) (Subst -> CoreExpr -> EtaInfo -> CoreExpr
go Subst
subst CoreExpr
e EtaInfo
eis)

    go Subst
subst (Cast CoreExpr
e Coercion
co) (EI [TyVar]
bs MCoercionR
mco)
      = Subst -> CoreExpr -> EtaInfo -> CoreExpr
go Subst
subst CoreExpr
e ([TyVar] -> MCoercionR -> EtaInfo
EI [TyVar]
bs (HasCallStack => Subst -> Coercion -> Coercion
Core.substCo Subst
subst Coercion
co Coercion -> MCoercionR -> MCoercionR
`mkTransMCoR` MCoercionR
mco))

    go Subst
subst (Case CoreExpr
e TyVar
b Type
ty [Alt TyVar]
alts) EtaInfo
eis
      = forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case (HasDebugCallStack => Subst -> CoreExpr -> CoreExpr
Core.substExprSC Subst
subst CoreExpr
e) TyVar
b1 Type
ty' [Alt TyVar]
alts'
      where
        (Subst
subst1, TyVar
b1) = Subst -> TyVar -> (Subst, TyVar)
Core.substBndr Subst
subst TyVar
b
        alts' :: [Alt TyVar]
alts' = forall a b. (a -> b) -> [a] -> [b]
map Alt TyVar -> Alt TyVar
subst_alt [Alt TyVar]
alts
        ty' :: Type
ty'   = Type -> EtaInfo -> Type
etaInfoAppTy (Subst -> Type -> Type
Core.substTy Subst
subst Type
ty) EtaInfo
eis
        subst_alt :: Alt TyVar -> Alt TyVar
subst_alt (Alt AltCon
con [TyVar]
bs CoreExpr
rhs) = forall b. AltCon -> [b] -> Expr b -> Alt b
Alt AltCon
con [TyVar]
bs' (Subst -> CoreExpr -> EtaInfo -> CoreExpr
go Subst
subst2 CoreExpr
rhs EtaInfo
eis)
                 where
                  (Subst
subst2,[TyVar]
bs') = Subst -> [TyVar] -> (Subst, [TyVar])
Core.substBndrs Subst
subst1 [TyVar]
bs

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

    -- Beta-reduction if possible, pushing any intervening casts past
    -- the argument. See Note [The EtaInfo mechansim]
    go Subst
subst (Lam TyVar
v CoreExpr
e) (EI (TyVar
b:[TyVar]
bs) MCoercionR
mco)
      | Just (CoreExpr
arg,MCoercionR
mco') <- MCoercionR -> CoreExpr -> Maybe (CoreExpr, MCoercionR)
pushMCoArg MCoercionR
mco (forall b. TyVar -> Expr b
varToCoreExpr TyVar
b)
      = Subst -> CoreExpr -> EtaInfo -> CoreExpr
go (Subst -> TyVar -> CoreExpr -> Subst
Core.extendSubst Subst
subst TyVar
v CoreExpr
arg) CoreExpr
e ([TyVar] -> MCoercionR -> EtaInfo
EI [TyVar]
bs MCoercionR
mco')

    -- Stop pushing down; just wrap the expression up
    go Subst
subst CoreExpr
e (EI [TyVar]
bs MCoercionR
mco) = HasDebugCallStack => Subst -> CoreExpr -> CoreExpr
Core.substExprSC Subst
subst CoreExpr
e
                             CoreExpr -> MCoercionR -> CoreExpr
`mkCastMCo` MCoercionR
mco
                             forall b. Expr b -> [TyVar] -> Expr b
`mkVarApps` [TyVar]
bs

--------------
etaInfoAppTy :: Type -> EtaInfo -> Type
-- If                    e :: ty
-- then   etaInfoApp e eis :: etaInfoApp ty eis
etaInfoAppTy :: Type -> EtaInfo -> Type
etaInfoAppTy Type
ty (EI [TyVar]
bs MCoercionR
mco)
  = SDoc -> Type -> [CoreExpr] -> Type
applyTypeToArgs (String -> SDoc
text String
"etaInfoAppTy") Type
ty1 (forall a b. (a -> b) -> [a] -> [b]
map forall b. TyVar -> Expr b
varToCoreExpr [TyVar]
bs)
  where
    ty1 :: Type
ty1 = case MCoercionR
mco of
             MCoercionR
MRefl  -> Type
ty
             MCo Coercion
co -> Coercion -> Type
coercionRKind Coercion
co

--------------
etaInfoAbs :: EtaInfo -> CoreExpr -> CoreExpr
-- See Note [The EtaInfo mechanism]
etaInfoAbs :: EtaInfo -> CoreExpr -> CoreExpr
etaInfoAbs (EI [TyVar]
bs MCoercionR
mco) CoreExpr
expr = (forall b. [b] -> Expr b -> Expr b
mkLams [TyVar]
bs CoreExpr
expr) CoreExpr -> MCoercionR -> CoreExpr
`mkCastMCo` MCoercionR -> MCoercionR
mkSymMCo MCoercionR
mco

--------------
-- | @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
  = Int -> [OneShotInfo] -> TCvSubst -> Type -> (InScopeSet, EtaInfo)
go Int
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)
       -> (InScopeSet, EtaInfo)
    -- (go [o1,..,on] subst ty) = (in_scope, EI [b1,..,bn] co)
    --    co :: subst(ty) ~ b1_ty -> ... -> bn_ty -> tr

    go :: Int -> [OneShotInfo] -> TCvSubst -> Type -> (InScopeSet, EtaInfo)
go Int
_ [] TCvSubst
subst Type
_       -- See Note [exprArity invariant]
       ----------- Done!  No more expansion needed
       = (TCvSubst -> InScopeSet
getTCvInScope TCvSubst
subst, [TyVar] -> MCoercionR -> EtaInfo
EI [] MCoercionR
MRefl)

    go Int
n oss :: [OneShotInfo]
oss@(OneShotInfo
one_shot:[OneShotInfo]
oss1) TCvSubst
subst Type
ty       -- See Note [exprArity invariant]
       ----------- Forall types  (forall a. ty)
       | Just (TyVar
tcv,Type
ty') <- Type -> Maybe (TyVar, Type)
splitForAllTyCoVar_maybe Type
ty
       , (TCvSubst
subst', TyVar
tcv') <- HasCallStack => TCvSubst -> TyVar -> (TCvSubst, TyVar)
Type.substVarBndr TCvSubst
subst TyVar
tcv
       , let oss' :: [OneShotInfo]
oss' | TyVar -> Bool
isTyVar TyVar
tcv = [OneShotInfo]
oss
                  | Bool
otherwise   = [OneShotInfo]
oss1
         -- A forall can bind a CoVar, in which case
         -- we consume one of the [OneShotInfo]
       , (InScopeSet
in_scope, EI [TyVar]
bs MCoercionR
mco) <- Int -> [OneShotInfo] -> TCvSubst -> Type -> (InScopeSet, EtaInfo)
go Int
n [OneShotInfo]
oss' TCvSubst
subst' Type
ty'
       = (InScopeSet
in_scope, [TyVar] -> MCoercionR -> EtaInfo
EI (TyVar
tcv' forall a. a -> [a] -> [a]
: [TyVar]
bs) (TyVar -> MCoercionR -> MCoercionR
mkHomoForAllMCo TyVar
tcv' MCoercionR
mco))

       ----------- 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 [Representation polymorphism invariants] in GHC.Core
          -- See also test case typecheck/should_run/EtaExpandLevPoly

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

       , let eta_id' :: TyVar
eta_id' = TyVar
eta_id TyVar -> OneShotInfo -> TyVar
`setIdOneShotInfo` OneShotInfo
one_shot
       , (InScopeSet
in_scope, EI [TyVar]
bs MCoercionR
mco) <- Int -> [OneShotInfo] -> TCvSubst -> Type -> (InScopeSet, EtaInfo)
go (Int
nforall a. Num a => a -> a -> a
+Int
1) [OneShotInfo]
oss1 TCvSubst
subst' Type
res_ty
       = (InScopeSet
in_scope, [TyVar] -> MCoercionR -> EtaInfo
EI (TyVar
eta_id' forall a. a -> [a] -> [a]
: [TyVar]
bs) (Scaled Type -> MCoercionR -> MCoercionR
mkFunResMCo (TyVar -> Scaled Type
idScaledType TyVar
eta_id') MCoercionR
mco))

       ----------- 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 (Coercion
co, Type
ty') <- Type -> Maybe (Coercion, Type)
topNormaliseNewType_maybe Type
ty
       , -- co :: ty ~ ty'
         let co' :: Coercion
co' = HasCallStack => TCvSubst -> Coercion -> Coercion
Type.substCo TCvSubst
subst Coercion
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)
       , (InScopeSet
in_scope, EI [TyVar]
bs MCoercionR
mco) <- Int -> [OneShotInfo] -> TCvSubst -> Type -> (InScopeSet, EtaInfo)
go Int
n [OneShotInfo]
oss TCvSubst
subst Type
ty'
         -- mco :: subst(ty') ~ b1_ty -> ... -> bn_ty -> tr
       = (InScopeSet
in_scope, [TyVar] -> MCoercionR -> EtaInfo
EI [TyVar]
bs (Coercion -> MCoercionR -> MCoercionR
mkTransMCoR Coercion
co' MCoercionR
mco))

       | Bool
otherwise       -- We have an expression of arity > 0,
                         -- but its type isn't a function, or a binder
                         -- is representation-polymorphic
       = WARN( True, (ppr orig_oss <+> ppr orig_ty) $$ ppr_orig_expr )
         (TCvSubst -> InScopeSet
getTCvInScope TCvSubst
subst, [TyVar] -> MCoercionR -> EtaInfo
EI [] MCoercionR
MRefl)
        -- 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


{- *********************************************************************
*                                                                      *
              The "push rules"
*                                                                      *
************************************************************************

Here we implement the "push rules" from FC papers:

* The push-argument rules, where we can move a coercion past an argument.
  We have
      (fun |> co) arg
  and we want to transform it to
    (fun arg') |> co'
  for some suitable co' and transformed arg'.

* The PushK rule for data constructors.  We have
       (K e1 .. en) |> co
  and we want to transform to
       (K e1' .. en')
  by pushing the coercion into the arguments
-}

pushCoArgs :: CoercionR -> [CoreArg] -> Maybe ([CoreArg], MCoercion)
pushCoArgs :: Coercion -> [CoreExpr] -> Maybe ([CoreExpr], MCoercionR)
pushCoArgs Coercion
co []         = forall (m :: * -> *) a. Monad m => a -> m a
return ([], Coercion -> MCoercionR
MCo Coercion
co)
pushCoArgs Coercion
co (CoreExpr
arg:[CoreExpr]
args) = do { (CoreExpr
arg',  MCoercionR
m_co1) <- Coercion -> CoreExpr -> Maybe (CoreExpr, MCoercionR)
pushCoArg  Coercion
co  CoreExpr
arg
                              ; case MCoercionR
m_co1 of
                                  MCo Coercion
co1 -> do { ([CoreExpr]
args', MCoercionR
m_co2) <- Coercion -> [CoreExpr] -> Maybe ([CoreExpr], MCoercionR)
pushCoArgs Coercion
co1 [CoreExpr]
args
                                                 ; forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr
arg'forall a. a -> [a] -> [a]
:[CoreExpr]
args', MCoercionR
m_co2) }
                                  MCoercionR
MRefl  -> forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr
arg'forall a. a -> [a] -> [a]
:[CoreExpr]
args, MCoercionR
MRefl) }

pushMCoArg :: MCoercionR -> CoreArg -> Maybe (CoreArg, MCoercion)
pushMCoArg :: MCoercionR -> CoreExpr -> Maybe (CoreExpr, MCoercionR)
pushMCoArg MCoercionR
MRefl    CoreExpr
arg = forall a. a -> Maybe a
Just (CoreExpr
arg, MCoercionR
MRefl)
pushMCoArg (MCo Coercion
co) CoreExpr
arg = Coercion -> CoreExpr -> Maybe (CoreExpr, MCoercionR)
pushCoArg Coercion
co CoreExpr
arg

pushCoArg :: CoercionR -> CoreArg -> Maybe (CoreArg, MCoercion)
-- We have (fun |> co) arg, and we want to transform it to
--         (fun arg) |> co
-- This may fail, e.g. if (fun :: N) where N is a newtype
-- C.f. simplCast in GHC.Core.Opt.Simplify
-- 'co' is always Representational
pushCoArg :: Coercion -> CoreExpr -> Maybe (CoreExpr, MCoercionR)
pushCoArg Coercion
co (Type Type
ty) = do { (Type
ty', MCoercionR
m_co') <- Coercion -> Type -> Maybe (Type, MCoercionR)
pushCoTyArg Coercion
co Type
ty
                            ; forall (m :: * -> *) a. Monad m => a -> m a
return (forall b. Type -> Expr b
Type Type
ty', MCoercionR
m_co') }
pushCoArg Coercion
co CoreExpr
val_arg   = do { (MCoercionR
arg_co, MCoercionR
m_co') <- Coercion -> Maybe (MCoercionR, MCoercionR)
pushCoValArg Coercion
co
                            ; forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr
val_arg CoreExpr -> MCoercionR -> CoreExpr
`mkCastMCo` MCoercionR
arg_co, MCoercionR
m_co') }

pushCoTyArg :: CoercionR -> Type -> Maybe (Type, MCoercionR)
-- We have (fun |> co) @ty
-- Push the coercion through to return
--         (fun @ty') |> co'
-- 'co' is always Representational
-- If the returned coercion is Nothing, then it would have been reflexive;
-- it's faster not to compute it, though.
pushCoTyArg :: Coercion -> Type -> Maybe (Type, MCoercionR)
pushCoTyArg Coercion
co Type
ty
  -- The following is inefficient - don't do `eqType` here, the coercion
  -- optimizer will take care of it. See #14737.
  -- -- | tyL `eqType` tyR
  -- -- = Just (ty, Nothing)

  | Coercion -> Bool
isReflCo Coercion
co
  = forall a. a -> Maybe a
Just (Type
ty, MCoercionR
MRefl)

  | Type -> Bool
isForAllTy_ty Type
tyL
  = ASSERT2( isForAllTy_ty tyR, ppr co $$ ppr ty )
    forall a. a -> Maybe a
Just (Type
ty Type -> Coercion -> Type
`mkCastTy` Coercion
co1, Coercion -> MCoercionR
MCo Coercion
co2)

  | Bool
otherwise
  = forall a. Maybe a
Nothing
  where
    Pair Type
tyL Type
tyR = Coercion -> Pair Type
coercionKind Coercion
co
       -- co :: tyL ~ tyR
       -- tyL = forall (a1 :: k1). ty1
       -- tyR = forall (a2 :: k2). ty2

    co1 :: Coercion
co1 = Coercion -> Coercion
mkSymCo (HasDebugCallStack => Role -> Int -> Coercion -> Coercion
mkNthCo Role
Nominal Int
0 Coercion
co)
       -- co1 :: k2 ~N k1
       -- Note that NthCo can extract a Nominal equality between the
       -- kinds of the types related by a coercion between forall-types.
       -- See the NthCo case in GHC.Core.Lint.

    co2 :: Coercion
co2 = Coercion -> Coercion -> Coercion
mkInstCo Coercion
co (Role -> Type -> Coercion -> Coercion
mkGReflLeftCo Role
Nominal Type
ty Coercion
co1)
        -- co2 :: ty1[ (ty|>co1)/a1 ] ~ ty2[ ty/a2 ]
        -- Arg of mkInstCo is always nominal, hence mkNomReflCo

pushCoValArg :: CoercionR -> Maybe (MCoercionR, MCoercionR)
-- We have (fun |> co) arg
-- Push the coercion through to return
--         (fun (arg |> co_arg)) |> co_res
-- 'co' is always Representational
-- If the second returned Coercion is actually Nothing, then no cast is necessary;
-- the returned coercion would have been reflexive.
pushCoValArg :: Coercion -> Maybe (MCoercionR, MCoercionR)
pushCoValArg Coercion
co
  -- The following is inefficient - don't do `eqType` here, the coercion
  -- optimizer will take care of it. See #14737.
  -- -- | tyL `eqType` tyR
  -- -- = Just (mkRepReflCo arg, Nothing)

  | Coercion -> Bool
isReflCo Coercion
co
  = forall a. a -> Maybe a
Just (MCoercionR
MRefl, MCoercionR
MRefl)

  | Type -> Bool
isFunTy Type
tyL
  , (Coercion
co_mult, Coercion
co1, Coercion
co2) <- HasDebugCallStack =>
Role -> Coercion -> (Coercion, Coercion, Coercion)
decomposeFunCo Role
Representational Coercion
co
  , Coercion -> Bool
isReflexiveCo Coercion
co_mult
    -- We can't push the coercion in the case where co_mult isn't reflexivity:
    -- it could be an unsafe axiom, and losing this information could yield
    -- ill-typed terms. For instance (fun x ::(1) Int -> (fun _ -> () |> co) x)
    -- with co :: (Int -> ()) ~ (Int %1 -> ()), would reduce to (fun x ::(1) Int
    -- -> (fun _ ::(Many) Int -> ()) x) which is ill-typed

              -- If   co  :: (tyL1 -> tyL2) ~ (tyR1 -> tyR2)
              -- then co1 :: tyL1 ~ tyR1
              --      co2 :: tyL2 ~ tyR2
  = ASSERT2( isFunTy tyR, ppr co $$ ppr arg )
    forall a. a -> Maybe a
Just (Coercion -> MCoercionR
coToMCo (Coercion -> Coercion
mkSymCo Coercion
co1), Coercion -> MCoercionR
coToMCo Coercion
co2)
    -- Critically, coToMCo to checks for ReflCo; the whole coercion may not
    -- be reflexive, but either of its components might be
    -- We could use isReflexiveCo, but it's not clear if the benefit
    -- is worth the cost, and it makes no difference in #18223

  | Bool
otherwise
  = forall a. Maybe a
Nothing
  where
    arg :: Type
arg = Type -> Type
funArgTy Type
tyR
    Pair Type
tyL Type
tyR = Coercion -> Pair Type
coercionKind Coercion
co

pushCoercionIntoLambda
    :: HasDebugCallStack => InScopeSet -> Var -> CoreExpr -> CoercionR -> Maybe (Var, CoreExpr)
-- This implements the Push rule from the paper on coercions
--    (\x. e) |> co
-- ===>
--    (\x'. e |> co')
pushCoercionIntoLambda :: HasDebugCallStack =>
InScopeSet
-> TyVar -> CoreExpr -> Coercion -> Maybe (TyVar, CoreExpr)
pushCoercionIntoLambda InScopeSet
in_scope TyVar
x CoreExpr
e Coercion
co
    | ASSERT(not (isTyVar x) && not (isCoVar x)) True
    , Pair Type
s1s2 Type
t1t2 <- Coercion -> Pair Type
coercionKind Coercion
co
    , Just (Type
_, Type
_s1,Type
_s2) <- Type -> Maybe (Type, Type, Type)
splitFunTy_maybe Type
s1s2
    , Just (Type
w1, Type
t1,Type
_t2) <- Type -> Maybe (Type, Type, Type)
splitFunTy_maybe Type
t1t2
    , (Coercion
co_mult, Coercion
co1, Coercion
co2) <- HasDebugCallStack =>
Role -> Coercion -> (Coercion, Coercion, Coercion)
decomposeFunCo Role
Representational Coercion
co
    , Coercion -> Bool
isReflexiveCo Coercion
co_mult
      -- We can't push the coercion in the case where co_mult isn't
      -- reflexivity. See pushCoValArg for more details.
    = let
          -- Should we optimize the coercions here?
          -- Otherwise they might not match too well
          x' :: TyVar
x' = TyVar
x TyVar -> Type -> TyVar
`setIdType` Type
t1 TyVar -> Type -> TyVar
`setIdMult` Type
w1
          in_scope' :: InScopeSet
in_scope' = InScopeSet
in_scope InScopeSet -> TyVar -> InScopeSet
`extendInScopeSet` TyVar
x'
          subst :: Subst
subst = Subst -> TyVar -> CoreExpr -> Subst
extendIdSubst (InScopeSet -> Subst
mkEmptySubst InScopeSet
in_scope')
                                TyVar
x
                                (CoreExpr -> Coercion -> CoreExpr
mkCast (forall b. TyVar -> Expr b
Var TyVar
x') Coercion
co1)
      in forall a. a -> Maybe a
Just (TyVar
x', HasDebugCallStack => Subst -> CoreExpr -> CoreExpr
substExpr Subst
subst CoreExpr
e CoreExpr -> Coercion -> CoreExpr
`mkCast` Coercion
co2)
    | Bool
otherwise
      -- See #21555 / #21577 for a case where this trace fired but the cause was benign
    = -- pprTrace "exprIsLambda_maybe: Unexpected lambda in case" (ppr (Lam x e))
      forall a. Maybe a
Nothing

pushCoDataCon :: DataCon -> [CoreExpr] -> Coercion
              -> Maybe (DataCon
                       , [Type]      -- Universal type args
                       , [CoreExpr]) -- All other args incl existentials
-- Implement the KPush reduction rule as described in "Down with kinds"
-- The transformation applies iff we have
--      (C e1 ... en) `cast` co
-- where co :: (T t1 .. tn) ~ to_ty
-- The left-hand one must be a T, because exprIsConApp returned True
-- but the right-hand one might not be.  (Though it usually will.)
pushCoDataCon :: DataCon
-> [CoreExpr] -> Coercion -> Maybe (DataCon, [Type], [CoreExpr])
pushCoDataCon DataCon
dc [CoreExpr]
dc_args Coercion
co
  | Coercion -> Bool
isReflCo Coercion
co Bool -> Bool -> Bool
|| Type
from_ty Type -> Type -> Bool
`eqType` Type
to_ty  -- try cheap test first
  , let ([CoreExpr]
univ_ty_args, [CoreExpr]
rest_args) = forall b a. [b] -> [a] -> ([a], [a])
splitAtList (DataCon -> [TyVar]
dataConUnivTyVars DataCon
dc) [CoreExpr]
dc_args
  = forall a. a -> Maybe a
Just (DataCon
dc, forall a b. (a -> b) -> [a] -> [b]
map CoreExpr -> Type
exprToType [CoreExpr]
univ_ty_args, [CoreExpr]
rest_args)

  | Just (TyCon
to_tc, [Type]
to_tc_arg_tys) <- HasDebugCallStack => Type -> Maybe (TyCon, [Type])
splitTyConApp_maybe Type
to_ty
  , TyCon
to_tc forall a. Eq a => a -> a -> Bool
== DataCon -> TyCon
dataConTyCon DataCon
dc
        -- These two tests can fail; we might see
        --      (C x y) `cast` (g :: T a ~ S [a]),
        -- where S is a type function.  In fact, exprIsConApp
        -- will probably not be called in such circumstances,
        -- but there's nothing wrong with it

  = let
        tc_arity :: Int
tc_arity       = TyCon -> Int
tyConArity TyCon
to_tc
        dc_univ_tyvars :: [TyVar]
dc_univ_tyvars = DataCon -> [TyVar]
dataConUnivTyVars DataCon
dc
        dc_ex_tcvars :: [TyVar]
dc_ex_tcvars   = DataCon -> [TyVar]
dataConExTyCoVars DataCon
dc
        arg_tys :: [Scaled Type]
arg_tys        = DataCon -> [Scaled Type]
dataConRepArgTys DataCon
dc

        non_univ_args :: [CoreExpr]
non_univ_args  = forall b a. [b] -> [a] -> [a]
dropList [TyVar]
dc_univ_tyvars [CoreExpr]
dc_args
        ([CoreExpr]
ex_args, [CoreExpr]
val_args) = forall b a. [b] -> [a] -> ([a], [a])
splitAtList [TyVar]
dc_ex_tcvars [CoreExpr]
non_univ_args

        -- Make the "Psi" from the paper
        omegas :: [Coercion]
omegas = Int -> Coercion -> [Role] -> [Coercion]
decomposeCo Int
tc_arity Coercion
co (TyCon -> [Role]
tyConRolesRepresentational TyCon
to_tc)
        (Type -> Coercion
psi_subst, [Type]
to_ex_arg_tys)
          = Role
-> [TyVar]
-> [Coercion]
-> [TyVar]
-> [Type]
-> (Type -> Coercion, [Type])
liftCoSubstWithEx Role
Representational
                              [TyVar]
dc_univ_tyvars
                              [Coercion]
omegas
                              [TyVar]
dc_ex_tcvars
                              (forall a b. (a -> b) -> [a] -> [b]
map CoreExpr -> Type
exprToType [CoreExpr]
ex_args)

          -- Cast the value arguments (which include dictionaries)
        new_val_args :: [CoreExpr]
new_val_args = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Type -> CoreExpr -> CoreExpr
cast_arg (forall a b. (a -> b) -> [a] -> [b]
map forall a. Scaled a -> a
scaledThing [Scaled Type]
arg_tys) [CoreExpr]
val_args
        cast_arg :: Type -> CoreExpr -> CoreExpr
cast_arg Type
arg_ty CoreExpr
arg = CoreExpr -> Coercion -> CoreExpr
mkCast CoreExpr
arg (Type -> Coercion
psi_subst Type
arg_ty)

        to_ex_args :: [CoreExpr]
to_ex_args = forall a b. (a -> b) -> [a] -> [b]
map forall b. Type -> Expr b
Type [Type]
to_ex_arg_tys

        dump_doc :: SDoc
dump_doc = [SDoc] -> SDoc
vcat [forall a. Outputable a => a -> SDoc
ppr DataCon
dc,      forall a. Outputable a => a -> SDoc
ppr [TyVar]
dc_univ_tyvars, forall a. Outputable a => a -> SDoc
ppr [TyVar]
dc_ex_tcvars,
                         forall a. Outputable a => a -> SDoc
ppr [Scaled Type]
arg_tys, forall a. Outputable a => a -> SDoc
ppr [CoreExpr]
dc_args,
                         forall a. Outputable a => a -> SDoc
ppr [CoreExpr]
ex_args, forall a. Outputable a => a -> SDoc
ppr [CoreExpr]
val_args, forall a. Outputable a => a -> SDoc
ppr Coercion
co, forall a. Outputable a => a -> SDoc
ppr Type
from_ty, forall a. Outputable a => a -> SDoc
ppr Type
to_ty, forall a. Outputable a => a -> SDoc
ppr TyCon
to_tc
                         , forall a. Outputable a => a -> SDoc
ppr forall a b. (a -> b) -> a -> b
$ TyCon -> [Type] -> Type
mkTyConApp TyCon
to_tc (forall a b. (a -> b) -> [a] -> [b]
map CoreExpr -> Type
exprToType forall a b. (a -> b) -> a -> b
$ forall b a. [b] -> [a] -> [a]
takeList [TyVar]
dc_univ_tyvars [CoreExpr]
dc_args) ]
    in
    ASSERT2( eqType from_ty (mkTyConApp to_tc (map exprToType $ takeList dc_univ_tyvars dc_args)), dump_doc )
    ASSERT2( equalLength val_args arg_tys, dump_doc )
    forall a. a -> Maybe a
Just (DataCon
dc, [Type]
to_tc_arg_tys, [CoreExpr]
to_ex_args forall a. [a] -> [a] -> [a]
++ [CoreExpr]
new_val_args)

  | Bool
otherwise
  = forall a. Maybe a
Nothing

  where
    Pair Type
from_ty Type
to_ty = Coercion -> Pair Type
coercionKind Coercion
co

collectBindersPushingCo :: CoreExpr -> ([Var], CoreExpr)
-- Collect lambda binders, pushing coercions inside if possible
-- E.g.   (\x.e) |> g         g :: <Int> -> blah
--        = (\x. e |> Nth 1 g)
--
-- That is,
--
-- collectBindersPushingCo ((\x.e) |> g) === ([x], e |> Nth 1 g)
collectBindersPushingCo :: CoreExpr -> ([TyVar], CoreExpr)
collectBindersPushingCo CoreExpr
e
  = [TyVar] -> CoreExpr -> ([TyVar], CoreExpr)
go [] CoreExpr
e
  where
    -- Peel off lambdas until we hit a cast.
    go :: [Var] -> CoreExpr -> ([Var], CoreExpr)
    -- The accumulator is in reverse order
    go :: [TyVar] -> CoreExpr -> ([TyVar], CoreExpr)
go [TyVar]
bs (Lam TyVar
b CoreExpr
e)   = [TyVar] -> CoreExpr -> ([TyVar], CoreExpr)
go (TyVar
bforall a. a -> [a] -> [a]
:[TyVar]
bs) CoreExpr
e
    go [TyVar]
bs (Cast CoreExpr
e Coercion
co) = [TyVar] -> CoreExpr -> Coercion -> ([TyVar], CoreExpr)
go_c [TyVar]
bs CoreExpr
e Coercion
co
    go [TyVar]
bs CoreExpr
e           = (forall a. [a] -> [a]
reverse [TyVar]
bs, CoreExpr
e)

    -- We are in a cast; peel off casts until we hit a lambda.
    go_c :: [Var] -> CoreExpr -> CoercionR -> ([Var], CoreExpr)
    -- (go_c bs e c) is same as (go bs e (e |> c))
    go_c :: [TyVar] -> CoreExpr -> Coercion -> ([TyVar], CoreExpr)
go_c [TyVar]
bs (Cast CoreExpr
e Coercion
co1) Coercion
co2 = [TyVar] -> CoreExpr -> Coercion -> ([TyVar], CoreExpr)
go_c [TyVar]
bs CoreExpr
e (Coercion
co1 Coercion -> Coercion -> Coercion
`mkTransCo` Coercion
co2)
    go_c [TyVar]
bs (Lam TyVar
b CoreExpr
e)    Coercion
co  = [TyVar] -> TyVar -> CoreExpr -> Coercion -> ([TyVar], CoreExpr)
go_lam [TyVar]
bs TyVar
b CoreExpr
e Coercion
co
    go_c [TyVar]
bs CoreExpr
e            Coercion
co  = (forall a. [a] -> [a]
reverse [TyVar]
bs, CoreExpr -> Coercion -> CoreExpr
mkCast CoreExpr
e Coercion
co)

    -- We are in a lambda under a cast; peel off lambdas and build a
    -- new coercion for the body.
    go_lam :: [Var] -> Var -> CoreExpr -> CoercionR -> ([Var], CoreExpr)
    -- (go_lam bs b e c) is same as (go_c bs (\b.e) c)
    go_lam :: [TyVar] -> TyVar -> CoreExpr -> Coercion -> ([TyVar], CoreExpr)
go_lam [TyVar]
bs TyVar
b CoreExpr
e Coercion
co
      | TyVar -> Bool
isTyVar TyVar
b
      , let Pair Type
tyL Type
tyR = Coercion -> Pair Type
coercionKind Coercion
co
      , ASSERT( isForAllTy_ty tyL )
        Type -> Bool
isForAllTy_ty Type
tyR
      , Coercion -> Bool
isReflCo (HasDebugCallStack => Role -> Int -> Coercion -> Coercion
mkNthCo Role
Nominal Int
0 Coercion
co)  -- See Note [collectBindersPushingCo]
      = [TyVar] -> CoreExpr -> Coercion -> ([TyVar], CoreExpr)
go_c (TyVar
bforall a. a -> [a] -> [a]
:[TyVar]
bs) CoreExpr
e (Coercion -> Coercion -> Coercion
mkInstCo Coercion
co (Type -> Coercion
mkNomReflCo (TyVar -> Type
mkTyVarTy TyVar
b)))

      | TyVar -> Bool
isCoVar TyVar
b
      , let Pair Type
tyL Type
tyR = Coercion -> Pair Type
coercionKind Coercion
co
      , ASSERT( isForAllTy_co tyL )
        Type -> Bool
isForAllTy_co Type
tyR
      , Coercion -> Bool
isReflCo (HasDebugCallStack => Role -> Int -> Coercion -> Coercion
mkNthCo Role
Nominal Int
0 Coercion
co)  -- See Note [collectBindersPushingCo]
      , let cov :: Coercion
cov = TyVar -> Coercion
mkCoVarCo TyVar
b
      = [TyVar] -> CoreExpr -> Coercion -> ([TyVar], CoreExpr)
go_c (TyVar
bforall a. a -> [a] -> [a]
:[TyVar]
bs) CoreExpr
e (Coercion -> Coercion -> Coercion
mkInstCo Coercion
co (Type -> Coercion
mkNomReflCo (Coercion -> Type
mkCoercionTy Coercion
cov)))

      | TyVar -> Bool
isId TyVar
b
      , let Pair Type
tyL Type
tyR = Coercion -> Pair Type
coercionKind Coercion
co
      , ASSERT( isFunTy tyL) isFunTy tyR
      , (Coercion
co_mult, Coercion
co_arg, Coercion
co_res) <- HasDebugCallStack =>
Role -> Coercion -> (Coercion, Coercion, Coercion)
decomposeFunCo Role
Representational Coercion
co
      , Coercion -> Bool
isReflCo Coercion
co_mult -- See Note [collectBindersPushingCo]
      , Coercion -> Bool
isReflCo Coercion
co_arg  -- See Note [collectBindersPushingCo]
      = [TyVar] -> CoreExpr -> Coercion -> ([TyVar], CoreExpr)
go_c (TyVar
bforall a. a -> [a] -> [a]
:[TyVar]
bs) CoreExpr
e Coercion
co_res

      | Bool
otherwise = (forall a. [a] -> [a]
reverse [TyVar]
bs, CoreExpr -> Coercion -> CoreExpr
mkCast (forall b. b -> Expr b -> Expr b
Lam TyVar
b CoreExpr
e) Coercion
co)

{-

Note [collectBindersPushingCo]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
We just look for coercions of form
   <type> # w -> blah
(and similarly for foralls) to keep this function simple.  We could do
more elaborate stuff, but it'd involve substitution etc.

-}

{- *********************************************************************
*                                                                      *
                Join points
*                                                                      *
********************************************************************* -}

-------------------
-- | 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 :: Int -> CoreExpr -> ([TyVar], CoreExpr)
etaExpandToJoinPoint Int
join_arity CoreExpr
expr
  = Int -> [TyVar] -> CoreExpr -> ([TyVar], CoreExpr)
go Int
join_arity [] CoreExpr
expr
  where
    go :: Int -> [TyVar] -> CoreExpr -> ([TyVar], CoreExpr)
go Int
0 [TyVar]
rev_bs CoreExpr
e         = (forall a. [a] -> [a]
reverse [TyVar]
rev_bs, CoreExpr
e)
    go Int
n [TyVar]
rev_bs (Lam TyVar
b CoreExpr
e) = Int -> [TyVar] -> CoreExpr -> ([TyVar], CoreExpr)
go (Int
nforall a. Num a => a -> a -> a
-Int
1) (TyVar
b forall a. a -> [a] -> [a]
: [TyVar]
rev_bs) CoreExpr
e
    go Int
n [TyVar]
rev_bs CoreExpr
e         = case Int -> CoreExpr -> ([TyVar], CoreExpr)
etaBodyForJoinPoint Int
n CoreExpr
e of
                              ([TyVar]
bs, CoreExpr
e') -> (forall a. [a] -> [a]
reverse [TyVar]
rev_bs forall a. [a] -> [a] -> [a]
++ [TyVar]
bs, CoreExpr
e')

etaExpandToJoinPointRule :: JoinArity -> CoreRule -> CoreRule
etaExpandToJoinPointRule :: Int -> CoreRule -> CoreRule
etaExpandToJoinPointRule Int
_ 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 Int
join_arity rule :: CoreRule
rule@(Rule { ru_bndrs :: CoreRule -> [TyVar]
ru_bndrs = [TyVar]
bndrs, ru_rhs :: CoreRule -> CoreExpr
ru_rhs = CoreExpr
rhs
                                               , ru_args :: CoreRule -> [CoreExpr]
ru_args  = [CoreExpr]
args })
  | Int
need_args forall a. Eq a => a -> a -> Bool
== Int
0
  = CoreRule
rule
  | Int
need_args forall a. Ord a => a -> a -> Bool
< Int
0
  = forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"etaExpandToJoinPointRule" (forall a. Outputable a => a -> SDoc
ppr Int
join_arity SDoc -> SDoc -> SDoc
$$ forall a. Outputable a => a -> SDoc
ppr CoreRule
rule)
  | Bool
otherwise
  = CoreRule
rule { ru_bndrs :: [TyVar]
ru_bndrs = [TyVar]
bndrs forall a. [a] -> [a] -> [a]
++ [TyVar]
new_bndrs, ru_args :: [CoreExpr]
ru_args = [CoreExpr]
args forall a. [a] -> [a] -> [a]
++ [CoreExpr]
new_args
         , ru_rhs :: CoreExpr
ru_rhs = CoreExpr
new_rhs }
  where
    need_args :: Int
need_args = Int
join_arity forall a. Num a => a -> a -> a
- forall (t :: * -> *) a. Foldable t => t a -> Int
length [CoreExpr]
args
    ([TyVar]
new_bndrs, CoreExpr
new_rhs) = Int -> CoreExpr -> ([TyVar], CoreExpr)
etaBodyForJoinPoint Int
need_args CoreExpr
rhs
    new_args :: [CoreExpr]
new_args = forall b. [TyVar] -> [Expr b]
varsToCoreExprs [TyVar]
new_bndrs

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

    init_subst :: CoreExpr -> TCvSubst
init_subst CoreExpr
e = InScopeSet -> TCvSubst
mkEmptyTCvSubst (VarSet -> InScopeSet
mkInScopeSet (CoreExpr -> VarSet
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 :: Int -> TCvSubst -> Scaled Type -> (TCvSubst, TyVar)
freshEtaId Int
n TCvSubst
subst Scaled Type
ty
      = (TCvSubst
subst', TyVar
eta_id')
      where
        Scaled Type
mult' Type
ty' = HasCallStack => TCvSubst -> Scaled Type -> Scaled Type
Type.substScaledTyUnchecked TCvSubst
subst Scaled Type
ty
        eta_id' :: TyVar
eta_id' = InScopeSet -> TyVar -> TyVar
uniqAway (TCvSubst -> InScopeSet
getTCvInScope TCvSubst
subst) forall a b. (a -> b) -> a -> b
$
                  FastString -> Unique -> Type -> Type -> TyVar
mkSysLocalOrCoVar (String -> FastString
fsLit String
"eta") (Int -> Unique
mkBuiltinUnique Int
n) Type
mult' Type
ty'
                  -- "OrCoVar" since this can be used to eta-expand
                  -- coercion abstractions
        subst' :: TCvSubst
subst'  = TCvSubst -> TyVar -> TCvSubst
extendTCvInScope TCvSubst
subst TyVar
eta_id'