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