compiler/coreSyn/CoreArity.lhs

   1 %
   2 % (c) The University of Glasgow 2006
   3 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
   4 %
   5
   6         Arity and ete expansion
   7
   8 \begin{code}
   9 -- | Arit and eta expansion
  10 module CoreArity (
  11         manifestArity, exprArity, exprBotStrictness_maybe,
  12         exprEtaExpandArity, CheapFun, etaExpand
  13     ) where
  14
  15 #include "HsVersions.h"
  16
  17 import CoreSyn
  18 import CoreFVs
  19 import CoreUtils
  20 import CoreSubst
  21 import Demand
  22 import Var
  23 import VarEnv
  24 import Id
  25 import Type
  26 import TyCon    ( isRecursiveTyCon, isClassTyCon )
  27 import Coercion
  28 import BasicTypes
  29 import Unique
  30 import Outputable
  31 import FastString
  32 import Pair
  33 \end{code}
  34
  35 %************************************************************************
  36 %*                                                                      *
  37               manifestArity and exprArity
  38 %*                                                                      *
  39 %************************************************************************
  40
  41 exprArity is a cheap-and-cheerful version of exprEtaExpandArity.
  42 It tells how many things the expression can be applied to before doing
  43 any work.  It doesn't look inside cases, lets, etc.  The idea is that
  44 exprEtaExpandArity will do the hard work, leaving something that's easy
  45 for exprArity to grapple with.  In particular, Simplify uses exprArity to
  46 compute the ArityInfo for the Id.
  47
  48 Originally I thought that it was enough just to look for top-level lambdas, but
  49 it isn't.  I've seen this
  50
  51         foo = PrelBase.timesInt
  52
  53 We want foo to get arity 2 even though the eta-expander will leave it
  54 unchanged, in the expectation that it'll be inlined.  But occasionally it
  55 isn't, because foo is blacklisted (used in a rule).
  56
  57 Similarly, see the ok_note check in exprEtaExpandArity.  So
  58         f = __inline_me (\x -> e)
  59 won't be eta-expanded.
  60
  61 And in any case it seems more robust to have exprArity be a bit more intelligent.
  62 But note that   (\x y z -> f x y z)
  63 should have arity 3, regardless of f's arity.
  64
  65 \begin{code}
  66 manifestArity :: CoreExpr -> Arity
  67 -- ^ manifestArity sees how many leading value lambdas there are
  68 manifestArity (Lam v e) | isId v        = 1 + manifestArity e
  69                         | otherwise     = manifestArity e
  70 manifestArity (Note n e) | notSccNote n = manifestArity e
  71 manifestArity (Cast e _)                = manifestArity e
  72 manifestArity _                         = 0
  73
  74 ---------------
  75 exprArity :: CoreExpr -> Arity
  76 -- ^ An approximate, fast, version of 'exprEtaExpandArity'
  77 exprArity e = go e
  78   where
  79     go (Var v)                     = idArity v
  80     go (Lam x e) | isId x          = go e + 1
  81                  | otherwise       = go e
  82     go (Note n e) | notSccNote n   = go e
  83     go (Cast e co)                 = go e `min` length (typeArity (pSnd (coercionKind co)))
  84                                         -- Note [exprArity invariant]
  85     go (App e (Type _))            = go e
  86     go (App f a) | exprIsTrivial a = (go f - 1) `max` 0
  87         -- See Note [exprArity for applications]
  88         -- NB: coercions count as a value argument
  89
  90     go _                           = 0
  91
  92
  93 ---------------
  94 typeArity :: Type -> [OneShot]
  95 -- How many value arrows are visible in the type?
  96 -- We look through foralls, and newtypes
  97 -- See Note [exprArity invariant]
  98 typeArity ty
  99   | Just (_, ty')  <- splitForAllTy_maybe ty
 100   = typeArity ty'
 101
 102   | Just (arg,res) <- splitFunTy_maybe ty
 103   = isStateHackType arg : typeArity res
 104
 105   | Just (tc,tys) <- splitTyConApp_maybe ty
 106   , Just (ty', _) <- instNewTyCon_maybe tc tys
 107   , not (isRecursiveTyCon tc)
 108   , not (isClassTyCon tc)       -- Do not eta-expand through newtype classes
 109                                 -- See Note [Newtype classes and eta expansion]
 110   = typeArity ty'
 111         -- Important to look through non-recursive newtypes, so that, eg
 112         --      (f x)   where f has arity 2, f :: Int -> IO ()
 113         -- Here we want to get arity 1 for the result!
 114
 115   | otherwise
 116   = []
 117
 118 ---------------
 119 exprBotStrictness_maybe :: CoreExpr -> Maybe (Arity, StrictSig)
 120 -- A cheap and cheerful function that identifies bottoming functions
 121 -- and gives them a suitable strictness signatures.  It's used during
 122 -- float-out
 123 exprBotStrictness_maybe e
 124   = case getBotArity (arityType is_cheap e) of
 125         Nothing -> Nothing
 126         Just ar -> Just (ar, mkStrictSig (mkTopDmdType (replicate ar topDmd) BotRes))
 127   where
 128     is_cheap _ _ = False  -- Irrelevant for this purpose
 129 \end{code}
 130
 131 Note [exprArity invariant]
 132 ~~~~~~~~~~~~~~~~~~~~~~~~~~
 133 exprArity has the following invariant:
 134
 135   * If typeArity (exprType e) = n,
 136     then manifestArity (etaExpand e n) = n
 137
 138     That is, etaExpand can always expand as much as typeArity says
 139     So the case analysis in etaExpand and in typeArity must match
 140
 141   * exprArity e <= typeArity (exprType e)
 142
 143   * Hence if (exprArity e) = n, then manifestArity (etaExpand e n) = n
 144
 145     That is, if exprArity says "the arity is n" then etaExpand really
 146     can get "n" manifest lambdas to the top.
 147
 148 Why is this important?  Because
 149   - In TidyPgm we use exprArity to fix the *final arity* of
 150     each top-level Id, and in
 151   - In CorePrep we use etaExpand on each rhs, so that the visible lambdas
 152     actually match that arity, which in turn means
 153     that the StgRhs has the right number of lambdas
 154
 155 An alternative would be to do the eta-expansion in TidyPgm, at least
 156 for top-level bindings, in which case we would not need the trim_arity
 157 in exprArity.  That is a less local change, so I'm going to leave it for today!
 158
 159 Note [Newtype classes and eta expansion]
 160 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 161 We have to be careful when eta-expanding through newtypes.  In general
 162 it's a good idea, but annoyingly it interacts badly with the class-op
 163 rule mechanism.  Consider
 164
 165    class C a where { op :: a -> a }
 166    instance C b => C [b] where
 167      op x = ...
 168
 169 These translate to
 170
 171    co :: forall a. (a->a) ~ C a
 172
 173    $copList :: C b -> [b] -> [b]
 174    $copList d x = ...
 175
 176    $dfList :: C b -> C [b]
 177    {-# DFunUnfolding = [$copList] #-}
 178    $dfList d = $copList d |> co@[b]
 179
 180 Now suppose we have:
 181
 182    dCInt :: C Int
 183
 184    blah :: [Int] -> [Int]
 185    blah = op ($dfList dCInt)
 186
 187 Now we want the built-in op/$dfList rule will fire to give
 188    blah = $copList dCInt
 189
 190 But with eta-expansion 'blah' might (and in Trac #3772, which is
 191 slightly more complicated, does) turn into
 192
 193    blah = op (\eta. ($dfList dCInt |> sym co) eta)
 194
 195 and now it is *much* harder for the op/$dfList rule to fire, becuase
 196 exprIsConApp_maybe won't hold of the argument to op.  I considered
 197 trying to *make* it hold, but it's tricky and I gave up.
 198
 199 The test simplCore/should_compile/T3722 is an excellent example.
 200
 201
 202 Note [exprArity for applications]
 203 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 204 When we come to an application we check that the arg is trivial.
 205    eg  f (fac x) does not have arity 2,
 206                  even if f has arity 3!
 207
 208 * We require that is trivial rather merely cheap.  Suppose f has arity 2.
 209   Then    f (Just y)
 210   has arity 0, because if we gave it arity 1 and then inlined f we'd get
 211           let v = Just y in \w. <f-body>
 212   which has arity 0.  And we try to maintain the invariant that we don't
 213   have arity decreases.
 214
 215 *  The `max 0` is important!  (\x y -> f x) has arity 2, even if f is
 216    unknown, hence arity 0
 217
 218
 219 %************************************************************************
 220 %*                                                                      *
 221            Computing the "arity" of an expression
 222 %*                                                                      *
 223 %************************************************************************
 224
 225 Note [Definition of arity]
 226 ~~~~~~~~~~~~~~~~~~~~~~~~~~
 227 The "arity" of an expression 'e' is n if
 228    applying 'e' to *fewer* than n *value* arguments
 229    converges rapidly
 230
 231 Or, to put it another way
 232
 233    there is no work lost in duplicating the partial
 234    application (e x1 .. x(n-1))
 235
 236 In the divegent case, no work is lost by duplicating because if the thing
 237 is evaluated once, that's the end of the program.
 238
 239 Or, to put it another way, in any context C
 240
 241    C[ (\x1 .. xn. e x1 .. xn) ]
 242          is as efficient as
 243    C[ e ]
 244
 245 It's all a bit more subtle than it looks:
 246
 247 Note [Arity of case expressions]
 248 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 249 We treat the arity of
 250         case x of p -> \s -> ...
 251 as 1 (or more) because for I/O ish things we really want to get that
 252 \s to the top.  We are prepared to evaluate x each time round the loop
 253 in order to get that.
 254
 255 This isn't really right in the presence of seq.  Consider
 256         f = \x -> case x of
 257                         True  -> \y -> x+y
 258                         False -> \y -> x-y
 259 Can we eta-expand here?  At first the answer looks like "yes of course", but
 260 consider
 261         (f bot) `seq` 1
 262 This should diverge!  But if we eta-expand, it won't.   Again, we ignore this
 263 "problem", because being scrupulous would lose an important transformation for
 264 many programs.
 265
 266 1.  Note [One-shot lambdas]
 267 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 268 Consider one-shot lambdas
 269                 let x = expensive in \y z -> E
 270 We want this to have arity 1 if the \y-abstraction is a 1-shot lambda.
 271
 272 3.  Note [Dealing with bottom]
 273 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 274 Consider
 275         f = \x -> error "foo"
 276 Here, arity 1 is fine.  But if it is
 277         f = \x -> case x of
 278                         True  -> error "foo"
 279                         False -> \y -> x+y
 280 then we want to get arity 2.  Technically, this isn't quite right, because
 281         (f True) `seq` 1
 282 should diverge, but it'll converge if we eta-expand f.  Nevertheless, we
 283 do so; it improves some programs significantly, and increasing convergence
 284 isn't a bad thing.  Hence the ABot/ATop in ArityType.
 285
 286 4. Note [Newtype arity]
 287 ~~~~~~~~~~~~~~~~~~~~~~~~
 288 Non-recursive newtypes are transparent, and should not get in the way.
 289 We do (currently) eta-expand recursive newtypes too.  So if we have, say
 290
 291         newtype T = MkT ([T] -> Int)
 292
 293 Suppose we have
 294         e = coerce T f
 295 where f has arity 1.  Then: etaExpandArity e = 1;
 296 that is, etaExpandArity looks through the coerce.
 297
 298 When we eta-expand e to arity 1: eta_expand 1 e T
 299 we want to get:                  coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
 300
 301   HOWEVER, note that if you use coerce bogusly you can ge
 302         coerce Int negate
 303   And since negate has arity 2, you might try to eta expand.  But you can't
 304   decopose Int to a function type.   Hence the final case in eta_expand.
 305
 306 Note [The state-transformer hack]
 307 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 308 Suppose we have
 309         f = e
 310 where e has arity n.  Then, if we know from the context that f has
 311 a usage type like
 312         t1 -> ... -> tn -1-> t(n+1) -1-> ... -1-> tm -> ...
 313 then we can expand the arity to m.  This usage type says that
 314 any application (x e1 .. en) will be applied to uniquely to (m-n) more args
 315 Consider f = \x. let y = <expensive>
 316                  in case x of
 317                       True  -> foo
 318                       False -> \(s:RealWorld) -> e
 319 where foo has arity 1.  Then we want the state hack to
 320 apply to foo too, so we can eta expand the case.
 321
 322 Then we expect that if f is applied to one arg, it'll be applied to two
 323 (that's the hack -- we don't really know, and sometimes it's false)
 324 See also Id.isOneShotBndr.
 325
 326 Note [State hack and bottoming functions]
 327 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 328 It's a terrible idea to use the state hack on a bottoming function.
 329 Here's what happens (Trac #2861):
 330
 331   f :: String -> IO T
 332   f = \p. error "..."
 333
 334 Eta-expand, using the state hack:
 335
 336   f = \p. (\s. ((error "...") |> g1) s) |> g2
 337   g1 :: IO T ~ (S -> (S,T))
 338   g2 :: (S -> (S,T)) ~ IO T
 339
 340 Extrude the g2
 341
 342   f' = \p. \s. ((error "...") |> g1) s
 343   f = f' |> (String -> g2)
 344
 345 Discard args for bottomming function
 346
 347   f' = \p. \s. ((error "...") |> g1 |> g3
 348   g3 :: (S -> (S,T)) ~ (S,T)
 349
 350 Extrude g1.g3
 351
 352   f'' = \p. \s. (error "...")
 353   f' = f'' |> (String -> S -> g1.g3)
 354
 355 And now we can repeat the whole loop.  Aargh!  The bug is in applying the
 356 state hack to a function which then swallows the argument.
 357
 358 This arose in another guise in Trac #3959.  Here we had
 359
 360      catch# (throw exn >> return ())
 361
 362 Note that (throw :: forall a e. Exn e => e -> a) is called with [a = IO ()].
 363 After inlining (>>) we get
 364
 365      catch# (\_. throw {IO ()} exn)
 366
 367 We must *not* eta-expand to
 368
 369      catch# (\_ _. throw {...} exn)
 370
 371 because 'catch#' expects to get a (# _,_ #) after applying its argument to
 372 a State#, not another function!
 373
 374 In short, we use the state hack to allow us to push let inside a lambda,
 375 but not to introduce a new lambda.
 376
 377
 378 Note [ArityType]
 379 ~~~~~~~~~~~~~~~~
 380 ArityType is the result of a compositional analysis on expressions,
 381 from which we can decide the real arity of the expression (extracted
 382 with function exprEtaExpandArity).
 383
 384 Here is what the fields mean. If an arbitrary expression 'f' has
 385 ArityType 'at', then
 386
 387  * If at = ABot n, then (f x1..xn) definitely diverges. Partial
 388    applications to fewer than n args may *or may not* diverge.
 389
 390    We allow ourselves to eta-expand bottoming functions, even
 391    if doing so may lose some `seq` sharing,
 392        let x = <expensive> in \y. error (g x y)
 393        ==> \y. let x = <expensive> in error (g x y)
 394
 395  * If at = ATop as, and n=length as,
 396    then expanding 'f' to (\x1..xn. f x1 .. xn) loses no sharing,
 397    assuming the calls of f respect the one-shot-ness of of
 398    its definition.
 399
 400    NB 'f' is an arbitary expression, eg (f = g e1 e2).  This 'f'
 401    can have ArityType as ATop, with length as > 0, only if e1 e2 are
 402    themselves.
 403
 404  * In both cases, f, (f x1), ... (f x1 ... f(n-1)) are definitely
 405    really functions, or bottom, but *not* casts from a data type, in
 406    at least one case branch.  (If it's a function in one case branch but
 407    an unsafe cast from a data type in another, the program is bogus.)
 408    So eta expansion is dynamically ok; see Note [State hack and
 409    bottoming functions], the part about catch#
 410
 411 Example:
 412       f = \x\y. let v = <expensive> in
 413           \s(one-shot) \t(one-shot). blah
 414       'f' has ArityType [ManyShot,ManyShot,OneShot,OneShot]
 415       The one-shot-ness means we can, in effect, push that
 416       'let' inside the \st.
 417
 418
 419 Suppose f = \xy. x+y
 420 Then  f             :: AT [False,False] ATop
 421       f v           :: AT [False]       ATop
 422       f <expensive> :: AT []            ATop
 423
 424 -------------------- Main arity code ----------------------------
 425 \begin{code}
 426 -- See Note [ArityType]
 427 data ArityType = ATop [OneShot] | ABot Arity
 428      -- There is always an explicit lambda
 429      -- to justify the [OneShot], or the Arity
 430
 431 type OneShot = Bool    -- False <=> Know nothing
 432                        -- True  <=> Can definitely float inside this lambda
 433                        -- The 'True' case can arise either because a binder
 434                        -- is marked one-shot, or because it's a state lambda
 435                        -- and we have the state hack on
 436
 437 vanillaArityType :: ArityType
 438 vanillaArityType = ATop []      -- Totally uninformative
 439
 440 -- ^ The Arity returned is the number of value args the [_$_]
 441 -- expression can be applied to without doing much work
 442 exprEtaExpandArity :: CheapFun -> CoreExpr -> Arity
 443 -- exprEtaExpandArity is used when eta expanding
 444 --      e  ==>  \xy -> e x y
 445 exprEtaExpandArity cheap_fun e
 446   = case (arityType cheap_fun e) of
 447       ATop (os:oss)
 448         | os || has_lam e -> 1 + length oss     -- Note [Eta expanding thunks]
 449         | otherwise       -> 0
 450       ATop []             -> 0
 451       ABot n              -> n
 452   where
 453     has_lam (Note _ e) = has_lam e
 454     has_lam (Lam b e)  = isId b || has_lam e
 455     has_lam _          = False
 456
 457 getBotArity :: ArityType -> Maybe Arity
 458 -- Arity of a divergent function
 459 getBotArity (ABot n) = Just n
 460 getBotArity _        = Nothing
 461 \end{code}
 462
 463 Note [Eta expanding thunks]
 464 ~~~~~~~~~~~~~~~~~~~~~~~~~~~
 465 When we see
 466      f = case y of p -> \x -> blah
 467 should we eta-expand it? Well, if 'x' is a one-shot state token
 468 then 'yes' because 'f' will only be applied once.  But otherwise
 469 we (conservatively) say no.  My main reason is to avoid expanding
 470 PAPSs
 471         f = g d  ==>  f = \x. g d x
 472 because that might in turn make g inline (if it has an inline pragma),
 473 which we might not want.  After all, INLINE pragmas say "inline only
 474 when saturate" so we don't want to be too gung-ho about saturating!
 475
 476 \begin{code}
 477 arityLam :: Id -> ArityType -> ArityType
 478 arityLam id (ATop as) = ATop (isOneShotBndr id : as)
 479 arityLam _  (ABot n)  = ABot (n+1)
 480
 481 floatIn :: Bool -> ArityType -> ArityType
 482 -- We have something like (let x = E in b),
 483 -- where b has the given arity type.
 484 floatIn _     (ABot n)  = ABot n
 485 floatIn True  (ATop as) = ATop as
 486 floatIn False (ATop as) = ATop (takeWhile id as)
 487    -- If E is not cheap, keep arity only for one-shots
 488
 489 arityApp :: ArityType -> Bool -> ArityType
 490 -- Processing (fun arg) where at is the ArityType of fun,
 491 -- Knock off an argument and behave like 'let'
 492 arityApp (ABot 0)      _     = ABot 0
 493 arityApp (ABot n)      _     = ABot (n-1)
 494 arityApp (ATop [])     _     = ATop []
 495 arityApp (ATop (_:as)) cheap = floatIn cheap (ATop as)
 496
 497 andArityType :: ArityType -> ArityType -> ArityType   -- Used for branches of a 'case'
 498 andArityType (ABot n1) (ABot n2)
 499   = ABot (n1 `min` n2)
 500 andArityType (ATop as)  (ABot _)  = ATop as
 501 andArityType (ABot _)   (ATop bs) = ATop bs
 502 andArityType (ATop as)  (ATop bs) = ATop (as `combine` bs)
 503   where      -- See Note [Combining case branches]
 504     combine (a:as) (b:bs) = (a && b) : combine as bs
 505     combine []     bs     = take_one_shots bs
 506     combine as     []     = take_one_shots as
 507
 508     take_one_shots [] = []
 509     take_one_shots (one_shot : as)
 510       | one_shot  = True : take_one_shots as
 511       | otherwise = []
 512 \end{code}
 513
 514 Note [Combining case branches]
 515 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 516 Consider
 517   go = \x. let z = go e0
 518                go2 = \x. case x of
 519                            True  -> z
 520                            False -> \s(one-shot). e1
 521            in go2 x
 522 We *really* want to eta-expand go and go2.
 523 When combining the barnches of the case we have
 524      ATop [] `andAT` ATop [True]
 525 and we want to get ATop [True].  But if the inner
 526 lambda wasn't one-shot we don't want to do this.
 527 (We need a proper arity analysis to justify that.)
 528
 529
 530 \begin{code}
 531 ---------------------------
 532 type CheapFun = CoreExpr -> Maybe Type -> Bool
 533         -- How to decide if an expression is cheap
 534         -- If the Maybe is Just, the type is the type
 535         -- of the expression; Nothing means "don't know"
 536
 537 arityType :: CheapFun -> CoreExpr -> ArityType
 538 arityType _ (Var v)
 539   | Just strict_sig <- idStrictness_maybe v
 540   , (ds, res) <- splitStrictSig strict_sig
 541   , let arity = length ds
 542   = if isBotRes res then ABot arity
 543                     else ATop (take arity one_shots)
 544   | otherwise
 545   = ATop (take (idArity v) one_shots)
 546   where
 547     one_shots :: [Bool]     -- One-shot-ness derived from the type
 548     one_shots = typeArity (idType v)
 549
 550         -- Lambdas; increase arity
 551 arityType cheap_fn (Lam x e)
 552   | isId x    = arityLam x (arityType cheap_fn e)
 553   | otherwise = arityType cheap_fn e
 554
 555         -- Applications; decrease arity, except for types
 556 arityType cheap_fn (App fun (Type _))
 557    = arityType cheap_fn fun
 558 arityType cheap_fn (App fun arg )
 559    = arityApp (arityType cheap_fn fun) (cheap_fn arg Nothing)
 560
 561         -- Case/Let; keep arity if either the expression is cheap
 562         -- or it's a 1-shot lambda
 563         -- The former is not really right for Haskell
 564         --      f x = case x of { (a,b) -> \y. e }
 565         --  ===>
 566         --      f x y = case x of { (a,b) -> e }
 567         -- The difference is observable using 'seq'
 568 arityType cheap_fn (Case scrut bndr _ alts)
 569   = floatIn (cheap_fn scrut (Just (idType bndr)))
 570             (foldr1 andArityType [arityType cheap_fn rhs | (_,_,rhs) <- alts])
 571
 572 arityType cheap_fn (Let b e)
 573   = floatIn (cheap_bind b) (arityType cheap_fn e)
 574   where
 575     cheap_bind (NonRec b e) = is_cheap (b,e)
 576     cheap_bind (Rec prs)    = all is_cheap prs
 577     is_cheap (b,e) = cheap_fn e (Just (idType b))
 578
 579 arityType cheap_fn (Note n e)
 580   | notSccNote n              = arityType cheap_fn e
 581 arityType cheap_fn (Cast e _) = arityType cheap_fn e
 582 arityType _           _       = vanillaArityType
 583 \end{code}
 584
 585
 586 %************************************************************************
 587 %*                                                                      *
 588               The main eta-expander
 589 %*                                                                      *
 590 %************************************************************************
 591
 592 We go for:
 593    f = \x1..xn -> N  ==>   f = \x1..xn y1..ym -> N y1..ym
 594                                  (n >= 0)
 595
 596 where (in both cases)
 597
 598         * The xi can include type variables
 599
 600         * The yi are all value variables
 601
 602         * N is a NORMAL FORM (i.e. no redexes anywhere)
 603           wanting a suitable number of extra args.
 604
 605 The biggest reason for doing this is for cases like
 606
 607         f = \x -> case x of
 608                     True  -> \y -> e1
 609                     False -> \y -> e2
 610
 611 Here we want to get the lambdas together.  A good exmaple is the nofib
 612 program fibheaps, which gets 25% more allocation if you don't do this
 613 eta-expansion.
 614
 615 We may have to sandwich some coerces between the lambdas
 616 to make the types work.   exprEtaExpandArity looks through coerces
 617 when computing arity; and etaExpand adds the coerces as necessary when
 618 actually computing the expansion.
 619
 620
 621 Note [No crap in eta-expanded code]
 622 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 623 The eta expander is careful not to introduce "crap".  In particular,
 624 given a CoreExpr satisfying the 'CpeRhs' invariant (in CorePrep), it
 625 returns a CoreExpr satisfying the same invariant. See Note [Eta
 626 expansion and the CorePrep invariants] in CorePrep.
 627
 628 This means the eta-expander has to do a bit of on-the-fly
 629 simplification but it's not too hard.  The alernative, of relying on
 630 a subsequent clean-up phase of the Simplifier to de-crapify the result,
 631 means you can't really use it in CorePrep, which is painful.
 632
 633 Note [Eta expansion and SCCs]
 634 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 635 Note that SCCs are not treated specially by etaExpand.  If we have
 636         etaExpand 2 (\x -> scc "foo" e)
 637         = (\xy -> (scc "foo" e) y)
 638 So the costs of evaluating 'e' (not 'e y') are attributed to "foo"
 639
 640 \begin{code}
 641 -- | @etaExpand n us e ty@ returns an expression with
 642 -- the same meaning as @e@, but with arity @n@.
 643 --
 644 -- Given:
 645 --
 646 -- > e' = etaExpand n us e ty
 647 --
 648 -- We should have that:
 649 --
 650 -- > ty = exprType e = exprType e'
 651 etaExpand :: Arity              -- ^ Result should have this number of value args
 652           -> CoreExpr           -- ^ Expression to expand
 653           -> CoreExpr
 654 -- etaExpand deals with for-alls. For example:
 655 --              etaExpand 1 E
 656 -- where  E :: forall a. a -> a
 657 -- would return
 658 --      (/\b. \y::a -> E b y)
 659 --
 660 -- It deals with coerces too, though they are now rare
 661 -- so perhaps the extra code isn't worth it
 662
 663 etaExpand n orig_expr
 664   = go n orig_expr
 665   where
 666       -- Strip off existing lambdas and casts
 667       -- Note [Eta expansion and SCCs]
 668     go 0 expr = expr
 669     go n (Lam v body) | isTyVar v = Lam v (go n     body)
 670                       | otherwise = Lam v (go (n-1) body)
 671     go n (Cast expr co) = Cast (go n expr) co
 672     go n expr           = -- pprTrace "ee" (vcat [ppr orig_expr, ppr expr, ppr etas]) $
 673                           etaInfoAbs etas (etaInfoApp subst' expr etas)
 674                         where
 675                             in_scope = mkInScopeSet (exprFreeVars expr)
 676                             (in_scope', etas) = mkEtaWW n orig_expr in_scope (exprType expr)
 677                             subst' = mkEmptySubst in_scope'
 678
 679                                 -- Wrapper    Unwrapper
 680 --------------
 681 data EtaInfo = EtaVar Var       -- /\a. [],   [] a
 682                                 -- \x.  [],   [] x
 683              | EtaCo Coercion   -- [] |> co,  [] |> (sym co)
 684
 685 instance Outputable EtaInfo where
 686    ppr (EtaVar v) = ptext (sLit "EtaVar") <+> ppr v
 687    ppr (EtaCo co) = ptext (sLit "EtaCo")  <+> ppr co
 688
 689 pushCoercion :: Coercion -> [EtaInfo] -> [EtaInfo]
 690 pushCoercion co1 (EtaCo co2 : eis)
 691   | isReflCo co = eis
 692   | otherwise   = EtaCo co : eis
 693   where
 694     co = co1 `mkTransCo` co2
 695
 696 pushCoercion co eis = EtaCo co : eis
 697
 698 --------------
 699 etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr
 700 etaInfoAbs []               expr = expr
 701 etaInfoAbs (EtaVar v : eis) expr = Lam v (etaInfoAbs eis expr)
 702 etaInfoAbs (EtaCo co : eis) expr = Cast (etaInfoAbs eis expr) (mkSymCo co)
 703
 704 --------------
 705 etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
 706 -- (etaInfoApp s e eis) returns something equivalent to
 707 --             ((substExpr s e) `appliedto` eis)
 708
 709 etaInfoApp subst (Lam v1 e) (EtaVar v2 : eis)
 710   = etaInfoApp (CoreSubst.extendSubstWithVar subst v1 v2) e eis
 711
 712 etaInfoApp subst (Cast e co1) eis
 713   = etaInfoApp subst e (pushCoercion co' eis)
 714   where
 715     co' = CoreSubst.substCo subst co1
 716
 717 etaInfoApp subst (Case e b _ alts) eis
 718   = Case (subst_expr subst e) b1 (coreAltsType alts') alts'
 719   where
 720     (subst1, b1) = substBndr subst b
 721     alts' = map subst_alt alts
 722     subst_alt (con, bs, rhs) = (con, bs', etaInfoApp subst2 rhs eis)
 723               where
 724                  (subst2,bs') = substBndrs subst1 bs
 725
 726 etaInfoApp subst (Let b e) eis
 727   = Let b' (etaInfoApp subst' e eis)
 728   where
 729     (subst', b') = subst_bind subst b
 730
 731 etaInfoApp subst (Note note e) eis
 732   = Note note (etaInfoApp subst e eis)
 733
 734 etaInfoApp subst e eis
 735   = go (subst_expr subst e) eis
 736   where
 737     go e []                  = e
 738     go e (EtaVar v    : eis) = go (App e (varToCoreExpr v)) eis
 739     go e (EtaCo co    : eis) = go (Cast e co) eis
 740
 741 --------------
 742 mkEtaWW :: Arity -> CoreExpr -> InScopeSet -> Type
 743         -> (InScopeSet, [EtaInfo])
 744         -- EtaInfo contains fresh variables,
 745         --   not free in the incoming CoreExpr
 746         -- Outgoing InScopeSet includes the EtaInfo vars
 747         --   and the original free vars
 748
 749 mkEtaWW orig_n orig_expr in_scope orig_ty
 750   = go orig_n empty_subst orig_ty []
 751   where
 752     empty_subst = TvSubst in_scope emptyTvSubstEnv
 753
 754     go n subst ty eis       -- See Note [exprArity invariant]
 755        | n == 0
 756        = (getTvInScope subst, reverse eis)
 757
 758        | Just (tv,ty') <- splitForAllTy_maybe ty
 759        , let (subst', tv') = Type.substTyVarBndr subst tv
 760            -- Avoid free vars of the original expression
 761        = go n subst' ty' (EtaVar tv' : eis)
 762
 763        | Just (arg_ty, res_ty) <- splitFunTy_maybe ty
 764        , let (subst', eta_id') = freshEtaId n subst arg_ty
 765            -- Avoid free vars of the original expression
 766        = go (n-1) subst' res_ty (EtaVar eta_id' : eis)
 767
 768        | Just(ty',co) <- splitNewTypeRepCo_maybe ty
 769        =        -- Given this:
 770                 --      newtype T = MkT ([T] -> Int)
 771                 -- Consider eta-expanding this
 772                 --      eta_expand 1 e T
 773                 -- We want to get
 774                 --      coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
 775          go n subst ty' (EtaCo co : eis)
 776
 777        | otherwise       -- We have an expression of arity > 0,
 778                          -- but its type isn't a function.
 779        = WARN( True, (ppr orig_n <+> ppr orig_ty) $$ ppr orig_expr )
 780          (getTvInScope subst, reverse eis)
 781         -- This *can* legitmately happen:
 782         -- e.g.  coerce Int (\x. x) Essentially the programmer is
 783         -- playing fast and loose with types (Happy does this a lot).
 784         -- So we simply decline to eta-expand.  Otherwise we'd end up
 785         -- with an explicit lambda having a non-function type
 786
 787
 788 --------------
 789 -- Avoiding unnecessary substitution; use short-cutting versions
 790
 791 subst_expr :: Subst -> CoreExpr -> CoreExpr
 792 subst_expr = substExprSC (text "CoreArity:substExpr")
 793
 794 subst_bind :: Subst -> CoreBind -> (Subst, CoreBind)
 795 subst_bind = substBindSC
 796
 797
 798 --------------
 799 freshEtaId :: Int -> TvSubst -> Type -> (TvSubst, Id)
 800 -- Make a fresh Id, with specified type (after applying substitution)
 801 -- It should be "fresh" in the sense that it's not in the in-scope set
 802 -- of the TvSubstEnv; and it should itself then be added to the in-scope
 803 -- set of the TvSubstEnv
 804 --
 805 -- The Int is just a reasonable starting point for generating a unique;
 806 -- it does not necessarily have to be unique itself.
 807 freshEtaId n subst ty
 808       = (subst', eta_id')
 809       where
 810         ty'     = Type.substTy subst ty
 811         eta_id' = uniqAway (getTvInScope subst) $
 812                   mkSysLocal (fsLit "eta") (mkBuiltinUnique n) ty'
 813         subst'  = extendTvInScope subst eta_id'
 814 \end{code}
 815