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) = 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
142 exprBotStrictness_maybe :: CoreExpr -> Maybe (Arity, StrictSig)
143 -- A cheap and cheerful function that identifies bottoming functions
144 -- and gives them a suitable strictness signatures. It's used during
146 exprBotStrictness_maybe e
147 = case arityType False e of
149 AT a ABot -> Just (a, mkStrictSig (mkTopDmdType (replicate a topDmd) BotRes))
152 Note [Definition of arity]
153 ~~~~~~~~~~~~~~~~~~~~~~~~~~
154 The "arity" of an expression 'e' is n if
155 applying 'e' to *fewer* than n *value* arguments
158 Or, to put it another way
160 there is no work lost in duplicating the partial
161 application (e x1 .. x(n-1))
163 In the divegent case, no work is lost by duplicating because if the thing
164 is evaluated once, that's the end of the program.
166 Or, to put it another way, in any context C
168 C[ (\x1 .. xn. e x1 .. xn) ]
173 It's all a bit more subtle than it looks:
175 Note [Arity of case expressions]
176 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
177 We treat the arity of
178 case x of p -> \s -> ...
179 as 1 (or more) because for I/O ish things we really want to get that
180 \s to the top. We are prepared to evaluate x each time round the loop
181 in order to get that.
183 This isn't really right in the presence of seq. Consider
187 Can we eta-expand here? At first the answer looks like "yes of course", but
190 This should diverge! But if we eta-expand, it won't. Again, we ignore this
191 "problem", because being scrupulous would lose an important transformation for
195 1. Note [One-shot lambdas]
196 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
197 Consider one-shot lambdas
198 let x = expensive in \y z -> E
199 We want this to have arity 1 if the \y-abstraction is a 1-shot lambda.
201 3. Note [Dealing with bottom]
202 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
204 f = \x -> error "foo"
205 Here, arity 1 is fine. But if it is
209 then we want to get arity 2. Technically, this isn't quite right, because
211 should diverge, but it'll converge if we eta-expand f. Nevertheless, we
212 do so; it improves some programs significantly, and increasing convergence
213 isn't a bad thing. Hence the ABot/ATop in ArityType.
216 4. Note [Newtype arity]
217 ~~~~~~~~~~~~~~~~~~~~~~~~
218 Non-recursive newtypes are transparent, and should not get in the way.
219 We do (currently) eta-expand recursive newtypes too. So if we have, say
221 newtype T = MkT ([T] -> Int)
225 where f has arity 1. Then: etaExpandArity e = 1;
226 that is, etaExpandArity looks through the coerce.
228 When we eta-expand e to arity 1: eta_expand 1 e T
229 we want to get: coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
231 HOWEVER, note that if you use coerce bogusly you can ge
233 And since negate has arity 2, you might try to eta expand. But you can't
234 decopose Int to a function type. Hence the final case in eta_expand.
236 Note [The state-transformer hack]
237 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
240 where e has arity n. Then, if we know from the context that f has
242 t1 -> ... -> tn -1-> t(n+1) -1-> ... -1-> tm -> ...
243 then we can expand the arity to m. This usage type says that
244 any application (x e1 .. en) will be applied to uniquely to (m-n) more args
245 Consider f = \x. let y = <expensive>
248 False -> \(s:RealWorld) -> e
249 where foo has arity 1. Then we want the state hack to
250 apply to foo too, so we can eta expand the case.
252 Then we expect that if f is applied to one arg, it'll be applied to two
253 (that's the hack -- we don't really know, and sometimes it's false)
254 See also Id.isOneShotBndr.
257 applyStateHack :: CoreExpr -> ArityType -> Arity
258 applyStateHack e (AT orig_arity is_bot)
259 | opt_NoStateHack = orig_arity
260 | ABot <- is_bot = orig_arity -- Note [State hack and bottoming functions]
261 | otherwise = go orig_ty orig_arity
262 where -- Note [The state-transformer hack]
264 go :: Type -> Arity -> Arity
265 go ty arity -- This case analysis should match that in eta_expand
266 | Just (_, ty') <- splitForAllTy_maybe ty = go ty' arity
268 | Just (tc,tys) <- splitTyConApp_maybe ty
269 , Just (ty', _) <- instNewTyCon_maybe tc tys
270 , not (isRecursiveTyCon tc) = go ty' arity
271 -- Important to look through non-recursive newtypes, so that, eg
272 -- (f x) where f has arity 2, f :: Int -> IO ()
273 -- Here we want to get arity 1 for the result!
275 | Just (arg,res) <- splitFunTy_maybe ty
276 , arity > 0 || isStateHackType arg = 1 + go res (arity-1)
278 = if arity > 0 then 1 + go res (arity-1)
279 else if isStateHackType arg then
280 pprTrace "applystatehack" (vcat [ppr orig_arity, ppr orig_ty,
281 ppr ty, ppr res, ppr e]) $
283 else WARN( arity > 0, ppr arity ) 0
285 | otherwise = WARN( arity > 0, ppr arity <+> ppr ty) 0
288 Note [State hack and bottoming functions]
289 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
290 It's a terrible idea to use the state hack on a bottoming function.
291 Here's what happens (Trac #2861):
296 Eta-expand, using the state hack:
298 f = \p. (\s. ((error "...") |> g1) s) |> g2
299 g1 :: IO T ~ (S -> (S,T))
300 g2 :: (S -> (S,T)) ~ IO T
304 f' = \p. \s. ((error "...") |> g1) s
305 f = f' |> (String -> g2)
307 Discard args for bottomming function
309 f' = \p. \s. ((error "...") |> g1 |> g3
310 g3 :: (S -> (S,T)) ~ (S,T)
314 f'' = \p. \s. (error "...")
315 f' = f'' |> (String -> S -> g1.g3)
317 And now we can repeat the whole loop. Aargh! The bug is in applying the
318 state hack to a function which then swallows the argument.
321 -------------------- Main arity code ----------------------------
323 -- If e has ArityType (AT n r), then the term 'e'
324 -- * Must be applied to at least n *value* args
325 -- before doing any significant work
326 -- * It will not diverge before being applied to n
328 -- * If 'r' is ABot, then it guarantees to diverge if
329 -- applied to n arguments (or more)
331 data ArityType = AT Arity ArityRes
332 data ArityRes = ATop -- Know nothing
335 vanillaArityType :: ArityType
336 vanillaArityType = AT 0 ATop -- Totally uninformative
338 incArity :: ArityType -> ArityType
339 incArity (AT a r) = AT (a+1) r
341 decArity :: ArityType -> ArityType
342 decArity (AT 0 r) = AT 0 r
343 decArity (AT a r) = AT (a-1) r
345 andArityType :: ArityType -> ArityType -> ArityType -- Used for branches of a 'case'
346 andArityType (AT a1 ATop) (AT a2 ATop) = AT (a1 `min` a2) ATop
347 andArityType (AT _ ABot) (AT a2 ATop) = AT a2 ATop
348 andArityType (AT a1 ATop) (AT _ ABot) = AT a1 ATop
349 andArityType (AT a1 ABot) (AT a2 ABot) = AT (a1 `max` a2) ABot
351 trimArity :: Bool -> ArityType -> ArityType
352 -- We have something like (let x = E in b), where b has the given
354 -- * If E is cheap we can push it inside as far as we like
355 -- * If b eventually diverges, we allow ourselves to push inside
356 -- arbitrarily, even though that is not quite right
357 trimArity _cheap (AT a ABot) = AT a ABot
358 trimArity True (AT a ATop) = AT a ATop
359 trimArity False (AT _ ATop) = AT 0 ATop -- Bale out
361 ---------------------------
362 arityType :: Bool -> CoreExpr -> ArityType
364 | Just strict_sig <- idStrictness_maybe v
365 , (ds, res) <- splitStrictSig strict_sig
367 = AT (length ds) ABot -- Function diverges
369 = AT (idArity v) ATop
371 -- Lambdas; increase arity
372 arityType dicts_cheap (Lam x e)
373 | isId x = incArity (arityType dicts_cheap e)
374 | otherwise = arityType dicts_cheap e
376 -- Applications; decrease arity
377 arityType dicts_cheap (App fun (Type _))
378 = arityType dicts_cheap fun
379 arityType dicts_cheap (App fun arg )
380 = trimArity (exprIsCheap arg) (decArity (arityType dicts_cheap fun))
382 -- Case/Let; keep arity if either the expression is cheap
383 -- or it's a 1-shot lambda
384 -- The former is not really right for Haskell
385 -- f x = case x of { (a,b) -> \y. e }
387 -- f x y = case x of { (a,b) -> e }
388 -- The difference is observable using 'seq'
389 arityType dicts_cheap (Case scrut _ _ alts)
390 = trimArity (exprIsCheap scrut)
391 (foldr1 andArityType [arityType dicts_cheap rhs | (_,_,rhs) <- alts])
393 arityType dicts_cheap (Let b e)
394 = trimArity (cheap_bind b) (arityType dicts_cheap e)
396 cheap_bind (NonRec b e) = is_cheap (b,e)
397 cheap_bind (Rec prs) = all is_cheap prs
398 is_cheap (b,e) = (dicts_cheap && isDictLikeTy (idType b))
400 -- If the experimental -fdicts-cheap flag is on, we eta-expand through
401 -- dictionary bindings. This improves arities. Thereby, it also
402 -- means that full laziness is less prone to floating out the
403 -- application of a function to its dictionary arguments, which
404 -- can thereby lose opportunities for fusion. Example:
405 -- foo :: Ord a => a -> ...
406 -- foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....
407 -- -- So foo has arity 1
409 -- f = \x. foo dInt $ bar x
411 -- The (foo DInt) is floated out, and makes ineffective a RULE
414 -- One could go further and make exprIsCheap reply True to any
415 -- dictionary-typed expression, but that's more work.
417 -- See Note [Dictionary-like types] in TcType.lhs for why we use
418 -- isDictLikeTy here rather than isDictTy
420 arityType dicts_cheap (Note _ e) = arityType dicts_cheap e
421 arityType dicts_cheap (Cast e _) = arityType dicts_cheap e
422 arityType _ _ = vanillaArityType
426 %************************************************************************
428 The main eta-expander
430 %************************************************************************
432 IMPORTANT NOTE: The eta expander is careful not to introduce "crap".
433 In particular, given a CoreExpr satisfying the 'CpeRhs' invariant (in
434 CorePrep), it returns a CoreExpr satisfying the same invariant. See
435 Note [Eta expansion and the CorePrep invariants] in CorePrep.
437 This means the eta-expander has to do a bit of on-the-fly
438 simplification but it's not too hard. The alernative, of relying on
439 a subsequent clean-up phase of the Simplifier to de-crapify the result,
440 means you can't really use it in CorePrep, which is painful.
442 Note [Eta expansion and SCCs]
443 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
444 Note that SCCs are not treated specially by etaExpand. If we have
445 etaExpand 2 (\x -> scc "foo" e)
446 = (\xy -> (scc "foo" e) y)
447 So the costs of evaluating 'e' (not 'e y') are attributed to "foo"
450 -- | @etaExpand n us e ty@ returns an expression with
451 -- the same meaning as @e@, but with arity @n@.
455 -- > e' = etaExpand n us e ty
457 -- We should have that:
459 -- > ty = exprType e = exprType e'
460 etaExpand :: Arity -- ^ Result should have this number of value args
461 -> CoreExpr -- ^ Expression to expand
463 -- etaExpand deals with for-alls. For example:
465 -- where E :: forall a. a -> a
467 -- (/\b. \y::a -> E b y)
469 -- It deals with coerces too, though they are now rare
470 -- so perhaps the extra code isn't worth it
472 etaExpand n orig_expr
473 | manifestArity orig_expr >= n = orig_expr -- The no-op case
477 -- Strip off existing lambdas
478 -- Note [Eta expansion and SCCs]
480 go n (Lam v body) | isTyVar v = Lam v (go n body)
481 | otherwise = Lam v (go (n-1) body)
482 go n (Cast expr co) = Cast (go n expr) co
483 go n expr = -- pprTrace "ee" (vcat [ppr orig_expr, ppr expr, ppr etas]) $
484 etaInfoAbs etas (etaInfoApp subst' expr etas)
486 in_scope = mkInScopeSet (exprFreeVars expr)
487 (in_scope', etas) = mkEtaWW n in_scope (exprType expr)
488 subst' = mkEmptySubst in_scope'
492 data EtaInfo = EtaVar Var -- /\a. [], [] a
494 | EtaCo Coercion -- [] |> co, [] |> (sym co)
496 instance Outputable EtaInfo where
497 ppr (EtaVar v) = ptext (sLit "EtaVar") <+> ppr v
498 ppr (EtaCo co) = ptext (sLit "EtaCo") <+> ppr co
500 pushCoercion :: Coercion -> [EtaInfo] -> [EtaInfo]
501 pushCoercion co1 (EtaCo co2 : eis)
502 | isIdentityCoercion co = eis
503 | otherwise = EtaCo co : eis
505 co = co1 `mkTransCoercion` co2
507 pushCoercion co eis = EtaCo co : eis
510 etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr
511 etaInfoAbs [] expr = expr
512 etaInfoAbs (EtaVar v : eis) expr = Lam v (etaInfoAbs eis expr)
513 etaInfoAbs (EtaCo co : eis) expr = Cast (etaInfoAbs eis expr) (mkSymCoercion co)
516 etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
517 -- (etaInfoApp s e eis) returns something equivalent to
518 -- ((substExpr s e) `appliedto` eis)
520 etaInfoApp subst (Lam v1 e) (EtaVar v2 : eis)
521 = etaInfoApp subst' e eis
523 subst' | isTyVar v1 = CoreSubst.extendTvSubst subst v1 (mkTyVarTy v2)
524 | otherwise = CoreSubst.extendIdSubst subst v1 (Var v2)
526 etaInfoApp subst (Cast e co1) eis
527 = etaInfoApp subst e (pushCoercion co' eis)
529 co' = CoreSubst.substTy subst co1
531 etaInfoApp subst (Case e b _ alts) eis
532 = Case (subst_expr subst e) b1 (coreAltsType alts') alts'
534 (subst1, b1) = substBndr subst b
535 alts' = map subst_alt alts
536 subst_alt (con, bs, rhs) = (con, bs', etaInfoApp subst2 rhs eis)
538 (subst2,bs') = substBndrs subst1 bs
540 etaInfoApp subst (Let b e) eis
541 = Let b' (etaInfoApp subst' e eis)
543 (subst', b') = subst_bind subst b
545 etaInfoApp subst (Note note e) eis
546 = Note note (etaInfoApp subst e eis)
548 etaInfoApp subst e eis
549 = go (subst_expr subst e) eis
552 go e (EtaVar v : eis) = go (App e (varToCoreExpr v)) eis
553 go e (EtaCo co : eis) = go (Cast e co) eis
556 mkEtaWW :: Arity -> InScopeSet -> Type
557 -> (InScopeSet, [EtaInfo])
558 -- EtaInfo contains fresh variables,
559 -- not free in the incoming CoreExpr
560 -- Outgoing InScopeSet includes the EtaInfo vars
561 -- and the original free vars
563 mkEtaWW n in_scope ty
564 = go n empty_subst ty []
566 empty_subst = mkTvSubst in_scope emptyTvSubstEnv
570 = (getTvInScope subst, reverse eis)
572 | Just (tv,ty') <- splitForAllTy_maybe ty
573 , let (subst', tv') = substTyVarBndr subst tv
574 -- Avoid free vars of the original expression
575 = go n subst' ty' (EtaVar tv' : eis)
577 | Just (arg_ty, res_ty) <- splitFunTy_maybe ty
578 , let (subst', eta_id') = freshEtaId n subst arg_ty
579 -- Avoid free vars of the original expression
580 = go (n-1) subst' res_ty (EtaVar eta_id' : eis)
582 | Just(ty',co) <- splitNewTypeRepCo_maybe ty
584 -- newtype T = MkT ([T] -> Int)
585 -- Consider eta-expanding this
588 -- coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
589 go n subst ty' (EtaCo (substTy subst co) : eis)
591 | otherwise -- We have an expression of arity > 0,
592 = (getTvInScope subst, reverse eis) -- but its type isn't a function.
593 -- This *can* legitmately happen:
594 -- e.g. coerce Int (\x. x) Essentially the programmer is
595 -- playing fast and loose with types (Happy does this a lot).
596 -- So we simply decline to eta-expand. Otherwise we'd end up
597 -- with an explicit lambda having a non-function type
601 -- Avoiding unnecessary substitution
603 subst_expr :: Subst -> CoreExpr -> CoreExpr
604 subst_expr s e | isEmptySubst s = e
605 | otherwise = substExpr s e
607 subst_bind :: Subst -> CoreBind -> (Subst, CoreBind)
608 subst_bind subst (NonRec b r)
609 = (subst', NonRec b' (subst_expr subst r))
611 (subst', b') = substBndr subst b
612 subst_bind subst (Rec prs)
613 = (subst', Rec (bs1 `zip` map (subst_expr subst') rhss))
615 (bs, rhss) = unzip prs
616 (subst', bs1) = substBndrs subst bs
620 freshEtaId :: Int -> TvSubst -> Type -> (TvSubst, Id)
621 -- Make a fresh Id, with specified type (after applying substitution)
622 -- It should be "fresh" in the sense that it's not in the in-scope set
623 -- of the TvSubstEnv; and it should itself then be added to the in-scope
624 -- set of the TvSubstEnv
626 -- The Int is just a reasonable starting point for generating a unique;
627 -- it does not necessarily have to be unique itself.
628 freshEtaId n subst ty
631 ty' = substTy subst ty
632 eta_id' = uniqAway (getTvInScope subst) $
633 mkSysLocal (fsLit "eta") (mkBuiltinUnique n) ty'
634 subst' = extendTvInScope subst eta_id'