{-
(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, mkTopArityType, expandableArityType
   , arityTypeArity, maxWithArity, idArityType

   -- ** Join points
   , etaExpandToJoinPoint, etaExpandToJoinPointRule

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

import GHC.Prelude

import GHC.Driver.Session ( DynFlags, GeneralFlag(..), gopt )

import GHC.Core
import GHC.Core.FVs
import GHC.Core.Utils
import GHC.Core.DataCon
import GHC.Core.TyCon     ( tyConArity )
import GHC.Core.TyCon.RecWalk     ( initRecTc, checkRecTc )
import GHC.Core.Predicate ( isDictTy, isCallStackPredTy )
import GHC.Core.Multiplicity

-- 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.Types.Demand
import GHC.Types.Var
import GHC.Types.Var.Env
import GHC.Types.Id
import GHC.Types.Var.Set
import GHC.Types.Basic
import GHC.Types.Tickish

import GHC.Builtin.Uniques
import GHC.Data.FastString
import GHC.Data.Pair

import GHC.Utils.Constants (debugIsOn)
import GHC.Utils.Outputable
import GHC.Utils.Panic
import GHC.Utils.Panic.Plain
import GHC.Utils.Trace
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 Int -> Int -> Int
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 (CoreTickish -> Bool
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 Int -> Int -> Int
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 Int -> Int -> Int
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 (CoreTickish -> Bool
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)
                                        -- See 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 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) Int -> Int -> Int
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 Int -> Int -> Int
forall a. Ord a => a -> a -> a
`min` [OneShotInfo] -> Int
forall a. [a] -> Int
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 OneShotInfo -> [OneShotInfo] -> [OneShotInfo]
forall a. a -> [a] -> [a]
: RecTcChecker -> Type -> [OneShotInfo]
go RecTcChecker
rec_nts Type
res

      | Just (TyCon
tc,[Type]
tys) <- (() :: Constraint) => Type -> Maybe (TyCon, [Type])
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 and products]
                                                -- 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, DmdSig)
-- 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, DmdSig)
exprBotStrictness_maybe CoreExpr
e
  = case ArityType -> Maybe Int
getBotArity (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
botStrictnessArityEnv CoreExpr
e) of
        Maybe Int
Nothing -> Maybe (Int, DmdSig)
forall a. Maybe a
Nothing
        Just Int
ar -> (Int, DmdSig) -> Maybe (Int, DmdSig)
forall a. a -> Maybe a
Just (Int
ar, Int -> DmdSig
sig Int
ar)
  where
    sig :: Int -> DmdSig
sig Int
ar = [Demand] -> Divergence -> DmdSig
mkClosedDmdSig (Int -> Demand -> [Demand]
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].
--
-- 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
(ArityType -> ArityType -> Bool)
-> (ArityType -> ArityType -> Bool) -> Eq ArityType
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ArityType -> ArityType -> Bool
== :: ArityType -> ArityType -> Bool
$c/= :: ArityType -> ArityType -> Bool
/= :: 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)
    | [OneShotInfo] -> Bool
