CoreArity.hs 42.6 KB
Newer Older
Austin Seipp's avatar
Austin Seipp committed
1 2 3 4
{-
(c) The University of Glasgow 2006
(c) The GRASP/AQUA Project, Glasgow University, 1992-1998

5

6
        Arity and eta expansion
Austin Seipp's avatar
Austin Seipp committed
7
-}
8

9
{-# LANGUAGE CPP #-}
10 11

-- | Arity and eta expansion
12
module CoreArity (
13
        manifestArity, joinRhsArity, exprArity, typeArity,
lukemaurer's avatar
lukemaurer committed
14
        exprEtaExpandArity, findRhsArity, CheapFun, etaExpand,
15 16
        etaExpandToJoinPoint, etaExpandToJoinPointRule,
        exprBotStrictness_maybe
17 18 19 20 21 22 23
    ) where

#include "HsVersions.h"

import CoreSyn
import CoreFVs
import CoreUtils
24
import CoreSubst
25
import Demand
26 27 28 29
import Var
import VarEnv
import Id
import Type
30
import TyCon    ( initRecTc, checkRecTc )
31 32 33
import Coercion
import BasicTypes
import Unique
ian@well-typed.com's avatar
ian@well-typed.com committed
34
import DynFlags ( DynFlags, GeneralFlag(..), gopt )
35 36
import Outputable
import FastString
37
import Pair
38
import Util     ( debugIsOn )
39

Austin Seipp's avatar
Austin Seipp committed
40 41 42
{-
************************************************************************
*                                                                      *
43
              manifestArity and exprArity
Austin Seipp's avatar
Austin Seipp committed
44 45
*                                                                      *
************************************************************************
46 47 48 49 50 51

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
52
compute the ArityInfo for the Id.
53 54 55 56

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

57
        foo = PrelBase.timesInt
58 59 60

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
61
isn't, because foo is blacklisted (used in a rule).
62

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

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

manifestArity :: CoreExpr -> Arity
73 74
-- ^ manifestArity sees how many leading value lambdas there are,
--   after looking through casts
75 76
manifestArity (Lam v e) | isId v        = 1 + manifestArity e
                        | otherwise     = manifestArity e
77
manifestArity (Tick t e) | not (tickishIsCode t) =  manifestArity e
78 79
manifestArity (Cast e _)                = manifestArity e
manifestArity _                         = 0
80

81 82 83 84 85 86 87 88
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 (Lam _ e) = 1 + joinRhsArity e
joinRhsArity _         = 0


89
---------------
90 91 92 93
exprArity :: CoreExpr -> Arity
-- ^ An approximate, fast, version of 'exprEtaExpandArity'
exprArity e = go e
  where
94 95 96
    go (Var v)                     = idArity v
    go (Lam x e) | isId x          = go e + 1
                 | otherwise       = go e
97
    go (Tick t e) | not (tickishIsCode t) = go e
98
    go (Cast e co)                 = trim_arity (go e) (pSnd (coercionKind co))
99
                                        -- Note [exprArity invariant]
100 101 102
    go (App e (Type _))            = go e
    go (App f a) | exprIsTrivial a = (go f - 1) `max` 0
        -- See Note [exprArity for applications]
103
        -- NB: coercions count as a value argument
104

105
    go _                           = 0
106

107 108
    trim_arity :: Arity -> Type -> Arity
    trim_arity arity ty = arity `min` length (typeArity ty)
109

110
---------------
111
typeArity :: Type -> [OneShotInfo]
112 113 114
-- How many value arrows are visible in the type?
-- We look through foralls, and newtypes
-- See Note [exprArity invariant]
115
typeArity ty
116 117
  = go initRecTc ty
  where
118
    go rec_nts ty
Simon Peyton Jones's avatar
Simon Peyton Jones committed
119 120 121 122 123
      | Just (_, ty')  <- splitForAllTy_maybe ty
      = go rec_nts ty'

      | Just (arg,res) <- splitFunTy_maybe ty
      = typeOneShot arg : go rec_nts res
Simon Peyton Jones's avatar
Simon Peyton Jones committed
124

125
      | Just (tc,tys) <- splitTyConApp_maybe ty
126 127 128
      , Just (ty', _) <- instNewTyCon_maybe tc tys
      , Just rec_nts' <- checkRecTc rec_nts tc  -- See Note [Expanding newtypes]
                                                -- in TyCon
129 130
--   , not (isClassTyCon tc)    -- Do not eta-expand through newtype classes
--                              -- See Note [Newtype classes and eta expansion]
131 132
--                              (no longer required)
      = go rec_nts' ty'
133 134 135
        -- 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!
136 137 138
        --
        -- AND through a layer of recursive newtypes
        -- e.g. newtype Stream m a b = Stream (m (Either b (a, Stream m a b)))
139

140 141
      | otherwise
      = []
142 143 144 145 146 147 148

---------------
exprBotStrictness_maybe :: CoreExpr -> Maybe (Arity, StrictSig)
-- A cheap and cheerful function that identifies bottoming functions
-- and gives them a suitable strictness signatures.  It's used during
-- float-out
exprBotStrictness_maybe e
149
  = case getBotArity (arityType env e) of
150 151
        Nothing -> Nothing
        Just ar -> Just (ar, sig ar)
152
  where
153
    env    = AE { ae_ped_bot = True, ae_cheap_fn = \ _ _ -> False }
154
    sig ar = mkClosedStrictSig (replicate ar topDmd) exnRes
155
                  -- For this purpose we can be very simple
156
                  -- exnRes is a bit less aggressive than botRes
157

Austin Seipp's avatar
Austin Seipp committed
158
{-
159 160 161 162
Note [exprArity invariant]
~~~~~~~~~~~~~~~~~~~~~~~~~~
exprArity has the following invariant:

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

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

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

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

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

176 177
Why is this important?  Because
  - In TidyPgm we use exprArity to fix the *final arity* of
178 179 180 181 182 183 184 185 186
    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 TidyPgm, 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!

187 188
Note [Newtype classes and eta expansion]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
189
    NB: this nasty special case is no longer required, because
190 191 192 193
    for newtype classes we don't use the class-op rule mechanism
    at all.  See Note [Single-method classes] in TcInstDcls. SLPJ May 2013

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

198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214
   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:

215
   dCInt :: C Int
216 217 218 219 220 221 222 223 224 225 226 227

   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 Trac #3772, which is
slightly more complicated, does) turn into

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

Gabor Greif's avatar
typos  
Gabor Greif committed
228
and now it is *much* harder for the op/$dfList rule to fire, because
229 230 231 232
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.
233
-------- End of old out of date comments, just for interest -----------
234 235


236 237 238
Note [exprArity for applications]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
When we come to an application we check that the arg is trivial.
239
   eg  f (fac x) does not have arity 2,
240 241 242 243 244 245 246 247 248 249 250 251 252
                 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


Austin Seipp's avatar
Austin Seipp committed
253 254
************************************************************************
*                                                                      *
255
           Computing the "arity" of an expression
Austin Seipp's avatar
Austin Seipp committed
256 257
*                                                                      *
************************************************************************
258

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

265
Or, to put it another way
266

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

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

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

275 276 277
   C[ (\x1 .. xn. e x1 .. xn) ]
         is as efficient as
   C[ e ]
278

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

281 282 283
Note [One-shot lambdas]
~~~~~~~~~~~~~~~~~~~~~~~
Consider one-shot lambdas
284
                let x = expensive in \y z -> E
285 286 287 288 289 290 291 292 293 294 295 296 297
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.
298

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

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

307
Consider also
308
        f = \x -> error "foo"
309
Here, arity 1 is fine.  But if it is
310 311 312
        f = \x -> case x of
                        True  -> error "foo"
                        False -> \y -> x+y
313
then we want to get arity 2.  Technically, this isn't quite right, because
314
        (f True) `seq` 1
315 316 317 318
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.

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

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

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

Simon Peyton Jones's avatar
Simon Peyton Jones committed
334
 (3) Do NOT move a lambda outside a case unless
335
     (a) The scrutinee is ok-for-speculation, or
Simon Peyton Jones's avatar
Simon Peyton Jones committed
336 337
     (b) more liberally: the scrutinee is cheap (e.g. a variable), and
         -fpedantic-bottoms is not enforced (see Trac #2915 for an example)
338 339

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

341 342
4. Note [Newtype arity]
~~~~~~~~~~~~~~~~~~~~~~~~
343 344 345
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

346
        newtype T = MkT ([T] -> Int)
347 348

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

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

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

361 362
Note [The state-transformer hack]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
363 364
Suppose we have
        f = e
365 366
where e has arity n.  Then, if we know from the context that f has
a usage type like
367
        t1 -> ... -> tn -1-> t(n+1) -1-> ... -1-> tm -> ...
368 369
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
370 371 372 373
Consider f = \x. let y = <expensive>
                 in case x of
                      True  -> foo
                      False -> \(s:RealWorld) -> e
374 375 376 377 378 379 380
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.

381 382 383 384 385 386 387 388 389
Note [State hack and bottoming functions]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
It's a terrible idea to use the state hack on a bottoming function.
Here's what happens (Trac #2861):

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

Eta-expand, using the state hack:
390

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

395
Extrude the g2
396

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

400
Discard args for bottomming function
401

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

405 406 407 408 409 410 411 412
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.

413 414 415 416 417
This arose in another guise in Trac #3959.  Here we had

     catch# (throw exn >> return ())

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

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

422
We must *not* eta-expand to
423 424 425 426

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

because 'catch#' expects to get a (# _,_ #) after applying its argument to
427
a State#, not another function!
428 429 430 431 432 433 434 435 436

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
437
with function exprEtaExpandArity).
438

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

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

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

450 451
 * If at = ATop as, and n=length as,
   then expanding 'f' to (\x1..xn. f x1 .. xn) loses no sharing,
Javran Cheng's avatar
Javran Cheng committed
452
   assuming the calls of f respect the one-shot-ness of
453
   its definition.
454

Herbert Valerio Riedel's avatar
Herbert Valerio Riedel committed
455
   NB 'f' is an arbitrary expression, eg (f = g e1 e2).  This 'f'
456
   can have ArityType as ATop, with length as > 0, only if e1 e2 are
457 458 459 460 461 462 463 464 465
   themselves.

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

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


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

-------------------- Main arity code ----------------------------
Austin Seipp's avatar
Austin Seipp committed
480 481
-}

482
-- See Note [ArityType]
483
data ArityType = ATop [OneShotInfo] | ABot Arity
484
     -- There is always an explicit lambda
485
     -- to justify the [OneShot], or the Arity
486

487
vanillaArityType :: ArityType
488
vanillaArityType = ATop []      -- Totally uninformative
489

490
-- ^ The Arity returned is the number of value args the
491
-- expression can be applied to without doing much work
492
exprEtaExpandArity :: DynFlags -> CoreExpr -> Arity
493
-- exprEtaExpandArity is used when eta expanding
494
--      e  ==>  \xy -> e x y
495
exprEtaExpandArity dflags e
496
  = case (arityType env e) of
497 498
      ATop oss -> length oss
      ABot n   -> n
499
  where
500
    env = AE { ae_cheap_fn = mk_cheap_fn dflags isCheapApp
ian@well-typed.com's avatar
ian@well-typed.com committed
501
             , ae_ped_bot  = gopt Opt_PedanticBottoms dflags }
502

503 504
getBotArity :: ArityType -> Maybe Arity
-- Arity of a divergent function
505 506
getBotArity (ABot n) = Just n
getBotArity _        = Nothing
507

508
mk_cheap_fn :: DynFlags -> CheapAppFun -> CheapFun
509
mk_cheap_fn dflags cheap_app
ian@well-typed.com's avatar
ian@well-typed.com committed
510
  | not (gopt Opt_DictsCheap dflags)
511
  = \e _     -> exprIsOk cheap_app e
512
  | otherwise
513
  = \e mb_ty -> exprIsOk cheap_app e
514 515 516
             || case mb_ty of
                  Nothing -> False
                  Just ty -> isDictLikeTy ty
517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542


----------------------
findRhsArity :: DynFlags -> Id -> CoreExpr -> Arity -> Arity
-- This implements the fixpoint loop for arity analysis
-- See Note [Arity analysis]
findRhsArity dflags bndr rhs old_arity
  = go (rhsEtaExpandArity dflags init_cheap_app rhs)
       -- We always call exprEtaExpandArity once, but usually
       -- that produces a result equal to old_arity, and then
       -- we stop right away (since arities should not decrease)
       -- Result: the common case is that there is just one iteration
  where
    init_cheap_app :: CheapAppFun
    init_cheap_app fn n_val_args
      | fn == bndr = True   -- On the first pass, this binder gets infinite arity
      | otherwise  = isCheapApp fn n_val_args

    go :: Arity -> Arity
    go cur_arity
      | cur_arity <= old_arity = cur_arity
      | new_arity == cur_arity = cur_arity
      | otherwise = ASSERT( new_arity < cur_arity )
#ifdef DEBUG
                    pprTrace "Exciting arity"
                       (vcat [ ppr bndr <+> ppr cur_arity <+> ppr new_arity
543
                                                    , ppr rhs])
544 545 546 547 548 549 550 551 552 553 554 555 556 557
#endif
                    go new_arity
      where
        new_arity = rhsEtaExpandArity dflags cheap_app rhs

        cheap_app :: CheapAppFun
        cheap_app fn n_val_args
          | fn == bndr = n_val_args < cur_arity
          | otherwise  = isCheapApp fn n_val_args

-- ^ The Arity returned is the number of value args the
-- expression can be applied to without doing much work
rhsEtaExpandArity :: DynFlags -> CheapAppFun -> CoreExpr -> Arity
-- exprEtaExpandArity is used when eta expanding
558
--      e  ==>  \xy -> e x y
559 560 561
rhsEtaExpandArity dflags cheap_app e
  = case (arityType env e) of
      ATop (os:oss)
562
        | isOneShotInfo os || has_lam e -> 1 + length oss
563 564
                                   -- Don't expand PAPs/thunks
                                   -- Note [Eta expanding thunks]
565 566 567 568
        | otherwise       -> 0
      ATop []             -> 0
      ABot n              -> n
  where
569
    env = AE { ae_cheap_fn = mk_cheap_fn dflags cheap_app
570 571 572 573 574
             , ae_ped_bot  = gopt Opt_PedanticBottoms dflags }

    has_lam (Tick _ e) = has_lam e
    has_lam (Lam b e)  = isId b || has_lam e
    has_lam _          = False
575

Austin Seipp's avatar
Austin Seipp committed
576
{-
577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596 597 598 599 600 601 602 603 604
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 Trac #4138.

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

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


605 606 607 608 609 610 611
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:
612
        foo :: Ord a => a -> ...
613
     foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....
614
        -- So foo has arity 1
615 616 617

     f = \x. foo dInt $ bar x

618
The (foo DInt) is floated out, and makes ineffective a RULE
619 620 621 622 623
     foo (bar x) = ...

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

624
See Note [Dictionary-like types] in TcType.hs for why we use
625 626
isDictLikeTy here rather than isDictTy

627 628
Note [Eta expanding thunks]
~~~~~~~~~~~~~~~~~~~~~~~~~~~
629 630 631 632 633
We don't eta-expand
   * Trivial RHSs     x = y
   * PAPs             x = map g
   * Thunks           f = case y of p -> \x -> blah

634 635
When we see
     f = case y of p -> \x -> blah
636
should we eta-expand it? Well, if 'x' is a one-shot state token
637 638 639
then 'yes' because 'f' will only be applied once.  But otherwise
we (conservatively) say no.  My main reason is to avoid expanding
PAPSs
640 641
        f = g d  ==>  f = \x. g d x
because that might in turn make g inline (if it has an inline pragma),
642
which we might not want.  After all, INLINE pragmas say "inline only
643
when saturated" so we don't want to be too gung-ho about saturating!
Austin Seipp's avatar
Austin Seipp committed
644
-}
645

646
arityLam :: Id -> ArityType -> ArityType
647
arityLam id (ATop as) = ATop (idStateHackOneShotInfo id : as)
648
arityLam _  (ABot n)  = ABot (n+1)
649 650

floatIn :: Bool -> ArityType -> ArityType
651 652
-- We have something like (let x = E in b),
-- where b has the given arity type.
653 654
floatIn _     (ABot n)  = ABot n
floatIn True  (ATop as) = ATop as
655
floatIn False (ATop as) = ATop (takeWhile isOneShotInfo as)
656
   -- If E is not cheap, keep arity only for one-shots
657

658
arityApp :: ArityType -> Bool -> ArityType
659
-- Processing (fun arg) where at is the ArityType of fun,
660
-- Knock off an argument and behave like 'let'
661 662 663 664
arityApp (ABot 0)      _     = ABot 0
arityApp (ABot n)      _     = ABot (n-1)
arityApp (ATop [])     _     = ATop []
arityApp (ATop (_:as)) cheap = floatIn cheap (ATop as)
665 666

andArityType :: ArityType -> ArityType -> ArityType   -- Used for branches of a 'case'
667
andArityType (ABot n1) (ABot n2)  = ABot (n1 `max` n2) -- Note [ABot branches: use max]
668 669 670
andArityType (ATop as)  (ABot _)  = ATop as
andArityType (ABot _)   (ATop bs) = ATop bs
andArityType (ATop as)  (ATop bs) = ATop (as `combine` bs)
671
  where      -- See Note [Combining case branches]
672 673 674
    combine (a:as) (b:bs) = (a `bestOneShot` b) : combine as bs
    combine []     bs     = takeWhile isOneShotInfo bs
    combine as     []     = takeWhile isOneShotInfo as
675

676 677 678 679 680 681 682 683 684
{- Note [ABot branches: use max]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider   case x of
             True  -> \x.  error "urk"
             False -> \xy. error "urk2"

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

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

700 701
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.
Austin Seipp's avatar
Austin Seipp committed
702
-}
703

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

710
data ArityEnv
711
  = AE { ae_cheap_fn :: CheapFun
712 713
       , ae_ped_bot  :: Bool       -- True <=> be pedantic about bottoms
  }
714

715 716 717 718
arityType :: ArityEnv -> CoreExpr -> ArityType

arityType env (Cast e co)
  = case arityType env e of
719 720 721 722 723
      ATop os -> ATop (take co_arity os)
      ABot n  -> ABot (n `min` co_arity)
  where
    co_arity = length (typeArity (pSnd (coercionKind co)))
    -- See Note [exprArity invariant] (2); must be true of
724 725 726
    -- arityType too, since that is how we compute the arity
    -- of variables, and they in turn affect result of exprArity
    -- Trac #5441 is a nice demo
727 728
    -- However, do make sure that ATop -> ATop and ABot -> ABot!
    --   Casts don't affect that part. Getting this wrong provoked #5475
729

730
arityType _ (Var v)
731
  | strict_sig <- idStrictness v
Joachim Breitner's avatar
Joachim Breitner committed
732
  , not $ isTopSig strict_sig
733
  , (ds, res) <- splitStrictSig strict_sig
734 735 736
  , let arity = length ds
  = if isBotRes res then ABot arity
                    else ATop (take arity one_shots)
737
  | otherwise
738
  = ATop (take (idArity v) one_shots)
739
  where
740
    one_shots :: [OneShotInfo]  -- One-shot-ness derived from the type
741
    one_shots = typeArity (idType v)
742

743
        -- Lambdas; increase arity
744
arityType env (Lam x e)
745
  | isId x    = arityLam x (arityType env e)
746
  | otherwise = arityType env e
747

748
        -- Applications; decrease arity, except for types
749 750 751
arityType env (App fun (Type _))
   = arityType env fun
arityType env (App fun arg )
752
   = arityApp (arityType env fun) (ae_cheap_fn env arg Nothing)
753

754 755 756 757 758 759 760 761
        -- 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'
        --
762
arityType env (Case scrut _ _ alts)
763
  | exprIsBottom scrut || null alts
764
  = ABot 0     -- Do not eta expand
765
               -- See Note [Dealing with bottom (1)]
766 767
  | otherwise
  = case alts_type of
768 769 770
     ABot n  | n>0       -> ATop []    -- Don't eta expand
             | otherwise -> ABot 0     -- if RHS is bottomming
                                       -- See Note [Dealing with bottom (2)]
771

Simon Peyton Jones's avatar
Simon Peyton Jones committed
772
     ATop as | not (ae_ped_bot env)    -- See Note [Dealing with bottom (3)]
773
             , ae_cheap_fn env scrut Nothing -> ATop as
Simon Peyton Jones's avatar
Simon Peyton Jones committed
774 775
             | exprOkForSpeculation scrut    -> ATop as
             | otherwise                     -> ATop (takeWhile isOneShotInfo as)
776
  where
777
    alts_type = foldr1 andArityType [arityType env rhs | (_,_,rhs) <- alts]
778

779
arityType env (Let b e)
780
  = floatIn (cheap_bind b) (arityType env e)
781 782 783
  where
    cheap_bind (NonRec b e) = is_cheap (b,e)
    cheap_bind (Rec prs)    = all is_cheap prs
784
    is_cheap (b,e) = ae_cheap_fn env e (Just (idType b))
785

786 787
arityType env (Tick t e)
  | not (tickishIsCode t)     = arityType env e
788

789
arityType _ _ = vanillaArityType
790

Austin Seipp's avatar
Austin Seipp committed
791
{-
792 793
%************************************************************************
%*                                                                      *
794
              The main eta-expander
795 796
%*                                                                      *
%************************************************************************
797

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

802
where (in both cases)
803

804
        * The xi can include type variables
805

806
        * The yi are all value variables
807

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

The biggest reason for doing this is for cases like

813 814 815
        f = \x -> case x of
                    True  -> \y -> e1
                    False -> \y -> e2
816

817
Here we want to get the lambdas together.  A good example is the nofib
818 819 820 821 822 823 824 825 826 827 828 829 830 831
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.
832 833

This means the eta-expander has to do a bit of on-the-fly
834
simplification but it's not too hard.  The alernative, of relying on
835 836 837
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.

838 839 840
Note [Eta expansion and SCCs]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Note that SCCs are not treated specially by etaExpand.  If we have
841 842
        etaExpand 2 (\x -> scc "foo" e)
        = (\xy -> (scc "foo" e) y)
843
So the costs of evaluating 'e' (not 'e y') are attributed to "foo"
Peter Wortmann's avatar
Peter Wortmann committed
844 845 846 847 848 849 850 851 852 853 854 855 856 857 858 859

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.

Austin Seipp's avatar
Austin Seipp committed
860
-}
861

862
-- | @etaExpand n e@ returns an expression with
863 864 865 866
-- the same meaning as @e@, but with arity @n@.
--
-- Given:
--
867
-- > e' = etaExpand n e
868 869 870 871
--
-- We should have that:
--
-- > ty = exprType e = exprType e'
872 873 874
etaExpand :: Arity              -- ^ Result should have this number of value args
          -> CoreExpr           -- ^ Expression to expand
          -> CoreExpr
875 876 877
-- etaExpand arity e = res
-- Then 'res' has at least 'arity' lambdas at the top
--
878
-- etaExpand deals with for-alls. For example:
879
--              etaExpand 1 E
880 881
-- where  E :: forall a. a -> a
-- would return
882
--      (/\b. \y::a -> E b y)
883 884 885 886 887 888 889
--
-- It deals with coerces too, though they are now rare
-- so perhaps the extra code isn't worth it

etaExpand n orig_expr
  = go n orig_expr
  where
890
      -- Strip off existing lambdas and casts
891
      -- Note [Eta expansion and SCCs]
892
    go 0 expr = expr
893
    go n (Lam v body) | isTyVar v = Lam v (go n     body)
894
                      | otherwise = Lam v (go (n-1) body)
Peter Wortmann's avatar
Peter Wortmann committed
895 896 897 898 899 900 901 902 903 904 905 906 907 908 909
    go n (Cast expr co)           = Cast (go n expr) co
    go n expr
      = -- pprTrace "ee" (vcat [ppr orig_expr, ppr expr, ppr etas]) $
        retick $ etaInfoAbs etas (etaInfoApp subst' sexpr etas)
      where
          in_scope = mkInScopeSet (exprFreeVars expr)
          (in_scope', etas) = mkEtaWW n orig_expr in_scope (exprType expr)
          subst' = mkEmptySubst in_scope'

          -- Find ticks behind type apps.
          -- See Note [Eta expansion and source notes]
          (expr', args) = collectArgs expr
          (ticks, expr'') = stripTicksTop tickishFloatable expr'
          sexpr = foldl App expr'' args
          retick expr = foldr mkTick expr ticks
910

911
                                -- Wrapper    Unwrapper
912
--------------