2 % (c) The University of Glasgow 2006
3 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
6 Arity and ete expansion
9 -- | Arit and eta expansion
11 manifestArity, exprArity,
12 exprEtaExpandArity, etaExpand
15 #include "HsVersions.h"
21 import TyCon ( isRecursiveTyCon )
22 import qualified CoreSubst
23 import CoreSubst ( Subst, substBndr, substBndrs, substExpr
24 , mkEmptySubst, isEmptySubst )
29 import TcType ( isDictLikeTy )
35 import StaticFlags ( opt_NoStateHack )
39 %************************************************************************
41 manifestArity and exprArity
43 %************************************************************************
45 exprArity is a cheap-and-cheerful version of exprEtaExpandArity.
46 It tells how many things the expression can be applied to before doing
47 any work. It doesn't look inside cases, lets, etc. The idea is that
48 exprEtaExpandArity will do the hard work, leaving something that's easy
49 for exprArity to grapple with. In particular, Simplify uses exprArity to
50 compute the ArityInfo for the Id.
52 Originally I thought that it was enough just to look for top-level lambdas, but
53 it isn't. I've seen this
55 foo = PrelBase.timesInt
57 We want foo to get arity 2 even though the eta-expander will leave it
58 unchanged, in the expectation that it'll be inlined. But occasionally it
59 isn't, because foo is blacklisted (used in a rule).
61 Similarly, see the ok_note check in exprEtaExpandArity. So
62 f = __inline_me (\x -> e)
63 won't be eta-expanded.
65 And in any case it seems more robust to have exprArity be a bit more intelligent.
66 But note that (\x y z -> f x y z)
67 should have arity 3, regardless of f's arity.
69 Note [exprArity invariant]
70 ~~~~~~~~~~~~~~~~~~~~~~~~~~
71 exprArity has the following invariant:
72 (exprArity e) = n, then manifestArity (etaExpand e n) = n
74 That is, if exprArity says "the arity is n" then etaExpand really can get
75 "n" manifest lambdas to the top.
77 Why is this important? Because
78 - In TidyPgm we use exprArity to fix the *final arity* of
79 each top-level Id, and in
80 - In CorePrep we use etaExpand on each rhs, so that the visible lambdas
81 actually match that arity, which in turn means
82 that the StgRhs has the right number of lambdas
84 An alternative would be to do the eta-expansion in TidyPgm, at least
85 for top-level bindings, in which case we would not need the trim_arity
86 in exprArity. That is a less local change, so I'm going to leave it for today!
90 manifestArity :: CoreExpr -> Arity
91 -- ^ manifestArity sees how many leading value lambdas there are
92 manifestArity (Lam v e) | isId v = 1 + manifestArity e
93 | otherwise = manifestArity e
94 manifestArity (Note _ e) = manifestArity e
95 manifestArity (Cast e _) = manifestArity e
98 exprArity :: CoreExpr -> Arity
99 -- ^ An approximate, fast, version of 'exprEtaExpandArity'
102 go (Var v) = idArity v
103 go (Lam x e) | isId x = go e + 1
106 go (Cast e co) = trim_arity (go e) 0 (snd (coercionKind co))
107 go (App e (Type _)) = go e
108 go (App f a) | exprIsCheap a = (go f - 1) `max` 0
109 -- NB: exprIsCheap a!
110 -- f (fac x) does not have arity 2,
111 -- even if f has arity 3!
112 -- NB: `max 0`! (\x y -> f x) has arity 2, even if f is
113 -- unknown, hence arity 0
116 -- Note [exprArity invariant]
119 | Just (_, ty') <- splitForAllTy_maybe ty = trim_arity n a ty'
120 | Just (_, ty') <- splitFunTy_maybe ty = trim_arity n (a+1) ty'
121 | Just (ty',_) <- splitNewTypeRepCo_maybe ty = trim_arity n a ty'
125 %************************************************************************
129 %************************************************************************
132 -- ^ The Arity returned is the number of value args the
133 -- expression can be applied to without doing much work
134 exprEtaExpandArity :: DynFlags -> CoreExpr -> Arity
135 -- exprEtaExpandArity is used when eta expanding
136 -- e ==> \xy -> e x y
137 exprEtaExpandArity dflags e
138 = applyStateHack e (arityType dicts_cheap e)
140 dicts_cheap = dopt Opt_DictsCheap dflags
143 Note [Definition of arity]
144 ~~~~~~~~~~~~~~~~~~~~~~~~~~
145 The "arity" of an expression 'e' is n if
146 applying 'e' to *fewer* than n *value* arguments
149 Or, to put it another way
151 there is no work lost in duplicating the partial
152 application (e x1 .. x(n-1))
154 In the divegent case, no work is lost by duplicating because if the thing
155 is evaluated once, that's the end of the program.
157 Or, to put it another way, in any context C
159 C[ (\x1 .. xn. e x1 .. xn) ]
164 It's all a bit more subtle than it looks:
166 Note [Arity of case expressions]
167 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
168 We treat the arity of
169 case x of p -> \s -> ...
170 as 1 (or more) because for I/O ish things we really want to get that
171 \s to the top. We are prepared to evaluate x each time round the loop
172 in order to get that.
174 This isn't really right in the presence of seq. Consider
178 Can we eta-expand here? At first the answer looks like "yes of course", but
181 This should diverge! But if we eta-expand, it won't. Again, we ignore this
182 "problem", because being scrupulous would lose an important transformation for
186 1. Note [One-shot lambdas]
187 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
188 Consider one-shot lambdas
189 let x = expensive in \y z -> E
190 We want this to have arity 1 if the \y-abstraction is a 1-shot lambda.
192 3. Note [Dealing with bottom]
193 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
195 f = \x -> error "foo"
196 Here, arity 1 is fine. But if it is
200 then we want to get arity 2. Technically, this isn't quite right, because
202 should diverge, but it'll converge if we eta-expand f. Nevertheless, we
203 do so; it improves some programs significantly, and increasing convergence
204 isn't a bad thing. Hence the ABot/ATop in ArityType.
207 4. Note [Newtype arity]
208 ~~~~~~~~~~~~~~~~~~~~~~~~
209 Non-recursive newtypes are transparent, and should not get in the way.
210 We do (currently) eta-expand recursive newtypes too. So if we have, say
212 newtype T = MkT ([T] -> Int)
216 where f has arity 1. Then: etaExpandArity e = 1;
217 that is, etaExpandArity looks through the coerce.
219 When we eta-expand e to arity 1: eta_expand 1 e T
220 we want to get: coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
222 HOWEVER, note that if you use coerce bogusly you can ge
224 And since negate has arity 2, you might try to eta expand. But you can't
225 decopose Int to a function type. Hence the final case in eta_expand.
227 Note [The state-transformer hack]
228 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
231 where e has arity n. Then, if we know from the context that f has
233 t1 -> ... -> tn -1-> t(n+1) -1-> ... -1-> tm -> ...
234 then we can expand the arity to m. This usage type says that
235 any application (x e1 .. en) will be applied to uniquely to (m-n) more args
236 Consider f = \x. let y = <expensive>
239 False -> \(s:RealWorld) -> e
240 where foo has arity 1. Then we want the state hack to
241 apply to foo too, so we can eta expand the case.
243 Then we expect that if f is applied to one arg, it'll be applied to two
244 (that's the hack -- we don't really know, and sometimes it's false)
245 See also Id.isOneShotBndr.
248 applyStateHack :: CoreExpr -> ArityType -> Arity
249 applyStateHack e (AT orig_arity is_bot)
250 | opt_NoStateHack = orig_arity
251 | ABot <- is_bot = orig_arity -- Note [State hack and bottoming functions]
252 | otherwise = go orig_ty orig_arity
253 where -- Note [The state-transformer hack]
255 go :: Type -> Arity -> Arity
256 go ty arity -- This case analysis should match that in eta_expand
257 | Just (_, ty') <- splitForAllTy_maybe ty = go ty' arity
259 | Just (tc,tys) <- splitTyConApp_maybe ty
260 , Just (ty', _) <- instNewTyCon_maybe tc tys
261 , not (isRecursiveTyCon tc) = go ty' arity
262 -- Important to look through non-recursive newtypes, so that, eg
263 -- (f x) where f has arity 2, f :: Int -> IO ()
264 -- Here we want to get arity 1 for the result!
266 | Just (arg,res) <- splitFunTy_maybe ty
267 , arity > 0 || isStateHackType arg = 1 + go res (arity-1)
269 = if arity > 0 then 1 + go res (arity-1)
270 else if isStateHackType arg then
271 pprTrace "applystatehack" (vcat [ppr orig_arity, ppr orig_ty,
272 ppr ty, ppr res, ppr e]) $
274 else WARN( arity > 0, ppr arity ) 0
276 | otherwise = WARN( arity > 0, ppr arity ) 0
279 Note [State hack and bottoming functions]
280 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
281 It's a terrible idea to use the state hack on a bottoming function.
282 Here's what happens (Trac #2861):
287 Eta-expand, using the state hack:
289 f = \p. (\s. ((error "...") |> g1) s) |> g2
290 g1 :: IO T ~ (S -> (S,T))
291 g2 :: (S -> (S,T)) ~ IO T
295 f' = \p. \s. ((error "...") |> g1) s
296 f = f' |> (String -> g2)
298 Discard args for bottomming function
300 f' = \p. \s. ((error "...") |> g1 |> g3
301 g3 :: (S -> (S,T)) ~ (S,T)
305 f'' = \p. \s. (error "...")
306 f' = f'' |> (String -> S -> g1.g3)
308 And now we can repeat the whole loop. Aargh! The bug is in applying the
309 state hack to a function which then swallows the argument.
312 -------------------- Main arity code ----------------------------
314 -- If e has ArityType (AT n r), then the term 'e'
315 -- * Must be applied to at least n *value* args
316 -- before doing any significant work
317 -- * It will not diverge before being applied to n
319 -- * If 'r' is ABot, then it guarantees to diverge if
320 -- applied to n arguments (or more)
322 data ArityType = AT Arity ArityRes
323 data ArityRes = ATop -- Know nothing
326 vanillaArityType :: ArityType
327 vanillaArityType = AT 0 ATop -- Totally uninformative
329 incArity :: ArityType -> ArityType
330 incArity (AT a r) = AT (a+1) r
332 decArity :: ArityType -> ArityType
333 decArity (AT 0 r) = AT 0 r
334 decArity (AT a r) = AT (a-1) r
336 andArityType :: ArityType -> ArityType -> ArityType -- Used for branches of a 'case'
337 andArityType (AT a1 ATop) (AT a2 ATop) = AT (a1 `min` a2) ATop
338 andArityType (AT _ ABot) (AT a2 ATop) = AT a2 ATop
339 andArityType (AT a1 ATop) (AT _ ABot) = AT a1 ATop
340 andArityType (AT a1 ABot) (AT a2 ABot) = AT (a1 `max` a2) ABot
342 trimArity :: Bool -> ArityType -> ArityType
343 -- We have something like (let x = E in b), where b has the given
345 -- * If E is cheap we can push it inside as far as we like
346 -- * If b eventually diverges, we allow ourselves to push inside
347 -- arbitrarily, even though that is not quite right
348 trimArity _cheap (AT a ABot) = AT a ABot
349 trimArity True (AT a ATop) = AT a ATop
350 trimArity False (AT _ ATop) = AT 0 ATop -- Bale out
352 ---------------------------
353 arityType :: Bool -> CoreExpr -> ArityType
355 | Just strict_sig <- idNewStrictness_maybe v
356 , (ds, res) <- splitStrictSig strict_sig
358 = AT (length ds) ABot -- Function diverges
360 = AT (idArity v) ATop
362 -- Lambdas; increase arity
363 arityType dicts_cheap (Lam x e)
364 | isId x = incArity (arityType dicts_cheap e)
365 | otherwise = arityType dicts_cheap e
367 -- Applications; decrease arity
368 arityType dicts_cheap (App fun (Type _))
369 = arityType dicts_cheap fun
370 arityType dicts_cheap (App fun arg )
371 = trimArity (exprIsCheap arg) (decArity (arityType dicts_cheap fun))
373 -- Case/Let; keep arity if either the expression is cheap
374 -- or it's a 1-shot lambda
375 -- The former is not really right for Haskell
376 -- f x = case x of { (a,b) -> \y. e }
378 -- f x y = case x of { (a,b) -> e }
379 -- The difference is observable using 'seq'
380 arityType dicts_cheap (Case scrut _ _ alts)
381 = trimArity (exprIsCheap scrut)
382 (foldr1 andArityType [arityType dicts_cheap rhs | (_,_,rhs) <- alts])
384 arityType dicts_cheap (Let b e)
385 = trimArity (cheap_bind b) (arityType dicts_cheap e)
387 cheap_bind (NonRec b e) = is_cheap (b,e)
388 cheap_bind (Rec prs) = all is_cheap prs
389 is_cheap (b,e) = (dicts_cheap && isDictLikeTy (idType b))
391 -- If the experimental -fdicts-cheap flag is on, we eta-expand through
392 -- dictionary bindings. This improves arities. Thereby, it also
393 -- means that full laziness is less prone to floating out the
394 -- application of a function to its dictionary arguments, which
395 -- can thereby lose opportunities for fusion. Example:
396 -- foo :: Ord a => a -> ...
397 -- foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....
398 -- -- So foo has arity 1
400 -- f = \x. foo dInt $ bar x
402 -- The (foo DInt) is floated out, and makes ineffective a RULE
405 -- One could go further and make exprIsCheap reply True to any
406 -- dictionary-typed expression, but that's more work.
408 -- See Note [Dictionary-like types] in TcType.lhs for why we use
409 -- isDictLikeTy here rather than isDictTy
411 arityType dicts_cheap (Note _ e) = arityType dicts_cheap e
412 arityType dicts_cheap (Cast e _) = arityType dicts_cheap e
413 arityType _ _ = vanillaArityType
417 %************************************************************************
419 The main eta-expander
421 %************************************************************************
423 IMPORTANT NOTE: The eta expander is careful not to introduce "crap".
424 In particular, given a CoreExpr satisfying the 'CpeRhs' invariant (in
425 CorePrep), it returns a CoreExpr satisfying the same invariant. See
426 Note [Eta expansion and the CorePrep invariants] in CorePrep.
428 This means the eta-expander has to do a bit of on-the-fly
429 simplification but it's not too hard. The alernative, of relying on
430 a subsequent clean-up phase of the Simplifier to de-crapify the result,
431 means you can't really use it in CorePrep, which is painful.
434 -- | @etaExpand n us e ty@ returns an expression with
435 -- the same meaning as @e@, but with arity @n@.
439 -- > e' = etaExpand n us e ty
441 -- We should have that:
443 -- > ty = exprType e = exprType e'
444 etaExpand :: Arity -- ^ Result should have this number of value args
445 -> CoreExpr -- ^ Expression to expand
447 -- Note that SCCs are not treated specially. If we have
448 -- etaExpand 2 (\x -> scc "foo" e)
449 -- = (\xy -> (scc "foo" e) y)
450 -- So the costs of evaluating 'e' (not 'e y') are attributed to "foo"
452 -- etaExpand deals with for-alls. For example:
454 -- where E :: forall a. a -> a
456 -- (/\b. \y::a -> E b y)
458 -- It deals with coerces too, though they are now rare
459 -- so perhaps the extra code isn't worth it
461 etaExpand n orig_expr
462 | manifestArity orig_expr >= n = orig_expr -- The no-op case
466 -- Strip off existing lambdas
467 -- Note [Eta expansion and SCCs]
469 go n (Lam v body) | isTyVar v = Lam v (go n body)
470 | otherwise = Lam v (go (n-1) body)
471 go n (Note InlineMe expr) = Note InlineMe (go n expr)
472 go n (Cast expr co) = Cast (go n expr) co
473 go n expr = -- pprTrace "ee" (vcat [ppr orig_expr, ppr expr, ppr etas]) $
474 etaInfoAbs etas (etaInfoApp subst' expr etas)
476 in_scope = mkInScopeSet (exprFreeVars expr)
477 (in_scope', etas) = mkEtaWW n in_scope (exprType expr)
478 subst' = mkEmptySubst in_scope'
482 data EtaInfo = EtaVar Var -- /\a. [], [] a
484 | EtaCo Coercion -- [] |> co, [] |> (sym co)
486 instance Outputable EtaInfo where
487 ppr (EtaVar v) = ptext (sLit "EtaVar") <+> ppr v
488 ppr (EtaCo co) = ptext (sLit "EtaCo") <+> ppr co
490 pushCoercion :: Coercion -> [EtaInfo] -> [EtaInfo]
491 pushCoercion co1 (EtaCo co2 : eis)
492 | isIdentityCoercion co = eis
493 | otherwise = EtaCo co : eis
495 co = co1 `mkTransCoercion` co2
497 pushCoercion co eis = EtaCo co : eis
500 etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr
501 etaInfoAbs [] expr = expr
502 etaInfoAbs (EtaVar v : eis) expr = Lam v (etaInfoAbs eis expr)
503 etaInfoAbs (EtaCo co : eis) expr = Cast (etaInfoAbs eis expr) (mkSymCoercion co)
506 etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
507 -- (etaInfoApp s e eis) returns something equivalent to
508 -- ((substExpr s e) `appliedto` eis)
510 etaInfoApp subst (Lam v1 e) (EtaVar v2 : eis)
511 = etaInfoApp subst' e eis
513 subst' | isTyVar v1 = CoreSubst.extendTvSubst subst v1 (mkTyVarTy v2)
514 | otherwise = CoreSubst.extendIdSubst subst v1 (Var v2)
516 etaInfoApp subst (Cast e co1) eis
517 = etaInfoApp subst e (pushCoercion co' eis)
519 co' = CoreSubst.substTy subst co1
521 etaInfoApp subst (Case e b _ alts) eis
522 = Case (subst_expr subst e) b1 (coreAltsType alts') alts'
524 (subst1, b1) = substBndr subst b
525 alts' = map subst_alt alts
526 subst_alt (con, bs, rhs) = (con, bs', etaInfoApp subst2 rhs eis)
528 (subst2,bs') = substBndrs subst1 bs
530 etaInfoApp subst (Let b e) eis
531 = Let b' (etaInfoApp subst' e eis)
533 (subst', b') = subst_bind subst b
535 etaInfoApp subst (Note note e) eis
536 = Note note (etaInfoApp subst e eis)
538 etaInfoApp subst e eis
539 = go (subst_expr subst e) eis
542 go e (EtaVar v : eis) = go (App e (varToCoreExpr v)) eis
543 go e (EtaCo co : eis) = go (Cast e co) eis
546 mkEtaWW :: Arity -> InScopeSet -> Type
547 -> (InScopeSet, [EtaInfo])
548 -- EtaInfo contains fresh variables,
549 -- not free in the incoming CoreExpr
550 -- Outgoing InScopeSet includes the EtaInfo vars
551 -- and the original free vars
553 mkEtaWW n in_scope ty
554 = go n empty_subst ty []
556 empty_subst = mkTvSubst in_scope emptyTvSubstEnv
560 = (getTvInScope subst, reverse eis)
562 | Just (tv,ty') <- splitForAllTy_maybe ty
563 , let (subst', tv') = substTyVarBndr subst tv
564 -- Avoid free vars of the original expression
565 = go n subst' ty' (EtaVar tv' : eis)
567 | Just (arg_ty, res_ty) <- splitFunTy_maybe ty
568 , let (subst', eta_id') = freshEtaId n subst arg_ty
569 -- Avoid free vars of the original expression
570 = go (n-1) subst' res_ty (EtaVar eta_id' : eis)
572 | Just(ty',co) <- splitNewTypeRepCo_maybe ty
574 -- newtype T = MkT ([T] -> Int)
575 -- Consider eta-expanding this
578 -- coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
579 go n subst ty' (EtaCo (substTy subst co) : eis)
581 | otherwise -- We have an expression of arity > 0,
582 = (getTvInScope subst, reverse eis) -- but its type isn't a function.
583 -- This *can* legitmately happen:
584 -- e.g. coerce Int (\x. x) Essentially the programmer is
585 -- playing fast and loose with types (Happy does this a lot).
586 -- So we simply decline to eta-expand. Otherwise we'd end up
587 -- with an explicit lambda having a non-function type
591 -- Avoiding unnecessary substitution
593 subst_expr :: Subst -> CoreExpr -> CoreExpr
594 subst_expr s e | isEmptySubst s = e
595 | otherwise = substExpr s e
597 subst_bind :: Subst -> CoreBind -> (Subst, CoreBind)
598 subst_bind subst (NonRec b r)
599 = (subst', NonRec b' (subst_expr subst r))
601 (subst', b') = substBndr subst b
602 subst_bind subst (Rec prs)
603 = (subst', Rec (bs1 `zip` map (subst_expr subst') rhss))
605 (bs, rhss) = unzip prs
606 (subst', bs1) = substBndrs subst bs
610 freshEtaId :: Int -> TvSubst -> Type -> (TvSubst, Id)
611 -- Make a fresh Id, with specified type (after applying substitution)
612 -- It should be "fresh" in the sense that it's not in the in-scope set
613 -- of the TvSubstEnv; and it should itself then be added to the in-scope
614 -- set of the TvSubstEnv
616 -- The Int is just a reasonable starting point for generating a unique;
617 -- it does not necessarily have to be unique itself.
618 freshEtaId n subst ty
621 ty' = substTy subst ty
622 eta_id' = uniqAway (getTvInScope subst) $
623 mkSysLocal (fsLit "eta") (mkBuiltinUnique n) ty'
624 subst' = extendTvInScope subst [eta_id']