never executed always true always false
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 {-# OPTIONS_GHC -Wno-incomplete-record-updates #-}
11
12 -- | Arity and eta expansion
13 module GHC.Core.Opt.Arity
14 ( manifestArity, joinRhsArity, exprArity, typeArity
15 , exprEtaExpandArity, findRhsArity
16 , etaExpand, etaExpandAT
17 , exprBotStrictness_maybe
18
19 -- ** ArityType
20 , ArityType(..), mkBotArityType, mkTopArityType, expandableArityType
21 , arityTypeArity, maxWithArity, idArityType
22
23 -- ** Join points
24 , etaExpandToJoinPoint, etaExpandToJoinPointRule
25
26 -- ** Coercions and casts
27 , pushCoArg, pushCoArgs, pushCoValArg, pushCoTyArg
28 , pushCoercionIntoLambda, pushCoDataCon, collectBindersPushingCo
29 )
30 where
31
32 import GHC.Prelude
33
34 import GHC.Driver.Session ( DynFlags, GeneralFlag(..), gopt )
35
36 import GHC.Core
37 import GHC.Core.FVs
38 import GHC.Core.Utils
39 import GHC.Core.DataCon
40 import GHC.Core.TyCon ( tyConArity )
41 import GHC.Core.TyCon.RecWalk ( initRecTc, checkRecTc )
42 import GHC.Core.Predicate ( isDictTy, isCallStackPredTy )
43 import GHC.Core.Multiplicity
44
45 -- We have two sorts of substitution:
46 -- GHC.Core.Subst.Subst, and GHC.Core.TyCo.TCvSubst
47 -- Both have substTy, substCo Hence need for qualification
48 import GHC.Core.Subst as Core
49 import GHC.Core.Type as Type
50 import GHC.Core.Coercion as Type
51
52 import GHC.Types.Demand
53 import GHC.Types.Var
54 import GHC.Types.Var.Env
55 import GHC.Types.Id
56 import GHC.Types.Var.Set
57 import GHC.Types.Basic
58 import GHC.Types.Tickish
59
60 import GHC.Builtin.Uniques
61 import GHC.Data.FastString
62 import GHC.Data.Pair
63
64 import GHC.Utils.Constants (debugIsOn)
65 import GHC.Utils.Outputable
66 import GHC.Utils.Panic
67 import GHC.Utils.Panic.Plain
68 import GHC.Utils.Trace
69 import GHC.Utils.Misc
70
71 {-
72 ************************************************************************
73 * *
74 manifestArity and exprArity
75 * *
76 ************************************************************************
77
78 exprArity is a cheap-and-cheerful version of exprEtaExpandArity.
79 It tells how many things the expression can be applied to before doing
80 any work. It doesn't look inside cases, lets, etc. The idea is that
81 exprEtaExpandArity will do the hard work, leaving something that's easy
82 for exprArity to grapple with. In particular, Simplify uses exprArity to
83 compute the ArityInfo for the Id.
84
85 Originally I thought that it was enough just to look for top-level lambdas, but
86 it isn't. I've seen this
87
88 foo = PrelBase.timesInt
89
90 We want foo to get arity 2 even though the eta-expander will leave it
91 unchanged, in the expectation that it'll be inlined. But occasionally it
92 isn't, because foo is blacklisted (used in a rule).
93
94 Similarly, see the ok_note check in exprEtaExpandArity. So
95 f = __inline_me (\x -> e)
96 won't be eta-expanded.
97
98 And in any case it seems more robust to have exprArity be a bit more intelligent.
99 But note that (\x y z -> f x y z)
100 should have arity 3, regardless of f's arity.
101 -}
102
103 manifestArity :: CoreExpr -> Arity
104 -- ^ manifestArity sees how many leading value lambdas there are,
105 -- after looking through casts
106 manifestArity (Lam v e) | isId v = 1 + manifestArity e
107 | otherwise = manifestArity e
108 manifestArity (Tick t e) | not (tickishIsCode t) = manifestArity e
109 manifestArity (Cast e _) = manifestArity e
110 manifestArity _ = 0
111
112 joinRhsArity :: CoreExpr -> JoinArity
113 -- Join points are supposed to have manifestly-visible
114 -- lambdas at the top: no ticks, no casts, nothing
115 -- Moreover, type lambdas count in JoinArity
116 joinRhsArity (Lam _ e) = 1 + joinRhsArity e
117 joinRhsArity _ = 0
118
119
120 ---------------
121 exprArity :: CoreExpr -> Arity
122 -- ^ An approximate, fast, version of 'exprEtaExpandArity'
123 exprArity e = go e
124 where
125 go (Var v) = idArity v
126 go (Lam x e) | isId x = go e + 1
127 | otherwise = go e
128 go (Tick t e) | not (tickishIsCode t) = go e
129 go (Cast e co) = trim_arity (go e) (coercionRKind co)
130 -- Note [exprArity invariant]
131 go (App e (Type _)) = go e
132 go (App f a) | exprIsTrivial a = (go f - 1) `max` 0
133 -- See Note [exprArity for applications]
134 -- NB: coercions count as a value argument
135
136 go _ = 0
137
138 trim_arity :: Arity -> Type -> Arity
139 trim_arity arity ty = arity `min` length (typeArity ty)
140
141 ---------------
142 typeArity :: Type -> [OneShotInfo]
143 -- How many value arrows are visible in the type?
144 -- We look through foralls, and newtypes
145 -- See Note [exprArity invariant]
146 typeArity ty
147 = go initRecTc ty
148 where
149 go rec_nts ty
150 | Just (_, ty') <- splitForAllTyCoVar_maybe ty
151 = go rec_nts ty'
152
153 | Just (_,arg,res) <- splitFunTy_maybe ty
154 = typeOneShot arg : go rec_nts res
155
156 | Just (tc,tys) <- splitTyConApp_maybe ty
157 , Just (ty', _) <- instNewTyCon_maybe tc tys
158 , Just rec_nts' <- checkRecTc rec_nts tc -- See Note [Expanding newtypes]
159 -- in GHC.Core.TyCon
160 -- , not (isClassTyCon tc) -- Do not eta-expand through newtype classes
161 -- -- See Note [Newtype classes and eta expansion]
162 -- (no longer required)
163 = go rec_nts' ty'
164 -- Important to look through non-recursive newtypes, so that, eg
165 -- (f x) where f has arity 2, f :: Int -> IO ()
166 -- Here we want to get arity 1 for the result!
167 --
168 -- AND through a layer of recursive newtypes
169 -- e.g. newtype Stream m a b = Stream (m (Either b (a, Stream m a b)))
170
171 | otherwise
172 = []
173
174 ---------------
175 exprBotStrictness_maybe :: CoreExpr -> Maybe (Arity, DmdSig)
176 -- A cheap and cheerful function that identifies bottoming functions
177 -- and gives them a suitable strictness signatures. It's used during
178 -- float-out
179 exprBotStrictness_maybe e
180 = case getBotArity (arityType botStrictnessArityEnv e) of
181 Nothing -> Nothing
182 Just ar -> Just (ar, sig ar)
183 where
184 sig ar = mkClosedDmdSig (replicate ar topDmd) botDiv
185
186 {-
187 Note [exprArity invariant]
188 ~~~~~~~~~~~~~~~~~~~~~~~~~~
189 exprArity has the following invariants:
190
191 (1) If typeArity (exprType e) = n,
192 then manifestArity (etaExpand e n) = n
193
194 That is, etaExpand can always expand as much as typeArity says
195 So the case analysis in etaExpand and in typeArity must match
196
197 (2) exprArity e <= typeArity (exprType e)
198
199 (3) Hence if (exprArity e) = n, then manifestArity (etaExpand e n) = n
200
201 That is, if exprArity says "the arity is n" then etaExpand really
202 can get "n" manifest lambdas to the top.
203
204 Why is this important? Because
205 - In GHC.Iface.Tidy we use exprArity to fix the *final arity* of
206 each top-level Id, and in
207 - In CorePrep we use etaExpand on each rhs, so that the visible lambdas
208 actually match that arity, which in turn means
209 that the StgRhs has the right number of lambdas
210
211 An alternative would be to do the eta-expansion in GHC.Iface.Tidy, at least
212 for top-level bindings, in which case we would not need the trim_arity
213 in exprArity. That is a less local change, so I'm going to leave it for today!
214
215 Note [Newtype classes and eta expansion]
216 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
217 NB: this nasty special case is no longer required, because
218 for newtype classes we don't use the class-op rule mechanism
219 at all. See Note [Single-method classes] in GHC.Tc.TyCl.Instance. SLPJ May 2013
220
221 -------- Old out of date comments, just for interest -----------
222 We have to be careful when eta-expanding through newtypes. In general
223 it's a good idea, but annoyingly it interacts badly with the class-op
224 rule mechanism. Consider
225
226 class C a where { op :: a -> a }
227 instance C b => C [b] where
228 op x = ...
229
230 These translate to
231
232 co :: forall a. (a->a) ~ C a
233
234 $copList :: C b -> [b] -> [b]
235 $copList d x = ...
236
237 $dfList :: C b -> C [b]
238 {-# DFunUnfolding = [$copList] #-}
239 $dfList d = $copList d |> co@[b]
240
241 Now suppose we have:
242
243 dCInt :: C Int
244
245 blah :: [Int] -> [Int]
246 blah = op ($dfList dCInt)
247
248 Now we want the built-in op/$dfList rule will fire to give
249 blah = $copList dCInt
250
251 But with eta-expansion 'blah' might (and in #3772, which is
252 slightly more complicated, does) turn into
253
254 blah = op (\eta. ($dfList dCInt |> sym co) eta)
255
256 and now it is *much* harder for the op/$dfList rule to fire, because
257 exprIsConApp_maybe won't hold of the argument to op. I considered
258 trying to *make* it hold, but it's tricky and I gave up.
259
260 The test simplCore/should_compile/T3722 is an excellent example.
261 -------- End of old out of date comments, just for interest -----------
262
263
264 Note [exprArity for applications]
265 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
266 When we come to an application we check that the arg is trivial.
267 eg f (fac x) does not have arity 2,
268 even if f has arity 3!
269
270 * We require that is trivial rather merely cheap. Suppose f has arity 2.
271 Then f (Just y)
272 has arity 0, because if we gave it arity 1 and then inlined f we'd get
273 let v = Just y in \w. <f-body>
274 which has arity 0. And we try to maintain the invariant that we don't
275 have arity decreases.
276
277 * The `max 0` is important! (\x y -> f x) has arity 2, even if f is
278 unknown, hence arity 0
279
280
281 ************************************************************************
282 * *
283 Computing the "arity" of an expression
284 * *
285 ************************************************************************
286
287 Note [Definition of arity]
288 ~~~~~~~~~~~~~~~~~~~~~~~~~~
289 The "arity" of an expression 'e' is n if
290 applying 'e' to *fewer* than n *value* arguments
291 converges rapidly
292
293 Or, to put it another way
294
295 there is no work lost in duplicating the partial
296 application (e x1 .. x(n-1))
297
298 In the divergent case, no work is lost by duplicating because if the thing
299 is evaluated once, that's the end of the program.
300
301 Or, to put it another way, in any context C
302
303 C[ (\x1 .. xn. e x1 .. xn) ]
304 is as efficient as
305 C[ e ]
306
307 It's all a bit more subtle than it looks:
308
309 Note [One-shot lambdas]
310 ~~~~~~~~~~~~~~~~~~~~~~~
311 Consider one-shot lambdas
312 let x = expensive in \y z -> E
313 We want this to have arity 1 if the \y-abstraction is a 1-shot lambda.
314
315 Note [Dealing with bottom]
316 ~~~~~~~~~~~~~~~~~~~~~~~~~~
317 A Big Deal with computing arities is expressions like
318
319 f = \x -> case x of
320 True -> \s -> e1
321 False -> \s -> e2
322
323 This happens all the time when f :: Bool -> IO ()
324 In this case we do eta-expand, in order to get that \s to the
325 top, and give f arity 2.
326
327 This isn't really right in the presence of seq. Consider
328 (f bot) `seq` 1
329
330 This should diverge! But if we eta-expand, it won't. We ignore this
331 "problem" (unless -fpedantic-bottoms is on), because being scrupulous
332 would lose an important transformation for many programs. (See
333 #5587 for an example.)
334
335 Consider also
336 f = \x -> error "foo"
337 Here, arity 1 is fine. But if it is
338 f = \x -> case x of
339 True -> error "foo"
340 False -> \y -> x+y
341 then we want to get arity 2. Technically, this isn't quite right, because
342 (f True) `seq` 1
343 should diverge, but it'll converge if we eta-expand f. Nevertheless, we
344 do so; it improves some programs significantly, and increasing convergence
345 isn't a bad thing. Hence the ABot/ATop in ArityType.
346
347 So these two transformations aren't always the Right Thing, and we
348 have several tickets reporting unexpected behaviour resulting from
349 this transformation. So we try to limit it as much as possible:
350
351 (1) Do NOT move a lambda outside a known-bottom case expression
352 case undefined of { (a,b) -> \y -> e }
353 This showed up in #5557
354
355 (2) Do NOT move a lambda outside a case unless
356 (a) The scrutinee is ok-for-speculation, or
357 (b) more liberally: the scrutinee is cheap (e.g. a variable), and
358 -fpedantic-bottoms is not enforced (see #2915 for an example)
359
360 Of course both (1) and (2) are readily defeated by disguising the bottoms.
361
362 4. Note [Newtype arity]
363 ~~~~~~~~~~~~~~~~~~~~~~~~
364 Non-recursive newtypes are transparent, and should not get in the way.
365 We do (currently) eta-expand recursive newtypes too. So if we have, say
366
367 newtype T = MkT ([T] -> Int)
368
369 Suppose we have
370 e = coerce T f
371 where f has arity 1. Then: etaExpandArity e = 1;
372 that is, etaExpandArity looks through the coerce.
373
374 When we eta-expand e to arity 1: eta_expand 1 e T
375 we want to get: coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
376
377 HOWEVER, note that if you use coerce bogusly you can ge
378 coerce Int negate
379 And since negate has arity 2, you might try to eta expand. But you can't
380 decompose Int to a function type. Hence the final case in eta_expand.
381
382 Note [The state-transformer hack]
383 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
384 Suppose we have
385 f = e
386 where e has arity n. Then, if we know from the context that f has
387 a usage type like
388 t1 -> ... -> tn -1-> t(n+1) -1-> ... -1-> tm -> ...
389 then we can expand the arity to m. This usage type says that
390 any application (x e1 .. en) will be applied to uniquely to (m-n) more args
391 Consider f = \x. let y = <expensive>
392 in case x of
393 True -> foo
394 False -> \(s:RealWorld) -> e
395 where foo has arity 1. Then we want the state hack to
396 apply to foo too, so we can eta expand the case.
397
398 Then we expect that if f is applied to one arg, it'll be applied to two
399 (that's the hack -- we don't really know, and sometimes it's false)
400 See also Id.isOneShotBndr.
401
402 Note [State hack and bottoming functions]
403 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
404 It's a terrible idea to use the state hack on a bottoming function.
405 Here's what happens (#2861):
406
407 f :: String -> IO T
408 f = \p. error "..."
409
410 Eta-expand, using the state hack:
411
412 f = \p. (\s. ((error "...") |> g1) s) |> g2
413 g1 :: IO T ~ (S -> (S,T))
414 g2 :: (S -> (S,T)) ~ IO T
415
416 Extrude the g2
417
418 f' = \p. \s. ((error "...") |> g1) s
419 f = f' |> (String -> g2)
420
421 Discard args for bottomming function
422
423 f' = \p. \s. ((error "...") |> g1 |> g3
424 g3 :: (S -> (S,T)) ~ (S,T)
425
426 Extrude g1.g3
427
428 f'' = \p. \s. (error "...")
429 f' = f'' |> (String -> S -> g1.g3)
430
431 And now we can repeat the whole loop. Aargh! The bug is in applying the
432 state hack to a function which then swallows the argument.
433
434 This arose in another guise in #3959. Here we had
435
436 catch# (throw exn >> return ())
437
438 Note that (throw :: forall a e. Exn e => e -> a) is called with [a = IO ()].
439 After inlining (>>) we get
440
441 catch# (\_. throw {IO ()} exn)
442
443 We must *not* eta-expand to
444
445 catch# (\_ _. throw {...} exn)
446
447 because 'catch#' expects to get a (# _,_ #) after applying its argument to
448 a State#, not another function!
449
450 In short, we use the state hack to allow us to push let inside a lambda,
451 but not to introduce a new lambda.
452
453
454 Note [ArityType]
455 ~~~~~~~~~~~~~~~~
456 ArityType is the result of a compositional analysis on expressions,
457 from which we can decide the real arity of the expression (extracted
458 with function exprEtaExpandArity).
459
460 We use the following notation:
461 at ::= \o1..on.div
462 div ::= T | x | ⊥
463 o ::= ? | 1
464 And omit the \. if n = 0. Examples:
465 \?11.T stands for @AT [NoOneShotInfo,OneShotLam,OneShotLam] topDiv@
466 ⊥ stands for @AT [] botDiv@
467 See the 'Outputable' instance for more information. It's pretty simple.
468
469 Here is what the fields mean. If an arbitrary expression 'f' has
470 ArityType 'at', then
471
472 * If @at = AT [o1,..,on] botDiv@ (notation: \o1..on.⊥), then @f x1..xn@
473 definitely diverges. Partial applications to fewer than n args may *or
474 may not* diverge.
475
476 We allow ourselves to eta-expand bottoming functions, even
477 if doing so may lose some `seq` sharing,
478 let x = <expensive> in \y. error (g x y)
479 ==> \y. let x = <expensive> in error (g x y)
480
481 * If @at = AT [o1,..,on] topDiv@ (notation: \o1..on.T), then expanding 'f'
482 to @\x1..xn. f x1..xn@ loses no sharing, assuming the calls of f respect
483 the one-shot-ness o1..on of its definition.
484
485 NB 'f' is an arbitrary expression, eg @f = g e1 e2@. This 'f' can have
486 arity type @AT oss _@, with @length oss > 0@, only if e1 e2 are themselves
487 cheap.
488
489 * In both cases, @f@, @f x1@, ... @f x1 ... x(n-1)@ are definitely
490 really functions, or bottom, but *not* casts from a data type, in
491 at least one case branch. (If it's a function in one case branch but
492 an unsafe cast from a data type in another, the program is bogus.)
493 So eta expansion is dynamically ok; see Note [State hack and
494 bottoming functions], the part about catch#
495
496 Example:
497 f = \x\y. let v = <expensive> in
498 \s(one-shot) \t(one-shot). blah
499 'f' has arity type \??11.T
500 The one-shot-ness means we can, in effect, push that
501 'let' inside the \st.
502
503
504 Suppose f = \xy. x+y
505 Then f :: \??.T
506 f v :: \?.T
507 f <expensive> :: T
508 -}
509
510
511 -- | The analysis lattice of arity analysis. It is isomorphic to
512 --
513 -- @
514 -- data ArityType'
515 -- = AEnd Divergence
516 -- | ALam OneShotInfo ArityType'
517 -- @
518 --
519 -- Which is easier to display the Hasse diagram for:
520 --
521 -- @
522 -- ALam OneShotLam at
523 -- |
524 -- AEnd topDiv
525 -- |
526 -- ALam NoOneShotInfo at
527 -- |
528 -- AEnd exnDiv
529 -- |
530 -- AEnd botDiv
531 -- @
532 --
533 -- where the @at@ fields of @ALam@ are inductively subject to the same order.
534 -- That is, @ALam os at1 < ALam os at2@ iff @at1 < at2@.
535 --
536 -- Why the strange Top element? See Note [Combining case branches].
537 --
538 -- We rely on this lattice structure for fixed-point iteration in
539 -- 'findRhsArity'. For the semantics of 'ArityType', see Note [ArityType].
540 data ArityType
541 = AT ![OneShotInfo] !Divergence
542 -- ^ @AT oss div@ means this value can safely be eta-expanded @length oss@
543 -- times, provided use sites respect the 'OneShotInfo's in @oss@.
544 -- A 'OneShotLam' annotation can come from two sources:
545 -- * The user annotated a lambda as one-shot with 'GHC.Exts.oneShot'
546 -- * It's from a lambda binder of a type affected by `-fstate-hack`.
547 -- See 'idStateHackOneShotInfo'.
548 -- In both cases, 'OneShotLam' should win over 'NoOneShotInfo', see
549 -- Note [Combining case branches].
550 --
551 -- If @div@ is dead-ending ('isDeadEndDiv'), then application to
552 -- @length os@ arguments will surely diverge, similar to the situation
553 -- with 'DmdType'.
554 deriving Eq
555
556 -- | This is the BNF of the generated output:
557 --
558 -- @
559 -- @
560 --
561 -- We format
562 -- @AT [o1,..,on] topDiv@ as @\o1..on.T@ and
563 -- @AT [o1,..,on] botDiv@ as @\o1..on.⊥@, respectively.
564 -- More concretely, @AT [NOI,OS,OS] topDiv@ is formatted as @\?11.T@.
565 -- If the one-shot info is empty, we omit the leading @\.@.
566 instance Outputable ArityType where
567 ppr (AT oss div)
568 | null oss = pp_div div
569 | otherwise = char '\\' <> hcat (map pp_os oss) <> dot <> pp_div div
570 where
571 pp_div Diverges = char '⊥'
572 pp_div ExnOrDiv = char 'x'
573 pp_div Dunno = char 'T'
574 pp_os OneShotLam = char '1'
575 pp_os NoOneShotInfo = char '?'
576
577 mkBotArityType :: [OneShotInfo] -> ArityType
578 mkBotArityType oss = AT oss botDiv
579
580 botArityType :: ArityType
581 botArityType = mkBotArityType []
582
583 mkTopArityType :: [OneShotInfo] -> ArityType
584 mkTopArityType oss = AT oss topDiv
585
586 topArityType :: ArityType
587 topArityType = mkTopArityType []
588
589 -- | The number of value args for the arity type
590 arityTypeArity :: ArityType -> Arity
591 arityTypeArity (AT oss _) = length oss
592
593 -- | True <=> eta-expansion will add at least one lambda
594 expandableArityType :: ArityType -> Bool
595 expandableArityType at = arityTypeArity at > 0
596
597 -- | See Note [Dead ends] in "GHC.Types.Demand".
598 -- Bottom implies a dead end.
599 isDeadEndArityType :: ArityType -> Bool
600 isDeadEndArityType (AT _ div) = isDeadEndDiv div
601
602 -- | Expand a non-bottoming arity type so that it has at least the given arity.
603 maxWithArity :: ArityType -> Arity -> ArityType
604 maxWithArity at@(AT oss div) !ar
605 | isDeadEndArityType at = at
606 | oss `lengthAtLeast` ar = at
607 | otherwise = AT (take ar $ oss ++ repeat NoOneShotInfo) div
608
609 -- | Trim an arity type so that it has at most the given arity.
610 -- Any excess 'OneShotInfo's are truncated to 'topDiv', even if they end in
611 -- 'ABot'.
612 minWithArity :: ArityType -> Arity -> ArityType
613 minWithArity at@(AT oss _) ar
614 | oss `lengthAtMost` ar = at
615 | otherwise = AT (take ar oss) topDiv
616
617 takeWhileOneShot :: ArityType -> ArityType
618 takeWhileOneShot (AT oss div)
619 | isDeadEndDiv div = AT (takeWhile isOneShotInfo oss) topDiv
620 | otherwise = AT (takeWhile isOneShotInfo oss) div
621
622 -- | The Arity returned is the number of value args the
623 -- expression can be applied to without doing much work
624 exprEtaExpandArity :: DynFlags -> CoreExpr -> ArityType
625 -- exprEtaExpandArity is used when eta expanding
626 -- e ==> \xy -> e x y
627 exprEtaExpandArity dflags e = arityType (etaExpandArityEnv dflags) e
628
629 getBotArity :: ArityType -> Maybe Arity
630 -- Arity of a divergent function
631 getBotArity (AT oss div)
632 | isDeadEndDiv div = Just $ length oss
633 | otherwise = Nothing
634
635 ----------------------
636 findRhsArity :: DynFlags -> Id -> CoreExpr -> Arity -> ArityType
637 -- This implements the fixpoint loop for arity analysis
638 -- See Note [Arity analysis]
639 -- If findRhsArity e = (n, is_bot) then
640 -- (a) any application of e to <n arguments will not do much work,
641 -- so it is safe to expand e ==> (\x1..xn. e x1 .. xn)
642 -- (b) if is_bot=True, then e applied to n args is guaranteed bottom
643 findRhsArity dflags bndr rhs old_arity
644 = go 0 botArityType
645 -- We always do one step, but usually that produces a result equal to
646 -- old_arity, and then we stop right away, because old_arity is assumed
647 -- to be sound. In other words, arities should never decrease.
648 -- Result: the common case is that there is just one iteration
649 where
650 go :: Int -> ArityType -> ArityType
651 go !n cur_at@(AT oss div)
652 | not (isDeadEndDiv div) -- the "stop right away" case
653 , length oss <= old_arity = cur_at -- from above
654 | next_at == cur_at = cur_at
655 | otherwise =
656 -- Warn if more than 2 iterations. Why 2? See Note [Exciting arity]
657 warnPprTrace (debugIsOn && n > 2)
658 (text "Exciting arity" $$ nest 2
659 ( ppr bndr <+> ppr cur_at <+> ppr next_at $$ ppr rhs)) $
660 go (n+1) next_at
661 where
662 next_at = step cur_at
663
664 step :: ArityType -> ArityType
665 step at = -- pprTrace "step" (ppr bndr <+> ppr at <+> ppr (arityType env rhs)) $
666 arityType env rhs
667 where
668 env = extendSigEnv (findRhsArityEnv dflags) bndr at
669
670
671 {-
672 Note [Arity analysis]
673 ~~~~~~~~~~~~~~~~~~~~~
674 The motivating example for arity analysis is this:
675
676 f = \x. let g = f (x+1)
677 in \y. ...g...
678
679 What arity does f have? Really it should have arity 2, but a naive
680 look at the RHS won't see that. You need a fixpoint analysis which
681 says it has arity "infinity" the first time round.
682
683 This example happens a lot; it first showed up in Andy Gill's thesis,
684 fifteen years ago! It also shows up in the code for 'rnf' on lists
685 in #4138.
686
687 We do the necessary, quite simple fixed-point iteration in 'findRhsArity',
688 which assumes for a single binding 'ABot' on the first run and iterates
689 until it finds a stable arity type. Two wrinkles
690
691 * We often have to ask (see the Case or Let case of 'arityType') whether some
692 expression is cheap. In the case of an application, that depends on the arity
693 of the application head! That's why we have our own version of 'exprIsCheap',
694 'myExprIsCheap', that will integrate the optimistic arity types we have on
695 f and g into the cheapness check.
696
697 * Consider this (#18793)
698
699 go = \ds. case ds of
700 [] -> id
701 (x:ys) -> let acc = go ys in
702 case blah of
703 True -> acc
704 False -> \ x1 -> acc (negate x1)
705
706 We must propagate go's optimistically large arity to @acc@, so that the
707 tail call to @acc@ in the True branch has sufficient arity. This is done
708 by the 'am_sigs' field in 'FindRhsArity', and 'lookupSigEnv' in the Var case
709 of 'arityType'.
710
711 Note [Exciting Arity]
712 ~~~~~~~~~~~~~~~~~~~~~
713 The fixed-point iteration in 'findRhsArity' stabilises very quickly in almost
714 all cases. To get notified of cases where we need an usual number of iterations,
715 we emit a warning in debug mode, so that we can investigate and make sure that
716 we really can't do better. It's a gross hack, but catches real bugs (#18870).
717
718 Now, which number is "unusual"? We pick n > 2. Here's a pretty common and
719 expected example that takes two iterations and would ruin the specificity
720 of the warning (from T18937):
721
722 f :: [Int] -> Int -> Int
723 f [] = id
724 f (x:xs) = let y = sum [0..x]
725 in \z -> f xs (y + z)
726
727 Fixed-point iteration starts with arity type ⊥ for f. After the first
728 iteration, we get arity type \??.T, e.g. arity 2, because we unconditionally
729 'floatIn' the let-binding (see its bottom case). After the second iteration,
730 we get arity type \?.T, e.g. arity 1, because now we are no longer allowed
731 to floatIn the non-cheap let-binding. Which is all perfectly benign, but
732 means we do two iterations (well, actually 3 'step's to detect we are stable)
733 and don't want to emit the warning.
734
735 Note [Eta expanding through dictionaries]
736 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
737 If the experimental -fdicts-cheap flag is on, we eta-expand through
738 dictionary bindings. This improves arities. Thereby, it also
739 means that full laziness is less prone to floating out the
740 application of a function to its dictionary arguments, which
741 can thereby lose opportunities for fusion. Example:
742 foo :: Ord a => a -> ...
743 foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....
744 -- So foo has arity 1
745
746 f = \x. foo dInt $ bar x
747
748 The (foo DInt) is floated out, and makes ineffective a RULE
749 foo (bar x) = ...
750
751 One could go further and make exprIsCheap reply True to any
752 dictionary-typed expression, but that's more work.
753 -}
754
755 arityLam :: Id -> ArityType -> ArityType
756 arityLam id (AT oss div) = AT (idStateHackOneShotInfo id : oss) div
757
758 floatIn :: Bool -> ArityType -> ArityType
759 -- We have something like (let x = E in b),
760 -- where b has the given arity type.
761 floatIn cheap at
762 | isDeadEndArityType at || cheap = at
763 -- If E is not cheap, keep arity only for one-shots
764 | otherwise = takeWhileOneShot at
765
766 arityApp :: ArityType -> Bool -> ArityType
767 -- Processing (fun arg) where at is the ArityType of fun,
768 -- Knock off an argument and behave like 'let'
769 arityApp (AT (_:oss) div) cheap = floatIn cheap (AT oss div)
770 arityApp at _ = at
771
772 -- | Least upper bound in the 'ArityType' lattice.
773 -- See the haddocks on 'ArityType' for the lattice.
774 --
775 -- Used for branches of a @case@.
776 andArityType :: ArityType -> ArityType -> ArityType
777 andArityType (AT (os1:oss1) div1) (AT (os2:oss2) div2)
778 | AT oss' div' <- andArityType (AT oss1 div1) (AT oss2 div2)
779 = AT ((os1 `bestOneShot` os2) : oss') div' -- See Note [Combining case branches]
780 andArityType (AT [] div1) at2
781 | isDeadEndDiv div1 = at2 -- Note [ABot branches: max arity wins]
782 | otherwise = takeWhileOneShot at2 -- See Note [Combining case branches]
783 andArityType at1 (AT [] div2)
784 | isDeadEndDiv div2 = at1 -- Note [ABot branches: max arity wins]
785 | otherwise = takeWhileOneShot at1 -- See Note [Combining case branches]
786
787 {- Note [ABot branches: max arity wins]
788 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
789 Consider case x of
790 True -> \x. error "urk"
791 False -> \xy. error "urk2"
792
793 Remember: \o1..on.⊥ means "if you apply to n args, it'll definitely diverge".
794 So we need \??.⊥ for the whole thing, the /max/ of both arities.
795
796 Note [Combining case branches]
797 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
798 Consider
799 go = \x. let z = go e0
800 go2 = \x. case x of
801 True -> z
802 False -> \s(one-shot). e1
803 in go2 x
804 We *really* want to respect the one-shot annotation provided by the
805 user and eta-expand go and go2.
806 When combining the branches of the case we have
807 T `andAT` \1.T
808 and we want to get \1.T.
809 But if the inner lambda wasn't one-shot (\?.T) we don't want to do this.
810 (We need a usage analysis to justify that.)
811
812 So we combine the best of the two branches, on the (slightly dodgy)
813 basis that if we know one branch is one-shot, then they all must be.
814 Surprisingly, this means that the one-shot arity type is effectively the top
815 element of the lattice.
816
817 Note [Arity trimming]
818 ~~~~~~~~~~~~~~~~~~~~~
819 Consider ((\x y. blah) |> co), where co :: (Int->Int->Int) ~ (Int -> F a) , and
820 F is some type family.
821
822 Because of Note [exprArity invariant], item (2), we must return with arity at
823 most 1, because typeArity (Int -> F a) = 1. So we have to trim the result of
824 calling arityType on (\x y. blah). Failing to do so, and hence breaking the
825 exprArity invariant, led to #5441.
826
827 How to trim? If we end in topDiv, it's easy. But we must take great care with
828 dead ends (i.e. botDiv). Suppose the expression was (\x y. error "urk"),
829 we'll get \??.⊥. We absolutely must not trim that to \?.⊥, because that
830 claims that ((\x y. error "urk") |> co) diverges when given one argument,
831 which it absolutely does not. And Bad Things happen if we think something
832 returns bottom when it doesn't (#16066).
833
834 So, if we need to trim a dead-ending arity type, switch (conservatively) to
835 topDiv.
836
837 Historical note: long ago, we unconditionally switched to topDiv when we
838 encountered a cast, but that is far too conservative: see #5475
839
840 Note [Eta expanding through CallStacks]
841 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
842 Just as it's good to eta-expand through dictionaries, so it is good to
843 do so through CallStacks. #20103 is a case in point, where we got
844 foo :: HasCallStack => Int -> Int
845 foo = \(d::CallStack). let d2 = pushCallStack blah d in
846 \(x:Int). blah
847
848 We really want to eta-expand this! #20103 is quite convincing!
849 We do this regardless of -fdicts-cheap; it's not really a dictionary.
850 -}
851
852 ---------------------------
853
854 -- | Each of the entry-points of the analyser ('arityType') has different
855 -- requirements. The entry-points are
856 --
857 -- 1. 'exprBotStrictness_maybe'
858 -- 2. 'exprEtaExpandArity'
859 -- 3. 'findRhsArity'
860 --
861 -- For each of the entry-points, there is a separate mode that governs
862 --
863 -- 1. How pedantic we are wrt. ⊥, in 'pedanticBottoms'.
864 -- 2. Whether we store arity signatures for non-recursive let-bindings,
865 -- accessed in 'extendSigEnv'/'lookupSigEnv'.
866 -- See Note [Arity analysis] why that's important.
867 -- 3. Which expressions we consider cheap to float inside a lambda,
868 -- in 'myExprIsCheap'.
869 data AnalysisMode
870 = BotStrictness
871 -- ^ Used during 'exprBotStrictness_maybe'.
872 | EtaExpandArity { am_ped_bot :: !Bool
873 , am_dicts_cheap :: !Bool }
874 -- ^ Used for finding an expression's eta-expanding arity quickly, without
875 -- fixed-point iteration ('exprEtaExpandArity').
876 | FindRhsArity { am_ped_bot :: !Bool
877 , am_dicts_cheap :: !Bool
878 , am_sigs :: !(IdEnv ArityType) }
879 -- ^ Used for regular, fixed-point arity analysis ('findRhsArity').
880 -- See Note [Arity analysis] for details about fixed-point iteration.
881 -- INVARIANT: Disjoint with 'ae_joins'.
882
883 data ArityEnv
884 = AE
885 { ae_mode :: !AnalysisMode
886 -- ^ The analysis mode. See 'AnalysisMode'.
887 , ae_joins :: !IdSet
888 -- ^ In-scope join points. See Note [Eta-expansion and join points]
889 -- INVARIANT: Disjoint with the domain of 'am_sigs' (if present).
890 }
891
892 -- | The @ArityEnv@ used by 'exprBotStrictness_maybe'. Pedantic about bottoms
893 -- and no application is ever considered cheap.
894 botStrictnessArityEnv :: ArityEnv
895 botStrictnessArityEnv = AE { ae_mode = BotStrictness, ae_joins = emptyVarSet }
896
897 -- | The @ArityEnv@ used by 'exprEtaExpandArity'.
898 etaExpandArityEnv :: DynFlags -> ArityEnv
899 etaExpandArityEnv dflags
900 = AE { ae_mode = EtaExpandArity { am_ped_bot = gopt Opt_PedanticBottoms dflags
901 , am_dicts_cheap = gopt Opt_DictsCheap dflags }
902 , ae_joins = emptyVarSet }
903
904 -- | The @ArityEnv@ used by 'findRhsArity'.
905 findRhsArityEnv :: DynFlags -> ArityEnv
906 findRhsArityEnv dflags
907 = AE { ae_mode = FindRhsArity { am_ped_bot = gopt Opt_PedanticBottoms dflags
908 , am_dicts_cheap = gopt Opt_DictsCheap dflags
909 , am_sigs = emptyVarEnv }
910 , ae_joins = emptyVarSet }
911
912 -- First some internal functions in snake_case for deleting in certain VarEnvs
913 -- of the ArityType. Don't call these; call delInScope* instead!
914
915 modifySigEnv :: (IdEnv ArityType -> IdEnv ArityType) -> ArityEnv -> ArityEnv
916 modifySigEnv f env@AE { ae_mode = am@FindRhsArity{am_sigs = sigs} } =
917 env { ae_mode = am { am_sigs = f sigs } }
918 modifySigEnv _ env = env
919 {-# INLINE modifySigEnv #-}
920
921 del_sig_env :: Id -> ArityEnv -> ArityEnv -- internal!
922 del_sig_env id = modifySigEnv (\sigs -> delVarEnv sigs id)
923 {-# INLINE del_sig_env #-}
924
925 del_sig_env_list :: [Id] -> ArityEnv -> ArityEnv -- internal!
926 del_sig_env_list ids = modifySigEnv (\sigs -> delVarEnvList sigs ids)
927 {-# INLINE del_sig_env_list #-}
928
929 del_join_env :: JoinId -> ArityEnv -> ArityEnv -- internal!
930 del_join_env id env@(AE { ae_joins = joins })
931 = env { ae_joins = delVarSet joins id }
932 {-# INLINE del_join_env #-}
933
934 del_join_env_list :: [JoinId] -> ArityEnv -> ArityEnv -- internal!
935 del_join_env_list ids env@(AE { ae_joins = joins })
936 = env { ae_joins = delVarSetList joins ids }
937 {-# INLINE del_join_env_list #-}
938
939 -- end of internal deletion functions
940
941 extendJoinEnv :: ArityEnv -> [JoinId] -> ArityEnv
942 extendJoinEnv env@(AE { ae_joins = joins }) join_ids
943 = del_sig_env_list join_ids
944 $ env { ae_joins = joins `extendVarSetList` join_ids }
945
946 extendSigEnv :: ArityEnv -> Id -> ArityType -> ArityEnv
947 extendSigEnv env id ar_ty
948 = del_join_env id (modifySigEnv (\sigs -> extendVarEnv sigs id ar_ty) env)
949
950 delInScope :: ArityEnv -> Id -> ArityEnv
951 delInScope env id = del_join_env id $ del_sig_env id env
952
953 delInScopeList :: ArityEnv -> [Id] -> ArityEnv
954 delInScopeList env ids = del_join_env_list ids $ del_sig_env_list ids env
955
956 lookupSigEnv :: ArityEnv -> Id -> Maybe ArityType
957 lookupSigEnv AE{ ae_mode = mode } id = case mode of
958 BotStrictness -> Nothing
959 EtaExpandArity{} -> Nothing
960 FindRhsArity{ am_sigs = sigs } -> lookupVarEnv sigs id
961
962 -- | Whether the analysis should be pedantic about bottoms.
963 -- 'exprBotStrictness_maybe' always is.
964 pedanticBottoms :: ArityEnv -> Bool
965 pedanticBottoms AE{ ae_mode = mode } = case mode of
966 BotStrictness -> True
967 EtaExpandArity{ am_ped_bot = ped_bot } -> ped_bot
968 FindRhsArity{ am_ped_bot = ped_bot } -> ped_bot
969
970 -- | A version of 'exprIsCheap' that considers results from arity analysis
971 -- and optionally the expression's type.
972 -- Under 'exprBotStrictness_maybe', no expressions are cheap.
973 myExprIsCheap :: ArityEnv -> CoreExpr -> Maybe Type -> Bool
974 myExprIsCheap AE{ae_mode = mode} e mb_ty = case mode of
975 BotStrictness -> False
976 _ -> cheap_dict || cheap_fun e
977 where
978 cheap_dict = case mb_ty of
979 Nothing -> False
980 Just ty -> (am_dicts_cheap mode && isDictTy ty)
981 || isCallStackPredTy ty
982 -- See Note [Eta expanding through dictionaries]
983 -- See Note [Eta expanding through CallStacks]
984
985 cheap_fun e = case mode of
986 #if __GLASGOW_HASKELL__ <= 900
987 BotStrictness -> panic "impossible"
988 #endif
989 EtaExpandArity{} -> exprIsCheap e
990 FindRhsArity{am_sigs = sigs} -> exprIsCheapX (myIsCheapApp sigs) e
991
992 -- | A version of 'isCheapApp' that considers results from arity analysis.
993 -- See Note [Arity analysis] for what's in the signature environment and why
994 -- it's important.
995 myIsCheapApp :: IdEnv ArityType -> CheapAppFun
996 myIsCheapApp sigs fn n_val_args = case lookupVarEnv sigs fn of
997 -- Nothing means not a local function, fall back to regular
998 -- 'GHC.Core.Utils.isCheapApp'
999 Nothing -> isCheapApp fn n_val_args
1000 -- @Just at@ means local function with @at@ as current ArityType.
1001 -- Roughly approximate what 'isCheapApp' is doing.
1002 Just (AT oss div)
1003 | isDeadEndDiv div -> True -- See Note [isCheapApp: bottoming functions] in GHC.Core.Utils
1004 | n_val_args < length oss -> True -- Essentially isWorkFreeApp
1005 | otherwise -> False
1006
1007 ----------------
1008 arityType :: ArityEnv -> CoreExpr -> ArityType
1009
1010 arityType env (Cast e co)
1011 = minWithArity (arityType env e) co_arity -- See Note [Arity trimming]
1012 where
1013 co_arity = length (typeArity (coercionRKind co))
1014 -- See Note [exprArity invariant] (2); must be true of
1015 -- arityType too, since that is how we compute the arity
1016 -- of variables, and they in turn affect result of exprArity
1017 -- #5441 is a nice demo
1018
1019 arityType env (Var v)
1020 | v `elemVarSet` ae_joins env
1021 = botArityType -- See Note [Eta-expansion and join points]
1022 | Just at <- lookupSigEnv env v -- Local binding
1023 = at
1024 | otherwise
1025 = idArityType v
1026
1027 -- Lambdas; increase arity
1028 arityType env (Lam x e)
1029 | isId x = arityLam x (arityType env' e)
1030 | otherwise = arityType env' e
1031 where
1032 env' = delInScope env x
1033
1034 -- Applications; decrease arity, except for types
1035 arityType env (App fun (Type _))
1036 = arityType env fun
1037 arityType env (App fun arg )
1038 = arityApp (arityType env fun) (myExprIsCheap env arg Nothing)
1039
1040 -- Case/Let; keep arity if either the expression is cheap
1041 -- or it's a 1-shot lambda
1042 -- The former is not really right for Haskell
1043 -- f x = case x of { (a,b) -> \y. e }
1044 -- ===>
1045 -- f x y = case x of { (a,b) -> e }
1046 -- The difference is observable using 'seq'
1047 --
1048 arityType env (Case scrut bndr _ alts)
1049 | exprIsDeadEnd scrut || null alts
1050 = botArityType -- Do not eta expand. See Note [Dealing with bottom (1)]
1051 | not (pedanticBottoms env) -- See Note [Dealing with bottom (2)]
1052 , myExprIsCheap env scrut (Just (idType bndr))
1053 = alts_type
1054 | exprOkForSpeculation scrut
1055 = alts_type
1056
1057 | otherwise -- In the remaining cases we may not push
1058 = takeWhileOneShot alts_type -- evaluation of the scrutinee in
1059 where
1060 env' = delInScope env bndr
1061 arity_type_alt (Alt _con bndrs rhs) = arityType (delInScopeList env' bndrs) rhs
1062 alts_type = foldr1 andArityType (map arity_type_alt alts)
1063
1064 arityType env (Let (NonRec j rhs) body)
1065 | Just join_arity <- isJoinId_maybe j
1066 , (_, rhs_body) <- collectNBinders join_arity rhs
1067 = -- See Note [Eta-expansion and join points]
1068 andArityType (arityType env rhs_body)
1069 (arityType env' body)
1070 where
1071 env' = extendJoinEnv env [j]
1072
1073 arityType env (Let (Rec pairs) body)
1074 | ((j,_):_) <- pairs
1075 , isJoinId j
1076 = -- See Note [Eta-expansion and join points]
1077 foldr (andArityType . do_one) (arityType env' body) pairs
1078 where
1079 env' = extendJoinEnv env (map fst pairs)
1080 do_one (j,rhs)
1081 | Just arity <- isJoinId_maybe j
1082 = arityType env' $ snd $ collectNBinders arity rhs
1083 | otherwise
1084 = pprPanic "arityType:joinrec" (ppr pairs)
1085
1086 arityType env (Let (NonRec b r) e)
1087 = floatIn cheap_rhs (arityType env' e)
1088 where
1089 cheap_rhs = myExprIsCheap env r (Just (idType b))
1090 env' = extendSigEnv env b (arityType env r)
1091
1092 arityType env (Let (Rec prs) e)
1093 = floatIn (all is_cheap prs) (arityType env' e)
1094 where
1095 env' = delInScopeList env (map fst prs)
1096 is_cheap (b,e) = myExprIsCheap env' e (Just (idType b))
1097
1098 arityType env (Tick t e)
1099 | not (tickishIsCode t) = arityType env e
1100
1101 arityType _ _ = topArityType
1102
1103 {- Note [Eta-expansion and join points]
1104 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1105 Consider this (#18328)
1106
1107 f x = join j y = case y of
1108 True -> \a. blah
1109 False -> \b. blah
1110 in case x of
1111 A -> j True
1112 B -> \c. blah
1113 C -> j False
1114
1115 and suppose the join point is too big to inline. Now, what is the
1116 arity of f? If we inlined the join point, we'd definitely say "arity
1117 2" because we are prepared to push case-scrutinisation inside a
1118 lambda. But currently the join point totally messes all that up,
1119 because (thought of as a vanilla let-binding) the arity pinned on 'j'
1120 is just 1.
1121
1122 Why don't we eta-expand j? Because of
1123 Note [Do not eta-expand join points] in GHC.Core.Opt.Simplify.Utils
1124
1125 Even if we don't eta-expand j, why is its arity only 1?
1126 See invariant 2b in Note [Invariants on join points] in GHC.Core.
1127
1128 So we do this:
1129
1130 * Treat the RHS of a join-point binding, /after/ stripping off
1131 join-arity lambda-binders, as very like the body of the let.
1132 More precisely, do andArityType with the arityType from the
1133 body of the let.
1134
1135 * Dually, when we come to a /call/ of a join point, just no-op
1136 by returning ABot, the bottom element of ArityType,
1137 which so that: bot `andArityType` x = x
1138
1139 * This works if the join point is bound in the expression we are
1140 taking the arityType of. But if it's bound further out, it makes
1141 no sense to say that (say) the arityType of (j False) is ABot.
1142 Bad things happen. So we keep track of the in-scope join-point Ids
1143 in ae_join.
1144
1145 This will make f, above, have arity 2. Then, we'll eta-expand it thus:
1146
1147 f x eta = (join j y = ... in case x of ...) eta
1148
1149 and the Simplify will automatically push that application of eta into
1150 the join points.
1151
1152 An alternative (roughly equivalent) idea would be to carry an
1153 environment mapping let-bound Ids to their ArityType.
1154 -}
1155
1156 idArityType :: Id -> ArityType
1157 idArityType v
1158 | strict_sig <- idDmdSig v
1159 , not $ isTopSig strict_sig
1160 , (ds, div) <- splitDmdSig strict_sig
1161 , let arity = length ds
1162 -- Every strictness signature admits an arity signature!
1163 = AT (take arity one_shots) div
1164 | otherwise
1165 = AT (take (idArity v) one_shots) topDiv
1166 where
1167 one_shots :: [OneShotInfo] -- One-shot-ness derived from the type
1168 one_shots = typeArity (idType v)
1169
1170 {-
1171 %************************************************************************
1172 %* *
1173 The main eta-expander
1174 %* *
1175 %************************************************************************
1176
1177 We go for:
1178 f = \x1..xn -> N ==> f = \x1..xn y1..ym -> N y1..ym
1179 (n >= 0)
1180
1181 where (in both cases)
1182
1183 * The xi can include type variables
1184
1185 * The yi are all value variables
1186
1187 * N is a NORMAL FORM (i.e. no redexes anywhere)
1188 wanting a suitable number of extra args.
1189
1190 The biggest reason for doing this is for cases like
1191
1192 f = \x -> case x of
1193 True -> \y -> e1
1194 False -> \y -> e2
1195
1196 Here we want to get the lambdas together. A good example is the nofib
1197 program fibheaps, which gets 25% more allocation if you don't do this
1198 eta-expansion.
1199
1200 We may have to sandwich some coerces between the lambdas
1201 to make the types work. exprEtaExpandArity looks through coerces
1202 when computing arity; and etaExpand adds the coerces as necessary when
1203 actually computing the expansion.
1204
1205 Note [No crap in eta-expanded code]
1206 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1207 The eta expander is careful not to introduce "crap". In particular,
1208 given a CoreExpr satisfying the 'CpeRhs' invariant (in CorePrep), it
1209 returns a CoreExpr satisfying the same invariant. See Note [Eta
1210 expansion and the CorePrep invariants] in CorePrep.
1211
1212 This means the eta-expander has to do a bit of on-the-fly
1213 simplification but it's not too hard. The alternative, of relying on
1214 a subsequent clean-up phase of the Simplifier to de-crapify the result,
1215 means you can't really use it in CorePrep, which is painful.
1216
1217 Note [Eta expansion for join points]
1218 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1219 The no-crap rule is very tiresome to guarantee when
1220 we have join points. Consider eta-expanding
1221 let j :: Int -> Int -> Bool
1222 j x = e
1223 in b
1224
1225 The simple way is
1226 \(y::Int). (let j x = e in b) y
1227
1228 The no-crap way is
1229 \(y::Int). let j' :: Int -> Bool
1230 j' x = e y
1231 in b[j'/j] y
1232 where I have written to stress that j's type has
1233 changed. Note that (of course!) we have to push the application
1234 inside the RHS of the join as well as into the body. AND if j
1235 has an unfolding we have to push it into there too. AND j might
1236 be recursive...
1237
1238 So for now I'm abandoning the no-crap rule in this case. I think
1239 that for the use in CorePrep it really doesn't matter; and if
1240 it does, then CoreToStg.myCollectArgs will fall over.
1241
1242 (Moreover, I think that casts can make the no-crap rule fail too.)
1243
1244 Note [Eta expansion and SCCs]
1245 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1246 Note that SCCs are not treated specially by etaExpand. If we have
1247 etaExpand 2 (\x -> scc "foo" e)
1248 = (\xy -> (scc "foo" e) y)
1249 So the costs of evaluating 'e' (not 'e y') are attributed to "foo"
1250
1251 Note [Eta expansion and source notes]
1252 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1253 CorePrep puts floatable ticks outside of value applications, but not
1254 type applications. As a result we might be trying to eta-expand an
1255 expression like
1256
1257 (src<...> v) @a
1258
1259 which we want to lead to code like
1260
1261 \x -> src<...> v @a x
1262
1263 This means that we need to look through type applications and be ready
1264 to re-add floats on the top.
1265
1266 Note [Eta expansion with ArityType]
1267 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1268 The etaExpandAT function takes an ArityType (not just an Arity) to
1269 guide eta-expansion. Why? Because we want to preserve one-shot info.
1270 Consider
1271 foo = \x. case x of
1272 True -> (\s{os}. blah) |> co
1273 False -> wubble
1274 We'll get an ArityType for foo of \?1.T.
1275
1276 Then we want to eta-expand to
1277 foo = \x. (\eta{os}. (case x of ...as before...) eta) |> some_co
1278
1279 That 'eta' binder is fresh, and we really want it to have the
1280 one-shot flag from the inner \s{os}. By expanding with the
1281 ArityType gotten from analysing the RHS, we achieve this neatly.
1282
1283 This makes a big difference to the one-shot monad trick;
1284 see Note [The one-shot state monad trick] in GHC.Utils.Monad.
1285 -}
1286
1287 -- | @etaExpand n e@ returns an expression with
1288 -- the same meaning as @e@, but with arity @n@.
1289 --
1290 -- Given:
1291 --
1292 -- > e' = etaExpand n e
1293 --
1294 -- We should have that:
1295 --
1296 -- > ty = exprType e = exprType e'
1297
1298 etaExpand :: Arity -> CoreExpr -> CoreExpr
1299 etaExpand n orig_expr
1300 = eta_expand in_scope (replicate n NoOneShotInfo) orig_expr
1301 where
1302 in_scope = {-#SCC "eta_expand:in-scopeX" #-}
1303 mkInScopeSet (exprFreeVars orig_expr)
1304
1305 etaExpandAT :: InScopeSet -> ArityType -> CoreExpr -> CoreExpr
1306 -- See Note [Eta expansion with ArityType]
1307 --
1308 -- We pass in the InScopeSet from the simplifier to avoid recomputing
1309 -- it here, which can be jolly expensive if the casts are big
1310 -- In #18223 it took 10% of compile time just to do the exprFreeVars!
1311 etaExpandAT in_scope (AT oss _) orig_expr
1312 = eta_expand in_scope oss orig_expr
1313
1314 -- etaExpand arity e = res
1315 -- Then 'res' has at least 'arity' lambdas at the top
1316 -- possibly with a cast wrapped around the outside
1317 -- See Note [Eta expansion with ArityType]
1318 --
1319 -- etaExpand deals with for-alls. For example:
1320 -- etaExpand 1 E
1321 -- where E :: forall a. a -> a
1322 -- would return
1323 -- (/\b. \y::a -> E b y)
1324
1325 eta_expand :: InScopeSet -> [OneShotInfo] -> CoreExpr -> CoreExpr
1326 eta_expand in_scope one_shots (Cast expr co)
1327 = Cast (eta_expand in_scope one_shots expr) co
1328
1329 eta_expand in_scope one_shots orig_expr
1330 = go in_scope one_shots [] orig_expr
1331 where
1332 -- Strip off existing lambdas and casts before handing off to mkEtaWW
1333 -- This is mainly to avoid spending time cloning binders and substituting
1334 -- when there is actually nothing to do. It's slightly awkward to deal
1335 -- with casts here, apart from the topmost one, and they are rare, so
1336 -- if we find one we just hand off to mkEtaWW anyway
1337 -- Note [Eta expansion and SCCs]
1338 go _ [] _ _ = orig_expr -- Already has the specified arity; no-op
1339
1340 go in_scope oss@(_:oss1) vs (Lam v body)
1341 | isTyVar v = go (in_scope `extendInScopeSet` v) oss (v:vs) body
1342 | otherwise = go (in_scope `extendInScopeSet` v) oss1 (v:vs) body
1343
1344 go in_scope oss rev_vs expr
1345 = -- pprTrace "ee" (vcat [ppr in_scope', ppr top_bndrs, ppr eis]) $
1346 retick $
1347 etaInfoAbs top_eis $
1348 etaInfoApp in_scope' sexpr eis
1349 where
1350 (in_scope', eis@(EI eta_bndrs mco))
1351 = mkEtaWW oss (ppr orig_expr) in_scope (exprType expr)
1352 top_bndrs = reverse rev_vs
1353 top_eis = EI (top_bndrs ++ eta_bndrs) (mkPiMCos top_bndrs mco)
1354
1355 -- Find ticks behind type apps.
1356 -- See Note [Eta expansion and source notes]
1357 -- I don't really understand this code SLPJ May 21
1358 (expr', args) = collectArgs expr
1359 (ticks, expr'') = stripTicksTop tickishFloatable expr'
1360 sexpr = mkApps expr'' args
1361 retick expr = foldr mkTick expr ticks
1362
1363 {- *********************************************************************
1364 * *
1365 The EtaInfo mechanism
1366 mkEtaWW, etaInfoAbs, etaInfoApp
1367 * *
1368 ********************************************************************* -}
1369
1370 {- Note [The EtaInfo mechanism]
1371 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1372 Suppose we have (e :: ty) and we want to eta-expand it to arity N.
1373 This what eta_expand does. We do it in two steps:
1374
1375 1. mkEtaWW: from 'ty' and 'N' build a EtaInfo which describes
1376 the shape of the expansion necessary to expand to arity N.
1377
1378 2. Build the term
1379 \ v1..vn. e v1 .. vn
1380 where those abstractions and applications are described by
1381 the same EtaInfo. Specifically we build the term
1382
1383 etaInfoAbs etas (etaInfoApp in_scope e etas)
1384
1385 where etas :: EtaInfo
1386 etaInfoAbs builds the lambdas
1387 etaInfoApp builds the applictions
1388
1389 Note that the /same/ EtaInfo drives both etaInfoAbs and etaInfoApp
1390
1391 To a first approximation EtaInfo is just [Var]. But
1392 casts complicate the question. If we have
1393 newtype N a = MkN (S -> a)
1394 axN :: N a ~ S -> a
1395 and
1396 e :: N (N Int)
1397 then the eta-expansion should look like
1398 (\(x::S) (y::S) -> e |> co x y) |> sym co
1399 where
1400 co :: N (N Int) ~ S -> S -> Int
1401 co = axN @(N Int) ; (S -> axN @Int)
1402
1403 We want to get one cast, at the top, to account for all those
1404 nested newtypes. This is expressed by the EtaInfo type:
1405
1406 data EtaInfo = EI [Var] MCoercionR
1407
1408 Note [Check for reflexive casts in eta expansion]
1409 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1410 It turns out that the casts created by teh above mechanism are often Refl.
1411 When casts are very deeply nested (as happens in #18223), the repetition
1412 of types can make the overall term very large. So there is a big
1413 payoff in cancelling out casts aggressively wherever possible.
1414 (See also Note [No crap in eta-expanded code].)
1415
1416 This matters particularly in etaInfoApp, where we
1417 * Do beta-reduction on the fly
1418 * Use getArg_maybe to get a cast out of the way,
1419 so that we can do beta reduction
1420 Together this makes a big difference. Consider when e is
1421 case x of
1422 True -> (\x -> e1) |> c1
1423 False -> (\p -> e2) |> c2
1424
1425 When we eta-expand this to arity 1, say, etaInfoAbs will wrap
1426 a (\eta) around the outside and use etaInfoApp to apply each
1427 alternative to 'eta'. We want to beta-reduce all that junk
1428 away.
1429
1430 #18223 was a dramatic example in which the intermediate term was
1431 grotesquely huge, even though the next Simplifier iteration squashed
1432 it. Better to kill it at birth.
1433
1434 The crucial spots in etaInfoApp are:
1435 * `checkReflexiveMCo` in the (Cast e co) case of `go`
1436 * `checkReflexiveMCo` in `pushCoArg`
1437 * Less important: checkReflexiveMCo in the final case of `go`
1438 Collectively these make a factor-of-5 difference to the total
1439 allocation of T18223, so take care if you change this stuff!
1440
1441 Example:
1442 newtype N = MkN (Y->Z)
1443 f :: X -> N
1444 f = \(x::X). ((\(y::Y). blah) |> fco)
1445
1446 where fco :: (Y->Z) ~ N
1447
1448 mkEtaWW makes an EtaInfo of (EI [(eta1:X), (eta2:Y)] eta_co
1449 where
1450 eta_co :: (X->N) ~ (X->Y->Z)
1451 eta_co = (<X> -> nco)
1452 nco :: N ~ (Y->Z) -- Comes from topNormaliseNewType_maybe
1453
1454 Now, when we push that eta_co inward in etaInfoApp:
1455 * In the (Cast e co) case, the 'fco' and 'nco' will meet, and
1456 should cancel.
1457 * When we meet the (\y.e) we want no cast on the y.
1458
1459 -}
1460
1461 --------------
1462 data EtaInfo = EI [Var] MCoercionR
1463
1464 -- (EI bs co) describes a particular eta-expansion, as follows:
1465 -- Abstraction: (\b1 b2 .. bn. []) |> sym co
1466 -- Application: ([] |> co) b1 b2 .. bn
1467 --
1468 -- e :: T co :: T ~ (t1 -> t2 -> .. -> tn -> tr)
1469 -- e = (\b1 b2 ... bn. (e |> co) b1 b2 .. bn) |> sym co
1470
1471 instance Outputable EtaInfo where
1472 ppr (EI vs mco) = text "EI" <+> ppr vs <+> parens (ppr mco)
1473
1474
1475 etaInfoApp :: InScopeSet -> CoreExpr -> EtaInfo -> CoreExpr
1476 -- (etaInfoApp s e (EI bs mco) returns something equivalent to
1477 -- ((substExpr s e) |> mco b1 .. bn)
1478 -- See Note [The EtaInfo mechanism]
1479 --
1480 -- NB: With very deeply nested casts, this function can be expensive
1481 -- In T18223, this function alone costs 15% of allocation, all
1482 -- spent in the calls to substExprSC and substBindSC
1483
1484 etaInfoApp in_scope expr eis
1485 = go (mkEmptySubst in_scope) expr eis
1486 where
1487 go :: Subst -> CoreExpr -> EtaInfo -> CoreExpr
1488 -- 'go' pushed down the eta-infos into the branch of a case
1489 -- and the body of a let; and does beta-reduction if possible
1490 -- go subst fun co [b1,..,bn] returns (subst(fun) |> co) b1 .. bn
1491 go subst (Tick t e) eis
1492 = Tick (substTickish subst t) (go subst e eis)
1493
1494 go subst (Cast e co) (EI bs mco)
1495 = go subst e (EI bs mco')
1496 where
1497 mco' = checkReflexiveMCo (Core.substCo subst co `mkTransMCoR` mco)
1498 -- See Note [Check for reflexive casts in eta-expansion]
1499
1500 go subst (Case e b ty alts) eis
1501 = Case (Core.substExprSC subst e) b1 ty' alts'
1502 where
1503 (subst1, b1) = Core.substBndr subst b
1504 alts' = map subst_alt alts
1505 ty' = etaInfoAppTy (Core.substTy subst ty) eis
1506 subst_alt (Alt con bs rhs) = Alt con bs' (go subst2 rhs eis)
1507 where
1508 (subst2,bs') = Core.substBndrs subst1 bs
1509
1510 go subst (Let b e) eis
1511 | not (isJoinBind b) -- See Note [Eta expansion for join points]
1512 = Let b' (go subst' e eis)
1513 where
1514 (subst', b') = Core.substBindSC subst b
1515
1516 -- Beta-reduction if possible, pushing any intervening casts past
1517 -- the argument. See Note [The EtaInfo mechansim]
1518 go subst (Lam v e) (EI (b:bs) mco)
1519 | Just (arg,mco') <- pushMCoArg mco (varToCoreExpr b)
1520 = go (Core.extendSubst subst v arg) e (EI bs mco')
1521
1522 -- Stop pushing down; just wrap the expression up
1523 -- See Note [Check for reflexive casts in eta expansion]
1524 go subst e (EI bs mco) = Core.substExprSC subst e
1525 `mkCastMCo` checkReflexiveMCo mco
1526 `mkVarApps` bs
1527
1528 --------------
1529 etaInfoAppTy :: Type -> EtaInfo -> Type
1530 -- If e :: ty
1531 -- then etaInfoApp e eis :: etaInfoApp ty eis
1532 etaInfoAppTy ty (EI bs mco)
1533 = applyTypeToArgs (text "etaInfoAppTy") ty1 (map varToCoreExpr bs)
1534 where
1535 ty1 = case mco of
1536 MRefl -> ty
1537 MCo co -> coercionRKind co
1538
1539 --------------
1540 etaInfoAbs :: EtaInfo -> CoreExpr -> CoreExpr
1541 -- See Note [The EtaInfo mechanism]
1542 etaInfoAbs (EI bs mco) expr = (mkLams bs expr) `mkCastMCo` mkSymMCo mco
1543
1544 --------------
1545 -- | @mkEtaWW n _ fvs ty@ will compute the 'EtaInfo' necessary for eta-expanding
1546 -- an expression @e :: ty@ to take @n@ value arguments, where @fvs@ are the
1547 -- free variables of @e@.
1548 --
1549 -- Note that this function is entirely unconcerned about cost centres and other
1550 -- semantically-irrelevant source annotations, so call sites must take care to
1551 -- preserve that info. See Note [Eta expansion and SCCs].
1552 mkEtaWW
1553 :: [OneShotInfo]
1554 -- ^ How many value arguments to eta-expand
1555 -> SDoc
1556 -- ^ The pretty-printed original expression, for warnings.
1557 -> InScopeSet
1558 -- ^ A super-set of the free vars of the expression to eta-expand.
1559 -> Type
1560 -> (InScopeSet, EtaInfo)
1561 -- ^ The variables in 'EtaInfo' are fresh wrt. to the incoming 'InScopeSet'.
1562 -- The outgoing 'InScopeSet' extends the incoming 'InScopeSet' with the
1563 -- fresh variables in 'EtaInfo'.
1564
1565 mkEtaWW orig_oss ppr_orig_expr in_scope orig_ty
1566 = go 0 orig_oss empty_subst orig_ty
1567 where
1568 empty_subst = mkEmptyTCvSubst in_scope
1569
1570 go :: Int -- For fresh names
1571 -> [OneShotInfo] -- Number of value args to expand to
1572 -> TCvSubst -> Type -- We are really looking at subst(ty)
1573 -> (InScopeSet, EtaInfo)
1574 -- (go [o1,..,on] subst ty) = (in_scope, EI [b1,..,bn] co)
1575 -- co :: subst(ty) ~ b1_ty -> ... -> bn_ty -> tr
1576
1577 go _ [] subst _ -- See Note [exprArity invariant]
1578 ----------- Done! No more expansion needed
1579 = (getTCvInScope subst, EI [] MRefl)
1580
1581 go n oss@(one_shot:oss1) subst ty -- See Note [exprArity invariant]
1582 ----------- Forall types (forall a. ty)
1583 | Just (tcv,ty') <- splitForAllTyCoVar_maybe ty
1584 , (subst', tcv') <- Type.substVarBndr subst tcv
1585 , let oss' | isTyVar tcv = oss
1586 | otherwise = oss1
1587 -- A forall can bind a CoVar, in which case
1588 -- we consume one of the [OneShotInfo]
1589 , (in_scope, EI bs mco) <- go n oss' subst' ty'
1590 = (in_scope, EI (tcv' : bs) (mkHomoForAllMCo tcv' mco))
1591
1592 ----------- Function types (t1 -> t2)
1593 | Just (mult, arg_ty, res_ty) <- splitFunTy_maybe ty
1594 , typeHasFixedRuntimeRep arg_ty
1595 -- See Note [Representation polymorphism invariants] in GHC.Core
1596 -- See also test case typecheck/should_run/EtaExpandLevPoly
1597
1598 , (subst', eta_id) <- freshEtaId n subst (Scaled mult arg_ty)
1599 -- Avoid free vars of the original expression
1600
1601 , let eta_id' = eta_id `setIdOneShotInfo` one_shot
1602 , (in_scope, EI bs mco) <- go (n+1) oss1 subst' res_ty
1603 = (in_scope, EI (eta_id' : bs) (mkFunResMCo (idScaledType eta_id') mco))
1604
1605 ----------- Newtypes
1606 -- Given this:
1607 -- newtype T = MkT ([T] -> Int)
1608 -- Consider eta-expanding this
1609 -- eta_expand 1 e T
1610 -- We want to get
1611 -- coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
1612 | Just (co, ty') <- topNormaliseNewType_maybe ty
1613 , -- co :: ty ~ ty'
1614 let co' = Type.substCo subst co
1615 -- Remember to apply the substitution to co (#16979)
1616 -- (or we could have applied to ty, but then
1617 -- we'd have had to zap it for the recursive call)
1618 , (in_scope, EI bs mco) <- go n oss subst ty'
1619 -- mco :: subst(ty') ~ b1_ty -> ... -> bn_ty -> tr
1620 = (in_scope, EI bs (mkTransMCoR co' mco))
1621
1622 | otherwise -- We have an expression of arity > 0,
1623 -- but its type isn't a function, or a binder
1624 -- does not have a fixed runtime representation
1625 = warnPprTrace True ((ppr orig_oss <+> ppr orig_ty) $$ ppr_orig_expr)
1626 (getTCvInScope subst, EI [] MRefl)
1627 -- This *can* legitimately happen:
1628 -- e.g. coerce Int (\x. x) Essentially the programmer is
1629 -- playing fast and loose with types (Happy does this a lot).
1630 -- So we simply decline to eta-expand. Otherwise we'd end up
1631 -- with an explicit lambda having a non-function type
1632
1633
1634 {- *********************************************************************
1635 * *
1636 The "push rules"
1637 * *
1638 ************************************************************************
1639
1640 Here we implement the "push rules" from FC papers:
1641
1642 * The push-argument rules, where we can move a coercion past an argument.
1643 We have
1644 (fun |> co) arg
1645 and we want to transform it to
1646 (fun arg') |> co'
1647 for some suitable co' and transformed arg'.
1648
1649 * The PushK rule for data constructors. We have
1650 (K e1 .. en) |> co
1651 and we want to transform to
1652 (K e1' .. en')
1653 by pushing the coercion into the arguments
1654 -}
1655
1656 pushCoArgs :: CoercionR -> [CoreArg] -> Maybe ([CoreArg], MCoercion)
1657 pushCoArgs co [] = return ([], MCo co)
1658 pushCoArgs co (arg:args) = do { (arg', m_co1) <- pushCoArg co arg
1659 ; case m_co1 of
1660 MCo co1 -> do { (args', m_co2) <- pushCoArgs co1 args
1661 ; return (arg':args', m_co2) }
1662 MRefl -> return (arg':args, MRefl) }
1663
1664 pushMCoArg :: MCoercionR -> CoreArg -> Maybe (CoreArg, MCoercion)
1665 pushMCoArg MRefl arg = Just (arg, MRefl)
1666 pushMCoArg (MCo co) arg = pushCoArg co arg
1667
1668 pushCoArg :: CoercionR -> CoreArg -> Maybe (CoreArg, MCoercion)
1669 -- We have (fun |> co) arg, and we want to transform it to
1670 -- (fun arg) |> co
1671 -- This may fail, e.g. if (fun :: N) where N is a newtype
1672 -- C.f. simplCast in GHC.Core.Opt.Simplify
1673 -- 'co' is always Representational
1674 pushCoArg co arg
1675 | Type ty <- arg
1676 = do { (ty', m_co') <- pushCoTyArg co ty
1677 ; return (Type ty', m_co') }
1678 | otherwise
1679 = do { (arg_mco, m_co') <- pushCoValArg co
1680 ; let arg_mco' = checkReflexiveMCo arg_mco
1681 -- checkReflexiveMCo: see Note [Check for reflexive casts in eta expansion]
1682 -- The coercion is very often (arg_co -> res_co), but without
1683 -- the argument coercion actually being ReflCo
1684 ; return (arg `mkCastMCo` arg_mco', m_co') }
1685
1686 pushCoTyArg :: CoercionR -> Type -> Maybe (Type, MCoercionR)
1687 -- We have (fun |> co) @ty
1688 -- Push the coercion through to return
1689 -- (fun @ty') |> co'
1690 -- 'co' is always Representational
1691 -- If the returned coercion is Nothing, then it would have been reflexive;
1692 -- it's faster not to compute it, though.
1693 pushCoTyArg co ty
1694 -- The following is inefficient - don't do `eqType` here, the coercion
1695 -- optimizer will take care of it. See #14737.
1696 -- -- | tyL `eqType` tyR
1697 -- -- = Just (ty, Nothing)
1698
1699 | isReflCo co
1700 = Just (ty, MRefl)
1701
1702 | isForAllTy_ty tyL
1703 = assertPpr (isForAllTy_ty tyR) (ppr co $$ ppr ty) $
1704 Just (ty `mkCastTy` co1, MCo co2)
1705
1706 | otherwise
1707 = Nothing
1708 where
1709 Pair tyL tyR = coercionKind co
1710 -- co :: tyL ~ tyR
1711 -- tyL = forall (a1 :: k1). ty1
1712 -- tyR = forall (a2 :: k2). ty2
1713
1714 co1 = mkSymCo (mkNthCo Nominal 0 co)
1715 -- co1 :: k2 ~N k1
1716 -- Note that NthCo can extract a Nominal equality between the
1717 -- kinds of the types related by a coercion between forall-types.
1718 -- See the NthCo case in GHC.Core.Lint.
1719
1720 co2 = mkInstCo co (mkGReflLeftCo Nominal ty co1)
1721 -- co2 :: ty1[ (ty|>co1)/a1 ] ~ ty2[ ty/a2 ]
1722 -- Arg of mkInstCo is always nominal, hence mkNomReflCo
1723
1724 pushCoValArg :: CoercionR -> Maybe (MCoercionR, MCoercionR)
1725 -- We have (fun |> co) arg
1726 -- Push the coercion through to return
1727 -- (fun (arg |> co_arg)) |> co_res
1728 -- 'co' is always Representational
1729 -- If the second returned Coercion is actually Nothing, then no cast is necessary;
1730 -- the returned coercion would have been reflexive.
1731 pushCoValArg co
1732 -- The following is inefficient - don't do `eqType` here, the coercion
1733 -- optimizer will take care of it. See #14737.
1734 -- -- | tyL `eqType` tyR
1735 -- -- = Just (mkRepReflCo arg, Nothing)
1736
1737 | isReflCo co
1738 = Just (MRefl, MRefl)
1739
1740 | isFunTy tyL
1741 , (co_mult, co1, co2) <- decomposeFunCo Representational co
1742 , isReflexiveCo co_mult
1743 -- We can't push the coercion in the case where co_mult isn't reflexivity:
1744 -- it could be an unsafe axiom, and losing this information could yield
1745 -- ill-typed terms. For instance (fun x ::(1) Int -> (fun _ -> () |> co) x)
1746 -- with co :: (Int -> ()) ~ (Int %1 -> ()), would reduce to (fun x ::(1) Int
1747 -- -> (fun _ ::(Many) Int -> ()) x) which is ill-typed
1748
1749 -- If co :: (tyL1 -> tyL2) ~ (tyR1 -> tyR2)
1750 -- then co1 :: tyL1 ~ tyR1
1751 -- co2 :: tyL2 ~ tyR2
1752 = assertPpr (isFunTy tyR) (ppr co $$ ppr arg) $
1753 Just (coToMCo (mkSymCo co1), coToMCo co2)
1754 -- Critically, coToMCo to checks for ReflCo; the whole coercion may not
1755 -- be reflexive, but either of its components might be
1756 -- We could use isReflexiveCo, but it's not clear if the benefit
1757 -- is worth the cost, and it makes no difference in #18223
1758
1759 | otherwise
1760 = Nothing
1761 where
1762 arg = funArgTy tyR
1763 Pair tyL tyR = coercionKind co
1764
1765 pushCoercionIntoLambda
1766 :: HasDebugCallStack => InScopeSet -> Var -> CoreExpr -> CoercionR -> Maybe (Var, CoreExpr)
1767 -- This implements the Push rule from the paper on coercions
1768 -- (\x. e) |> co
1769 -- ===>
1770 -- (\x'. e |> co')
1771 pushCoercionIntoLambda in_scope x e co
1772 | assert (not (isTyVar x) && not (isCoVar x)) True
1773 , Pair s1s2 t1t2 <- coercionKind co
1774 , Just (_, _s1,_s2) <- splitFunTy_maybe s1s2
1775 , Just (w1, t1,_t2) <- splitFunTy_maybe t1t2
1776 , (co_mult, co1, co2) <- decomposeFunCo Representational co
1777 , isReflexiveCo co_mult
1778 -- We can't push the coercion in the case where co_mult isn't
1779 -- reflexivity. See pushCoValArg for more details.
1780 = let
1781 -- Should we optimize the coercions here?
1782 -- Otherwise they might not match too well
1783 x' = x `setIdType` t1 `setIdMult` w1
1784 in_scope' = in_scope `extendInScopeSet` x'
1785 subst = extendIdSubst (mkEmptySubst in_scope')
1786 x
1787 (mkCast (Var x') co1)
1788 in Just (x', substExpr subst e `mkCast` co2)
1789 | otherwise
1790 = pprTrace "exprIsLambda_maybe: Unexpected lambda in case" (ppr (Lam x e))
1791 Nothing
1792
1793 pushCoDataCon :: DataCon -> [CoreExpr] -> Coercion
1794 -> Maybe (DataCon
1795 , [Type] -- Universal type args
1796 , [CoreExpr]) -- All other args incl existentials
1797 -- Implement the KPush reduction rule as described in "Down with kinds"
1798 -- The transformation applies iff we have
1799 -- (C e1 ... en) `cast` co
1800 -- where co :: (T t1 .. tn) ~ to_ty
1801 -- The left-hand one must be a T, because exprIsConApp returned True
1802 -- but the right-hand one might not be. (Though it usually will.)
1803 pushCoDataCon dc dc_args co
1804 | isReflCo co || from_ty `eqType` to_ty -- try cheap test first
1805 , let (univ_ty_args, rest_args) = splitAtList (dataConUnivTyVars dc) dc_args
1806 = Just (dc, map exprToType univ_ty_args, rest_args)
1807
1808 | Just (to_tc, to_tc_arg_tys) <- splitTyConApp_maybe to_ty
1809 , to_tc == dataConTyCon dc
1810 -- These two tests can fail; we might see
1811 -- (C x y) `cast` (g :: T a ~ S [a]),
1812 -- where S is a type function. In fact, exprIsConApp
1813 -- will probably not be called in such circumstances,
1814 -- but there's nothing wrong with it
1815
1816 = let
1817 tc_arity = tyConArity to_tc
1818 dc_univ_tyvars = dataConUnivTyVars dc
1819 dc_ex_tcvars = dataConExTyCoVars dc
1820 arg_tys = dataConRepArgTys dc
1821
1822 non_univ_args = dropList dc_univ_tyvars dc_args
1823 (ex_args, val_args) = splitAtList dc_ex_tcvars non_univ_args
1824
1825 -- Make the "Psi" from the paper
1826 omegas = decomposeCo tc_arity co (tyConRolesRepresentational to_tc)
1827 (psi_subst, to_ex_arg_tys)
1828 = liftCoSubstWithEx Representational
1829 dc_univ_tyvars
1830 omegas
1831 dc_ex_tcvars
1832 (map exprToType ex_args)
1833
1834 -- Cast the value arguments (which include dictionaries)
1835 new_val_args = zipWith cast_arg (map scaledThing arg_tys) val_args
1836 cast_arg arg_ty arg = mkCast arg (psi_subst arg_ty)
1837
1838 to_ex_args = map Type to_ex_arg_tys
1839
1840 dump_doc = vcat [ppr dc, ppr dc_univ_tyvars, ppr dc_ex_tcvars,
1841 ppr arg_tys, ppr dc_args,
1842 ppr ex_args, ppr val_args, ppr co, ppr from_ty, ppr to_ty, ppr to_tc
1843 , ppr $ mkTyConApp to_tc (map exprToType $ takeList dc_univ_tyvars dc_args) ]
1844 in
1845 assertPpr (eqType from_ty (mkTyConApp to_tc (map exprToType $ takeList dc_univ_tyvars dc_args))) dump_doc $
1846 assertPpr (equalLength val_args arg_tys) dump_doc $
1847 Just (dc, to_tc_arg_tys, to_ex_args ++ new_val_args)
1848
1849 | otherwise
1850 = Nothing
1851
1852 where
1853 Pair from_ty to_ty = coercionKind co
1854
1855 collectBindersPushingCo :: CoreExpr -> ([Var], CoreExpr)
1856 -- Collect lambda binders, pushing coercions inside if possible
1857 -- E.g. (\x.e) |> g g :: <Int> -> blah
1858 -- = (\x. e |> Nth 1 g)
1859 --
1860 -- That is,
1861 --
1862 -- collectBindersPushingCo ((\x.e) |> g) === ([x], e |> Nth 1 g)
1863 collectBindersPushingCo e
1864 = go [] e
1865 where
1866 -- Peel off lambdas until we hit a cast.
1867 go :: [Var] -> CoreExpr -> ([Var], CoreExpr)
1868 -- The accumulator is in reverse order
1869 go bs (Lam b e) = go (b:bs) e
1870 go bs (Cast e co) = go_c bs e co
1871 go bs e = (reverse bs, e)
1872
1873 -- We are in a cast; peel off casts until we hit a lambda.
1874 go_c :: [Var] -> CoreExpr -> CoercionR -> ([Var], CoreExpr)
1875 -- (go_c bs e c) is same as (go bs e (e |> c))
1876 go_c bs (Cast e co1) co2 = go_c bs e (co1 `mkTransCo` co2)
1877 go_c bs (Lam b e) co = go_lam bs b e co
1878 go_c bs e co = (reverse bs, mkCast e co)
1879
1880 -- We are in a lambda under a cast; peel off lambdas and build a
1881 -- new coercion for the body.
1882 go_lam :: [Var] -> Var -> CoreExpr -> CoercionR -> ([Var], CoreExpr)
1883 -- (go_lam bs b e c) is same as (go_c bs (\b.e) c)
1884 go_lam bs b e co
1885 | isTyVar b
1886 , let Pair tyL tyR = coercionKind co
1887 , assert (isForAllTy_ty tyL) $
1888 isForAllTy_ty tyR
1889 , isReflCo (mkNthCo Nominal 0 co) -- See Note [collectBindersPushingCo]
1890 = go_c (b:bs) e (mkInstCo co (mkNomReflCo (mkTyVarTy b)))
1891
1892 | isCoVar b
1893 , let Pair tyL tyR = coercionKind co
1894 , assert (isForAllTy_co tyL) $
1895 isForAllTy_co tyR
1896 , isReflCo (mkNthCo Nominal 0 co) -- See Note [collectBindersPushingCo]
1897 , let cov = mkCoVarCo b
1898 = go_c (b:bs) e (mkInstCo co (mkNomReflCo (mkCoercionTy cov)))
1899
1900 | isId b
1901 , let Pair tyL tyR = coercionKind co
1902 , assert (isFunTy tyL) $ isFunTy tyR
1903 , (co_mult, co_arg, co_res) <- decomposeFunCo Representational co
1904 , isReflCo co_mult -- See Note [collectBindersPushingCo]
1905 , isReflCo co_arg -- See Note [collectBindersPushingCo]
1906 = go_c (b:bs) e co_res
1907
1908 | otherwise = (reverse bs, mkCast (Lam b e) co)
1909
1910 {-
1911
1912 Note [collectBindersPushingCo]
1913 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1914 We just look for coercions of form
1915 <type> # w -> blah
1916 (and similarly for foralls) to keep this function simple. We could do
1917 more elaborate stuff, but it'd involve substitution etc.
1918
1919 -}
1920
1921 {- *********************************************************************
1922 * *
1923 Join points
1924 * *
1925 ********************************************************************* -}
1926
1927 -------------------
1928 -- | Split an expression into the given number of binders and a body,
1929 -- eta-expanding if necessary. Counts value *and* type binders.
1930 etaExpandToJoinPoint :: JoinArity -> CoreExpr -> ([CoreBndr], CoreExpr)
1931 etaExpandToJoinPoint join_arity expr
1932 = go join_arity [] expr
1933 where
1934 go 0 rev_bs e = (reverse rev_bs, e)
1935 go n rev_bs (Lam b e) = go (n-1) (b : rev_bs) e
1936 go n rev_bs e = case etaBodyForJoinPoint n e of
1937 (bs, e') -> (reverse rev_bs ++ bs, e')
1938
1939 etaExpandToJoinPointRule :: JoinArity -> CoreRule -> CoreRule
1940 etaExpandToJoinPointRule _ rule@(BuiltinRule {})
1941 = warnPprTrace True (sep [text "Can't eta-expand built-in rule:", ppr rule])
1942 -- How did a local binding get a built-in rule anyway? Probably a plugin.
1943 rule
1944 etaExpandToJoinPointRule join_arity rule@(Rule { ru_bndrs = bndrs, ru_rhs = rhs
1945 , ru_args = args })
1946 | need_args == 0
1947 = rule
1948 | need_args < 0
1949 = pprPanic "etaExpandToJoinPointRule" (ppr join_arity $$ ppr rule)
1950 | otherwise
1951 = rule { ru_bndrs = bndrs ++ new_bndrs, ru_args = args ++ new_args
1952 , ru_rhs = new_rhs }
1953 where
1954 need_args = join_arity - length args
1955 (new_bndrs, new_rhs) = etaBodyForJoinPoint need_args rhs
1956 new_args = varsToCoreExprs new_bndrs
1957
1958 -- Adds as many binders as asked for; assumes expr is not a lambda
1959 etaBodyForJoinPoint :: Int -> CoreExpr -> ([CoreBndr], CoreExpr)
1960 etaBodyForJoinPoint need_args body
1961 = go need_args (exprType body) (init_subst body) [] body
1962 where
1963 go 0 _ _ rev_bs e
1964 = (reverse rev_bs, e)
1965 go n ty subst rev_bs e
1966 | Just (tv, res_ty) <- splitForAllTyCoVar_maybe ty
1967 , let (subst', tv') = substVarBndr subst tv
1968 = go (n-1) res_ty subst' (tv' : rev_bs) (e `App` varToCoreExpr tv')
1969 | Just (mult, arg_ty, res_ty) <- splitFunTy_maybe ty
1970 , let (subst', b) = freshEtaId n subst (Scaled mult arg_ty)
1971 = go (n-1) res_ty subst' (b : rev_bs) (e `App` Var b)
1972 | otherwise
1973 = pprPanic "etaBodyForJoinPoint" $ int need_args $$
1974 ppr body $$ ppr (exprType body)
1975
1976 init_subst e = mkEmptyTCvSubst (mkInScopeSet (exprFreeVars e))
1977
1978
1979
1980 --------------
1981 freshEtaId :: Int -> TCvSubst -> Scaled Type -> (TCvSubst, Id)
1982 -- Make a fresh Id, with specified type (after applying substitution)
1983 -- It should be "fresh" in the sense that it's not in the in-scope set
1984 -- of the TvSubstEnv; and it should itself then be added to the in-scope
1985 -- set of the TvSubstEnv
1986 --
1987 -- The Int is just a reasonable starting point for generating a unique;
1988 -- it does not necessarily have to be unique itself.
1989 freshEtaId n subst ty
1990 = (subst', eta_id')
1991 where
1992 Scaled mult' ty' = Type.substScaledTyUnchecked subst ty
1993 eta_id' = uniqAway (getTCvInScope subst) $
1994 mkSysLocalOrCoVar (fsLit "eta") (mkBuiltinUnique n) mult' ty'
1995 -- "OrCoVar" since this can be used to eta-expand
1996 -- coercion abstractions
1997 subst' = extendTCvInScope subst eta_id'