forall a. [a] -> Bool
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 ((OneShotInfo -> SDoc) -> [OneShotInfo] -> [SDoc]
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 []

mkTopArityType :: [OneShotInfo] -> ArityType
mkTopArityType :: [OneShotInfo] -> ArityType
mkTopArityType [OneShotInfo]
oss = [OneShotInfo] -> Divergence -> ArityType
AT [OneShotInfo]
oss Divergence
topDiv

topArityType :: ArityType
topArityType :: ArityType
topArityType = [OneShotInfo] -> ArityType
mkTopArityType []

-- | The number of value args for the arity type
arityTypeArity :: ArityType -> Arity
arityTypeArity :: ArityType -> Int
arityTypeArity (AT [OneShotInfo]
oss Divergence
_) = [OneShotInfo] -> Int
forall a. [a] -> Int
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 Int -> Int -> Bool
forall a. Ord 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 [OneShotInfo] -> Int -> Bool
forall a. [a] -> Int -> Bool
`lengthAtLeast` Int
ar   = ArityType
at
  | Bool
otherwise                = [OneShotInfo] -> Divergence -> ArityType
AT (Int -> [OneShotInfo] -> [OneShotInfo]
forall a. Int -> [a] -> [a]
take Int
ar ([OneShotInfo] -> [OneShotInfo]) -> [OneShotInfo] -> [OneShotInfo]
forall a b. (a -> b) -> a -> b
$ [OneShotInfo]
oss [OneShotInfo] -> [OneShotInfo] -> [OneShotInfo]
forall a. [a] -> [a] -> [a]
++ OneShotInfo -> [OneShotInfo]
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 [OneShotInfo] -> Int -> Bool
forall a. [a] -> Int -> Bool
`lengthAtMost` Int
ar = ArityType
at
  | Bool
otherwise             = [OneShotInfo] -> Divergence -> ArityType
AT (Int -> [OneShotInfo] -> [OneShotInfo]
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 ((OneShotInfo -> Bool) -> [OneShotInfo] -> [OneShotInfo]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile OneShotInfo -> Bool
isOneShotInfo [OneShotInfo]
oss) Divergence
topDiv
  | Bool
otherwise        = [OneShotInfo] -> Divergence -> ArityType
AT ((OneShotInfo -> Bool) -> [OneShotInfo] -> [OneShotInfo]
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 = ArityEnv -> CoreExpr -> ArityType
arityType (DynFlags -> ArityEnv
etaExpandArityEnv 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 = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int) -> Int -> Maybe Int
forall a b. (a -> b) -> a -> b
$ [OneShotInfo] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [OneShotInfo]
oss
  | Bool
otherwise        = Maybe Int
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
      , [OneShotInfo] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [OneShotInfo]
oss Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
old_arity = ArityType
cur_at -- from above
      | ArityType
next_at ArityType -> ArityType -> Bool
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]
         Bool -> String -> SDoc -> ArityType -> ArityType
forall a. HasCallStack => Bool -> String -> SDoc -> a -> a
warnPprTrace (Bool
debugIsOn Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
2)
            String
"Exciting arity"
            (Int -> SDoc -> SDoc
nest Int
2 (TyVar -> SDoc
forall a. Outputable a => a -> SDoc
ppr TyVar
bndr SDoc -> SDoc -> SDoc
<+> ArityType -> SDoc
forall a. Outputable a => a -> SDoc
ppr ArityType
cur_at SDoc -> SDoc -> SDoc
<+> ArityType -> SDoc
forall a. Outputable a => a -> SDoc
ppr ArityType
next_at SDoc -> SDoc -> SDoc
$$ CoreExpr -> SDoc
forall a. Outputable a => a -> SDoc
ppr CoreExpr
rhs)) (ArityType -> ArityType) -> ArityType -> ArityType
forall a b. (a -> b) -> a -> b
$
            Int -> ArityType -> ArityType
go (Int
nInt -> Int -> Int
forall 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)) $
              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 OneShotInfo -> [OneShotInfo] -> [OneShotInfo]
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 :: ArityType -> ArityType -> ArityType
andArityType :: ArityType -> ArityType -> ArityType
andArityType (AT (OneShotInfo
os1:[OneShotInfo]
oss1) Divergence
div1) (AT (OneShotInfo
os2:[OneShotInfo]
oss2) Divergence
div2)
  | AT [OneShotInfo]
oss' Divergence
div' <- ArityType -> ArityType -> ArityType
andArityType ([OneShotInfo] -> Divergence -> ArityType
AT [OneShotInfo]
oss1 Divergence
div1) ([OneShotInfo] -> Divergence -> ArityType
AT [OneShotInfo]
oss2 Divergence
div2)
  = [OneShotInfo] -> Divergence -> ArityType
AT ((OneShotInfo
os1 OneShotInfo -> OneShotInfo -> OneShotInfo
`bestOneShot` OneShotInfo
os2) OneShotInfo -> [OneShotInfo] -> [OneShotInfo]
forall a. a -> [a] -> [a]
: [OneShotInfo]
oss') Divergence
div' -- See Note [Combining case branches]
andArityType (AT []         Divergence
div1) ArityType
at2
  | Divergence -> Bool
isDeadEndDiv Divergence
div1 = ArityType
at2                  -- Note [ABot branches: max arity wins]
  | Bool
otherwise         = ArityType -> ArityType
takeWhileOneShot ArityType
at2 -- See Note [Combining case branches]
andArityType ArityType
at1                  (AT []         Divergence
div2)
  | Divergence -> Bool
isDeadEndDiv Divergence
div2 = ArityType
at1                  -- Note [ABot branches: max arity wins]
  | Bool
otherwise         = ArityType -> ArityType
takeWhileOneShot ArityType
at1 -- See Note [Combining case branches]

{- 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]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider
  go = \x. let z = go e0
               go2 = \x. case x of
                           True  -> z
                           False -> \s(one-shot). e1
           in go2 x
We *really* want to 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

Note [Eta expanding through CallStacks]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Just as it's good to eta-expand through dictionaries, so it is good to
do so through CallStacks.  #20103 is a case in point, where we got
  foo :: HasCallStack => Int -> Int
  foo = \(d::CallStack). let d2 = pushCallStack blah d in
        \(x:Int). blah

We really want to eta-expand this!  #20103 is quite convincing!
We do this regardless of -fdicts-cheap; it's not really a dictionary.
-}

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

-- | 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'.
  , ArityEnv -> IdSet
ae_joins  :: !IdSet
  -- ^ In-scope join points. See Note [Eta-expansion and join points]
  --   INVARIANT: Disjoint with the domain of 'am_sigs' (if present).
  }

-- | 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, ae_joins :: IdSet
ae_joins = IdSet
emptyVarSet }

-- | The @ArityEnv@ used by 'exprEtaExpandArity'.
etaExpandArityEnv :: DynFlags -> ArityEnv
etaExpandArityEnv :: DynFlags -> ArityEnv
etaExpandArityEnv DynFlags
dflags
  = AE { ae_mode :: AnalysisMode
ae_mode  = EtaExpandArity { 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 }
       , ae_joins :: IdSet
ae_joins = IdSet
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 = IdEnv ArityType
forall a. VarEnv a
emptyVarEnv }
       , ae_joins :: IdSet
ae_joins = IdSet
emptyVarSet }

