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