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"
26 import TyCon ( isRecursiveTyCon )
27 import TcType ( isDictLikeTy )
33 import StaticFlags ( opt_NoStateHack )
37 %************************************************************************
39 manifestArity and exprArity
41 %************************************************************************
43 exprArity is a cheap-and-cheerful version of exprEtaExpandArity.
44 It tells how many things the expression can be applied to before doing
45 any work. It doesn't look inside cases, lets, etc. The idea is that
46 exprEtaExpandArity will do the hard work, leaving something that's easy
47 for exprArity to grapple with. In particular, Simplify uses exprArity to
48 compute the ArityInfo for the Id.
50 Originally I thought that it was enough just to look for top-level lambdas, but
51 it isn't. I've seen this
53 foo = PrelBase.timesInt
55 We want foo to get arity 2 even though the eta-expander will leave it
56 unchanged, in the expectation that it'll be inlined. But occasionally it
57 isn't, because foo is blacklisted (used in a rule).
59 Similarly, see the ok_note check in exprEtaExpandArity. So
60 f = __inline_me (\x -> e)
61 won't be eta-expanded.
63 And in any case it seems more robust to have exprArity be a bit more intelligent.
64 But note that (\x y z -> f x y z)
65 should have arity 3, regardless of f's arity.
67 Note [exprArity invariant]
68 ~~~~~~~~~~~~~~~~~~~~~~~~~~
69 exprArity has the following invariant:
70 (exprArity e) = n, then manifestArity (etaExpand e n) = n
72 That is, if exprArity says "the arity is n" then etaExpand really can get
73 "n" manifest lambdas to the top.
75 Why is this important? Because
76 - In TidyPgm we use exprArity to fix the *final arity* of
77 each top-level Id, and in
78 - In CorePrep we use etaExpand on each rhs, so that the visible lambdas
79 actually match that arity, which in turn means
80 that the StgRhs has the right number of lambdas
82 An alternative would be to do the eta-expansion in TidyPgm, at least
83 for top-level bindings, in which case we would not need the trim_arity
84 in exprArity. That is a less local change, so I'm going to leave it for today!
88 manifestArity :: CoreExpr -> Arity
89 -- ^ manifestArity sees how many leading value lambdas there are
90 manifestArity (Lam v e) | isId v = 1 + manifestArity e
91 | otherwise = manifestArity e
92 manifestArity (Note _ e) = manifestArity e
93 manifestArity (Cast e _) = manifestArity e
96 exprArity :: CoreExpr -> Arity
97 -- ^ An approximate, fast, version of 'exprEtaExpandArity'
100 go (Var v) = idArity v
101 go (Lam x e) | isId x = go e + 1
104 go (Cast e co) = go e `min` typeArity (snd (coercionKind co))
105 -- Note [exprArity invariant]
106 go (App e (Type _)) = go e
107 go (App f a) | exprIsTrivial a = (go f - 1) `max` 0
108 -- See Note [exprArity for applications]
112 Note [exprArity for applications]
113 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
114 When we come to an application we check that the arg is trivial.
115 eg f (fac x) does not have arity 2,
116 even if f has arity 3!
118 * We require that is trivial rather merely cheap. Suppose f has arity 2.
120 has arity 0, because if we gave it arity 1 and then inlined f we'd get
121 let v = Just y in \w. <f-body>
122 which has arity 0. And we try to maintain the invariant that we don't
123 have arity decreases.
125 * The `max 0` is important! (\x y -> f x) has arity 2, even if f is
126 unknown, hence arity 0
129 %************************************************************************
133 %************************************************************************
136 -- ^ The Arity returned is the number of value args the
137 -- expression can be applied to without doing much work
138 exprEtaExpandArity :: DynFlags -> CoreExpr -> Arity
139 -- exprEtaExpandArity is used when eta expanding
140 -- e ==> \xy -> e x y
141 exprEtaExpandArity dflags e
142 = applyStateHack e (arityType dicts_cheap e)
144 dicts_cheap = dopt Opt_DictsCheap dflags
146 exprBotStrictness_maybe :: CoreExpr -> Maybe (Arity, StrictSig)
147 -- A cheap and cheerful function that identifies bottoming functions
148 -- and gives them a suitable strictness signatures. It's used during
150 exprBotStrictness_maybe e
151 = case arityType False e of
153 AT a ABot -> Just (a, mkStrictSig (mkTopDmdType (replicate a topDmd) BotRes))
156 Note [Definition of arity]
157 ~~~~~~~~~~~~~~~~~~~~~~~~~~
158 The "arity" of an expression 'e' is n if
159 applying 'e' to *fewer* than n *value* arguments
162 Or, to put it another way
164 there is no work lost in duplicating the partial
165 application (e x1 .. x(n-1))
167 In the divegent case, no work is lost by duplicating because if the thing
168 is evaluated once, that's the end of the program.
170 Or, to put it another way, in any context C
172 C[ (\x1 .. xn. e x1 .. xn) ]
176 It's all a bit more subtle than it looks:
178 Note [Arity of case expressions]
179 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
180 We treat the arity of
181 case x of p -> \s -> ...
182 as 1 (or more) because for I/O ish things we really want to get that
183 \s to the top. We are prepared to evaluate x each time round the loop
184 in order to get that.
186 This isn't really right in the presence of seq. Consider
190 Can we eta-expand here? At first the answer looks like "yes of course", but
193 This should diverge! But if we eta-expand, it won't. Again, we ignore this
194 "problem", because being scrupulous would lose an important transformation for
197 1. Note [One-shot lambdas]
198 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
199 Consider one-shot lambdas
200 let x = expensive in \y z -> E
201 We want this to have arity 1 if the \y-abstraction is a 1-shot lambda.
203 3. Note [Dealing with bottom]
204 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
206 f = \x -> error "foo"
207 Here, arity 1 is fine. But if it is
211 then we want to get arity 2. Technically, this isn't quite right, because
213 should diverge, but it'll converge if we eta-expand f. Nevertheless, we
214 do so; it improves some programs significantly, and increasing convergence
215 isn't a bad thing. Hence the ABot/ATop in ArityType.
217 4. Note [Newtype arity]
218 ~~~~~~~~~~~~~~~~~~~~~~~~
219 Non-recursive newtypes are transparent, and should not get in the way.
220 We do (currently) eta-expand recursive newtypes too. So if we have, say
222 newtype T = MkT ([T] -> Int)
226 where f has arity 1. Then: etaExpandArity e = 1;
227 that is, etaExpandArity looks through the coerce.
229 When we eta-expand e to arity 1: eta_expand 1 e T
230 we want to get: coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
232 HOWEVER, note that if you use coerce bogusly you can ge
234 And since negate has arity 2, you might try to eta expand. But you can't
235 decopose Int to a function type. Hence the final case in eta_expand.
238 applyStateHack :: CoreExpr -> ArityType -> Arity
239 applyStateHack e (AT orig_arity is_bot)
240 | opt_NoStateHack = orig_arity
241 | ABot <- is_bot = orig_arity -- Note [State hack and bottoming functions]
242 | otherwise = go orig_ty orig_arity
243 where -- Note [The state-transformer hack]
245 go :: Type -> Arity -> Arity
246 go ty arity -- This case analysis should match that in eta_expand
247 | Just (_, ty') <- splitForAllTy_maybe ty = go ty' arity
248 | Just (arg,res) <- splitFunTy_maybe ty
249 , arity > 0 || isStateHackType arg = 1 + go res (arity-1)
251 -- See Note [trimCast]
252 | Just (tc,tys) <- splitTyConApp_maybe ty
253 , Just (ty', _) <- instNewTyCon_maybe tc tys
254 , not (isRecursiveTyCon tc) = go ty' arity
255 -- Important to look through non-recursive newtypes, so that, eg
256 -- (f x) where f has arity 2, f :: Int -> IO ()
257 -- Here we want to get arity 1 for the result!
261 = if arity > 0 then 1 + go res (arity-1)
262 else if isStateHackType arg then
263 pprTrace "applystatehack" (vcat [ppr orig_arity, ppr orig_ty,
264 ppr ty, ppr res, ppr e]) $
266 else WARN( arity > 0, ppr arity ) 0
268 | otherwise = WARN( arity > 0, ppr arity <+> ppr ty) 0
271 Note [The state-transformer hack]
272 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
275 where e has arity n. Then, if we know from the context that f has
277 t1 -> ... -> tn -1-> t(n+1) -1-> ... -1-> tm -> ...
278 then we can expand the arity to m. This usage type says that
279 any application (x e1 .. en) will be applied to uniquely to (m-n) more args
280 Consider f = \x. let y = <expensive>
283 False -> \(s:RealWorld) -> e
284 where foo has arity 1. Then we want the state hack to
285 apply to foo too, so we can eta expand the case.
287 Then we expect that if f is applied to one arg, it'll be applied to two
288 (that's the hack -- we don't really know, and sometimes it's false)
289 See also Id.isOneShotBndr.
291 Note [State hack and bottoming functions]
292 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
293 It's a terrible idea to use the state hack on a bottoming function.
294 Here's what happens (Trac #2861):
299 Eta-expand, using the state hack:
301 f = \p. (\s. ((error "...") |> g1) s) |> g2
302 g1 :: IO T ~ (S -> (S,T))
303 g2 :: (S -> (S,T)) ~ IO T
307 f' = \p. \s. ((error "...") |> g1) s
308 f = f' |> (String -> g2)
310 Discard args for bottomming function
312 f' = \p. \s. ((error "...") |> g1 |> g3
313 g3 :: (S -> (S,T)) ~ (S,T)
317 f'' = \p. \s. (error "...")
318 f' = f'' |> (String -> S -> g1.g3)
320 And now we can repeat the whole loop. Aargh! The bug is in applying the
321 state hack to a function which then swallows the argument.
324 -------------------- Main arity code ----------------------------
326 -- If e has ArityType (AT n r), then the term 'e'
327 -- * Must be applied to at least n *value* args
328 -- before doing any significant work
329 -- * It will not diverge before being applied to n
331 -- * If 'r' is ABot, then it guarantees to diverge if
332 -- applied to n arguments (or more)
334 data ArityType = AT Arity ArityRes
335 data ArityRes = ATop -- Know nothing
338 vanillaArityType :: ArityType
339 vanillaArityType = AT 0 ATop -- Totally uninformative
341 incArity :: ArityType -> ArityType
342 incArity (AT a r) = AT (a+1) r
344 decArity :: ArityType -> ArityType
345 decArity (AT 0 r) = AT 0 r
346 decArity (AT a r) = AT (a-1) r
348 andArityType :: ArityType -> ArityType -> ArityType -- Used for branches of a 'case'
349 andArityType (AT a1 ATop) (AT a2 ATop) = AT (a1 `min` a2) ATop
350 andArityType (AT _ ABot) (AT a2 ATop) = AT a2 ATop
351 andArityType (AT a1 ATop) (AT _ ABot) = AT a1 ATop
352 andArityType (AT a1 ABot) (AT a2 ABot) = AT (a1 `max` a2) ABot
354 ---------------------------
355 trimCast :: Coercion -> ArityType -> ArityType
356 -- Trim the arity to be no more than allowed by the
357 -- arrows in ty2, where co :: ty1~ty2
360 {- Omitting for now Note [trimCast]
361 trimCast co at@(AT ar _)
362 | ar > co_arity = AT co_arity ATop
365 (_, ty2) = coercionKind co
366 co_arity = typeArity ty2
372 When you try putting trimCast back in, comment out the snippets
373 flagged by the other references to Note [trimCast]
376 ---------------------------
377 trimArity :: Bool -> ArityType -> ArityType
378 -- We have something like (let x = E in b), where b has the given
380 -- * If E is cheap we can push it inside as far as we like
381 -- * If b eventually diverges, we allow ourselves to push inside
382 -- arbitrarily, even though that is not quite right
383 trimArity _cheap (AT a ABot) = AT a ABot
384 trimArity True (AT a ATop) = AT a ATop
385 trimArity False (AT _ ATop) = AT 0 ATop -- Bale out
387 ---------------------------
388 arityType :: Bool -> CoreExpr -> ArityType
390 | Just strict_sig <- idStrictness_maybe v
391 , (ds, res) <- splitStrictSig strict_sig
393 = AT (length ds) ABot -- Function diverges
395 = AT (idArity v) ATop
397 -- Lambdas; increase arity
398 arityType dicts_cheap (Lam x e)
399 | isId x = incArity (arityType dicts_cheap e)
400 | otherwise = arityType dicts_cheap e
402 -- Applications; decrease arity
403 arityType dicts_cheap (App fun (Type _))
404 = arityType dicts_cheap fun
405 arityType dicts_cheap (App fun arg )
406 = trimArity (exprIsCheap arg) (decArity (arityType dicts_cheap fun))
408 -- Case/Let; keep arity if either the expression is cheap
409 -- or it's a 1-shot lambda
410 -- The former is not really right for Haskell
411 -- f x = case x of { (a,b) -> \y. e }
413 -- f x y = case x of { (a,b) -> e }
414 -- The difference is observable using 'seq'
415 arityType dicts_cheap (Case scrut _ _ alts)
416 = trimArity (exprIsCheap scrut)
417 (foldr1 andArityType [arityType dicts_cheap rhs | (_,_,rhs) <- alts])
419 arityType dicts_cheap (Let b e)
420 = trimArity (cheap_bind b) (arityType dicts_cheap e)
422 cheap_bind (NonRec b e) = is_cheap (b,e)
423 cheap_bind (Rec prs) = all is_cheap prs
424 is_cheap (b,e) = (dicts_cheap && isDictLikeTy (idType b))
426 -- If the experimental -fdicts-cheap flag is on, we eta-expand through
427 -- dictionary bindings. This improves arities. Thereby, it also
428 -- means that full laziness is less prone to floating out the
429 -- application of a function to its dictionary arguments, which
430 -- can thereby lose opportunities for fusion. Example:
431 -- foo :: Ord a => a -> ...
432 -- foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....
433 -- -- So foo has arity 1
435 -- f = \x. foo dInt $ bar x
437 -- The (foo DInt) is floated out, and makes ineffective a RULE
440 -- One could go further and make exprIsCheap reply True to any
441 -- dictionary-typed expression, but that's more work.
443 -- See Note [Dictionary-like types] in TcType.lhs for why we use
444 -- isDictLikeTy here rather than isDictTy
446 arityType dicts_cheap (Note _ e) = arityType dicts_cheap e
447 arityType dicts_cheap (Cast e co) = trimCast co (arityType dicts_cheap e)
448 arityType _ _ = vanillaArityType
452 %************************************************************************
454 The main eta-expander
456 %************************************************************************
458 IMPORTANT NOTE: The eta expander is careful not to introduce "crap".
459 In particular, given a CoreExpr satisfying the 'CpeRhs' invariant (in
460 CorePrep), it returns a CoreExpr satisfying the same invariant. See
461 Note [Eta expansion and the CorePrep invariants] in CorePrep.
463 This means the eta-expander has to do a bit of on-the-fly
464 simplification but it's not too hard. The alernative, of relying on
465 a subsequent clean-up phase of the Simplifier to de-crapify the result,
466 means you can't really use it in CorePrep, which is painful.
468 Note [Eta expansion and SCCs]
469 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
470 Note that SCCs are not treated specially by etaExpand. If we have
471 etaExpand 2 (\x -> scc "foo" e)
472 = (\xy -> (scc "foo" e) y)
473 So the costs of evaluating 'e' (not 'e y') are attributed to "foo"
476 -- | @etaExpand n us e ty@ returns an expression with
477 -- the same meaning as @e@, but with arity @n@.
481 -- > e' = etaExpand n us e ty
483 -- We should have that:
485 -- > ty = exprType e = exprType e'
486 etaExpand :: Arity -- ^ Result should have this number of value args
487 -> CoreExpr -- ^ Expression to expand
489 -- etaExpand deals with for-alls. For example:
491 -- where E :: forall a. a -> a
493 -- (/\b. \y::a -> E b y)
495 -- It deals with coerces too, though they are now rare
496 -- so perhaps the extra code isn't worth it
498 etaExpand n orig_expr
501 -- Strip off existing lambdas and casts
502 -- Note [Eta expansion and SCCs]
504 go n (Lam v body) | isTyVar v = Lam v (go n body)
505 | otherwise = Lam v (go (n-1) body)
506 go n (Cast expr co) = Cast (go n expr) co
507 go n expr = -- pprTrace "ee" (vcat [ppr orig_expr, ppr expr, ppr etas]) $
508 etaInfoAbs etas (etaInfoApp subst' expr etas)
510 in_scope = mkInScopeSet (exprFreeVars expr)
511 (in_scope', etas) = mkEtaWW n in_scope (exprType expr)
512 subst' = mkEmptySubst in_scope'
516 data EtaInfo = EtaVar Var -- /\a. [], [] a
518 | EtaCo Coercion -- [] |> co, [] |> (sym co)
520 instance Outputable EtaInfo where
521 ppr (EtaVar v) = ptext (sLit "EtaVar") <+> ppr v
522 ppr (EtaCo co) = ptext (sLit "EtaCo") <+> ppr co
524 pushCoercion :: Coercion -> [EtaInfo] -> [EtaInfo]
525 pushCoercion co1 (EtaCo co2 : eis)
526 | isIdentityCoercion co = eis
527 | otherwise = EtaCo co : eis
529 co = co1 `mkTransCoercion` co2
531 pushCoercion co eis = EtaCo co : eis
534 etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr
535 etaInfoAbs [] expr = expr
536 etaInfoAbs (EtaVar v : eis) expr = Lam v (etaInfoAbs eis expr)
537 etaInfoAbs (EtaCo co : eis) expr = Cast (etaInfoAbs eis expr) (mkSymCoercion co)
540 etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
541 -- (etaInfoApp s e eis) returns something equivalent to
542 -- ((substExpr s e) `appliedto` eis)
544 etaInfoApp subst (Lam v1 e) (EtaVar v2 : eis)
545 = etaInfoApp subst' e eis
547 subst' | isTyVar v1 = CoreSubst.extendTvSubst subst v1 (mkTyVarTy v2)
548 | otherwise = CoreSubst.extendIdSubst subst v1 (Var v2)
550 etaInfoApp subst (Cast e co1) eis
551 = etaInfoApp subst e (pushCoercion co' eis)
553 co' = CoreSubst.substTy subst co1
555 etaInfoApp subst (Case e b _ alts) eis
556 = Case (subst_expr subst e) b1 (coreAltsType alts') alts'
558 (subst1, b1) = substBndr subst b
559 alts' = map subst_alt alts
560 subst_alt (con, bs, rhs) = (con, bs', etaInfoApp subst2 rhs eis)
562 (subst2,bs') = substBndrs subst1 bs
564 etaInfoApp subst (Let b e) eis
565 = Let b' (etaInfoApp subst' e eis)
567 (subst', b') = subst_bind subst b
569 etaInfoApp subst (Note note e) eis
570 = Note note (etaInfoApp subst e eis)
572 etaInfoApp subst e eis
573 = go (subst_expr subst e) eis
576 go e (EtaVar v : eis) = go (App e (varToCoreExpr v)) eis
577 go e (EtaCo co : eis) = go (Cast e co) eis
580 mkEtaWW :: Arity -> InScopeSet -> Type
581 -> (InScopeSet, [EtaInfo])
582 -- EtaInfo contains fresh variables,
583 -- not free in the incoming CoreExpr
584 -- Outgoing InScopeSet includes the EtaInfo vars
585 -- and the original free vars
587 mkEtaWW orig_n in_scope orig_ty
588 = go orig_n empty_subst orig_ty []
590 empty_subst = mkTvSubst in_scope emptyTvSubstEnv
594 = (getTvInScope subst, reverse eis)
596 | Just (tv,ty') <- splitForAllTy_maybe ty
597 , let (subst', tv') = substTyVarBndr subst tv
598 -- Avoid free vars of the original expression
599 = go n subst' ty' (EtaVar tv' : eis)
601 | Just (arg_ty, res_ty) <- splitFunTy_maybe ty
602 , let (subst', eta_id') = freshEtaId n subst arg_ty
603 -- Avoid free vars of the original expression
604 = go (n-1) subst' res_ty (EtaVar eta_id' : eis)
606 -- See Note [trimCast]
607 | Just(ty',co) <- splitNewTypeRepCo_maybe ty
609 -- newtype T = MkT ([T] -> Int)
610 -- Consider eta-expanding this
613 -- coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
614 go n subst ty' (EtaCo (Type.substTy subst co) : eis)
617 | otherwise -- We have an expression of arity > 0,
618 = WARN( True, ppr orig_n <+> ppr orig_ty )
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; use short-cutting versions
630 subst_expr :: Subst -> CoreExpr -> CoreExpr
631 subst_expr = substExprSC (text "CoreArity:substExpr")
633 subst_bind :: Subst -> CoreBind -> (Subst, CoreBind)
634 subst_bind = substBindSC
638 freshEtaId :: Int -> TvSubst -> Type -> (TvSubst, Id)
639 -- Make a fresh Id, with specified type (after applying substitution)
640 -- It should be "fresh" in the sense that it's not in the in-scope set
641 -- of the TvSubstEnv; and it should itself then be added to the in-scope
642 -- set of the TvSubstEnv
644 -- The Int is just a reasonable starting point for generating a unique;
645 -- it does not necessarily have to be unique itself.
646 freshEtaId n subst ty
649 ty' = Type.substTy subst ty
650 eta_id' = uniqAway (getTvInScope subst) $
651 mkSysLocal (fsLit "eta") (mkBuiltinUnique n) ty'
652 subst' = extendTvInScope subst eta_id'