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