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, exprBotStrictness_maybe,
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) = go e `min` typeArity (snd (coercionKind co))
107 -- Note [exprArity invariant]
108 go (App e (Type _)) = go e
109 go (App f a) | exprIsTrivial a = (go f - 1) `max` 0
110 -- See Note [exprArity for applications]
114 Note [exprArity for applications]
115 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
116 When we come to an application we check that the arg is trivial.
117 eg f (fac x) does not have arity 2,
118 even if f has arity 3!
120 * We require that is trivial rather merely cheap. Suppose f has arity 2.
122 has arity 0, because if we gave it arity 1 and then inlined f we'd get
123 let v = Just y in \w. <f-body>
124 which has arity 0. And we try to maintain the invariant that we don't
125 have arity decreases.
127 * The `max 0` is important! (\x y -> f x) has arity 2, even if f is
128 unknown, hence arity 0
131 %************************************************************************
135 %************************************************************************
138 -- ^ The Arity returned is the number of value args the
139 -- expression can be applied to without doing much work
140 exprEtaExpandArity :: DynFlags -> CoreExpr -> Arity
141 -- exprEtaExpandArity is used when eta expanding
142 -- e ==> \xy -> e x y
143 exprEtaExpandArity dflags e
144 = applyStateHack e (arityType dicts_cheap e)
146 dicts_cheap = dopt Opt_DictsCheap dflags
148 exprBotStrictness_maybe :: CoreExpr -> Maybe (Arity, StrictSig)
149 -- A cheap and cheerful function that identifies bottoming functions
150 -- and gives them a suitable strictness signatures. It's used during
152 exprBotStrictness_maybe e
153 = case arityType False e of
155 AT a ABot -> Just (a, mkStrictSig (mkTopDmdType (replicate a topDmd) BotRes))
158 Note [Definition of arity]
159 ~~~~~~~~~~~~~~~~~~~~~~~~~~
160 The "arity" of an expression 'e' is n if
161 applying 'e' to *fewer* than n *value* arguments
164 Or, to put it another way
166 there is no work lost in duplicating the partial
167 application (e x1 .. x(n-1))
169 In the divegent case, no work is lost by duplicating because if the thing
170 is evaluated once, that's the end of the program.
172 Or, to put it another way, in any context C
174 C[ (\x1 .. xn. e x1 .. xn) ]
178 It's all a bit more subtle than it looks:
180 Note [Arity of case expressions]
181 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
182 We treat the arity of
183 case x of p -> \s -> ...
184 as 1 (or more) because for I/O ish things we really want to get that
185 \s to the top. We are prepared to evaluate x each time round the loop
186 in order to get that.
188 This isn't really right in the presence of seq. Consider
192 Can we eta-expand here? At first the answer looks like "yes of course", but
195 This should diverge! But if we eta-expand, it won't. Again, we ignore this
196 "problem", because being scrupulous would lose an important transformation for
199 1. Note [One-shot lambdas]
200 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
201 Consider one-shot lambdas
202 let x = expensive in \y z -> E
203 We want this to have arity 1 if the \y-abstraction is a 1-shot lambda.
205 3. Note [Dealing with bottom]
206 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
208 f = \x -> error "foo"
209 Here, arity 1 is fine. But if it is
213 then we want to get arity 2. Technically, this isn't quite right, because
215 should diverge, but it'll converge if we eta-expand f. Nevertheless, we
216 do so; it improves some programs significantly, and increasing convergence
217 isn't a bad thing. Hence the ABot/ATop in ArityType.
219 4. Note [Newtype arity]
220 ~~~~~~~~~~~~~~~~~~~~~~~~
221 Non-recursive newtypes are transparent, and should not get in the way.
222 We do (currently) eta-expand recursive newtypes too. So if we have, say
224 newtype T = MkT ([T] -> Int)
228 where f has arity 1. Then: etaExpandArity e = 1;
229 that is, etaExpandArity looks through the coerce.
231 When we eta-expand e to arity 1: eta_expand 1 e T
232 we want to get: coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
234 HOWEVER, note that if you use coerce bogusly you can ge
236 And since negate has arity 2, you might try to eta expand. But you can't
237 decopose Int to a function type. Hence the final case in eta_expand.
240 applyStateHack :: CoreExpr -> ArityType -> Arity
241 applyStateHack e (AT orig_arity is_bot)
242 | opt_NoStateHack = orig_arity
243 | ABot <- is_bot = orig_arity -- Note [State hack and bottoming functions]
244 | otherwise = go orig_ty orig_arity
245 where -- Note [The state-transformer hack]
247 go :: Type -> Arity -> Arity
248 go ty arity -- This case analysis should match that in eta_expand
249 | Just (_, ty') <- splitForAllTy_maybe ty = go ty' arity
250 | Just (arg,res) <- splitFunTy_maybe ty
251 , arity > 0 || isStateHackType arg = 1 + go res (arity-1)
253 -- See Note [trimCast]
254 | Just (tc,tys) <- splitTyConApp_maybe ty
255 , Just (ty', _) <- instNewTyCon_maybe tc tys
256 , not (isRecursiveTyCon tc) = go ty' arity
257 -- Important to look through non-recursive newtypes, so that, eg
258 -- (f x) where f has arity 2, f :: Int -> IO ()
259 -- Here we want to get arity 1 for the result!
263 = if arity > 0 then 1 + go res (arity-1)
264 else if isStateHackType arg then
265 pprTrace "applystatehack" (vcat [ppr orig_arity, ppr orig_ty,
266 ppr ty, ppr res, ppr e]) $
268 else WARN( arity > 0, ppr arity ) 0
270 | otherwise = WARN( arity > 0, ppr arity <+> ppr ty) 0
273 Note [The state-transformer hack]
274 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
277 where e has arity n. Then, if we know from the context that f has
279 t1 -> ... -> tn -1-> t(n+1) -1-> ... -1-> tm -> ...
280 then we can expand the arity to m. This usage type says that
281 any application (x e1 .. en) will be applied to uniquely to (m-n) more args
282 Consider f = \x. let y = <expensive>
285 False -> \(s:RealWorld) -> e
286 where foo has arity 1. Then we want the state hack to
287 apply to foo too, so we can eta expand the case.
289 Then we expect that if f is applied to one arg, it'll be applied to two
290 (that's the hack -- we don't really know, and sometimes it's false)
291 See also Id.isOneShotBndr.
293 Note [State hack and bottoming functions]
294 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
295 It's a terrible idea to use the state hack on a bottoming function.
296 Here's what happens (Trac #2861):
301 Eta-expand, using the state hack:
303 f = \p. (\s. ((error "...") |> g1) s) |> g2
304 g1 :: IO T ~ (S -> (S,T))
305 g2 :: (S -> (S,T)) ~ IO T
309 f' = \p. \s. ((error "...") |> g1) s
310 f = f' |> (String -> g2)
312 Discard args for bottomming function
314 f' = \p. \s. ((error "...") |> g1 |> g3
315 g3 :: (S -> (S,T)) ~ (S,T)
319 f'' = \p. \s. (error "...")
320 f' = f'' |> (String -> S -> g1.g3)
322 And now we can repeat the whole loop. Aargh! The bug is in applying the
323 state hack to a function which then swallows the argument.
326 -------------------- Main arity code ----------------------------
328 -- If e has ArityType (AT n r), then the term 'e'
329 -- * Must be applied to at least n *value* args
330 -- before doing any significant work
331 -- * It will not diverge before being applied to n
333 -- * If 'r' is ABot, then it guarantees to diverge if
334 -- applied to n arguments (or more)
336 data ArityType = AT Arity ArityRes
337 data ArityRes = ATop -- Know nothing
340 vanillaArityType :: ArityType
341 vanillaArityType = AT 0 ATop -- Totally uninformative
343 incArity :: ArityType -> ArityType
344 incArity (AT a r) = AT (a+1) r
346 decArity :: ArityType -> ArityType
347 decArity (AT 0 r) = AT 0 r
348 decArity (AT a r) = AT (a-1) r
350 andArityType :: ArityType -> ArityType -> ArityType -- Used for branches of a 'case'
351 andArityType (AT a1 ATop) (AT a2 ATop) = AT (a1 `min` a2) ATop
352 andArityType (AT _ ABot) (AT a2 ATop) = AT a2 ATop
353 andArityType (AT a1 ATop) (AT _ ABot) = AT a1 ATop
354 andArityType (AT a1 ABot) (AT a2 ABot) = AT (a1 `max` a2) ABot
356 ---------------------------
357 trimCast :: Coercion -> ArityType -> ArityType
358 -- Trim the arity to be no more than allowed by the
359 -- arrows in ty2, where co :: ty1~ty2
362 {- Omitting for now Note [trimCast]
363 trimCast co at@(AT ar _)
364 | ar > co_arity = AT co_arity ATop
367 (_, ty2) = coercionKind co
368 co_arity = typeArity ty2
374 When you try putting trimCast back in, comment out the snippets
375 flagged by the other references to Note [trimCast]
378 ---------------------------
379 trimArity :: Bool -> ArityType -> ArityType
380 -- We have something like (let x = E in b), where b has the given
382 -- * If E is cheap we can push it inside as far as we like
383 -- * If b eventually diverges, we allow ourselves to push inside
384 -- arbitrarily, even though that is not quite right
385 trimArity _cheap (AT a ABot) = AT a ABot
386 trimArity True (AT a ATop) = AT a ATop
387 trimArity False (AT _ ATop) = AT 0 ATop -- Bale out
389 ---------------------------
390 arityType :: Bool -> CoreExpr -> ArityType
392 | Just strict_sig <- idStrictness_maybe v
393 , (ds, res) <- splitStrictSig strict_sig
395 = AT (length ds) ABot -- Function diverges
397 = AT (idArity v) ATop
399 -- Lambdas; increase arity
400 arityType dicts_cheap (Lam x e)
401 | isId x = incArity (arityType dicts_cheap e)
402 | otherwise = arityType dicts_cheap e
404 -- Applications; decrease arity
405 arityType dicts_cheap (App fun (Type _))
406 = arityType dicts_cheap fun
407 arityType dicts_cheap (App fun arg )
408 = trimArity (exprIsCheap arg) (decArity (arityType dicts_cheap fun))
410 -- Case/Let; keep arity if either the expression is cheap
411 -- or it's a 1-shot lambda
412 -- The former is not really right for Haskell
413 -- f x = case x of { (a,b) -> \y. e }
415 -- f x y = case x of { (a,b) -> e }
416 -- The difference is observable using 'seq'
417 arityType dicts_cheap (Case scrut _ _ alts)
418 = trimArity (exprIsCheap scrut)
419 (foldr1 andArityType [arityType dicts_cheap rhs | (_,_,rhs) <- alts])
421 arityType dicts_cheap (Let b e)
422 = trimArity (cheap_bind b) (arityType dicts_cheap e)
424 cheap_bind (NonRec b e) = is_cheap (b,e)
425 cheap_bind (Rec prs) = all is_cheap prs
426 is_cheap (b,e) = (dicts_cheap && isDictLikeTy (idType b))
428 -- If the experimental -fdicts-cheap flag is on, we eta-expand through
429 -- dictionary bindings. This improves arities. Thereby, it also
430 -- means that full laziness is less prone to floating out the
431 -- application of a function to its dictionary arguments, which
432 -- can thereby lose opportunities for fusion. Example:
433 -- foo :: Ord a => a -> ...
434 -- foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....
435 -- -- So foo has arity 1
437 -- f = \x. foo dInt $ bar x
439 -- The (foo DInt) is floated out, and makes ineffective a RULE
442 -- One could go further and make exprIsCheap reply True to any
443 -- dictionary-typed expression, but that's more work.
445 -- See Note [Dictionary-like types] in TcType.lhs for why we use
446 -- isDictLikeTy here rather than isDictTy
448 arityType dicts_cheap (Note _ e) = arityType dicts_cheap e
449 arityType dicts_cheap (Cast e co) = trimCast co (arityType dicts_cheap e)
450 arityType _ _ = vanillaArityType
454 %************************************************************************
456 The main eta-expander
458 %************************************************************************
460 IMPORTANT NOTE: The eta expander is careful not to introduce "crap".
461 In particular, given a CoreExpr satisfying the 'CpeRhs' invariant (in
462 CorePrep), it returns a CoreExpr satisfying the same invariant. See
463 Note [Eta expansion and the CorePrep invariants] in CorePrep.
465 This means the eta-expander has to do a bit of on-the-fly
466 simplification but it's not too hard. The alernative, of relying on
467 a subsequent clean-up phase of the Simplifier to de-crapify the result,
468 means you can't really use it in CorePrep, which is painful.
470 Note [Eta expansion and SCCs]
471 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
472 Note that SCCs are not treated specially by etaExpand. If we have
473 etaExpand 2 (\x -> scc "foo" e)
474 = (\xy -> (scc "foo" e) y)
475 So the costs of evaluating 'e' (not 'e y') are attributed to "foo"
478 -- | @etaExpand n us e ty@ returns an expression with
479 -- the same meaning as @e@, but with arity @n@.
483 -- > e' = etaExpand n us e ty
485 -- We should have that:
487 -- > ty = exprType e = exprType e'
488 etaExpand :: Arity -- ^ Result should have this number of value args
489 -> CoreExpr -- ^ Expression to expand
491 -- etaExpand deals with for-alls. For example:
493 -- where E :: forall a. a -> a
495 -- (/\b. \y::a -> E b y)
497 -- It deals with coerces too, though they are now rare
498 -- so perhaps the extra code isn't worth it
500 etaExpand n orig_expr
503 -- Strip off existing lambdas and casts
504 -- Note [Eta expansion and SCCs]
506 go n (Lam v body) | isTyVar v = Lam v (go n body)
507 | otherwise = Lam v (go (n-1) body)
508 go n (Cast expr co) = Cast (go n expr) co
509 go n expr = -- pprTrace "ee" (vcat [ppr orig_expr, ppr expr, ppr etas]) $
510 etaInfoAbs etas (etaInfoApp subst' expr etas)
512 in_scope = mkInScopeSet (exprFreeVars expr)
513 (in_scope', etas) = mkEtaWW n in_scope (exprType expr)
514 subst' = mkEmptySubst in_scope'
518 data EtaInfo = EtaVar Var -- /\a. [], [] a
520 | EtaCo Coercion -- [] |> co, [] |> (sym co)
522 instance Outputable EtaInfo where
523 ppr (EtaVar v) = ptext (sLit "EtaVar") <+> ppr v
524 ppr (EtaCo co) = ptext (sLit "EtaCo") <+> ppr co
526 pushCoercion :: Coercion -> [EtaInfo] -> [EtaInfo]
527 pushCoercion co1 (EtaCo co2 : eis)
528 | isIdentityCoercion co = eis
529 | otherwise = EtaCo co : eis
531 co = co1 `mkTransCoercion` co2
533 pushCoercion co eis = EtaCo co : eis
536 etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr
537 etaInfoAbs [] expr = expr
538 etaInfoAbs (EtaVar v : eis) expr = Lam v (etaInfoAbs eis expr)
539 etaInfoAbs (EtaCo co : eis) expr = Cast (etaInfoAbs eis expr) (mkSymCoercion co)
542 etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
543 -- (etaInfoApp s e eis) returns something equivalent to
544 -- ((substExpr s e) `appliedto` eis)
546 etaInfoApp subst (Lam v1 e) (EtaVar v2 : eis)
547 = etaInfoApp subst' e eis
549 subst' | isTyVar v1 = CoreSubst.extendTvSubst subst v1 (mkTyVarTy v2)
550 | otherwise = CoreSubst.extendIdSubst subst v1 (Var v2)
552 etaInfoApp subst (Cast e co1) eis
553 = etaInfoApp subst e (pushCoercion co' eis)
555 co' = CoreSubst.substTy subst co1
557 etaInfoApp subst (Case e b _ alts) eis
558 = Case (subst_expr subst e) b1 (coreAltsType alts') alts'
560 (subst1, b1) = substBndr subst b
561 alts' = map subst_alt alts
562 subst_alt (con, bs, rhs) = (con, bs', etaInfoApp subst2 rhs eis)
564 (subst2,bs') = substBndrs subst1 bs
566 etaInfoApp subst (Let b e) eis
567 = Let b' (etaInfoApp subst' e eis)
569 (subst', b') = subst_bind subst b
571 etaInfoApp subst (Note note e) eis
572 = Note note (etaInfoApp subst e eis)
574 etaInfoApp subst e eis
575 = go (subst_expr subst e) eis
578 go e (EtaVar v : eis) = go (App e (varToCoreExpr v)) eis
579 go e (EtaCo co : eis) = go (Cast e co) eis
582 mkEtaWW :: Arity -> InScopeSet -> Type
583 -> (InScopeSet, [EtaInfo])
584 -- EtaInfo contains fresh variables,
585 -- not free in the incoming CoreExpr
586 -- Outgoing InScopeSet includes the EtaInfo vars
587 -- and the original free vars
589 mkEtaWW orig_n in_scope orig_ty
590 = go orig_n empty_subst orig_ty []
592 empty_subst = mkTvSubst in_scope emptyTvSubstEnv
596 = (getTvInScope subst, reverse eis)
598 | Just (tv,ty') <- splitForAllTy_maybe ty
599 , let (subst', tv') = substTyVarBndr subst tv
600 -- Avoid free vars of the original expression
601 = go n subst' ty' (EtaVar tv' : eis)
603 | Just (arg_ty, res_ty) <- splitFunTy_maybe ty
604 , let (subst', eta_id') = freshEtaId n subst arg_ty
605 -- Avoid free vars of the original expression
606 = go (n-1) subst' res_ty (EtaVar eta_id' : eis)
608 -- See Note [trimCast]
609 | Just(ty',co) <- splitNewTypeRepCo_maybe ty
611 -- newtype T = MkT ([T] -> Int)
612 -- Consider eta-expanding this
615 -- coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
616 go n subst ty' (EtaCo (substTy subst co) : eis)
618 | otherwise -- We have an expression of arity > 0,
619 = (getTvInScope subst, reverse eis) -- but its type isn't a function.
620 -- This *can* legitmately happen:
621 -- e.g. coerce Int (\x. x) Essentially the programmer is
622 -- playing fast and loose with types (Happy does this a lot).
623 -- So we simply decline to eta-expand. Otherwise we'd end up
624 -- with an explicit lambda having a non-function type
628 -- Avoiding unnecessary substitution
630 subst_expr :: Subst -> CoreExpr -> CoreExpr
631 subst_expr s e | isEmptySubst s = e
632 | otherwise = substExpr s e
634 subst_bind :: Subst -> CoreBind -> (Subst, CoreBind)
635 subst_bind subst (NonRec b r)
636 = (subst', NonRec b' (subst_expr subst r))
638 (subst', b') = substBndr subst b
639 subst_bind subst (Rec prs)
640 = (subst', Rec (bs1 `zip` map (subst_expr subst') rhss))
642 (bs, rhss) = unzip prs
643 (subst', bs1) = substBndrs subst bs
647 freshEtaId :: Int -> TvSubst -> Type -> (TvSubst, Id)
648 -- Make a fresh Id, with specified type (after applying substitution)
649 -- It should be "fresh" in the sense that it's not in the in-scope set
650 -- of the TvSubstEnv; and it should itself then be added to the in-scope
651 -- set of the TvSubstEnv
653 -- The Int is just a reasonable starting point for generating a unique;
654 -- it does not necessarily have to be unique itself.
655 freshEtaId n subst ty
658 ty' = substTy subst ty
659 eta_id' = uniqAway (getTvInScope subst) $
660 mkSysLocal (fsLit "eta") (mkBuiltinUnique n) ty'
661 subst' = extendTvInScope subst eta_id'