812f12ca8375913ef9da08495984dacde5eb3452
[ghc.git] / compiler / coreSyn / CoreArity.hs
1 {-
2 (c) The University of Glasgow 2006
3 (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
4
5
6 Arity and eta expansion
7 -}
8
9 {-# LANGUAGE CPP #-}
10
11 -- | Arity and eta expansion
12 module CoreArity (
13 manifestArity, exprArity, typeArity, exprBotStrictness_maybe,
14 exprEtaExpandArity, findRhsArity, CheapFun, etaExpand
15 ) where
16
17 #include "HsVersions.h"
18
19 import CoreSyn
20 import CoreFVs
21 import CoreUtils
22 import CoreSubst
23 import Demand
24 import Var
25 import VarEnv
26 import Id
27 import Type
28 import TyCon ( initRecTc, checkRecTc )
29 import Coercion
30 import BasicTypes
31 import Unique
32 import DynFlags ( DynFlags, GeneralFlag(..), gopt )
33 import Outputable
34 import FastString
35 import Pair
36 import Util ( debugIsOn )
37
38 {-
39 ************************************************************************
40 * *
41 manifestArity and exprArity
42 * *
43 ************************************************************************
44
45 exprArity is a cheap-and-cheerful version of exprEtaExpandArity.
46 It tells how many things the expression can be applied to before doing
47 any work. It doesn't look inside cases, lets, etc. The idea is that
48 exprEtaExpandArity will do the hard work, leaving something that's easy
49 for exprArity to grapple with. In particular, Simplify uses exprArity to
50 compute the ArityInfo for the Id.
51
52 Originally I thought that it was enough just to look for top-level lambdas, but
53 it isn't. I've seen this
54
55 foo = PrelBase.timesInt
56
57 We want foo to get arity 2 even though the eta-expander will leave it
58 unchanged, in the expectation that it'll be inlined. But occasionally it
59 isn't, because foo is blacklisted (used in a rule).
60
61 Similarly, see the ok_note check in exprEtaExpandArity. So
62 f = __inline_me (\x -> e)
63 won't be eta-expanded.
64
65 And in any case it seems more robust to have exprArity be a bit more intelligent.
66 But note that (\x y z -> f x y z)
67 should have arity 3, regardless of f's arity.
68 -}
69
70 manifestArity :: CoreExpr -> Arity
71 -- ^ manifestArity sees how many leading value lambdas there are,
72 -- after looking through casts
73 manifestArity (Lam v e) | isId v = 1 + manifestArity e
74 | otherwise = manifestArity e
75 manifestArity (Tick t e) | not (tickishIsCode t) = manifestArity e
76 manifestArity (Cast e _) = manifestArity e
77 manifestArity _ = 0
78
79 ---------------
80 exprArity :: CoreExpr -> Arity
81 -- ^ An approximate, fast, version of 'exprEtaExpandArity'
82 exprArity e = go e
83 where
84 go (Var v) = idArity v
85 go (Lam x e) | isId x = go e + 1
86 | otherwise = go e
87 go (Tick t e) | not (tickishIsCode t) = go e
88 go (Cast e co) = trim_arity (go e) (pSnd (coercionKind co))
89 -- Note [exprArity invariant]
90 go (App e (Type _)) = go e
91 go (App f a) | exprIsTrivial a = (go f - 1) `max` 0
92 -- See Note [exprArity for applications]
93 -- NB: coercions count as a value argument
94
95 go _ = 0
96
97 trim_arity :: Arity -> Type -> Arity
98 trim_arity arity ty = arity `min` length (typeArity ty)
99
100 ---------------
101 typeArity :: Type -> [OneShotInfo]
102 -- How many value arrows are visible in the type?
103 -- We look through foralls, and newtypes
104 -- See Note [exprArity invariant]
105 typeArity ty
106 = go initRecTc ty
107 where
108 go rec_nts ty
109 | Just (bndr, ty') <- splitPiTy_maybe ty
110 = if isIdLikeBinder bndr
111 then typeOneShot (binderType bndr) : go rec_nts ty'
112 else go rec_nts ty'
113
114 | Just (tc,tys) <- splitTyConApp_maybe ty
115 , Just (ty', _) <- instNewTyCon_maybe tc tys
116 , Just rec_nts' <- checkRecTc rec_nts tc -- See Note [Expanding newtypes]
117 -- in TyCon
118 -- , not (isClassTyCon tc) -- Do not eta-expand through newtype classes
119 -- -- See Note [Newtype classes and eta expansion]
120 -- (no longer required)
121 = go rec_nts' ty'
122 -- Important to look through non-recursive newtypes, so that, eg
123 -- (f x) where f has arity 2, f :: Int -> IO ()
124 -- Here we want to get arity 1 for the result!
125 --
126 -- AND through a layer of recursive newtypes
127 -- e.g. newtype Stream m a b = Stream (m (Either b (a, Stream m a b)))
128
129 | otherwise
130 = []
131
132 ---------------
133 exprBotStrictness_maybe :: CoreExpr -> Maybe (Arity, StrictSig)
134 -- A cheap and cheerful function that identifies bottoming functions
135 -- and gives them a suitable strictness signatures. It's used during
136 -- float-out
137 exprBotStrictness_maybe e
138 = case getBotArity (arityType env e) of
139 Nothing -> Nothing
140 Just ar -> Just (ar, sig ar)
141 where
142 env = AE { ae_ped_bot = True, ae_cheap_fn = \ _ _ -> False }
143 sig ar = mkClosedStrictSig (replicate ar topDmd) exnRes
144 -- For this purpose we can be very simple
145 -- exnRes is a bit less aggressive than botRes
146
147 {-
148 Note [exprArity invariant]
149 ~~~~~~~~~~~~~~~~~~~~~~~~~~
150 exprArity has the following invariant:
151
152 (1) If typeArity (exprType e) = n,
153 then manifestArity (etaExpand e n) = n
154
155 That is, etaExpand can always expand as much as typeArity says
156 So the case analysis in etaExpand and in typeArity must match
157
158 (2) exprArity e <= typeArity (exprType e)
159
160 (3) Hence if (exprArity e) = n, then manifestArity (etaExpand e n) = n
161
162 That is, if exprArity says "the arity is n" then etaExpand really
163 can get "n" manifest lambdas to the top.
164
165 Why is this important? Because
166 - In TidyPgm we use exprArity to fix the *final arity* of
167 each top-level Id, and in
168 - In CorePrep we use etaExpand on each rhs, so that the visible lambdas
169 actually match that arity, which in turn means
170 that the StgRhs has the right number of lambdas
171
172 An alternative would be to do the eta-expansion in TidyPgm, at least
173 for top-level bindings, in which case we would not need the trim_arity
174 in exprArity. That is a less local change, so I'm going to leave it for today!
175
176 Note [Newtype classes and eta expansion]
177 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
178 NB: this nasty special case is no longer required, because
179 for newtype classes we don't use the class-op rule mechanism
180 at all. See Note [Single-method classes] in TcInstDcls. SLPJ May 2013
181
182 -------- Old out of date comments, just for interest -----------
183 We have to be careful when eta-expanding through newtypes. In general
184 it's a good idea, but annoyingly it interacts badly with the class-op
185 rule mechanism. Consider
186
187 class C a where { op :: a -> a }
188 instance C b => C [b] where
189 op x = ...
190
191 These translate to
192
193 co :: forall a. (a->a) ~ C a
194
195 $copList :: C b -> [b] -> [b]
196 $copList d x = ...
197
198 $dfList :: C b -> C [b]
199 {-# DFunUnfolding = [$copList] #-}
200 $dfList d = $copList d |> co@[b]
201
202 Now suppose we have:
203
204 dCInt :: C Int
205
206 blah :: [Int] -> [Int]
207 blah = op ($dfList dCInt)
208
209 Now we want the built-in op/$dfList rule will fire to give
210 blah = $copList dCInt
211
212 But with eta-expansion 'blah' might (and in Trac #3772, which is
213 slightly more complicated, does) turn into
214
215 blah = op (\eta. ($dfList dCInt |> sym co) eta)
216
217 and now it is *much* harder for the op/$dfList rule to fire, because
218 exprIsConApp_maybe won't hold of the argument to op. I considered
219 trying to *make* it hold, but it's tricky and I gave up.
220
221 The test simplCore/should_compile/T3722 is an excellent example.
222 -------- End of old out of date comments, just for interest -----------
223
224
225 Note [exprArity for applications]
226 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
227 When we come to an application we check that the arg is trivial.
228 eg f (fac x) does not have arity 2,
229 even if f has arity 3!
230
231 * We require that is trivial rather merely cheap. Suppose f has arity 2.
232 Then f (Just y)
233 has arity 0, because if we gave it arity 1 and then inlined f we'd get
234 let v = Just y in \w. <f-body>
235 which has arity 0. And we try to maintain the invariant that we don't
236 have arity decreases.
237
238 * The `max 0` is important! (\x y -> f x) has arity 2, even if f is
239 unknown, hence arity 0
240
241
242 ************************************************************************
243 * *
244 Computing the "arity" of an expression
245 * *
246 ************************************************************************
247
248 Note [Definition of arity]
249 ~~~~~~~~~~~~~~~~~~~~~~~~~~
250 The "arity" of an expression 'e' is n if
251 applying 'e' to *fewer* than n *value* arguments
252 converges rapidly
253
254 Or, to put it another way
255
256 there is no work lost in duplicating the partial
257 application (e x1 .. x(n-1))
258
259 In the divegent case, no work is lost by duplicating because if the thing
260 is evaluated once, that's the end of the program.
261
262 Or, to put it another way, in any context C
263
264 C[ (\x1 .. xn. e x1 .. xn) ]
265 is as efficient as
266 C[ e ]
267
268 It's all a bit more subtle than it looks:
269
270 Note [One-shot lambdas]
271 ~~~~~~~~~~~~~~~~~~~~~~~
272 Consider one-shot lambdas
273 let x = expensive in \y z -> E
274 We want this to have arity 1 if the \y-abstraction is a 1-shot lambda.
275
276 Note [Dealing with bottom]
277 ~~~~~~~~~~~~~~~~~~~~~~~~~~
278 A Big Deal with computing arities is expressions like
279
280 f = \x -> case x of
281 True -> \s -> e1
282 False -> \s -> e2
283
284 This happens all the time when f :: Bool -> IO ()
285 In this case we do eta-expand, in order to get that \s to the
286 top, and give f arity 2.
287
288 This isn't really right in the presence of seq. Consider
289 (f bot) `seq` 1
290
291 This should diverge! But if we eta-expand, it won't. We ignore this
292 "problem" (unless -fpedantic-bottoms is on), because being scrupulous
293 would lose an important transformation for many programs. (See
294 Trac #5587 for an example.)
295
296 Consider also
297 f = \x -> error "foo"
298 Here, arity 1 is fine. But if it is
299 f = \x -> case x of
300 True -> error "foo"
301 False -> \y -> x+y
302 then we want to get arity 2. Technically, this isn't quite right, because
303 (f True) `seq` 1
304 should diverge, but it'll converge if we eta-expand f. Nevertheless, we
305 do so; it improves some programs significantly, and increasing convergence
306 isn't a bad thing. Hence the ABot/ATop in ArityType.
307
308 So these two transformations aren't always the Right Thing, and we
309 have several tickets reporting unexpected bahaviour resulting from
310 this transformation. So we try to limit it as much as possible:
311
312 (1) Do NOT move a lambda outside a known-bottom case expression
313 case undefined of { (a,b) -> \y -> e }
314 This showed up in Trac #5557
315
316 (2) Do NOT move a lambda outside a case if all the branches of
317 the case are known to return bottom.
318 case x of { (a,b) -> \y -> error "urk" }
319 This case is less important, but the idea is that if the fn is
320 going to diverge eventually anyway then getting the best arity
321 isn't an issue, so we might as well play safe
322
323 (3) Do NOT move a lambda outside a case unless
324 (a) The scrutinee is ok-for-speculation, or
325 (b) more liberally: the scrutinee is cheap (e.g. a variable), and
326 -fpedantic-bottoms is not enforced (see Trac #2915 for an example)
327
328 Of course both (1) and (2) are readily defeated by disguising the bottoms.
329
330 4. Note [Newtype arity]
331 ~~~~~~~~~~~~~~~~~~~~~~~~
332 Non-recursive newtypes are transparent, and should not get in the way.
333 We do (currently) eta-expand recursive newtypes too. So if we have, say
334
335 newtype T = MkT ([T] -> Int)
336
337 Suppose we have
338 e = coerce T f
339 where f has arity 1. Then: etaExpandArity e = 1;
340 that is, etaExpandArity looks through the coerce.
341
342 When we eta-expand e to arity 1: eta_expand 1 e T
343 we want to get: coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
344
345 HOWEVER, note that if you use coerce bogusly you can ge
346 coerce Int negate
347 And since negate has arity 2, you might try to eta expand. But you can't
348 decopose Int to a function type. Hence the final case in eta_expand.
349
350 Note [The state-transformer hack]
351 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
352 Suppose we have
353 f = e
354 where e has arity n. Then, if we know from the context that f has
355 a usage type like
356 t1 -> ... -> tn -1-> t(n+1) -1-> ... -1-> tm -> ...
357 then we can expand the arity to m. This usage type says that
358 any application (x e1 .. en) will be applied to uniquely to (m-n) more args
359 Consider f = \x. let y = <expensive>
360 in case x of
361 True -> foo
362 False -> \(s:RealWorld) -> e
363 where foo has arity 1. Then we want the state hack to
364 apply to foo too, so we can eta expand the case.
365
366 Then we expect that if f is applied to one arg, it'll be applied to two
367 (that's the hack -- we don't really know, and sometimes it's false)
368 See also Id.isOneShotBndr.
369
370 Note [State hack and bottoming functions]
371 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
372 It's a terrible idea to use the state hack on a bottoming function.
373 Here's what happens (Trac #2861):
374
375 f :: String -> IO T
376 f = \p. error "..."
377
378 Eta-expand, using the state hack:
379
380 f = \p. (\s. ((error "...") |> g1) s) |> g2
381 g1 :: IO T ~ (S -> (S,T))
382 g2 :: (S -> (S,T)) ~ IO T
383
384 Extrude the g2
385
386 f' = \p. \s. ((error "...") |> g1) s
387 f = f' |> (String -> g2)
388
389 Discard args for bottomming function
390
391 f' = \p. \s. ((error "...") |> g1 |> g3
392 g3 :: (S -> (S,T)) ~ (S,T)
393
394 Extrude g1.g3
395
396 f'' = \p. \s. (error "...")
397 f' = f'' |> (String -> S -> g1.g3)
398
399 And now we can repeat the whole loop. Aargh! The bug is in applying the
400 state hack to a function which then swallows the argument.
401
402 This arose in another guise in Trac #3959. Here we had
403
404 catch# (throw exn >> return ())
405
406 Note that (throw :: forall a e. Exn e => e -> a) is called with [a = IO ()].
407 After inlining (>>) we get
408
409 catch# (\_. throw {IO ()} exn)
410
411 We must *not* eta-expand to
412
413 catch# (\_ _. throw {...} exn)
414
415 because 'catch#' expects to get a (# _,_ #) after applying its argument to
416 a State#, not another function!
417
418 In short, we use the state hack to allow us to push let inside a lambda,
419 but not to introduce a new lambda.
420
421
422 Note [ArityType]
423 ~~~~~~~~~~~~~~~~
424 ArityType is the result of a compositional analysis on expressions,
425 from which we can decide the real arity of the expression (extracted
426 with function exprEtaExpandArity).
427
428 Here is what the fields mean. If an arbitrary expression 'f' has
429 ArityType 'at', then
430
431 * If at = ABot n, then (f x1..xn) definitely diverges. Partial
432 applications to fewer than n args may *or may not* diverge.
433
434 We allow ourselves to eta-expand bottoming functions, even
435 if doing so may lose some `seq` sharing,
436 let x = <expensive> in \y. error (g x y)
437 ==> \y. let x = <expensive> in error (g x y)
438
439 * If at = ATop as, and n=length as,
440 then expanding 'f' to (\x1..xn. f x1 .. xn) loses no sharing,
441 assuming the calls of f respect the one-shot-ness of
442 its definition.
443
444 NB 'f' is an arbitrary expression, eg (f = g e1 e2). This 'f'
445 can have ArityType as ATop, with length as > 0, only if e1 e2 are
446 themselves.
447
448 * In both cases, f, (f x1), ... (f x1 ... f(n-1)) are definitely
449 really functions, or bottom, but *not* casts from a data type, in
450 at least one case branch. (If it's a function in one case branch but
451 an unsafe cast from a data type in another, the program is bogus.)
452 So eta expansion is dynamically ok; see Note [State hack and
453 bottoming functions], the part about catch#
454
455 Example:
456 f = \x\y. let v = <expensive> in
457 \s(one-shot) \t(one-shot). blah
458 'f' has ArityType [ManyShot,ManyShot,OneShot,OneShot]
459 The one-shot-ness means we can, in effect, push that
460 'let' inside the \st.
461
462
463 Suppose f = \xy. x+y
464 Then f :: AT [False,False] ATop
465 f v :: AT [False] ATop
466 f <expensive> :: AT [] ATop
467
468 -------------------- Main arity code ----------------------------
469 -}
470
471 -- See Note [ArityType]
472 data ArityType = ATop [OneShotInfo] | ABot Arity
473 -- There is always an explicit lambda
474 -- to justify the [OneShot], or the Arity
475
476 vanillaArityType :: ArityType
477 vanillaArityType = ATop [] -- Totally uninformative
478
479 -- ^ The Arity returned is the number of value args the
480 -- expression can be applied to without doing much work
481 exprEtaExpandArity :: DynFlags -> CoreExpr -> Arity
482 -- exprEtaExpandArity is used when eta expanding
483 -- e ==> \xy -> e x y
484 exprEtaExpandArity dflags e
485 = case (arityType env e) of
486 ATop oss -> length oss
487 ABot n -> n
488 where
489 env = AE { ae_cheap_fn = mk_cheap_fn dflags isCheapApp
490 , ae_ped_bot = gopt Opt_PedanticBottoms dflags }
491
492 getBotArity :: ArityType -> Maybe Arity
493 -- Arity of a divergent function
494 getBotArity (ABot n) = Just n
495 getBotArity _ = Nothing
496
497 mk_cheap_fn :: DynFlags -> CheapAppFun -> CheapFun
498 mk_cheap_fn dflags cheap_app
499 | not (gopt Opt_DictsCheap dflags)
500 = \e _ -> exprIsCheap' cheap_app e
501 | otherwise
502 = \e mb_ty -> exprIsCheap' cheap_app e
503 || case mb_ty of
504 Nothing -> False
505 Just ty -> isDictLikeTy ty
506
507
508 ----------------------
509 findRhsArity :: DynFlags -> Id -> CoreExpr -> Arity -> Arity
510 -- This implements the fixpoint loop for arity analysis
511 -- See Note [Arity analysis]
512 findRhsArity dflags bndr rhs old_arity
513 = go (rhsEtaExpandArity dflags init_cheap_app rhs)
514 -- We always call exprEtaExpandArity once, but usually
515 -- that produces a result equal to old_arity, and then
516 -- we stop right away (since arities should not decrease)
517 -- Result: the common case is that there is just one iteration
518 where
519 init_cheap_app :: CheapAppFun
520 init_cheap_app fn n_val_args
521 | fn == bndr = True -- On the first pass, this binder gets infinite arity
522 | otherwise = isCheapApp fn n_val_args
523
524 go :: Arity -> Arity
525 go cur_arity
526 | cur_arity <= old_arity = cur_arity
527 | new_arity == cur_arity = cur_arity
528 | otherwise = ASSERT( new_arity < cur_arity )
529 #ifdef DEBUG
530 pprTrace "Exciting arity"
531 (vcat [ ppr bndr <+> ppr cur_arity <+> ppr new_arity
532 , ppr rhs])
533 #endif
534 go new_arity
535 where
536 new_arity = rhsEtaExpandArity dflags cheap_app rhs
537
538 cheap_app :: CheapAppFun
539 cheap_app fn n_val_args
540 | fn == bndr = n_val_args < cur_arity
541 | otherwise = isCheapApp fn n_val_args
542
543 -- ^ The Arity returned is the number of value args the
544 -- expression can be applied to without doing much work
545 rhsEtaExpandArity :: DynFlags -> CheapAppFun -> CoreExpr -> Arity
546 -- exprEtaExpandArity is used when eta expanding
547 -- e ==> \xy -> e x y
548 rhsEtaExpandArity dflags cheap_app e
549 = case (arityType env e) of
550 ATop (os:oss)
551 | isOneShotInfo os || has_lam e -> 1 + length oss
552 -- Don't expand PAPs/thunks
553 -- Note [Eta expanding thunks]
554 | otherwise -> 0
555 ATop [] -> 0
556 ABot n -> n
557 where
558 env = AE { ae_cheap_fn = mk_cheap_fn dflags cheap_app
559 , ae_ped_bot = gopt Opt_PedanticBottoms dflags }
560
561 has_lam (Tick _ e) = has_lam e
562 has_lam (Lam b e) = isId b || has_lam e
563 has_lam _ = False
564
565 {-
566 Note [Arity analysis]
567 ~~~~~~~~~~~~~~~~~~~~~
568 The motivating example for arity analysis is this:
569
570 f = \x. let g = f (x+1)
571 in \y. ...g...
572
573 What arity does f have? Really it should have arity 2, but a naive
574 look at the RHS won't see that. You need a fixpoint analysis which
575 says it has arity "infinity" the first time round.
576
577 This example happens a lot; it first showed up in Andy Gill's thesis,
578 fifteen years ago! It also shows up in the code for 'rnf' on lists
579 in Trac #4138.
580
581 The analysis is easy to achieve because exprEtaExpandArity takes an
582 argument
583 type CheapFun = CoreExpr -> Maybe Type -> Bool
584 used to decide if an expression is cheap enough to push inside a
585 lambda. And exprIsCheap' in turn takes an argument
586 type CheapAppFun = Id -> Int -> Bool
587 which tells when an application is cheap. This makes it easy to
588 write the analysis loop.
589
590 The analysis is cheap-and-cheerful because it doesn't deal with
591 mutual recursion. But the self-recursive case is the important one.
592
593
594 Note [Eta expanding through dictionaries]
595 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
596 If the experimental -fdicts-cheap flag is on, we eta-expand through
597 dictionary bindings. This improves arities. Thereby, it also
598 means that full laziness is less prone to floating out the
599 application of a function to its dictionary arguments, which
600 can thereby lose opportunities for fusion. Example:
601 foo :: Ord a => a -> ...
602 foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....
603 -- So foo has arity 1
604
605 f = \x. foo dInt $ bar x
606
607 The (foo DInt) is floated out, and makes ineffective a RULE
608 foo (bar x) = ...
609
610 One could go further and make exprIsCheap reply True to any
611 dictionary-typed expression, but that's more work.
612
613 See Note [Dictionary-like types] in TcType.hs for why we use
614 isDictLikeTy here rather than isDictTy
615
616 Note [Eta expanding thunks]
617 ~~~~~~~~~~~~~~~~~~~~~~~~~~~
618 We don't eta-expand
619 * Trivial RHSs x = y
620 * PAPs x = map g
621 * Thunks f = case y of p -> \x -> blah
622
623 When we see
624 f = case y of p -> \x -> blah
625 should we eta-expand it? Well, if 'x' is a one-shot state token
626 then 'yes' because 'f' will only be applied once. But otherwise
627 we (conservatively) say no. My main reason is to avoid expanding
628 PAPSs
629 f = g d ==> f = \x. g d x
630 because that might in turn make g inline (if it has an inline pragma),
631 which we might not want. After all, INLINE pragmas say "inline only
632 when saturated" so we don't want to be too gung-ho about saturating!
633 -}
634
635 arityLam :: Id -> ArityType -> ArityType
636 arityLam id (ATop as) = ATop (idStateHackOneShotInfo id : as)
637 arityLam _ (ABot n) = ABot (n+1)
638
639 floatIn :: Bool -> ArityType -> ArityType
640 -- We have something like (let x = E in b),
641 -- where b has the given arity type.
642 floatIn _ (ABot n) = ABot n
643 floatIn True (ATop as) = ATop as
644 floatIn False (ATop as) = ATop (takeWhile isOneShotInfo as)
645 -- If E is not cheap, keep arity only for one-shots
646
647 arityApp :: ArityType -> Bool -> ArityType
648 -- Processing (fun arg) where at is the ArityType of fun,
649 -- Knock off an argument and behave like 'let'
650 arityApp (ABot 0) _ = ABot 0
651 arityApp (ABot n) _ = ABot (n-1)
652 arityApp (ATop []) _ = ATop []
653 arityApp (ATop (_:as)) cheap = floatIn cheap (ATop as)
654
655 andArityType :: ArityType -> ArityType -> ArityType -- Used for branches of a 'case'
656 andArityType (ABot n1) (ABot n2)
657 = ABot (n1 `min` n2)
658 andArityType (ATop as) (ABot _) = ATop as
659 andArityType (ABot _) (ATop bs) = ATop bs
660 andArityType (ATop as) (ATop bs) = ATop (as `combine` bs)
661 where -- See Note [Combining case branches]
662 combine (a:as) (b:bs) = (a `bestOneShot` b) : combine as bs
663 combine [] bs = takeWhile isOneShotInfo bs
664 combine as [] = takeWhile isOneShotInfo as
665
666 {-
667 Note [Combining case branches]
668 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
669 Consider
670 go = \x. let z = go e0
671 go2 = \x. case x of
672 True -> z
673 False -> \s(one-shot). e1
674 in go2 x
675 We *really* want to eta-expand go and go2.
676 When combining the barnches of the case we have
677 ATop [] `andAT` ATop [OneShotLam]
678 and we want to get ATop [OneShotLam]. But if the inner
679 lambda wasn't one-shot we don't want to do this.
680 (We need a proper arity analysis to justify that.)
681
682 So we combine the best of the two branches, on the (slightly dodgy)
683 basis that if we know one branch is one-shot, then they all must be.
684 -}
685
686 ---------------------------
687 type CheapFun = CoreExpr -> Maybe Type -> Bool
688 -- How to decide if an expression is cheap
689 -- If the Maybe is Just, the type is the type
690 -- of the expression; Nothing means "don't know"
691
692 data ArityEnv
693 = AE { ae_cheap_fn :: CheapFun
694 , ae_ped_bot :: Bool -- True <=> be pedantic about bottoms
695 }
696
697 arityType :: ArityEnv -> CoreExpr -> ArityType
698
699 arityType env (Cast e co)
700 = case arityType env e of
701 ATop os -> ATop (take co_arity os)
702 ABot n -> ABot (n `min` co_arity)
703 where
704 co_arity = length (typeArity (pSnd (coercionKind co)))
705 -- See Note [exprArity invariant] (2); must be true of
706 -- arityType too, since that is how we compute the arity
707 -- of variables, and they in turn affect result of exprArity
708 -- Trac #5441 is a nice demo
709 -- However, do make sure that ATop -> ATop and ABot -> ABot!
710 -- Casts don't affect that part. Getting this wrong provoked #5475
711
712 arityType _ (Var v)
713 | strict_sig <- idStrictness v
714 , not $ isTopSig strict_sig
715 , (ds, res) <- splitStrictSig strict_sig
716 , let arity = length ds
717 = if isBotRes res then ABot arity
718 else ATop (take arity one_shots)
719 | otherwise
720 = ATop (take (idArity v) one_shots)
721 where
722 one_shots :: [OneShotInfo] -- One-shot-ness derived from the type
723 one_shots = typeArity (idType v)
724
725 -- Lambdas; increase arity
726 arityType env (Lam x e)
727 | isId x = arityLam x (arityType env e)
728 | otherwise = arityType env e
729
730 -- Applications; decrease arity, except for types
731 arityType env (App fun (Type _))
732 = arityType env fun
733 arityType env (App fun arg )
734 = arityApp (arityType env fun) (ae_cheap_fn env arg Nothing)
735
736 -- Case/Let; keep arity if either the expression is cheap
737 -- or it's a 1-shot lambda
738 -- The former is not really right for Haskell
739 -- f x = case x of { (a,b) -> \y. e }
740 -- ===>
741 -- f x y = case x of { (a,b) -> e }
742 -- The difference is observable using 'seq'
743 --
744 arityType env (Case scrut _ _ alts)
745 | exprIsBottom scrut || null alts
746 = ABot 0 -- Do not eta expand
747 -- See Note [Dealing with bottom (1)]
748 | otherwise
749 = case alts_type of
750 ABot n | n>0 -> ATop [] -- Don't eta expand
751 | otherwise -> ABot 0 -- if RHS is bottomming
752 -- See Note [Dealing with bottom (2)]
753
754 ATop as | not (ae_ped_bot env) -- See Note [Dealing with bottom (3)]
755 , ae_cheap_fn env scrut Nothing -> ATop as
756 | exprOkForSpeculation scrut -> ATop as
757 | otherwise -> ATop (takeWhile isOneShotInfo as)
758 where
759 alts_type = foldr1 andArityType [arityType env rhs | (_,_,rhs) <- alts]
760
761 arityType env (Let b e)
762 = floatIn (cheap_bind b) (arityType env e)
763 where
764 cheap_bind (NonRec b e) = is_cheap (b,e)
765 cheap_bind (Rec prs) = all is_cheap prs
766 is_cheap (b,e) = ae_cheap_fn env e (Just (idType b))
767
768 arityType env (Tick t e)
769 | not (tickishIsCode t) = arityType env e
770
771 arityType _ _ = vanillaArityType
772
773 {-
774 %************************************************************************
775 %* *
776 The main eta-expander
777 %* *
778 %************************************************************************
779
780 We go for:
781 f = \x1..xn -> N ==> f = \x1..xn y1..ym -> N y1..ym
782 (n >= 0)
783
784 where (in both cases)
785
786 * The xi can include type variables
787
788 * The yi are all value variables
789
790 * N is a NORMAL FORM (i.e. no redexes anywhere)
791 wanting a suitable number of extra args.
792
793 The biggest reason for doing this is for cases like
794
795 f = \x -> case x of
796 True -> \y -> e1
797 False -> \y -> e2
798
799 Here we want to get the lambdas together. A good example is the nofib
800 program fibheaps, which gets 25% more allocation if you don't do this
801 eta-expansion.
802
803 We may have to sandwich some coerces between the lambdas
804 to make the types work. exprEtaExpandArity looks through coerces
805 when computing arity; and etaExpand adds the coerces as necessary when
806 actually computing the expansion.
807
808 Note [No crap in eta-expanded code]
809 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
810 The eta expander is careful not to introduce "crap". In particular,
811 given a CoreExpr satisfying the 'CpeRhs' invariant (in CorePrep), it
812 returns a CoreExpr satisfying the same invariant. See Note [Eta
813 expansion and the CorePrep invariants] in CorePrep.
814
815 This means the eta-expander has to do a bit of on-the-fly
816 simplification but it's not too hard. The alernative, of relying on
817 a subsequent clean-up phase of the Simplifier to de-crapify the result,
818 means you can't really use it in CorePrep, which is painful.
819
820 Note [Eta expansion and SCCs]
821 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
822 Note that SCCs are not treated specially by etaExpand. If we have
823 etaExpand 2 (\x -> scc "foo" e)
824 = (\xy -> (scc "foo" e) y)
825 So the costs of evaluating 'e' (not 'e y') are attributed to "foo"
826
827 Note [Eta expansion and source notes]
828 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
829 CorePrep puts floatable ticks outside of value applications, but not
830 type applications. As a result we might be trying to eta-expand an
831 expression like
832
833 (src<...> v) @a
834
835 which we want to lead to code like
836
837 \x -> src<...> v @a x
838
839 This means that we need to look through type applications and be ready
840 to re-add floats on the top.
841
842 -}
843
844 -- | @etaExpand n e@ returns an expression with
845 -- the same meaning as @e@, but with arity @n@.
846 --
847 -- Given:
848 --
849 -- > e' = etaExpand n e
850 --
851 -- We should have that:
852 --
853 -- > ty = exprType e = exprType e'
854 etaExpand :: Arity -- ^ Result should have this number of value args
855 -> CoreExpr -- ^ Expression to expand
856 -> CoreExpr
857 -- etaExpand deals with for-alls. For example:
858 -- etaExpand 1 E
859 -- where E :: forall a. a -> a
860 -- would return
861 -- (/\b. \y::a -> E b y)
862 --
863 -- It deals with coerces too, though they are now rare
864 -- so perhaps the extra code isn't worth it
865
866 etaExpand n orig_expr
867 = go n orig_expr
868 where
869 -- Strip off existing lambdas and casts
870 -- Note [Eta expansion and SCCs]
871 go 0 expr = expr
872 go n (Lam v body) | isTyVar v = Lam v (go n body)
873 | otherwise = Lam v (go (n-1) body)
874 go n (Cast expr co) = Cast (go n expr) co
875 go n expr
876 = -- pprTrace "ee" (vcat [ppr orig_expr, ppr expr, ppr etas]) $
877 retick $ etaInfoAbs etas (etaInfoApp subst' sexpr etas)
878 where
879 in_scope = mkInScopeSet (exprFreeVars expr)
880 (in_scope', etas) = mkEtaWW n orig_expr in_scope (exprType expr)
881 subst' = mkEmptySubst in_scope'
882
883 -- Find ticks behind type apps.
884 -- See Note [Eta expansion and source notes]
885 (expr', args) = collectArgs expr
886 (ticks, expr'') = stripTicksTop tickishFloatable expr'
887 sexpr = foldl App expr'' args
888 retick expr = foldr mkTick expr ticks
889
890 -- Wrapper Unwrapper
891 --------------
892 data EtaInfo = EtaVar Var -- /\a. [], [] a
893 -- \x. [], [] x
894 | EtaCo Coercion -- [] |> co, [] |> (sym co)
895
896 instance Outputable EtaInfo where
897 ppr (EtaVar v) = text "EtaVar" <+> ppr v
898 ppr (EtaCo co) = text "EtaCo" <+> ppr co
899
900 pushCoercion :: Coercion -> [EtaInfo] -> [EtaInfo]
901 pushCoercion co1 (EtaCo co2 : eis)
902 | isReflCo co = eis
903 | otherwise = EtaCo co : eis
904 where
905 co = co1 `mkTransCo` co2
906
907 pushCoercion co eis = EtaCo co : eis
908
909 --------------
910 etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr
911 etaInfoAbs [] expr = expr
912 etaInfoAbs (EtaVar v : eis) expr = Lam v (etaInfoAbs eis expr)
913 etaInfoAbs (EtaCo co : eis) expr = Cast (etaInfoAbs eis expr) (mkSymCo co)
914
915 --------------
916 etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
917 -- (etaInfoApp s e eis) returns something equivalent to
918 -- ((substExpr s e) `appliedto` eis)
919
920 etaInfoApp subst (Lam v1 e) (EtaVar v2 : eis)
921 = etaInfoApp (CoreSubst.extendSubstWithVar subst v1 v2) e eis
922
923 etaInfoApp subst (Cast e co1) eis
924 = etaInfoApp subst e (pushCoercion co' eis)
925 where
926 co' = CoreSubst.substCo subst co1
927
928 etaInfoApp subst (Case e b ty alts) eis
929 = Case (subst_expr subst e) b1 (mk_alts_ty (CoreSubst.substTy subst ty) eis) alts'
930 where
931 (subst1, b1) = substBndr subst b
932 alts' = map subst_alt alts
933 subst_alt (con, bs, rhs) = (con, bs', etaInfoApp subst2 rhs eis)
934 where
935 (subst2,bs') = substBndrs subst1 bs
936
937 mk_alts_ty ty [] = ty
938 mk_alts_ty ty (EtaVar v : eis) = mk_alts_ty (applyTypeToArg ty (varToCoreExpr v)) eis
939 mk_alts_ty _ (EtaCo co : eis) = mk_alts_ty (pSnd (coercionKind co)) eis
940
941 etaInfoApp subst (Let b e) eis
942 = Let b' (etaInfoApp subst' e eis)
943 where
944 (subst', b') = subst_bind subst b
945
946 etaInfoApp subst (Tick t e) eis
947 = Tick (substTickish subst t) (etaInfoApp subst e eis)
948
949 etaInfoApp subst e eis
950 = go (subst_expr subst e) eis
951 where
952 go e [] = e
953 go e (EtaVar v : eis) = go (App e (varToCoreExpr v)) eis
954 go e (EtaCo co : eis) = go (Cast e co) eis
955
956 --------------
957 mkEtaWW :: Arity -> CoreExpr -> InScopeSet -> Type
958 -> (InScopeSet, [EtaInfo])
959 -- EtaInfo contains fresh variables,
960 -- not free in the incoming CoreExpr
961 -- Outgoing InScopeSet includes the EtaInfo vars
962 -- and the original free vars
963
964 mkEtaWW orig_n orig_expr in_scope orig_ty
965 = go orig_n empty_subst orig_ty []
966 where
967 empty_subst = mkEmptyTCvSubst in_scope
968
969 go n subst ty eis -- See Note [exprArity invariant]
970 | n == 0
971 = (getTCvInScope subst, reverse eis)
972
973 | Just (bndr,ty') <- splitPiTy_maybe ty
974 = let ((subst', eta_id'), new_n) = caseBinder bndr
975 (\tv -> (Type.substTyVarBndr subst tv, n))
976 (\arg_ty -> (freshEtaVar n subst arg_ty, n-1))
977 in
978 -- Avoid free vars of the original expression
979 go new_n subst' ty' (EtaVar eta_id' : eis)
980
981 | Just (co, ty') <- topNormaliseNewType_maybe ty
982 = -- Given this:
983 -- newtype T = MkT ([T] -> Int)
984 -- Consider eta-expanding this
985 -- eta_expand 1 e T
986 -- We want to get
987 -- coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
988 go n subst ty' (EtaCo co : eis)
989
990 | otherwise -- We have an expression of arity > 0,
991 -- but its type isn't a function.
992 = WARN( True, (ppr orig_n <+> ppr orig_ty) $$ ppr orig_expr )
993 (getTCvInScope subst, reverse eis)
994 -- This *can* legitmately happen:
995 -- e.g. coerce Int (\x. x) Essentially the programmer is
996 -- playing fast and loose with types (Happy does this a lot).
997 -- So we simply decline to eta-expand. Otherwise we'd end up
998 -- with an explicit lambda having a non-function type
999
1000
1001 --------------
1002 -- Avoiding unnecessary substitution; use short-cutting versions
1003
1004 subst_expr :: Subst -> CoreExpr -> CoreExpr
1005 subst_expr = substExprSC (text "CoreArity:substExpr")
1006
1007 subst_bind :: Subst -> CoreBind -> (Subst, CoreBind)
1008 subst_bind = substBindSC
1009
1010
1011 --------------
1012 freshEtaVar :: Int -> TCvSubst -> Type -> (TCvSubst, Var)
1013 -- Make a fresh Id, with specified type (after applying substitution)
1014 -- It should be "fresh" in the sense that it's not in the in-scope set
1015 -- of the TvSubstEnv; and it should itself then be added to the in-scope
1016 -- set of the TvSubstEnv
1017 --
1018 -- The Int is just a reasonable starting point for generating a unique;
1019 -- it does not necessarily have to be unique itself.
1020 freshEtaVar n subst ty
1021 = (subst', eta_id')
1022 where
1023 ty' = Type.substTy subst ty
1024 eta_id' = uniqAway (getTCvInScope subst) $
1025 mkSysLocalOrCoVar (fsLit "eta") (mkBuiltinUnique n) ty'
1026 subst' = extendTCvInScope subst eta_id'