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