CoreArity.hs 35.5 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 14
        manifestArity, exprArity, typeArity, exprBotStrictness_maybe,
        exprEtaExpandArity, findRhsArity, CheapFun, etaExpand
15 16 17 18 19 20 21
    ) where

#include "HsVersions.h"

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

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

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

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

55
        foo = PrelBase.timesInt
56 57 58

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

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

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

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

79
---------------
80 81 82 83
exprArity :: CoreExpr -> Arity
-- ^ An approximate, fast, version of 'exprEtaExpandArity'
exprArity e = go e
  where
84 85 86
    go (Var v)                     = idArity v
    go (Lam x e) | isId x          = go e + 1
                 | otherwise       = go e
87
    go (Tick t e) | not (tickishIsCode t) = go e
88
    go (Cast e co)                 = trim_arity (go e) (pSnd (coercionKind co))
89
                                        -- Note [exprArity invariant]
90 91 92
    go (App e (Type _))            = go e
    go (App f a) | exprIsTrivial a = (go f - 1) `max` 0
        -- See Note [exprArity for applications]
93
        -- NB: coercions count as a value argument
94

95
    go _                           = 0
96

97 98
    trim_arity :: Arity -> Type -> Arity
    trim_arity arity ty = arity `min` length (typeArity ty)
99

100
---------------
101
typeArity :: Type -> [OneShotInfo]
102 103 104
-- How many value arrows are visible in the type?
-- We look through foralls, and newtypes
-- See Note [exprArity invariant]
105
typeArity ty
106 107
  = go initRecTc ty
  where
108
    go rec_nts ty
