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