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