109 110 111 112
      | Just (bndr, ty')  <- splitPiTy_maybe ty
      = if isIdLikeBinder bndr
        then typeOneShot (binderType bndr) : go rec_nts ty'
        else go rec_nts ty'
Simon Peyton Jones's avatar
Simon Peyton Jones committed
113

114
      | Just (tc,tys) <- splitTyConApp_maybe ty
115 116 117
      , Just (ty', _) <- instNewTyCon_maybe tc tys
      , Just rec_nts' <- checkRecTc rec_nts tc  -- See Note [Expanding newtypes]
                                                -- in TyCon
118 119
--   , not (isClassTyCon tc)    -- Do not eta-expand through newtype classes
--                              -- See Note [Newtype classes and eta expansion]
120 121
--                              (no longer required)
      = go rec_nts' ty'
122 123 124
        -- 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!
125 126 127
        --
        -- AND through a layer of recursive newtypes
        -- e.g. newtype Stream m a b = Stream (m (Either b (a, Stream m a b)))
128

129 130
      | otherwise
      = []
131 132 133 134 135 136 137

---------------
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
138
  = case getBotArity (arityType env e) of
139 140
        Nothing -> Nothing
        Just ar -> Just (ar, sig ar)
141
  where
142
    env    = AE { ae_ped_bot = True, ae_cheap_fn = \ _ _ -> False }
143
    sig ar = mkClosedStrictSig (replicate ar topDmd) botRes
144
                  -- For this purpose we can be very simple
145

Austin Seipp's avatar
Austin Seipp committed
146
{-
147 148 149 150
Note [exprArity invariant]
~~~~~~~~~~~~~~~~~~~~~~~~~~
exprArity has the following invariant:

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

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

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

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

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

164 165
Why is this important?  Because
  - In TidyPgm we use exprArity to fix the *final arity* of
166 167 168 169 170 171 172 173 174
    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!

175 176
Note [Newtype classes and eta expansion]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
177
    NB: this nasty special case is no longer required, because
178 179 180 181
    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 -----------
182
We have to be careful when eta-expanding through newtypes.  In general
183
it's a good idea, but annoyingly it interacts badly with the class-op
184
rule mechanism.  Consider
185

186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202
   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:

203
   dCInt :: C Int
204 205 206 207 208 209 210 211 212 213 214 215

   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
216
and now it is *much* harder for the op/$dfList rule to fire, because
217 218 219 220
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.
221
-------- End of old out of date comments, just for interest -----------
222 223


224 225 226
Note [exprArity for applications]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
When we come to an application we check that the arg is trivial.
227
   eg  f (fac x) does not have arity 2,
228 229 230 231 232 233 234 235 236 237 238 239 240
                 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
241 242
************************************************************************
*                                                                      *
243
           Computing the "arity" of an expression
Austin Seipp's avatar
Austin Seipp committed
244 245
*                                                                      *
************************************************************************
246

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

253
Or, to put it another way
254

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

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

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

263 264 265
   C[ (\x1 .. xn. e x1 .. xn) ]
         is as efficient as
   C[ e ]
266

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

269 270 271
Note [One-shot lambdas]
~~~~~~~~~~~~~~~~~~~~~~~
Consider one-shot lambdas
272
                let x = expensive in \y z -> E
273 274 275 276 277 278 279 280 281 282 283 284 285
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.
286

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

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

295
Consider also
296
        f = \x -> error "foo"
297
Here, arity 1 is fine.  But if it is
298 299 300
        f = \x -> case x of
                        True  -> error "foo"
                        False -> \y -> x+y
301
then we want to get arity 2.  Technically, this isn't quite right, because
302
        (f True) `seq` 1
303 304 305 306
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.

307 308 309
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:
310

311 312 313
 (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
314

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

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

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

329 330
4. Note [Newtype arity]
~~~~~~~~~~~~~~~~~~~~~~~~
331 332 333
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

334
        newtype T = MkT ([T] -> Int)
335 336

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

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

344
  HOWEVER, note that if you use coerce bogusly you can ge
345
        coerce Int negate
346 347
  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.
348

349 350
Note [The state-transformer hack]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
351 352
Suppose we have
        f = e
353 354
where e has arity n.  Then, if we know from the context that f has
a usage type like
355
        t1 -> ... -> tn -1-> t(n+1) -1-> ... -1-> tm -> ...
356 357
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
358 359 360 361
Consider f = \x. let y = <expensive>
                 in case x of
                      True  -> foo
                      False -> \(s:RealWorld) -> e
362 363 364 365 366 367 368
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.

369 370 371 372 373 374 375 376 377
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:
378

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

383
Extrude the g2
384

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

388
Discard args for bottomming function
389

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

393 394 395 396 397 398 399 400
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.

401 402 403 404 405
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 ()].
406
After inlining (>>) we get
407 408 409

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

410
We must *not* eta-expand to
411 412 413 414

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

because 'catch#' expects to get a (# _,_ #) after applying its argument to
415
a State#, not another function!
416 417 418 419 420 421 422 423 424

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
425
with function exprEtaExpandArity).
426

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

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

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

438 439
 * 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
440
   assuming the calls of f respect the one-shot-ness of
441
   its definition.
442

Herbert Valerio Riedel's avatar
Herbert Valerio Riedel committed
443
   NB 'f' is an arbitrary expression, eg (f = g e1 e2).  This 'f'
444
   can have ArityType as ATop, with length as > 0, only if e1 e2 are
445 446 447 448 449 450 451 452 453
   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#

454 455
Example:
      f = \x\y. let v = <expensive> in
456 457 458 459
          \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.
460 461 462


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

-------------------- Main arity code ----------------------------
Austin Seipp's avatar
Austin Seipp committed
468 469
-}

470
-- See Note [ArityType]
471
data ArityType = ATop [OneShotInfo] | ABot Arity
472
     -- There is always an explicit lambda
473
     -- to justify the [OneShot], or the Arity
474

475
vanillaArityType :: ArityType
476
vanillaArityType = ATop []      -- Totally uninformative
477

478
-- ^ The Arity returned is the number of value args the
479
-- expression can be applied to without doing much work
480
exprEtaExpandArity :: DynFlags -> CoreExpr -> Arity
481
-- exprEtaExpandArity is used when eta expanding
482
--      e  ==>  \xy -> e x y
483
exprEtaExpandArity dflags e
484
  = case (arityType env e) of
485 486
      ATop oss -> length oss
      ABot n   -> n
487
  where
488
    env = AE { ae_cheap_fn = mk_cheap_fn dflags isCheapApp
ian@well-typed.com's avatar
ian@well-typed.com committed
489
             , ae_ped_bot  = gopt Opt_PedanticBottoms dflags }
490

491 492
getBotArity :: ArityType -> Maybe Arity
-- Arity of a divergent function
493 494
getBotArity (ABot n) = Just n
getBotArity _        = Nothing
495

496
mk_cheap_fn :: DynFlags -> CheapAppFun -> CheapFun
497
mk_cheap_fn dflags cheap_app
ian@well-typed.com's avatar
ian@well-typed.com committed
498
  | not (gopt Opt_DictsCheap dflags)
499 500 501 502 503 504
  = \e _     -> exprIsCheap' cheap_app e
  | otherwise
  = \e mb_ty -> exprIsCheap' cheap_app e
             || case mb_ty of
                  Nothing -> False
                  Just ty -> isDictLikeTy ty
505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530


----------------------
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
531
                                                    , ppr rhs])
532 533 534 535 536 537 538 539 540 541 542 543 544 545
#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
546
--      e  ==>  \xy -> e x y
547 548 549
rhsEtaExpandArity dflags cheap_app e
  = case (arityType env e) of
      ATop (os:oss)
