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