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