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 \begin{code}
  67 manifestArity :: CoreExpr -> Arity
  68 -- ^ manifestArity sees how many leading value lambdas there are
  69 manifestArity (Lam v e) | isId v        = 1 + manifestArity e
  70                         | otherwise     = manifestArity e
  71 manifestArity (Note n e) | notSccNote n = manifestArity e
  72 manifestArity (Cast e _)                = manifestArity e
  73 manifestArity _                         = 0
  74
  75 ---------------
  76 exprArity :: CoreExpr -> Arity
  77 -- ^ An approximate, fast, version of 'exprEtaExpandArity'
  78 exprArity e = go e
  79   where
  80     go (Var v)                     = idArity v
  81     go (Lam x e) | isId x          = go e + 1
  82                  | otherwise       = go e
  83     go (Note n e) | notSccNote n   = go e
  84     go (Cast e co)                 = go e `min` length (typeArity (snd (coercionKind co)))
  85                                         -- Note [exprArity invariant]
  86     go (App e (Type _))            = go e
  87     go (App f a) | exprIsTrivial a = (go f - 1) `max` 0
  88         -- See Note [exprArity for applications]
  89     go _                           = 0
  90
  91
  92 ---------------
  93 typeArity :: Type -> [OneShot]
  94 -- How many value arrows are visible in the type?
  95 -- We look through foralls, and newtypes
  96 -- See Note [exprArity invariant]
  97 typeArity ty
  98   | Just (_, ty')  <- splitForAllTy_maybe ty
  99   = typeArity ty'
 100
 101   | Just (arg,res) <- splitFunTy_maybe ty
 102   = isStateHackType arg : typeArity res
 103
 104   | Just (tc,tys) <- splitTyConApp_maybe ty
 105   , Just (ty', _) <- instNewTyCon_maybe tc tys
 106   , not (isRecursiveTyCon tc)
 107   , not (isClassTyCon tc)       -- Do not eta-expand through newtype classes
 108                                 -- See Note [Newtype classes and eta expansion]
 109   = typeArity ty'
 110         -- Important to look through non-recursive newtypes, so that, eg
 111         --      (f x)   where f has arity 2, f :: Int -> IO ()
 112         -- Here we want to get arity 1 for the result!
 113
 114   | otherwise
 115   = []
 116
 117 ---------------
 118 exprBotStrictness_maybe :: CoreExpr -> Maybe (Arity, StrictSig)
 119 -- A cheap and cheerful function that identifies bottoming functions
 120 -- and gives them a suitable strictness signatures.  It's used during
 121 -- float-out
 122 exprBotStrictness_maybe e
 123   = case getBotArity (arityType False e) of
 124         Nothing -> Nothing
 125         Just ar -> Just (ar, mkStrictSig (mkTopDmdType (replicate ar topDmd) BotRes))
 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 :: DynFlags -> CoreExpr -> Arity
 440 -- exprEtaExpandArity is used when eta expanding
 441 --      e  ==>  \xy -> e x y
 442 exprEtaExpandArity dflags e
 443   = case (arityType dicts_cheap 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     dicts_cheap = dopt Opt_DictsCheap dflags
 451     has_lam (Note _ e) = has_lam e
 452     has_lam (Lam b e)  = isId b || has_lam e
 453     has_lam _          = False
 454
 455 getBotArity :: ArityType -> Maybe Arity
 456 -- Arity of a divergent function
 457 getBotArity (ABot n) = Just n
 458 getBotArity _        = Nothing
 459 \end{code}
 460
 461 Note [Eta expanding thunks]
 462 ~~~~~~~~~~~~~~~~~~~~~~~~~~~
 463 When we see
 464      f = case y of p -> \x -> blah
 465 should we eta-expand it? Well, if 'x' is a one-shot state token
 466 then 'yes' because 'f' will only be applied once.  But otherwise
 467 we (conservatively) say no.  My main reason is to avoid expanding
 468 PAPSs
 469         f = g d  ==>  f = \x. g d x
 470 because that might in turn make g inline (if it has an inline pragma),
 471 which we might not want.  After all, INLINE pragmas say "inline only
 472 when saturate" so we don't want to be too gung-ho about saturating!
 473
 474 \begin{code}
 475 arityLam :: Id -> ArityType -> ArityType
 476 arityLam id (ATop as) = ATop (isOneShotBndr id : as)
 477 arityLam _  (ABot n)  = ABot (n+1)
 478
 479 floatIn :: Bool -> ArityType -> ArityType
 480 -- We have something like (let x = E in b),
 481 -- where b has the given arity type.
 482 floatIn _     (ABot n)  = ABot n
 483 floatIn True  (ATop as) = ATop as
 484 floatIn False (ATop as) = ATop (takeWhile id as)
 485    -- If E is not cheap, keep arity only for one-shots
 486
 487 arityApp :: ArityType -> CoreExpr -> ArityType
 488 -- Processing (fun arg) where at is the ArityType of fun,
 489 -- Knock off an argument and behave like 'let'
 490 arityApp (ABot 0)      _   = ABot 0
 491 arityApp (ABot n)      _   = ABot (n-1)
 492 arityApp (ATop [])     _   = ATop []
 493 arityApp (ATop (_:as)) arg = floatIn (exprIsCheap arg) (ATop as)
 494
 495 andArityType :: ArityType -> ArityType -> ArityType   -- Used for branches of a 'case'
 496 andArityType (ABot n1) (ABot n2)
 497   = ABot (n1 `min` n2)
 498 andArityType (ATop as)  (ABot _)  = ATop as
 499 andArityType (ABot _)   (ATop bs) = ATop bs
 500 andArityType (ATop as)  (ATop bs) = ATop (as `combine` bs)
 501   where      -- See Note [Combining case branches]
 502     combine (a:as) (b:bs) = (a && b) : combine as bs
 503     combine []     bs     = take_one_shots bs
 504     combine as     []     = take_one_shots as
 505
 506     take_one_shots [] = []
 507     take_one_shots (one_shot : as)
 508       | one_shot  = True : take_one_shots as
 509       | otherwise = []
 510 \end{code}
 511
 512 Note [Combining case branches]
 513 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 514 Consider
 515   go = \x. let z = go e0
 516                go2 = \x. case x of
 517                            True  -> z
 518                            False -> \s(one-shot). e1
 519            in go2 x
 520 We *really* want to eta-expand go and go2.
 521 When combining the barnches of the case we have
 522      ATop [] `andAT` ATop [True]
 523 and we want to get ATop [True].  But if the inner
 524 lambda wasn't one-shot we don't want to do this.
 525 (We need a proper arity analysis to justify that.)
 526
 527
 528 \begin{code}
 529 ---------------------------
 530 arityType :: Bool -> CoreExpr -> ArityType
 531 arityType _ (Var v)
 532   | Just strict_sig <- idStrictness_maybe v
 533   , (ds, res) <- splitStrictSig strict_sig
 534   , let arity = length ds
 535   = if isBotRes res then ABot arity
 536                     else ATop (take arity one_shots)
 537   | otherwise
 538   = ATop (take (idArity v) one_shots)
 539   where
 540     one_shots :: [Bool]     -- One-shot-ness derived from the type
 541     one_shots = typeArity (idType v)
 542
 543         -- Lambdas; increase arity
 544 arityType dicts_cheap (Lam x e)
 545   | isId x    = arityLam x (arityType dicts_cheap e)
 546   | otherwise = arityType dicts_cheap e
 547
 548         -- Applications; decrease arity
 549 arityType dicts_cheap (App fun (Type _))
 550    = arityType dicts_cheap fun
 551 arityType dicts_cheap (App fun arg )
 552    = arityApp (arityType dicts_cheap fun) arg
 553
 554         -- Case/Let; keep arity if either the expression is cheap
 555         -- or it's a 1-shot lambda
 556         -- The former is not really right for Haskell
 557         --      f x = case x of { (a,b) -> \y. e }
 558         --  ===>
 559         --      f x y = case x of { (a,b) -> e }
 560         -- The difference is observable using 'seq'
 561 arityType dicts_cheap (Case scrut _ _ alts)
 562   = floatIn (exprIsCheap scrut)
 563               (foldr1 andArityType [arityType dicts_cheap rhs | (_,_,rhs) <- alts])
 564
 565 arityType dicts_cheap (Let b e)
 566   = floatIn (cheap_bind b) (arityType dicts_cheap e)
 567   where
 568     cheap_bind (NonRec b e) = is_cheap (b,e)
 569     cheap_bind (Rec prs)    = all is_cheap prs
 570     is_cheap (b,e) = (dicts_cheap && isDictLikeTy (idType b))
 571                    || exprIsCheap e
 572         -- If the experimental -fdicts-cheap flag is on, we eta-expand through
 573         -- dictionary bindings.  This improves arities. Thereby, it also
 574         -- means that full laziness is less prone to floating out the
 575         -- application of a function to its dictionary arguments, which
 576         -- can thereby lose opportunities for fusion.  Example:
 577         --      foo :: Ord a => a -> ...
 578         --      foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....
 579         --              -- So foo has arity 1
 580         --
 581         --      f = \x. foo dInt $ bar x
 582         --
 583         -- The (foo DInt) is floated out, and makes ineffective a RULE
 584         --      foo (bar x) = ...
 585         --
 586         -- One could go further and make exprIsCheap reply True to any
 587         -- dictionary-typed expression, but that's more work.
 588         --
 589         -- See Note [Dictionary-like types] in TcType.lhs for why we use
 590         -- isDictLikeTy here rather than isDictTy
 591
 592 arityType dicts_cheap (Note n e)
 593   | notSccNote n                 = arityType dicts_cheap e
 594 arityType dicts_cheap (Cast e _) = arityType dicts_cheap e
 595 arityType _           _          = vanillaArityType
 596 \end{code}
 597
 598
 599 %************************************************************************
 600 %*                                                                      *
 601               The main eta-expander
 602 %*                                                                      *
 603 %************************************************************************
 604
 605 We go for:
 606    f = \x1..xn -> N  ==>   f = \x1..xn y1..ym -> N y1..ym
 607                                  (n >= 0)
 608
 609 where (in both cases)
 610
 611         * The xi can include type variables
 612
 613         * The yi are all value variables
 614
 615         * N is a NORMAL FORM (i.e. no redexes anywhere)
 616           wanting a suitable number of extra args.
 617
 618 The biggest reason for doing this is for cases like
 619
 620         f = \x -> case x of
 621                     True  -> \y -> e1
 622                     False -> \y -> e2
 623
 624 Here we want to get the lambdas together.  A good exmaple is the nofib
 625 program fibheaps, which gets 25% more allocation if you don't do this
 626 eta-expansion.
 627
 628 We may have to sandwich some coerces between the lambdas
 629 to make the types work.   exprEtaExpandArity looks through coerces
 630 when computing arity; and etaExpand adds the coerces as necessary when
 631 actually computing the expansion.
 632
 633
 634 Note [No crap in eta-expanded code]
 635 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 636 The eta expander is careful not to introduce "crap".  In particular,
 637 given a CoreExpr satisfying the 'CpeRhs' invariant (in CorePrep), it
 638 returns a CoreExpr satisfying the same invariant. See Note [Eta
 639 expansion and the CorePrep invariants] in CorePrep.
 640
 641 This means the eta-expander has to do a bit of on-the-fly
 642 simplification but it's not too hard.  The alernative, of relying on
 643 a subsequent clean-up phase of the Simplifier to de-crapify the result,
 644 means you can't really use it in CorePrep, which is painful.
 645
 646 Note [Eta expansion and SCCs]
 647 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 648 Note that SCCs are not treated specially by etaExpand.  If we have
 649         etaExpand 2 (\x -> scc "foo" e)
 650         = (\xy -> (scc "foo" e) y)
 651 So the costs of evaluating 'e' (not 'e y') are attributed to "foo"
 652
 653 \begin{code}
 654 -- | @etaExpand n us e ty@ returns an expression with
 655 -- the same meaning as @e@, but with arity @n@.
 656 --
 657 -- Given:
 658 --
 659 -- > e' = etaExpand n us e ty
 660 --
 661 -- We should have that:
 662 --
 663 -- > ty = exprType e = exprType e'
 664 etaExpand :: Arity              -- ^ Result should have this number of value args
 665           -> CoreExpr           -- ^ Expression to expand
 666           -> CoreExpr
 667 -- etaExpand deals with for-alls. For example:
 668 --              etaExpand 1 E
 669 -- where  E :: forall a. a -> a
 670 -- would return
 671 --      (/\b. \y::a -> E b y)
 672 --
 673 -- It deals with coerces too, though they are now rare
 674 -- so perhaps the extra code isn't worth it
 675
 676 etaExpand n orig_expr
 677   = go n orig_expr
 678   where
 679       -- Strip off existing lambdas and casts
 680       -- Note [Eta expansion and SCCs]
 681     go 0 expr = expr
 682     go n (Lam v body) | isTyCoVar v = Lam v (go n     body)
 683                       | otherwise   = Lam v (go (n-1) body)
 684     go n (Cast expr co) = Cast (go n expr) co
 685     go n expr           = -- pprTrace "ee" (vcat [ppr orig_expr, ppr expr, ppr etas]) $
 686                           etaInfoAbs etas (etaInfoApp subst' expr etas)
 687                         where
 688                             in_scope = mkInScopeSet (exprFreeVars expr)
 689                             (in_scope', etas) = mkEtaWW n in_scope (exprType expr)
 690                             subst' = mkEmptySubst in_scope'
 691
 692                                 -- Wrapper    Unwrapper
 693 --------------
 694 data EtaInfo = EtaVar Var       -- /\a. [],   [] a
 695                                 -- \x.  [],   [] x
 696              | EtaCo Coercion   -- [] |> co,  [] |> (sym co)
 697
 698 instance Outputable EtaInfo where
 699    ppr (EtaVar v) = ptext (sLit "EtaVar") <+> ppr v
 700    ppr (EtaCo co) = ptext (sLit "EtaCo")  <+> ppr co
 701
 702 pushCoercion :: Coercion -> [EtaInfo] -> [EtaInfo]
 703 pushCoercion co1 (EtaCo co2 : eis)
 704   | isIdentityCoercion co = eis
 705   | otherwise             = EtaCo co : eis
 706   where
 707     co = co1 `mkTransCoercion` co2
 708
 709 pushCoercion co eis = EtaCo co : eis
 710
 711 --------------
 712 etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr
 713 etaInfoAbs []               expr = expr
 714 etaInfoAbs (EtaVar v : eis) expr = Lam v (etaInfoAbs eis expr)
 715 etaInfoAbs (EtaCo co : eis) expr = Cast (etaInfoAbs eis expr) (mkSymCoercion co)
 716
 717 --------------
 718 etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
 719 -- (etaInfoApp s e eis) returns something equivalent to
 720 --             ((substExpr s e) `appliedto` eis)
 721
 722 etaInfoApp subst (Lam v1 e) (EtaVar v2 : eis)
 723   = etaInfoApp subst' e eis
 724   where
 725     subst' | isTyCoVar v1 = CoreSubst.extendTvSubst subst v1 (mkTyVarTy v2)
 726            | otherwise  = CoreSubst.extendIdSubst subst v1 (Var v2)
 727
 728 etaInfoApp subst (Cast e co1) eis
 729   = etaInfoApp subst e (pushCoercion co' eis)
 730   where
 731     co' = CoreSubst.substTy subst co1
 732
 733 etaInfoApp subst (Case e b _ alts) eis
 734   = Case (subst_expr subst e) b1 (coreAltsType alts') alts'
 735   where
 736     (subst1, b1) = substBndr subst b
 737     alts' = map subst_alt alts
 738     subst_alt (con, bs, rhs) = (con, bs', etaInfoApp subst2 rhs eis)
 739               where
 740                  (subst2,bs') = substBndrs subst1 bs
 741
 742 etaInfoApp subst (Let b e) eis
 743   = Let b' (etaInfoApp subst' e eis)
 744   where
 745     (subst', b') = subst_bind subst b
 746
 747 etaInfoApp subst (Note note e) eis
 748   = Note note (etaInfoApp subst e eis)
 749
 750 etaInfoApp subst e eis
 751   = go (subst_expr subst e) eis
 752   where
 753     go e []                  = e
 754     go e (EtaVar v    : eis) = go (App e (varToCoreExpr v)) eis
 755     go e (EtaCo co    : eis) = go (Cast e co) eis
 756
 757 --------------
 758 mkEtaWW :: Arity -> InScopeSet -> Type
 759         -> (InScopeSet, [EtaInfo])
 760         -- EtaInfo contains fresh variables,
 761         --   not free in the incoming CoreExpr
 762         -- Outgoing InScopeSet includes the EtaInfo vars
 763         --   and the original free vars
 764
 765 mkEtaWW orig_n in_scope orig_ty
 766   = go orig_n empty_subst orig_ty []
 767   where
 768     empty_subst = mkTvSubst in_scope emptyTvSubstEnv
 769
 770     go n subst ty eis       -- See Note [exprArity invariant]
 771        | n == 0
 772        = (getTvInScope subst, reverse eis)
 773
 774        | Just (tv,ty') <- splitForAllTy_maybe ty
 775        , let (subst', tv') = substTyVarBndr subst tv
 776            -- Avoid free vars of the original expression
 777        = go n subst' ty' (EtaVar tv' : eis)
 778
 779        | Just (arg_ty, res_ty) <- splitFunTy_maybe ty
 780        , let (subst', eta_id') = freshEtaId n subst arg_ty
 781            -- Avoid free vars of the original expression
 782        = go (n-1) subst' res_ty (EtaVar eta_id' : eis)
 783
 784        | Just(ty',co) <- splitNewTypeRepCo_maybe ty
 785        =        -- Given this:
 786                 --      newtype T = MkT ([T] -> Int)
 787                 -- Consider eta-expanding this
 788                 --      eta_expand 1 e T
 789                 -- We want to get
 790                 --      coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
 791          go n subst ty' (EtaCo (Type.substTy subst co) : eis)
 792
 793        | otherwise       -- We have an expression of arity > 0,
 794                          -- but its type isn't a function.
 795        = WARN( True, ppr orig_n <+> ppr orig_ty )
 796          (getTvInScope subst, reverse eis)
 797         -- This *can* legitmately happen:
 798         -- e.g.  coerce Int (\x. x) Essentially the programmer is
 799         -- playing fast and loose with types (Happy does this a lot).
 800         -- So we simply decline to eta-expand.  Otherwise we'd end up
 801         -- with an explicit lambda having a non-function type
 802
 803
 804 --------------
 805 -- Avoiding unnecessary substitution; use short-cutting versions
 806
 807 subst_expr :: Subst -> CoreExpr -> CoreExpr
 808 subst_expr = substExprSC (text "CoreArity:substExpr")
 809
 810 subst_bind :: Subst -> CoreBind -> (Subst, CoreBind)
 811 subst_bind = substBindSC
 812
 813
 814 --------------
 815 freshEtaId :: Int -> TvSubst -> Type -> (TvSubst, Id)
 816 -- Make a fresh Id, with specified type (after applying substitution)
 817 -- It should be "fresh" in the sense that it's not in the in-scope set
 818 -- of the TvSubstEnv; and it should itself then be added to the in-scope
 819 -- set of the TvSubstEnv
 820 --
 821 -- The Int is just a reasonable starting point for generating a unique;
 822 -- it does not necessarily have to be unique itself.
 823 freshEtaId n subst ty
 824       = (subst', eta_id')
 825       where
 826         ty'     = Type.substTy subst ty
 827         eta_id' = uniqAway (getTvInScope subst) $
 828                   mkSysLocal (fsLit "eta") (mkBuiltinUnique n) ty'
 829         subst'  = extendTvInScope subst eta_id'
 830 \end{code}
 831