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