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