-- 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 -> IdEnv ArityType -> TyVar -> IdEnv ArityType
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 -> IdEnv ArityType -> [TyVar] -> IdEnv ArityType
forall a. VarEnv a -> [TyVar] -> VarEnv a
delVarEnvList IdEnv ArityType
sigs [TyVar]
ids)
{-# INLINE del_sig_env_list #-}

del_join_env :: JoinId -> ArityEnv -> ArityEnv -- internal!
del_join_env :: TyVar -> ArityEnv -> ArityEnv
del_join_env TyVar
id env :: ArityEnv
env@(AE { ae_joins :: ArityEnv -> IdSet
ae_joins = IdSet
joins })
  = ArityEnv
env { ae_joins :: IdSet
ae_joins = IdSet -> TyVar -> IdSet
delVarSet IdSet
joins TyVar
id }
{-# INLINE del_join_env #-}

del_join_env_list :: [JoinId] -> ArityEnv -> ArityEnv -- internal!
del_join_env_list :: [TyVar] -> ArityEnv -> ArityEnv
del_join_env_list [TyVar]
ids env :: ArityEnv
env@(AE { ae_joins :: ArityEnv -> IdSet
ae_joins = IdSet
joins })
  = ArityEnv
env { ae_joins :: IdSet
ae_joins = IdSet -> [TyVar] -> IdSet
delVarSetList IdSet
joins [TyVar]
ids }
{-# INLINE del_join_env_list #-}

-- end of internal deletion functions

extendJoinEnv :: ArityEnv -> [JoinId] -> ArityEnv
extendJoinEnv :: ArityEnv -> [TyVar] -> ArityEnv
extendJoinEnv env :: ArityEnv
env@(AE { ae_joins :: ArityEnv -> IdSet
ae_joins = IdSet
joins }) [TyVar]
join_ids
  = [TyVar] -> ArityEnv -> ArityEnv
del_sig_env_list [TyVar]
join_ids
  (ArityEnv -> ArityEnv) -> ArityEnv -> ArityEnv
forall a b. (a -> b) -> a -> b
$ ArityEnv
env { ae_joins :: IdSet
ae_joins = IdSet
joins IdSet -> [TyVar] -> IdSet
`extendVarSetList` [TyVar]
join_ids }

extendSigEnv :: ArityEnv -> Id -> ArityType -> ArityEnv
extendSigEnv :: ArityEnv -> TyVar -> ArityType -> ArityEnv
extendSigEnv ArityEnv
env TyVar
id ArityType
ar_ty
  = TyVar -> ArityEnv -> ArityEnv
del_join_env TyVar
id ((IdEnv ArityType -> IdEnv ArityType) -> ArityEnv -> ArityEnv
modifySigEnv (\IdEnv ArityType
sigs -> IdEnv ArityType -> TyVar -> ArityType -> IdEnv ArityType
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_join_env TyVar
id (ArityEnv -> ArityEnv) -> ArityEnv -> ArityEnv
forall a b. (a -> b) -> a -> b
$ 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_join_env_list [TyVar]
ids (ArityEnv -> ArityEnv) -> ArityEnv -> ArityEnv
forall a b. (a -> b) -> a -> b
$ [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                  -> Maybe ArityType
forall a. Maybe a
Nothing
  EtaExpandArity{}               -> Maybe ArityType
forall a. Maybe a
Nothing
  FindRhsArity{ am_sigs :: AnalysisMode -> IdEnv ArityType
am_sigs = IdEnv ArityType
sigs } -> IdEnv ArityType -> TyVar -> Maybe ArityType
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 = case Maybe Type
mb_ty of
                     Maybe Type
Nothing -> Bool
False
                     Just Type
ty -> (AnalysisMode -> Bool
am_dicts_cheap AnalysisMode
mode Bool -> Bool -> Bool
&& Type -> Bool
isDictTy Type
ty)
                                Bool -> Bool -> Bool
|| Type -> Bool
isCallStackPredTy Type
ty
        -- See Note [Eta expanding through dictionaries]
        -- See Note [Eta expanding through CallStacks]

      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 IdEnv ArityType -> TyVar -> Maybe ArityType
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 ArityType.
  -- 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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [OneShotInfo] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [OneShotInfo]
oss -> Bool
True -- Essentially isWorkFreeApp
    | Bool
otherwise -> Bool
False

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

arityType :: ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env (Cast CoreExpr
e Coercion
co)
  = ArityType -> Int -> ArityType
minWithArity (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
e) Int
co_arity -- See Note [Arity trimming]
  where
    co_arity :: Int
co_arity = [OneShotInfo] -> Int
forall a. [a] -> Int
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)
  | TyVar
v TyVar -> IdSet -> Bool
`elemVarSet` ArityEnv -> IdSet
ae_joins ArityEnv
env
  = ArityType
botArityType  -- See Note [Eta-expansion and join points]
  | Just ArityType
at <- ArityEnv -> TyVar -> Maybe ArityType
lookupSigEnv ArityEnv
env TyVar
v -- Local binding
  = ArityType
at
  | Bool
otherwise
  = 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 (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env' CoreExpr
e)
  | Bool
otherwise = 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
_))
   = ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
fun
arityType ArityEnv
env (App CoreExpr
fun CoreExpr
arg )
   = ArityType -> Bool -> ArityType
arityApp (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
fun) (ArityEnv -> CoreExpr -> Maybe Type -> Bool
myExprIsCheap ArityEnv
env CoreExpr
arg Maybe Type
forall a. Maybe a
Nothing)

        -- Case/Let; keep arity if either the expression is cheap
        -- or it's a 1-shot lambda
        -- The former is not really right for Haskell
        --      f x = case x of { (a,b) -> \y. e }
        --  ===>
        --      f x y = case x of { (a,b) -> e }
        -- The difference is observable using 'seq'
        --
arityType ArityEnv
env (Case CoreExpr
scrut TyVar
bndr Type
_ [Alt TyVar]
alts)
  | CoreExpr -> Bool
exprIsDeadEnd CoreExpr
scrut Bool -> Bool -> Bool
|| [Alt TyVar] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Alt TyVar]
alts
  = ArityType
botArityType    -- Do not eta expand. See (1) in Note [Dealing with bottom]
  | Bool -> Bool
not (ArityEnv -> Bool
pedanticBottoms ArityEnv
env)  -- See (2) in Note [Dealing with bottom]
  , ArityEnv -> CoreExpr -> Maybe Type -> Bool
myExprIsCheap ArityEnv
env CoreExpr
scrut (Type -> Maybe Type
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) = ArityEnv -> CoreExpr -> ArityType
arityType (ArityEnv -> [TyVar] -> ArityEnv
delInScopeList ArityEnv
env' [TyVar]
bndrs) CoreExpr
rhs
    alts_type :: ArityType
alts_type = (ArityType -> ArityType -> ArityType) -> [ArityType] -> ArityType
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldr1 ArityType -> ArityType -> ArityType
andArityType ((Alt TyVar -> ArityType) -> [Alt TyVar] -> [ArityType]
forall a b. (a -> b) -> [a] -> [b]
map Alt TyVar -> ArityType
arity_type_alt [Alt TyVar]
alts)

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

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

arityType ArityEnv
env (Let (NonRec TyVar
b CoreExpr
r) CoreExpr
e)
  = Bool -> ArityType -> ArityType
floatIn Bool
cheap_rhs (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env' CoreExpr
e)
  where
    cheap_rhs :: Bool
cheap_rhs = ArityEnv -> CoreExpr -> Maybe Type -> Bool
myExprIsCheap ArityEnv
env CoreExpr
r (Type -> Maybe Type
forall a. a -> Maybe a
Just (TyVar -> Type
idType TyVar
b))
    env' :: ArityEnv
env'      = ArityEnv -> TyVar -> ArityType -> ArityEnv
extendSigEnv ArityEnv
env TyVar
b (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env CoreExpr
r)

arityType ArityEnv
env (Let (Rec [(TyVar, CoreExpr)]
prs) CoreExpr
e)
  = Bool -> ArityType -> ArityType
floatIn (((TyVar, CoreExpr) -> Bool) -> [(TyVar, CoreExpr)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (TyVar, CoreExpr) -> Bool
is_cheap [(TyVar, CoreExpr)]
prs) (ArityEnv -> CoreExpr -> ArityType
arityType ArityEnv
env' CoreExpr
e)
  where
    env' :: ArityEnv
env'           = ArityEnv -> [TyVar] -> ArityEnv
delInScopeList ArityEnv
env (((TyVar, CoreExpr) -> TyVar) -> [(TyVar, CoreExpr)] -> [TyVar]
forall a b. (a -> b) -> [a] -> [b]
map (TyVar, CoreExpr) -> TyVar
forall a b. (a, b) -> a
fst [(TyVar, CoreExpr)]
prs)
    is_cheap :: (TyVar, CoreExpr) -> Bool
is_cheap (TyVar
b,CoreExpr
e) = ArityEnv -> CoreExpr -> Maybe Type -> Bool
myExprIsCheap ArityEnv
env' CoreExpr
e (Type -> Maybe Type
forall a. a -> Maybe a
Just (TyVar -> Type
idType TyVar
b))

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

arityType ArityEnv
_ CoreExpr
_ = ArityType
topArityType

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

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

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

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

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

So we do this:

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

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

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

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

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

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

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

idArityType :: Id -> ArityType
idArityType :: TyVar -> ArityType
idArityType TyVar
v
  | DmdSig
strict_sig <- TyVar -> DmdSig
idDmdSig TyVar
v
  , Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ DmdSig -> Bool
isTopSig DmdSig
strict_sig
  , ([Demand]
ds, Divergence
div) <- DmdSig -> ([Demand], Divergence)
splitDmdSig DmdSig
strict_sig
  , let arity :: Int
arity = [Demand] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Demand]
ds
  -- Every strictness signature admits an arity signature!
  = [OneShotInfo] -> Divergence -> ArityType
AT (Int -> [OneShotInfo] -> [OneShotInfo]
forall a. Int -> [a] -> [a]
take Int
arity [OneShotInfo]
one_shots) Divergence
div
  | Bool
otherwise
  = [OneShotInfo] -> Divergence -> ArityType
AT (Int -> [OneShotInfo] -> [OneShotInfo]
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
etaExpand :: Int -> CoreExpr -> CoreExpr
etaExpand Int
n CoreExpr
orig_expr
  = InScopeSet -> [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand InScopeSet
in_scope (Int -> OneShotInfo -> [OneShotInfo]
forall a. Int -> a -> [a]
replicate Int
n OneShotInfo
NoOneShotInfo) CoreExpr
orig_expr
  where
    in_scope :: InScopeSet
in_scope = {-#SCC "eta_expand:in-scopeX" #-}
               IdSet -> InScopeSet
mkInScopeSet (CoreExpr -> IdSet
exprFreeVars CoreExpr
orig_expr)

etaExpandAT :: InScopeSet -> ArityType -> CoreExpr -> CoreExpr
-- See Note [Eta expansion with ArityType]
--
-- We pass in the InScopeSet from the simplifier to avoid recomputing
-- it here, which can be jolly expensive if the casts are big
-- In #18223 it took 10% of compile time just to do the exprFreeVars!
etaExpandAT :: InScopeSet -> ArityType -> CoreExpr -> CoreExpr
etaExpandAT InScopeSet
in_scope (AT [OneShotInfo]
oss Divergence
_) CoreExpr
orig_expr
  = InScopeSet -> [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand InScopeSet
in_scope [OneShotInfo]
oss CoreExpr
orig_expr

-- 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 :: InScopeSet -> [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand :: InScopeSet -> [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand InScopeSet
in_scope [OneShotInfo]
one_shots (Cast CoreExpr
expr Coercion
co)
  = CoreExpr -> Coercion -> CoreExpr
forall b. Expr b -> Coercion -> Expr b
Cast (InScopeSet -> [OneShotInfo] -> CoreExpr -> CoreExpr
eta_expand InScopeSet
in_scope [OneShotInfo]
one_shots CoreExpr
expr) Coercion
co

eta_expand InScopeSet
in_scope [OneShotInfo]
one_shots CoreExpr
orig_expr
  = InScopeSet -> [OneShotInfo] -> [TyVar] -> CoreExpr -> CoreExpr
go InScopeSet
in_scope [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 :: InScopeSet -> [OneShotInfo] -> [TyVar] -> CoreExpr -> CoreExpr
go InScopeSet
_ [] [TyVar]
_ CoreExpr
_ = CoreExpr
orig_expr  -- Already has the specified arity; no-op

    go InScopeSet
in_scope oss :: [OneShotInfo]
oss@(OneShotInfo
_:[OneShotInfo]
oss1) [TyVar]
vs (Lam TyVar
v CoreExpr
body)
      | TyVar -> Bool
isTyVar TyVar
v = InScopeSet -> [OneShotInfo] -> [TyVar] -> CoreExpr -> CoreExpr
go (InScopeSet
in_scope InScopeSet -> TyVar -> InScopeSet
`extendInScopeSet` TyVar
v) [OneShotInfo]
oss  (TyVar
vTyVar -> [TyVar] -> [TyVar]
forall a. a -> [a] -> [a]
:[TyVar]
vs) CoreExpr
body
      | Bool
otherwise = InScopeSet -> [OneShotInfo] -> [TyVar] -> CoreExpr -> CoreExpr
go (InScopeSet
in_scope InScopeSet -> TyVar -> InScopeSet
`extendInScopeSet` TyVar
v) [OneShotInfo]
oss1 (TyVar
vTyVar -> [TyVar] -> [TyVar]
forall a. a -> [a] -> [a]
:[TyVar]
vs) CoreExpr
body

    go InScopeSet
in_scope [OneShotInfo]
oss [TyVar]
rev_vs CoreExpr
expr
      = -- pprTrace "ee" (vcat [ppr in_scope', ppr top_bndrs, ppr eis]) $
        CoreExpr -> CoreExpr
retick (CoreExpr -> CoreExpr) -> CoreExpr -> CoreExpr
forall a b. (a -> b) -> a -> b
$
        EtaInfo -> CoreExpr -> CoreExpr
etaInfoAbs EtaInfo
top_eis (CoreExpr -> CoreExpr) -> CoreExpr -> CoreExpr
forall a b. (a -> b) -> a -> b
$
        InScopeSet -> CoreExpr -> EtaInfo -> CoreExpr
etaInfoApp InScopeSet
in_scope' CoreExpr
sexpr EtaInfo
eis
      where
          (InScopeSet
in_scope', eis :: EtaInfo
eis@(EI [TyVar]
eta_bndrs MCoercionR
mco))
                    = [OneShotInfo]
-> SDoc -> InScopeSet -> Type -> (InScopeSet, EtaInfo)
mkEtaWW [OneShotInfo]
oss (CoreExpr -> SDoc
forall a. Outputable a => a -> SDoc
ppr CoreExpr
orig_expr) InScopeSet
in_scope ((() :: Constraint) => CoreExpr -> Type
CoreExpr -> Type
exprType CoreExpr
expr)
          top_bndrs :: [TyVar]
top_bndrs = [TyVar] -> [TyVar]
forall a. [a] -> [a]
reverse [TyVar]
rev_vs
          top_eis :: EtaInfo
top_eis   = [TyVar] -> MCoercionR -> EtaInfo
EI ([TyVar]
top_bndrs [TyVar] -> [TyVar] -> [TyVar]
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) = CoreExpr -> (CoreExpr, [CoreExpr])
forall b. Expr b -> (Expr b, [Expr b])
collectArgs CoreExpr
expr
          ([CoreTickish]
ticks, CoreExpr
expr'') = (CoreTickish -> Bool) -> CoreExpr -> ([CoreTickish], CoreExpr)
forall b.
(CoreTickish -> Bool) -> Expr b -> ([CoreTickish], Expr b)
stripTicksTop CoreTickish -> Bool
forall (pass :: TickishPass). GenTickish pass -> Bool
tickishFloatable CoreExpr
expr'
          sexpr :: CoreExpr
sexpr = CoreExpr -> [CoreExpr] -> CoreExpr
forall b. Expr b -> [Expr b] -> Expr b
mkApps CoreExpr
expr'' [CoreExpr]
args
          retick :: CoreExpr -> CoreExpr
retick CoreExpr
expr = (CoreTickish -> CoreExpr -> CoreExpr)
-> CoreExpr -> [CoreTickish] -> CoreExpr
forall a b. (a -> b -> b) -> b -> [a] -> b
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)
     axN :: N a  ~  S -> a
and
   e :: N (N Int)
then the eta-expansion should look like
   (\(x::S) (y::S) -> e |> co x y) |> sym co
where
  co :: N (N Int) ~ S -> S -> Int
  co = axN @(N Int) ; (S -> axN @Int)

We want to get one cast, at the top, to account for all those
nested newtypes. This is expressed by the EtaInfo type:

   data EtaInfo = EI [Var] MCoercionR

Note [Check for reflexive casts in eta expansion]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
It turns out that the casts created by teh above mechanism are often Refl.
When casts are 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 particularly in etaInfoApp, 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.

The crucial spots in etaInfoApp are:
* `checkReflexiveMCo` in the (Cast e co) case of `go`
* `checkReflexiveMCo` in `pushCoArg`
* Less important: checkReflexiveMCo in the final case of `go`
Collectively these make a factor-of-5 difference to the total
allocation of T18223, so take care if you change this stuff!

Example:
   newtype N = MkN (Y->Z)
   f :: X -> N
   f = \(x::X). ((\(y::Y). blah) |> fco)

where fco :: (Y->Z) ~ N

mkEtaWW makes an EtaInfo of (EI [(eta1:X), (eta2:Y)] eta_co
  where
    eta_co :: (X->N) ~ (X->Y->Z)
    eta_co =  (<X> -> nco)
    nco :: N ~ (Y->Z)  -- Comes from topNormaliseNewType_maybe

Now, when we push that eta_co inward in etaInfoApp:
* In the (Cast e co) case, the 'fco' and 'nco' will meet, and
  should cancel.
* When we meet the (\y.e) we want no cast on the y.

-}

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

-- (EI bs co) describes a particular eta-expansion, as follows:
--  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

instance Outputable EtaInfo where
  ppr :: EtaInfo -> SDoc
ppr (EI [TyVar]
vs MCoercionR
mco) = String -> SDoc
text String
"EI" SDoc -> SDoc -> SDoc
<+> [TyVar] -> SDoc
forall a. Outputable a => a -> SDoc
ppr [TyVar]
vs SDoc -> SDoc -> SDoc
<+> SDoc -> SDoc
parens (MCoercionR -> SDoc
forall a. Outputable a => a -> SDoc
ppr MCoercionR
mco)


etaInfoApp :: InScopeSet -> CoreExpr -> EtaInfo -> CoreExpr
-- (etaInfoApp s e (EI bs mco) returns something equivalent to
--             ((substExpr s e) |> mco b1 .. bn)
-- See Note [The EtaInfo mechanism]
--
-- NB: With very deeply nested casts, this function can be expensive
--     In T18223, this function alone costs 15% of allocation, all
--     spent in the calls to substExprSC and substBindSC

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
      = CoreTickish -> CoreExpr -> CoreExpr
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 MCoercionR
mco')
      where
        mco' :: MCoercionR
mco' = MCoercionR -> MCoercionR
checkReflexiveMCo (HasCallStack => Subst -> Coercion -> Coercion
Subst -> Coercion -> Coercion
Core.substCo Subst
subst Coercion
co Coercion -> MCoercionR -> MCoercionR
`mkTransMCoR` MCoercionR
mco)
               -- See Note [Check for reflexive casts in eta expansion]

    go Subst
subst (Case CoreExpr
e TyVar
b Type
ty [Alt TyVar]
alts) EtaInfo
eis
      = CoreExpr -> TyVar -> Type -> [Alt TyVar] -> CoreExpr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case ((() :: Constraint) => Subst -> CoreExpr -> CoreExpr
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' = (Alt TyVar -> Alt TyVar) -> [Alt TyVar] -> [Alt TyVar]
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) = AltCon -> [TyVar] -> CoreExpr -> Alt TyVar
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]
      = Bind TyVar -> CoreExpr -> CoreExpr
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') = (() :: Constraint) => Subst -> Bind TyVar -> (Subst, Bind TyVar)
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 mechanism]
    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 (TyVar -> CoreExpr
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
    -- See Note [Check for reflexive casts in eta expansion]
    go Subst
subst CoreExpr
e (EI [TyVar]
bs MCoercionR
mco) = (() :: Constraint) => Subst -> CoreExpr -> CoreExpr
Subst -> CoreExpr -> CoreExpr
Core.substExprSC Subst
subst CoreExpr
e
                             CoreExpr -> MCoercionR -> CoreExpr
`mkCastMCo` MCoercionR -> MCoercionR
checkReflexiveMCo MCoercionR
mco
                             CoreExpr -> [TyVar] -> CoreExpr
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)
  = (() :: Constraint) => SDoc -> Type -> [CoreExpr] -> Type
SDoc -> Type -> [CoreExpr] -> Type
applyTypeToArgs (String -> SDoc
text String
"etaInfoAppTy") Type
ty1 ((TyVar -> CoreExpr) -> [TyVar] -> [CoreExpr]
forall a b. (a -> b) -> [a] -> [b]
map TyVar -> CoreExpr
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 = ([TyVar] -> CoreExpr -> CoreExpr
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') <- (() :: Constraint) => TCvSubst -> TyVar -> (TCvSubst, TyVar)
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' TyVar -> [TyVar] -> [TyVar]
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
       , Type -> Bool
typeHasFixedRuntimeRep 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 (Type -> Type -> Scaled Type
forall a. Type -> a -> Scaled a
Scaled Type
mult Type
arg_ty)
          -- Avoid free vars of the original expression

       , let eta_id' :: 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
nInt -> Int -> Int
forall 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' TyVar -> [TyVar] -> [TyVar]
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' = (() :: Constraint) => TCvSubst -> Coercion -> Coercion
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
                         -- does not have a fixed runtime representation
       = Bool
-> String -> SDoc -> (InScopeSet, EtaInfo) -> (InScopeSet, EtaInfo)
forall a. HasCallStack => Bool -> String -> SDoc -> a -> a
warnPprTrace Bool
True String
"mkEtaWW" (([OneShotInfo] -> SDoc
forall a. Outputable a => a -> SDoc
ppr [OneShotInfo]
orig_oss SDoc -> SDoc -> SDoc
<+> Type -> SDoc
forall a. Outputable a => a -> SDoc
ppr Type
orig_ty) SDoc -> SDoc -> SDoc
$$ SDoc
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 []         = ([CoreExpr], MCoercionR) -> Maybe ([CoreExpr], MCoercionR)
forall a. a -> Maybe a
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
                                                 ; ([CoreExpr], MCoercionR) -> Maybe ([CoreExpr], MCoercionR)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr
arg'CoreExpr -> [CoreExpr] -> [CoreExpr]
forall a. a -> [a] -> [a]
:[CoreExpr]
args', MCoercionR
m_co2) }
                                  MCoercionR
MRefl  -> ([CoreExpr], MCoercionR) -> Maybe ([CoreExpr], MCoercionR)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr
arg'CoreExpr -> [CoreExpr] -> [CoreExpr]
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 = (CoreExpr, MCoercionR) -> Maybe (CoreExpr, MCoercionR)
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 CoreExpr
arg
  | Type Type
ty <- CoreExpr
arg
  = do { (Type
ty', MCoercionR
m_co') <- Coercion -> Type -> Maybe (Type, MCoercionR)
pushCoTyArg Coercion
co Type
ty
       ; (CoreExpr, MCoercionR) -> Maybe (CoreExpr, MCoercionR)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> CoreExpr
forall b. Type -> Expr b
Type Type
ty', MCoercionR
m_co') }
  | Bool
otherwise
  = do { (MCoercionR
arg_mco, MCoercionR
m_co') <- Coercion -> Maybe (MCoercionR, MCoercionR)
pushCoValArg Coercion
co
       ; let arg_mco' :: MCoercionR
arg_mco' = MCoercionR -> MCoercionR
checkReflexiveMCo MCoercionR
arg_mco
             -- checkReflexiveMCo: see Note [Check for reflexive casts in eta expansion]
             -- The coercion is very often (arg_co -> res_co), but without
             -- the argument coercion actually being ReflCo
       ; (CoreExpr, MCoercionR) -> Maybe (CoreExpr, MCoercionR)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreExpr
arg CoreExpr -> MCoercionR -> CoreExpr
`mkCastMCo` MCoercionR
arg_mco', 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
  = (Type, MCoercionR) -> Maybe (Type, MCoercionR)
forall a. a -> Maybe a
Just (Type
ty, MCoercionR
MRefl)

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

  | Bool
otherwise
  = Maybe (Type, MCoercionR)
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 ((() :: Constraint) => Role -> Int -> Coercion -> Coercion
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
  = (MCoercionR, MCoercionR) -> Maybe (MCoercionR, MCoercionR)
forall a. a -> Maybe a
Just (MCoercionR
MRefl, MCoercionR
MRefl)

  | Type -> Bool
isFunTy Type
tyL
  , (Coercion
co_mult, Coercion
co1, Coercion
co2) <- (() :: Constraint) =>
Role -> Coercion -> (Coercion, Coercion, Coercion)
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
  = Bool
-> SDoc
-> Maybe (MCoercionR, MCoercionR)
-> Maybe (MCoercionR, MCoercionR)
forall a. HasCallStack => Bool -> SDoc -> a -> a
assertPpr (Type -> Bool
isFunTy Type
tyR) (Coercion -> SDoc
forall a. Outputable a => a -> SDoc
ppr Coercion
co SDoc -> SDoc -> SDoc
$$ Type -> SDoc
forall a. Outputable a => a -> SDoc
ppr Type
arg) (Maybe (MCoercionR, MCoercionR) -> Maybe (MCoercionR, MCoercionR))
-> Maybe (MCoercionR, MCoercionR) -> Maybe (MCoercionR, MCoercionR)
forall a b. (a -> b) -> a -> b
$
    (MCoercionR, MCoercionR) -> Maybe (MCoercionR, MCoercionR)
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
  = Maybe (MCoercionR, MCoercionR)
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 :: (() :: Constraint) =>
InScopeSet
-> TyVar -> CoreExpr -> Coercion -> Maybe (TyVar, CoreExpr)
pushCoercionIntoLambda InScopeSet
in_scope TyVar
x CoreExpr
e Coercion
co
    | Bool -> Bool -> Bool
forall a. HasCallStack => Bool -> a -> a
assert (Bool -> Bool
not (TyVar -> Bool
isTyVar TyVar
x) Bool -> Bool -> Bool
&& Bool -> Bool
not (TyVar -> Bool
isCoVar TyVar
x)) Bool
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) <- (() :: Constraint) =>
Role -> Coercion -> (Coercion, Coercion, Coercion)
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 (TyVar -> CoreExpr
forall b. TyVar -> Expr b
Var TyVar
x') (Coercion -> Coercion
mkSymCo Coercion
co1))
            -- We substitute x' for x, except we need to preserve types.
            -- The types are as follows:
            --   x :: s1,  x' :: t1,  co1 :: s1 ~# t1,
            -- so we extend the substitution with x |-> (x' |> sym co1).
      in (TyVar, CoreExpr) -> Maybe (TyVar, CoreExpr)
forall a. a -> Maybe a
Just (TyVar
x', (() :: Constraint) => Subst -> CoreExpr -> CoreExpr
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))
      Maybe (TyVar, CoreExpr)
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) = [TyVar] -> [CoreExpr] -> ([CoreExpr], [CoreExpr])
forall b a. [b] -> [a] -> ([a], [a])
splitAtList (DataCon -> [TyVar]
dataConUnivTyVars DataCon
dc) [CoreExpr]
dc_args
  = (DataCon, [Type], [CoreExpr])
-> Maybe (DataCon, [Type], [CoreExpr])
forall a. a -> Maybe a
Just (DataCon
dc, (CoreExpr -> Type) -> [CoreExpr] -> [Type]
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) <- (() :: Constraint) => Type -> Maybe (TyCon, [Type])
Type -> Maybe (TyCon, [Type])
splitTyConApp_maybe Type
to_ty
  , TyCon
to_tc TyCon -> TyCon -> Bool
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  = [TyVar] -> [CoreExpr] -> [CoreExpr]
forall b a. [b] -> [a] -> [a]
dropList [TyVar]
dc_univ_tyvars [CoreExpr]
dc_args
        ([CoreExpr]
ex_args, [CoreExpr]
val_args) = [TyVar] -> [CoreExpr] -> ([CoreExpr], [CoreExpr])
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
                              ((CoreExpr -> Type) -> [CoreExpr] -> [Type]
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 = (Type -> CoreExpr -> CoreExpr)
-> [Type] -> [CoreExpr] -> [CoreExpr]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Type -> CoreExpr -> CoreExpr
cast_arg ((Scaled Type -> Type) -> [Scaled Type] -> [Type]
forall a b. (a -> b) -> [a] -> [b]
map Scaled Type -> Type
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 = (Type -> CoreExpr) -> [Type] -> [CoreExpr]
forall a b. (a -> b) -> [a] -> [b]
map Type -> CoreExpr
forall b. Type -> Expr b
Type [Type]
to_ex_arg_tys

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

  | Bool
otherwise
  = Maybe (DataCon, [Type], [CoreExpr])
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
bTyVar -> [TyVar] -> [TyVar]
forall 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           = ([TyVar] -> [TyVar]
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  = ([TyVar] -> [TyVar]
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
      , Bool -> Bool -> Bool
forall a. HasCallStack => Bool -> a -> a
assert (Type -> Bool
isForAllTy_ty Type
tyL) (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$
        Type -> Bool
isForAllTy_ty Type
tyR
      , Coercion -> Bool
isReflCo ((() :: Constraint) => Role -> Int -> Coercion -> Coercion
Role -> Int -> Coercion -> Coercion
mkNthCo Role
Nominal Int
0 Coercion
co)  -- See Note [collectBindersPushingCo]
      = [TyVar] -> CoreExpr -> Coercion -> ([TyVar], CoreExpr)
go_c (TyVar
bTyVar -> [TyVar] -> [TyVar]
forall 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
      , Bool -> Bool -> Bool
forall a. HasCallStack => Bool -> a -> a
assert (Type -> Bool
isForAllTy_co Type
tyL) (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$
        Type -> Bool
isForAllTy_co Type
tyR
      , Coercion -> Bool
isReflCo ((() :: Constraint) => Role -> Int -> Coercion -> Coercion
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
bTyVar -> [TyVar] -> [TyVar]
forall 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
      , Bool -> Bool -> Bool
forall a. HasCallStack => Bool -> a -> a
assert (Type -> Bool
isFunTy Type
tyL) (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Type -> Bool
isFunTy Type
tyR
      , (Coercion
co_mult, Coercion
co_arg, Coercion
co_res) <- (() :: Constraint) =>
Role -> Coercion -> (Coercion, Coercion, Coercion)
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
bTyVar -> [TyVar] -> [TyVar]
forall a. a -> [a] -> [a]
:[TyVar]
bs) CoreExpr
e Coercion
co_res

      | Bool
otherwise = ([TyVar] -> [TyVar]
forall a. [a] -> [a]
reverse [TyVar]
bs, CoreExpr -> Coercion -> CoreExpr
mkCast (TyVar -> CoreExpr -> CoreExpr
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         = ([TyVar] -> [TyVar]
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
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (TyVar
b TyVar -> [TyVar] -> [TyVar]
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') -> ([TyVar] -> [TyVar]
forall a. [a] -> [a]
reverse [TyVar]
rev_bs [TyVar] -> [TyVar] -> [TyVar]
forall a. [a] -> [a] -> [a]
++ [TyVar]
bs, CoreExpr
e')

etaExpandToJoinPointRule :: JoinArity -> CoreRule -> CoreRule
etaExpandToJoinPointRule :: Int -> CoreRule -> CoreRule
etaExpandToJoinPointRule Int
_ rule :: CoreRule
rule@(BuiltinRule {})
  = Bool -> String -> SDoc -> CoreRule -> CoreRule
forall a. HasCallStack => Bool -> String -> SDoc -> a -> a
warnPprTrace Bool
True String
"Can't eta-expand built-in rule:" (CoreRule -> SDoc
forall a. Outputable a => a -> SDoc
ppr CoreRule
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 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
  = CoreRule
rule
  | Int
need_args Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
  = String -> SDoc -> CoreRule
forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"etaExpandToJoinPointRule" (Int -> SDoc
forall a. Outputable a => a -> SDoc
ppr Int
join_arity SDoc -> SDoc -> SDoc
$$ CoreRule -> SDoc
forall a. Outputable a => a -> SDoc
ppr CoreRule
rule)
  | Bool
otherwise
  = CoreRule
rule { ru_bndrs :: [TyVar]
ru_bndrs = [TyVar]
bndrs [TyVar] -> [TyVar] -> [TyVar]
forall a. [a] -> [a] -> [a]
++ [TyVar]
new_bndrs, ru_args :: [CoreExpr]
ru_args = [CoreExpr]
args [CoreExpr] -> [CoreExpr] -> [CoreExpr]
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 Int -> Int -> Int
forall a. Num a => a -> a -> a
- [CoreExpr] -> Int
forall a. [a] -> Int
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 = [TyVar] -> [CoreExpr]
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 ((() :: Constraint) => CoreExpr -> Type
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
      = ([TyVar] -> [TyVar]
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') = (() :: Constraint) => TCvSubst -> TyVar -> (TCvSubst, TyVar)
TCvSubst -> TyVar -> (TCvSubst, TyVar)
substVarBndr TCvSubst
subst TyVar
tv
      = Int
-> Type -> TCvSubst -> [TyVar] -> CoreExpr -> ([TyVar], CoreExpr)
go (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Type
res_ty TCvSubst
subst' (TyVar
tv' TyVar -> [TyVar] -> [TyVar]
forall a. a -> [a] -> [a]
: [TyVar]
rev_bs) (CoreExpr
e CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
`App` TyVar -> CoreExpr
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 (Type -> Type -> Scaled Type
forall a. Type -> a -> Scaled a
Scaled Type
mult Type
arg_ty)
      = Int
-> Type -> TCvSubst -> [TyVar] -> CoreExpr -> ([TyVar], CoreExpr)
go (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Type
res_ty TCvSubst
subst' (TyVar
b TyVar -> [TyVar] -> [TyVar]
forall a. a -> [a] -> [a]
: [TyVar]
rev_bs) (CoreExpr
e CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
`App` TyVar -> CoreExpr
forall b. TyVar -> Expr b
Var TyVar
b)
      | Bool
otherwise
      = String -> SDoc -> ([TyVar], CoreExpr)
forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"etaBodyForJoinPoint" (SDoc -> ([TyVar], CoreExpr)) -> SDoc -> ([TyVar], CoreExpr)
forall a b. (a -> b) -> a -> b
$ Int -> SDoc
int Int
need_args SDoc -> SDoc -> SDoc
$$
                                         CoreExpr -> SDoc
forall a. Outputable a => a -> SDoc
ppr CoreExpr
body SDoc -> SDoc -> SDoc
$$ Type -> SDoc
forall a. Outputable a => a -> SDoc
ppr ((() :: Constraint) => CoreExpr -> Type
CoreExpr -> Type
exprType CoreExpr
body)

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



--------------
freshEtaId :: Int -> TCvSubst -> Scaled Type -> (TCvSubst, Id)
-- Make a fresh Id, with specified type (after applying substitution)
-- It should be "fresh" in the sense that it's not in the in-scope set
-- of the TvSubstEnv; and it should itself then be added to the in-scope
-- set of the TvSubstEnv
--
-- The Int is just a reasonable starting point for generating a unique;
-- it does not necessarily have to be unique itself.
freshEtaId :: 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' = (() :: Constraint) => TCvSubst -> Scaled Type -> Scaled Type
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) (TyVar -> TyVar) -> TyVar -> TyVar
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'