550
        | isOneShotInfo os || has_lam e -> 1 + length oss
551 552
                                   -- Don't expand PAPs/thunks
                                   -- Note [Eta expanding thunks]
553 554 555 556
        | otherwise       -> 0
      ATop []             -> 0
      ABot n              -> n
  where
557
    env = AE { ae_cheap_fn = mk_cheap_fn dflags cheap_app
558 559 560 561 562
             , 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
563

Austin Seipp's avatar
Austin Seipp committed
564
{-
565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592
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.


593 594 595 596 597 598 599
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:
600
        foo :: Ord a => a -> ...
601
     foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....
602
        -- So foo has arity 1
603 604 605

     f = \x. foo dInt $ bar x

606
The (foo DInt) is floated out, and makes ineffective a RULE
607 608 609 610 611
     foo (bar x) = ...

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

612
See Note [Dictionary-like types] in TcType.hs for why we use
613 614
isDictLikeTy here rather than isDictTy

615 616
Note [Eta expanding thunks]
~~~~~~~~~~~~~~~~~~~~~~~~~~~
617 618 619 620 621
We don't eta-expand
   * Trivial RHSs     x = y
   * PAPs             x = map g
   * Thunks           f = case y of p -> \x -> blah

622 623
When we see
     f = case y of p -> \x -> blah
624
should we eta-expand it? Well, if 'x' is a one-shot state token
625 626 627
then 'yes' because 'f' will only be applied once.  But otherwise
we (conservatively) say no.  My main reason is to avoid expanding
PAPSs
628 629
        f = g d  ==>  f = \x. g d x
because that might in turn make g inline (if it has an inline pragma),
630
which we might not want.  After all, INLINE pragmas say "inline only
631
when saturated" so we don't want to be too gung-ho about saturating!
Austin Seipp's avatar
Austin Seipp committed
632
-}
633

634
arityLam :: Id -> ArityType -> ArityType
635
arityLam id (ATop as) = ATop (idOneShotInfo id : as)
636
arityLam _  (ABot n)  = ABot (n+1)
637 638

floatIn :: Bool -> ArityType -> ArityType
639 640
-- We have something like (let x = E in b),
-- where b has the given arity type.
641 642
floatIn _     (ABot n)  = ABot n
floatIn True  (ATop as) = ATop as
643
floatIn False (ATop as) = ATop (takeWhile isOneShotInfo as)
644
   -- If E is not cheap, keep arity only for one-shots
645

646
arityApp :: ArityType -> Bool -> ArityType
647
-- Processing (fun arg) where at is the ArityType of fun,
648
-- Knock off an argument and behave like 'let'
649 650 651 652
arityApp (ABot 0)      _     = ABot 0
arityApp (ABot n)      _     = ABot (n-1)
arityApp (ATop [])     _     = ATop []
arityApp (ATop (_:as)) cheap = floatIn cheap (ATop as)
653 654

andArityType :: ArityType -> ArityType -> ArityType   -- Used for branches of a 'case'
655
andArityType (ABot n1) (ABot n2)
656 657 658 659
  = ABot (n1 `min` n2)
andArityType (ATop as)  (ABot _)  = ATop as
andArityType (ABot _)   (ATop bs) = ATop bs
andArityType (ATop as)  (ATop bs) = ATop (as `combine` bs)
660
  where      -- See Note [Combining case branches]
661 662 663
    combine (a:as) (b:bs) = (a `bestOneShot` b) : combine as bs
    combine []     bs     = takeWhile isOneShotInfo bs
    combine as     []     = takeWhile isOneShotInfo as
664

Austin Seipp's avatar
Austin Seipp committed
665
{-
666 667
Note [Combining case branches]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
668
Consider
669 670 671 672 673
  go = \x. let z = go e0
               go2 = \x. case x of
                           True  -> z
                           False -> \s(one-shot). e1
           in go2 x
674
We *really* want to eta-expand go and go2.
675
When combining the barnches of the case we have
676 677
     ATop [] `andAT` ATop [OneShotLam]
and we want to get ATop [OneShotLam].  But if the inner
678 679 680
lambda wasn't one-shot we don't want to do this.
(We need a proper arity analysis to justify that.)

681 682
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
683
-}
684

685
---------------------------
686
type CheapFun = CoreExpr -> Maybe Type -> Bool
687 688 689
        -- 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"
690

691
data ArityEnv
692
  = AE { ae_cheap_fn :: CheapFun
693 694
       , ae_ped_bot  :: Bool       -- True <=> be pedantic about bottoms
  }
695

696 697 698 699
arityType :: ArityEnv -> CoreExpr -> ArityType

arityType env (Cast e co)
  = case arityType env e of
700 701 702 703 704
      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
705 706 707
    -- 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
708 709
    -- However, do make sure that ATop -> ATop and ABot -> ABot!
    --   Casts don't affect that part. Getting this wrong provoked #5475
710

711
arityType _ (Var v)
712
  | strict_sig <- idStrictness v
713
  , not $ isNopSig strict_sig
714
  , (ds, res) <- splitStrictSig strict_sig
715 716 717
  , let arity = length ds
  = if isBotRes res then ABot arity
                    else ATop (take arity one_shots)
718
  | otherwise
719
  = ATop (take (idArity v) one_shots)
720
  where
721
    one_shots :: [OneShotInfo]  -- One-shot-ness derived from the type
722
    one_shots = typeArity (idType v)
723

724
        -- Lambdas; increase arity
725
arityType env (Lam x e)
726
  | isId x    = arityLam x (arityType env e)
727
  | otherwise = arityType env e
728

729
        -- Applications; decrease arity, except for types
730 731 732
arityType env (App fun (Type _))
   = arityType env fun
arityType env (App fun arg )
733
   = arityApp (arityType env fun) (ae_cheap_fn env arg Nothing)
734

735 736 737 738 739 740 741 742
        -- 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'
        --
743
arityType env (Case scrut _ _ alts)
744
  | exprIsBottom scrut || null alts
745
  = ABot 0     -- Do not eta expand
746
               -- See Note [Dealing with bottom (1)]
747 748
  | otherwise
  = case alts_type of
749 750 751
     ABot n  | n>0       -> ATop []    -- Don't eta expand
             | otherwise -> ABot 0     -- if RHS is bottomming
                                       -- See Note [Dealing with bottom (2)]
752

Simon Peyton Jones's avatar
Simon Peyton Jones committed
753
     ATop as | not (ae_ped_bot env)    -- See Note [Dealing with bottom (3)]
754
             , ae_cheap_fn env scrut Nothing -> ATop as
Simon Peyton Jones's avatar
Simon Peyton Jones committed
755 756
             | exprOkForSpeculation scrut    -> ATop as
             | otherwise                     -> ATop (takeWhile isOneShotInfo as)
757
  where
758
    alts_type = foldr1 andArityType [arityType env rhs | (_,_,rhs) <- alts]
759

760
arityType env (Let b e)
761
  = floatIn (cheap_bind b) (arityType env e)
762 763 764
  where
    cheap_bind (NonRec b e) = is_cheap (b,e)
    cheap_bind (Rec prs)    = all is_cheap prs
765
    is_cheap (b,e) = ae_cheap_fn env e (Just (idType b))
766

767 768
arityType env (Tick t e)
  | not (tickishIsCode t)     = arityType env e
769

770
arityType _ _ = vanillaArityType
771

Austin Seipp's avatar
Austin Seipp committed
772
{-
773 774
%************************************************************************
%*                                                                      *
775
              The main eta-expander
776 777
%*                                                                      *
%************************************************************************
778

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

783
where (in both cases)
784

785
        * The xi can include type variables
786

787
        * The yi are all value variables
788

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

The biggest reason for doing this is for cases like

794 795 796
        f = \x -> case x of
                    True  -> \y -> e1
                    False -> \y -> e2
797

798
Here we want to get the lambdas together.  A good example is the nofib
799 800 801 802 803 804 805 806 807 808 809 810 811 812
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.
813 814

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

819 820 821
Note [Eta expansion and SCCs]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Note that SCCs are not treated specially by etaExpand.  If we have
822 823
        etaExpand 2 (\x -> scc "foo" e)
        = (\xy -> (scc "foo" e) y)
824
So the costs of evaluating 'e' (not 'e y') are attributed to "foo"
Peter Wortmann's avatar
Peter Wortmann committed
825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840

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
841
-}
842

843 844 845 846 847 848 849 850 851 852
-- | @etaExpand n us e ty@ returns an expression with
-- the same meaning as @e@, but with arity @n@.
--
-- Given:
--
-- > e' = etaExpand n us e ty
--
-- We should have that:
--
-- > ty = exprType e = exprType e'
853 854 855
etaExpand :: Arity              -- ^ Result should have this number of value args
          -> CoreExpr           -- ^ Expression to expand
          -> CoreExpr
856
-- etaExpand deals with for-alls. For example:
857
--              etaExpand 1 E
858 859
-- where  E :: forall a. a -> a
-- would return
860
--      (/\b. \y::a -> E b y)
861 862 863 864 865 866 867
--
-- 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
868
      -- Strip off existing lambdas and casts
869
      -- Note [Eta expansion and SCCs]
870
    go 0 expr = expr
871
    go n (Lam v body) | isTyVar v = Lam v (go n     body)