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, 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 TcType   ( isDictLikeTy )
  28 import Coercion
  29 import BasicTypes
  30 import Unique
  31 import Outputable
  32 import DynFlags
  33 import FastString
  34 \end{code}
  35
  36 %************************************************************************
  37 %*                                                                      *
  38               manifestArity and exprArity
  39 %*                                                                      *
  40 %************************************************************************
  41
  42 exprArity is a cheap-and-cheerful version of exprEtaExpandArity.
  43 It tells how many things the expression can be applied to before doing
  44 any work.  It doesn't look inside cases, lets, etc.  The idea is that
  45 exprEtaExpandArity will do the hard work, leaving something that's easy
  46 for exprArity to grapple with.  In particular, Simplify uses exprArity to
  47 compute the ArityInfo for the Id.
  48
  49 Originally I thought that it was enough just to look for top-level lambdas, but
  50 it isn't.  I've seen this
  51
  52         foo = PrelBase.timesInt
  53
  54 We want foo to get arity 2 even though the eta-expander will leave it
  55 unchanged, in the expectation that it'll be inlined.  But occasionally it
  56 isn't, because foo is blacklisted (used in a rule).
  57
  58 Similarly, see the ok_note check in exprEtaExpandArity.  So
  59         f = __inline_me (\x -> e)
  60 won't be eta-expanded.
  61
  62 And in any case it seems more robust to have exprArity be a bit more intelligent.
  63 But note that   (\x y z -> f x y z)
  64 should have arity 3, regardless of f's arity.
  65
  66 Note [exprArity invariant]
  67 ~~~~~~~~~~~~~~~~~~~~~~~~~~
  68 exprArity has the following invariant:
  69
  70   * If typeArity (exprType e) = n,
  71     then manifestArity (etaExpand e n) = n
  72
  73     That is, etaExpand can always expand as much as typeArity says
  74     So the case analysis in etaExpand and in typeArity must match
  75
  76   * exprArity e <= typeArity (exprType e)
  77
  78   * Hence if (exprArity e) = n, then manifestArity (etaExpand e n) = n
  79
  80     That is, if exprArity says "the arity is n" then etaExpand really
  81     can get "n" manifest lambdas to the top.
  82
  83 Why is this important?  Because
  84   - In TidyPgm we use exprArity to fix the *final arity* of
  85     each top-level Id, and in
  86   - In CorePrep we use etaExpand on each rhs, so that the visible lambdas
  87     actually match that arity, which in turn means
  88     that the StgRhs has the right number of lambdas
  89
  90 An alternative would be to do the eta-expansion in TidyPgm, at least
  91 for top-level bindings, in which case we would not need the trim_arity
  92 in exprArity.  That is a less local change, so I'm going to leave it for today!
  93
  94
  95 \begin{code}
  96 manifestArity :: CoreExpr -> Arity
  97 -- ^ manifestArity sees how many leading value lambdas there are
  98 manifestArity (Lam v e) | isId v    = 1 + manifestArity e
  99                         | otherwise = manifestArity e
 100 manifestArity (Note _ e)            = manifestArity e
 101 manifestArity (Cast e _)            = manifestArity e
 102 manifestArity _                     = 0
 103
 104 exprArity :: CoreExpr -> Arity
 105 -- ^ An approximate, fast, version of 'exprEtaExpandArity'
 106 exprArity e = go e
 107   where
 108     go (Var v)                     = idArity v
 109     go (Lam x e) | isId x          = go e + 1
 110                  | otherwise       = go e
 111     go (Note _ e)                  = go e
 112     go (Cast e co)                 = go e `min` length (typeArity (snd (coercionKind co)))
 113                                         -- Note [exprArity invariant]
 114     go (App e (Type _))            = go e
 115     go (App f a) | exprIsTrivial a = (go f - 1) `max` 0
 116         -- See Note [exprArity for applications]
 117     go _                           = 0
 118
 119
 120 typeArity :: Type -> [OneShot]
 121 -- How many value arrows are visible in the type?
 122 -- We look through foralls, and newtypes
 123 -- See Note [exprArity invariant]
 124 typeArity ty
 125   | Just (_, ty')  <- splitForAllTy_maybe ty
 126   = typeArity ty'
 127
 128   | Just (arg,res) <- splitFunTy_maybe ty
 129   = isStateHackType arg : typeArity res
 130
 131   | Just (tc,tys) <- splitTyConApp_maybe ty
 132   , Just (ty', _) <- instNewTyCon_maybe tc tys
 133   , not (isRecursiveTyCon tc)
 134   , not (isClassTyCon tc)       -- Do not eta-expand through newtype classes
 135                                 -- See Note [Newtype classes and eta expansion]
 136   = typeArity ty'
 137         -- Important to look through non-recursive newtypes, so that, eg
 138         --      (f x)   where f has arity 2, f :: Int -> IO ()
 139         -- Here we want to get arity 1 for the result!
 140
 141   | otherwise
 142   = []
 143 \end{code}
 144
 145 Note [Newtype classes and eta expansion]
 146 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 147 We have to be careful when eta-expanding through newtypes.  In general
 148 it's a good idea, but annoyingly it interacts badly with the class-op
 149 rule mechanism.  Consider
 150
 151    class C a where { op :: a -> a }
 152    instance C b => C [b] where
 153      op x = ...
 154
 155 These translate to
 156
 157    co :: forall a. (a->a) ~ C a
 158
 159    $copList :: C b -> [b] -> [b]
 160    $copList d x = ...
 161
 162    $dfList :: C b -> C [b]
 163    {-# DFunUnfolding = [$copList] #-}
 164    $dfList d = $copList d |> co@[b]
 165
 166 Now suppose we have:
 167
 168    dCInt :: C Int
 169
 170    blah :: [Int] -> [Int]
 171    blah = op ($dfList dCInt)
 172
 173 Now we want the built-in op/$dfList rule will fire to give
 174    blah = $copList dCInt
 175
 176 But with eta-expansion 'blah' might (and in Trac #3772, which is
 177 slightly more complicated, does) turn into
 178
 179    blah = op (\eta. ($dfList dCInt |> sym co) eta)
 180
 181 and now it is *much* harder for the op/$dfList rule to fire, becuase
 182 exprIsConApp_maybe won't hold of the argument to op.  I considered
 183 trying to *make* it hold, but it's tricky and I gave up.
 184
 185 The test simplCore/should_compile/T3722 is an excellent example.
 186
 187
 188 Note [exprArity for applications]
 189 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 190 When we come to an application we check that the arg is trivial.
 191    eg  f (fac x) does not have arity 2,
 192                  even if f has arity 3!
 193
 194 * We require that is trivial rather merely cheap.  Suppose f has arity 2.
 195   Then    f (Just y)
 196   has arity 0, because if we gave it arity 1 and then inlined f we'd get
 197           let v = Just y in \w. <f-body>
 198   which has arity 0.  And we try to maintain the invariant that we don't
 199   have arity decreases.
 200
 201 *  The `max 0` is important!  (\x y -> f x) has arity 2, even if f is
 202    unknown, hence arity 0
 203
 204
 205 %************************************************************************
 206 %*                                                                      *
 207            Eta expansion
 208 %*                                                                      *
 209 %************************************************************************
 210
 211 \begin{code}
 212 exprBotStrictness_maybe :: CoreExpr -> Maybe (Arity, StrictSig)
 213 -- A cheap and cheerful function that identifies bottoming functions
 214 -- and gives them a suitable strictness signatures.  It's used during
 215 -- float-out
 216 exprBotStrictness_maybe e
 217   = case getBotArity (arityType False e) of
 218         Nothing -> Nothing
 219         Just ar -> Just (ar, mkStrictSig (mkTopDmdType (replicate ar topDmd) BotRes))
 220 \end{code}
 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 getArity).
 380
 381 Here is what the fields mean. If e has ArityType
 382      (AT as r), where n = length as,
 383 then
 384
 385  * If r is ABot then (e x1..xn) definitely diverges
 386    Partial applications may or may not diverge
 387
 388  * If r is ACheap then (e x1..x(n-1)) is cheap,
 389    including any nested sub-expressions inside e
 390    (say e is (f e1 e2) then e1,e2 are cheap too)
 391
 392  * e, (e x1), ... (e x1 ... x(n-1)) are definitely really
 393    functions, or bottom, not casts from a data type
 394    So eta expansion is dynamically ok;
 395     see Note [State hack and bottoming functions],
 396     the part about catch#
 397
 398 We regard ABot as stronger than ACheap; ie if ABot holds
 399 we don't bother about ACheap
 400
 401 Suppose f = \xy. x+y
 402 Then  f             :: AT [False,False] ACheap
 403       f v           :: AT [False]       ACheap
 404       f <expensive> :: AT [False]       ATop
 405 Note the ArityRes flag tells whether the whole expression is cheap.
 406 Note also that having a non-empty 'as' doesn't mean it has that
 407 arity; see (f <expensive>) which does not have arity 1!
 408
 409 The key function getArity extracts the arity (which in turn guides
 410 eta-expansion) from ArityType.
 411   * If the term is cheap or diverges we can certainly eta expand it
 412       e.g.   (f x)   where x has arity 2
 413
 414   * If its a function whose first arg is one-shot (probably via the
 415     state hack) we can eta expand it
 416       e.g.   (getChar <expensive>)
 417
 418 -------------------- Main arity code ----------------------------
 419 \begin{code}
 420 -- See Note [ArityType]
 421 data ArityType = AT [OneShot] ArityRes
 422      -- There is always an explicit lambda
 423      -- to justify the [OneShot]
 424
 425 type OneShot = Bool    -- False <=> Know nothing
 426                        -- True  <=> Can definitely float inside this lambda
 427                        -- The 'True' case can arise either because a binder
 428                        -- is marked one-shot, or because it's a state lambda
 429                        -- and we have the state hack on
 430
 431 data ArityRes  = ATop | ACheap | ABot
 432
 433 vanillaArityType :: ArityType
 434 vanillaArityType = AT [] ATop   -- Totally uninformative
 435
 436 -- ^ The Arity returned is the number of value args the [_$_]
 437 -- expression can be applied to without doing much work
 438 exprEtaExpandArity :: DynFlags -> CoreExpr -> Arity
 439 -- exprEtaExpandArity is used when eta expanding
 440 --      e  ==>  \xy -> e x y
 441 exprEtaExpandArity dflags e
 442   = case (arityType dicts_cheap e) of
 443       AT (a:as) res | want_eta a res -> 1 + length as
 444       _                              -> 0
 445   where
 446     want_eta one_shot ATop   = one_shot
 447     want_eta _        _      = True
 448
 449     dicts_cheap = dopt Opt_DictsCheap dflags
 450
 451 getBotArity :: ArityType -> Maybe Arity
 452 -- Arity of a divergent function
 453 getBotArity (AT as ABot) = Just (length as)
 454 getBotArity _            = Nothing
 455
 456 arityLam :: Id -> ArityType -> ArityType
 457 arityLam id (AT as r) = AT (isOneShotBndr id : as) r
 458
 459 floatIn :: Bool -> ArityType -> ArityType
 460 -- We have something like (let x = E in b),
 461 -- where b has the given arity type.
 462 floatIn c (AT as r) = AT as (extendArityRes r c)
 463
 464 arityApp :: ArityType -> CoreExpr -> ArityType
 465 -- Processing (fun arg) where at is the ArityType of fun,
 466 arityApp (AT [] r)     arg = AT [] (extendArityRes r (exprIsCheap arg))
 467 arityApp (AT (_:as) r) arg = AT as (extendArityRes r (exprIsCheap arg))
 468
 469 extendArityRes :: ArityRes -> Bool -> ArityRes
 470 extendArityRes ABot   _    = ABot
 471 extendArityRes ACheap True = ACheap
 472 extendArityRes _      _    = ATop
 473
 474 andArityType :: ArityType -> ArityType -> ArityType   -- Used for branches of a 'case'
 475 andArityType (AT as1 r1) (AT as2 r2)
 476   = AT (go_as as1 as2) (go_r r1 r2)
 477   where
 478     go_r ABot ABot     = ABot
 479     go_r ABot ACheap   = ACheap
 480     go_r ACheap ABot   = ACheap
 481     go_r ACheap ACheap = ACheap
 482     go_r _    _        = ATop
 483
 484     go_as (os1:as1) (os2:as2) = (os1 || os2) : go_as as1 as2
 485     go_as []        as2       = as2
 486     go_as as1       []        = as1
 487 \end{code}
 488
 489
 490 \begin{code}
 491 ---------------------------
 492 arityType :: Bool -> CoreExpr -> ArityType
 493 arityType _ (Var v)
 494   | Just strict_sig <- idStrictness_maybe v
 495   , (ds, res) <- splitStrictSig strict_sig
 496   = mk_arity (length ds) res
 497   | otherwise
 498   = mk_arity (idArity v) TopRes
 499
 500   where
 501     mk_arity id_arity res
 502       | isBotRes res = AT (take id_arity one_shots) ABot
 503       | id_arity>0   = AT (take id_arity one_shots) ACheap
 504       | otherwise    = AT []                        ATop
 505
 506     one_shots = typeArity (idType v)
 507
 508         -- Lambdas; increase arity
 509 arityType dicts_cheap (Lam x e)
 510   | isId x    = arityLam x (arityType dicts_cheap e)
 511   | otherwise = arityType dicts_cheap e
 512
 513         -- Applications; decrease arity
 514 arityType dicts_cheap (App fun (Type _))
 515    = arityType dicts_cheap fun
 516 arityType dicts_cheap (App fun arg )
 517    = arityApp (arityType dicts_cheap fun) arg
 518
 519         -- Case/Let; keep arity if either the expression is cheap
 520         -- or it's a 1-shot lambda
 521         -- The former is not really right for Haskell
 522         --      f x = case x of { (a,b) -> \y. e }
 523         --  ===>
 524         --      f x y = case x of { (a,b) -> e }
 525         -- The difference is observable using 'seq'
 526 arityType dicts_cheap (Case scrut _ _ alts)
 527   = floatIn (exprIsCheap scrut)
 528               (foldr1 andArityType [arityType dicts_cheap rhs | (_,_,rhs) <- alts])
 529
 530 arityType dicts_cheap (Let b e)
 531   = floatIn (cheap_bind b) (arityType dicts_cheap e)
 532   where
 533     cheap_bind (NonRec b e) = is_cheap (b,e)
 534     cheap_bind (Rec prs)    = all is_cheap prs
 535     is_cheap (b,e) = (dicts_cheap && isDictLikeTy (idType b))
 536                    || exprIsCheap e
 537         -- If the experimental -fdicts-cheap flag is on, we eta-expand through
 538         -- dictionary bindings.  This improves arities. Thereby, it also
 539         -- means that full laziness is less prone to floating out the
 540         -- application of a function to its dictionary arguments, which
 541         -- can thereby lose opportunities for fusion.  Example:
 542         --      foo :: Ord a => a -> ...
 543         --      foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....
 544         --              -- So foo has arity 1
 545         --
 546         --      f = \x. foo dInt $ bar x
 547         --
 548         -- The (foo DInt) is floated out, and makes ineffective a RULE
 549         --      foo (bar x) = ...
 550         --
 551         -- One could go further and make exprIsCheap reply True to any
 552         -- dictionary-typed expression, but that's more work.
 553         --
 554         -- See Note [Dictionary-like types] in TcType.lhs for why we use
 555         -- isDictLikeTy here rather than isDictTy
 556
 557 arityType dicts_cheap (Note _ e) = arityType dicts_cheap e
 558 arityType dicts_cheap (Cast e _) = arityType dicts_cheap e
 559 arityType _           _          = vanillaArityType
 560 \end{code}
 561
 562
 563 %************************************************************************
 564 %*                                                                      *
 565               The main eta-expander
 566 %*                                                                      *
 567 %************************************************************************
 568
 569 IMPORTANT NOTE: The eta expander is careful not to introduce "crap".
 570 In particular, given a CoreExpr satisfying the 'CpeRhs' invariant (in
 571 CorePrep), it returns a CoreExpr satisfying the same invariant. See
 572 Note [Eta expansion and the CorePrep invariants] in CorePrep.
 573
 574 This means the eta-expander has to do a bit of on-the-fly
 575 simplification but it's not too hard.  The alernative, of relying on
 576 a subsequent clean-up phase of the Simplifier to de-crapify the result,
 577 means you can't really use it in CorePrep, which is painful.
 578
 579 Note [Eta expansion and SCCs]
 580 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 581 Note that SCCs are not treated specially by etaExpand.  If we have
 582         etaExpand 2 (\x -> scc "foo" e)
 583         = (\xy -> (scc "foo" e) y)
 584 So the costs of evaluating 'e' (not 'e y') are attributed to "foo"
 585
 586 \begin{code}
 587 -- | @etaExpand n us e ty@ returns an expression with
 588 -- the same meaning as @e@, but with arity @n@.
 589 --
 590 -- Given:
 591 --
 592 -- > e' = etaExpand n us e ty
 593 --
 594 -- We should have that:
 595 --
 596 -- > ty = exprType e = exprType e'
 597 etaExpand :: Arity              -- ^ Result should have this number of value args
 598           -> CoreExpr           -- ^ Expression to expand
 599           -> CoreExpr
 600 -- etaExpand deals with for-alls. For example:
 601 --              etaExpand 1 E
 602 -- where  E :: forall a. a -> a
 603 -- would return
 604 --      (/\b. \y::a -> E b y)
 605 --
 606 -- It deals with coerces too, though they are now rare
 607 -- so perhaps the extra code isn't worth it
 608
 609 etaExpand n orig_expr
 610   = go n orig_expr
 611   where
 612       -- Strip off existing lambdas and casts
 613       -- Note [Eta expansion and SCCs]
 614     go 0 expr = expr
 615     go n (Lam v body) | isTyCoVar v = Lam v (go n     body)
 616                       | otherwise = Lam v (go (n-1) body)
 617     go n (Cast expr co) = Cast (go n expr) co
 618     go n expr           = -- pprTrace "ee" (vcat [ppr orig_expr, ppr expr, ppr etas]) $
 619                           etaInfoAbs etas (etaInfoApp subst' expr etas)
 620                         where
 621                             in_scope = mkInScopeSet (exprFreeVars expr)
 622                             (in_scope', etas) = mkEtaWW n in_scope (exprType expr)
 623                             subst' = mkEmptySubst in_scope'
 624
 625                                 -- Wrapper    Unwrapper
 626 --------------
 627 data EtaInfo = EtaVar Var       -- /\a. [],   [] a
 628                                 -- \x.  [],   [] x
 629              | EtaCo Coercion   -- [] |> co,  [] |> (sym co)
 630
 631 instance Outputable EtaInfo where
 632    ppr (EtaVar v) = ptext (sLit "EtaVar") <+> ppr v
 633    ppr (EtaCo co) = ptext (sLit "EtaCo")  <+> ppr co
 634
 635 pushCoercion :: Coercion -> [EtaInfo] -> [EtaInfo]
 636 pushCoercion co1 (EtaCo co2 : eis)
 637   | isIdentityCoercion co = eis
 638   | otherwise             = EtaCo co : eis
 639   where
 640     co = co1 `mkTransCoercion` co2
 641
 642 pushCoercion co eis = EtaCo co : eis
 643
 644 --------------
 645 etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr
 646 etaInfoAbs []               expr = expr
 647 etaInfoAbs (EtaVar v : eis) expr = Lam v (etaInfoAbs eis expr)
 648 etaInfoAbs (EtaCo co : eis) expr = Cast (etaInfoAbs eis expr) (mkSymCoercion co)
 649
 650 --------------
 651 etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
 652 -- (etaInfoApp s e eis) returns something equivalent to
 653 --             ((substExpr s e) `appliedto` eis)
 654
 655 etaInfoApp subst (Lam v1 e) (EtaVar v2 : eis)
 656   = etaInfoApp subst' e eis
 657   where
 658     subst' | isTyCoVar v1 = CoreSubst.extendTvSubst subst v1 (mkTyVarTy v2)
 659            | otherwise  = CoreSubst.extendIdSubst subst v1 (Var v2)
 660
 661 etaInfoApp subst (Cast e co1) eis
 662   = etaInfoApp subst e (pushCoercion co' eis)
 663   where
 664     co' = CoreSubst.substTy subst co1
 665
 666 etaInfoApp subst (Case e b _ alts) eis
 667   = Case (subst_expr subst e) b1 (coreAltsType alts') alts'
 668   where
 669     (subst1, b1) = substBndr subst b
 670     alts' = map subst_alt alts
 671     subst_alt (con, bs, rhs) = (con, bs', etaInfoApp subst2 rhs eis)
 672               where
 673                  (subst2,bs') = substBndrs subst1 bs
 674
 675 etaInfoApp subst (Let b e) eis
 676   = Let b' (etaInfoApp subst' e eis)
 677   where
 678     (subst', b') = subst_bind subst b
 679
 680 etaInfoApp subst (Note note e) eis
 681   = Note note (etaInfoApp subst e eis)
 682
 683 etaInfoApp subst e eis
 684   = go (subst_expr subst e) eis
 685   where
 686     go e []                  = e
 687     go e (EtaVar v    : eis) = go (App e (varToCoreExpr v)) eis
 688     go e (EtaCo co    : eis) = go (Cast e co) eis
 689
 690 --------------
 691 mkEtaWW :: Arity -> InScopeSet -> Type
 692         -> (InScopeSet, [EtaInfo])
 693         -- EtaInfo contains fresh variables,
 694         --   not free in the incoming CoreExpr
 695         -- Outgoing InScopeSet includes the EtaInfo vars
 696         --   and the original free vars
 697
 698 mkEtaWW orig_n in_scope orig_ty
 699   = go orig_n empty_subst orig_ty []
 700   where
 701     empty_subst = mkTvSubst in_scope emptyTvSubstEnv
 702
 703     go n subst ty eis       -- See Note [exprArity invariant]
 704        | n == 0
 705        = (getTvInScope subst, reverse eis)
 706
 707        | Just (tv,ty') <- splitForAllTy_maybe ty
 708        , let (subst', tv') = substTyVarBndr subst tv
 709            -- Avoid free vars of the original expression
 710        = go n subst' ty' (EtaVar tv' : eis)
 711
 712        | Just (arg_ty, res_ty) <- splitFunTy_maybe ty
 713        , let (subst', eta_id') = freshEtaId n subst arg_ty
 714            -- Avoid free vars of the original expression
 715        = go (n-1) subst' res_ty (EtaVar eta_id' : eis)
 716
 717        | Just(ty',co) <- splitNewTypeRepCo_maybe ty
 718        =        -- Given this:
 719                 --      newtype T = MkT ([T] -> Int)
 720                 -- Consider eta-expanding this
 721                 --      eta_expand 1 e T
 722                 -- We want to get
 723                 --      coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
 724          go n subst ty' (EtaCo (Type.substTy subst co) : eis)
 725
 726        | otherwise                         -- We have an expression of arity > 0,
 727        = WARN( True, ppr orig_n <+> ppr orig_ty )
 728          (getTvInScope subst, reverse eis) -- but its type isn't a function.
 729         -- This *can* legitmately happen:
 730         -- e.g.  coerce Int (\x. x) Essentially the programmer is
 731         -- playing fast and loose with types (Happy does this a lot).
 732         -- So we simply decline to eta-expand.  Otherwise we'd end up
 733         -- with an explicit lambda having a non-function type
 734
 735
 736 --------------
 737 -- Avoiding unnecessary substitution; use short-cutting versions
 738
 739 subst_expr :: Subst -> CoreExpr -> CoreExpr
 740 subst_expr = substExprSC (text "CoreArity:substExpr")
 741
 742 subst_bind :: Subst -> CoreBind -> (Subst, CoreBind)
 743 subst_bind = substBindSC
 744
 745
 746 --------------
 747 freshEtaId :: Int -> TvSubst -> Type -> (TvSubst, Id)
 748 -- Make a fresh Id, with specified type (after applying substitution)
 749 -- It should be "fresh" in the sense that it's not in the in-scope set
 750 -- of the TvSubstEnv; and it should itself then be added to the in-scope
 751 -- set of the TvSubstEnv
 752 --
 753 -- The Int is just a reasonable starting point for generating a unique;
 754 -- it does not necessarily have to be unique itself.
 755 freshEtaId n subst ty
 756       = (subst', eta_id')
 757       where
 758         ty'     = Type.substTy subst ty
 759         eta_id' = uniqAway (getTvInScope subst) $
 760                   mkSysLocal (fsLit "eta") (mkBuiltinUnique n) ty'
 761         subst'  = extendTvInScope subst eta_id'
 762 \end{code}
 763