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