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 import GHC.Exts -- For `xori`
42 %************************************************************************
44 manifestArity and exprArity
46 %************************************************************************
48 exprArity is a cheap-and-cheerful version of exprEtaExpandArity.
49 It tells how many things the expression can be applied to before doing
50 any work. It doesn't look inside cases, lets, etc. The idea is that
51 exprEtaExpandArity will do the hard work, leaving something that's easy
52 for exprArity to grapple with. In particular, Simplify uses exprArity to
53 compute the ArityInfo for the Id.
55 Originally I thought that it was enough just to look for top-level lambdas, but
56 it isn't. I've seen this
58 foo = PrelBase.timesInt
60 We want foo to get arity 2 even though the eta-expander will leave it
61 unchanged, in the expectation that it'll be inlined. But occasionally it
62 isn't, because foo is blacklisted (used in a rule).
64 Similarly, see the ok_note check in exprEtaExpandArity. So
65 f = __inline_me (\x -> e)
66 won't be eta-expanded.
68 And in any case it seems more robust to have exprArity be a bit more intelligent.
69 But note that (\x y z -> f x y z)
70 should have arity 3, regardless of f's arity.
72 Note [exprArity invariant]
73 ~~~~~~~~~~~~~~~~~~~~~~~~~~
74 exprArity has the following invariant:
75 (exprArity e) = n, then manifestArity (etaExpand e n) = n
77 That is, if exprArity says "the arity is n" then etaExpand really can get
78 "n" manifest lambdas to the top.
80 Why is this important? Because
81 - In TidyPgm we use exprArity to fix the *final arity* of
82 each top-level Id, and in
83 - In CorePrep we use etaExpand on each rhs, so that the visible lambdas
84 actually match that arity, which in turn means
85 that the StgRhs has the right number of lambdas
87 An alternative would be to do the eta-expansion in TidyPgm, at least
88 for top-level bindings, in which case we would not need the trim_arity
89 in exprArity. That is a less local change, so I'm going to leave it for today!
93 manifestArity :: CoreExpr -> Arity
94 -- ^ manifestArity sees how many leading value lambdas there are
95 manifestArity (Lam v e) | isId v = 1 + manifestArity e
96 | otherwise = manifestArity e
97 manifestArity (Note _ e) = manifestArity e
98 manifestArity (Cast e _) = manifestArity e
101 exprArity :: CoreExpr -> Arity
102 -- ^ An approximate, fast, version of 'exprEtaExpandArity'
105 go (Var v) = idArity v
106 go (Lam x e) | isId x = go e + 1
109 go (Cast e co) = trim_arity (go e) 0 (snd (coercionKind co))
110 go (App e (Type _)) = go e
111 go (App f a) | exprIsCheap a = (go f - 1) `max` 0
112 -- NB: exprIsCheap a!
113 -- f (fac x) does not have arity 2,
114 -- even if f has arity 3!
115 -- NB: `max 0`! (\x y -> f x) has arity 2, even if f is
116 -- unknown, hence arity 0
119 -- Note [exprArity invariant]
122 | Just (_, ty') <- splitForAllTy_maybe ty = trim_arity n a ty'
123 | Just (_, ty') <- splitFunTy_maybe ty = trim_arity n (a+1) ty'
124 | Just (ty',_) <- splitNewTypeRepCo_maybe ty = trim_arity n a ty'
128 %************************************************************************
132 %************************************************************************
135 -- ^ The Arity returned is the number of value args the
136 -- expression can be applied to without doing much work
137 exprEtaExpandArity :: DynFlags -> CoreExpr -> Arity
138 -- exprEtaExpandArity is used when eta expanding
139 -- e ==> \xy -> e x y
140 exprEtaExpandArity dflags e
141 = applyStateHack e (arityType dicts_cheap e)
143 dicts_cheap = dopt Opt_DictsCheap dflags
146 Note [Definition of arity]
147 ~~~~~~~~~~~~~~~~~~~~~~~~~~
148 The "arity" of an expression 'e' is n if
149 applying 'e' to *fewer* than n *value* arguments
152 Or, to put it another way
154 there is no work lost in duplicating the partial
155 application (e x1 .. x(n-1))
157 In the divegent case, no work is lost by duplicating because if the thing
158 is evaluated once, that's the end of the program.
160 Or, to put it another way, in any context C
162 C[ (\x1 .. xn. e x1 .. xn) ]
167 It's all a bit more subtle than it looks:
169 Note [Arity of case expressions]
170 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
171 We treat the arity of
172 case x of p -> \s -> ...
173 as 1 (or more) because for I/O ish things we really want to get that
174 \s to the top. We are prepared to evaluate x each time round the loop
175 in order to get that.
177 This isn't really right in the presence of seq. Consider
181 Can we eta-expand here? At first the answer looks like "yes of course", but
184 This should diverge! But if we eta-expand, it won't. Again, we ignore this
185 "problem", because being scrupulous would lose an important transformation for
189 1. Note [One-shot lambdas]
190 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
191 Consider one-shot lambdas
192 let x = expensive in \y z -> E
193 We want this to have arity 1 if the \y-abstraction is a 1-shot lambda.
195 3. Note [Dealing with bottom]
196 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
198 f = \x -> error "foo"
199 Here, arity 1 is fine. But if it is
203 then we want to get arity 2. Technically, this isn't quite right, because
205 should diverge, but it'll converge if we eta-expand f. Nevertheless, we
206 do so; it improves some programs significantly, and increasing convergence
207 isn't a bad thing. Hence the ABot/ATop in ArityType.
210 4. Note [Newtype arity]
211 ~~~~~~~~~~~~~~~~~~~~~~~~
212 Non-recursive newtypes are transparent, and should not get in the way.
213 We do (currently) eta-expand recursive newtypes too. So if we have, say
215 newtype T = MkT ([T] -> Int)
219 where f has arity 1. Then: etaExpandArity e = 1;
220 that is, etaExpandArity looks through the coerce.
222 When we eta-expand e to arity 1: eta_expand 1 e T
223 we want to get: coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
225 HOWEVER, note that if you use coerce bogusly you can ge
227 And since negate has arity 2, you might try to eta expand. But you can't
228 decopose Int to a function type. Hence the final case in eta_expand.
230 Note [The state-transformer hack]
231 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
234 where e has arity n. Then, if we know from the context that f has
236 t1 -> ... -> tn -1-> t(n+1) -1-> ... -1-> tm -> ...
237 then we can expand the arity to m. This usage type says that
238 any application (x e1 .. en) will be applied to uniquely to (m-n) more args
239 Consider f = \x. let y = <expensive>
242 False -> \(s:RealWorld) -> e
243 where foo has arity 1. Then we want the state hack to
244 apply to foo too, so we can eta expand the case.
246 Then we expect that if f is applied to one arg, it'll be applied to two
247 (that's the hack -- we don't really know, and sometimes it's false)
248 See also Id.isOneShotBndr.
251 applyStateHack :: CoreExpr -> ArityType -> Arity
252 applyStateHack e (AT orig_arity is_bot)
253 | opt_NoStateHack = orig_arity
254 | ABot <- is_bot = orig_arity -- Note [State hack and bottoming functions]
255 | otherwise = go orig_ty orig_arity
256 where -- Note [The state-transformer hack]
258 go :: Type -> Arity -> Arity
259 go ty arity -- This case analysis should match that in eta_expand
260 | Just (_, ty') <- splitForAllTy_maybe ty = go ty' arity
262 | Just (tc,tys) <- splitTyConApp_maybe ty
263 , Just (ty', _) <- instNewTyCon_maybe tc tys
264 , not (isRecursiveTyCon tc) = go ty' arity
265 -- Important to look through non-recursive newtypes, so that, eg
266 -- (f x) where f has arity 2, f :: Int -> IO ()
267 -- Here we want to get arity 1 for the result!
269 | Just (arg,res) <- splitFunTy_maybe ty
270 , arity > 0 || isStateHackType arg = 1 + go res (arity-1)
272 = if arity > 0 then 1 + go res (arity-1)
273 else if isStateHackType arg then
274 pprTrace "applystatehack" (vcat [ppr orig_arity, ppr orig_ty,
275 ppr ty, ppr res, ppr e]) $
277 else WARN( arity > 0, ppr arity ) 0
279 | otherwise = WARN( arity > 0, ppr arity ) 0
282 Note [State hack and bottoming functions]
283 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
284 It's a terrible idea to use the state hack on a bottoming function.
285 Here's what happens (Trac #2861):
290 Eta-expand, using the state hack:
292 f = \p. (\s. ((error "...") |> g1) s) |> g2
293 g1 :: IO T ~ (S -> (S,T))
294 g2 :: (S -> (S,T)) ~ IO T
298 f' = \p. \s. ((error "...") |> g1) s
299 f = f' |> (String -> g2)
301 Discard args for bottomming function
303 f' = \p. \s. ((error "...") |> g1 |> g3
304 g3 :: (S -> (S,T)) ~ (S,T)
308 f'' = \p. \s. (error "...")
309 f' = f'' |> (String -> S -> g1.g3)
311 And now we can repeat the whole loop. Aargh! The bug is in applying the
312 state hack to a function which then swallows the argument.
315 -------------------- Main arity code ----------------------------
317 -- If e has ArityType (AT n r), then the term 'e'
318 -- * Must be applied to at least n *value* args
319 -- before doing any significant work
320 -- * It will not diverge before being applied to n
322 -- * If 'r' is ABot, then it guarantees to diverge if
323 -- applied to n arguments (or more)
325 data ArityType = AT Arity ArityRes
326 data ArityRes = ATop -- Know nothing
329 vanillaArityType :: ArityType
330 vanillaArityType = AT 0 ATop -- Totally uninformative
332 incArity :: ArityType -> ArityType
333 incArity (AT a r) = AT (a+1) r
335 decArity :: ArityType -> ArityType
336 decArity (AT 0 r) = AT 0 r
337 decArity (AT a r) = AT (a-1) r
339 andArityType :: ArityType -> ArityType -> ArityType -- Used for branches of a 'case'
340 andArityType (AT a1 ATop) (AT a2 ATop) = AT (a1 `min` a2) ATop
341 andArityType (AT _ ABot) (AT a2 ATop) = AT a2 ATop
342 andArityType (AT a1 ATop) (AT _ ABot) = AT a1 ATop
343 andArityType (AT a1 ABot) (AT a2 ABot) = AT (a1 `max` a2) ABot
345 trimArity :: Bool -> ArityType -> ArityType
346 -- We have something like (let x = E in b), where b has the given
348 -- * If E is cheap we can push it inside as far as we like
349 -- * If b eventually diverges, we allow ourselves to push inside
350 -- arbitrarily, even though that is not quite right
351 trimArity _cheap (AT a ABot) = AT a ABot
352 trimArity True (AT a ATop) = AT a ATop
353 trimArity False (AT _ ATop) = AT 0 ATop -- Bale out
355 ---------------------------
356 arityType :: Bool -> CoreExpr -> ArityType
358 | Just strict_sig <- idNewStrictness_maybe v
359 , (ds, res) <- splitStrictSig strict_sig
361 = AT (length ds) ABot -- Function diverges
363 = AT (idArity v) ATop
365 -- Lambdas; increase arity
366 arityType dicts_cheap (Lam x e)
367 | isId x = incArity (arityType dicts_cheap e)
368 | otherwise = arityType dicts_cheap e
370 -- Applications; decrease arity
371 arityType dicts_cheap (App fun (Type _))
372 = arityType dicts_cheap fun
373 arityType dicts_cheap (App fun arg )
374 = trimArity (exprIsCheap arg) (decArity (arityType dicts_cheap fun))
376 -- Case/Let; keep arity if either the expression is cheap
377 -- or it's a 1-shot lambda
378 -- The former is not really right for Haskell
379 -- f x = case x of { (a,b) -> \y. e }
381 -- f x y = case x of { (a,b) -> e }
382 -- The difference is observable using 'seq'
383 arityType dicts_cheap (Case scrut _ _ alts)
384 = trimArity (exprIsCheap scrut)
385 (foldr1 andArityType [arityType dicts_cheap rhs | (_,_,rhs) <- alts])
387 arityType dicts_cheap (Let b e)
388 = trimArity (cheap_bind b) (arityType dicts_cheap e)
390 cheap_bind (NonRec b e) = is_cheap (b,e)
391 cheap_bind (Rec prs) = all is_cheap prs
392 is_cheap (b,e) = (dicts_cheap && isDictLikeTy (idType b))
394 -- If the experimental -fdicts-cheap flag is on, we eta-expand through
395 -- dictionary bindings. This improves arities. Thereby, it also
396 -- means that full laziness is less prone to floating out the
397 -- application of a function to its dictionary arguments, which
398 -- can thereby lose opportunities for fusion. Example:
399 -- foo :: Ord a => a -> ...
400 -- foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....
401 -- -- So foo has arity 1
403 -- f = \x. foo dInt $ bar x
405 -- The (foo DInt) is floated out, and makes ineffective a RULE
408 -- One could go further and make exprIsCheap reply True to any
409 -- dictionary-typed expression, but that's more work.
411 -- See Note [Dictionary-like types] in TcType.lhs for why we use
412 -- isDictLikeTy here rather than isDictTy
414 arityType dicts_cheap (Note _ e) = arityType dicts_cheap e
415 arityType dicts_cheap (Cast e _) = arityType dicts_cheap e
416 arityType _ _ = vanillaArityType
420 %************************************************************************
422 The main eta-expander
424 %************************************************************************
426 IMPORTANT NOTE: The eta expander is careful not to introduce "crap".
427 In particular, given a CoreExpr satisfying the 'CpeRhs' invariant (in
428 CorePrep), it returns a CoreExpr satisfying the same invariant. See
429 Note [Eta expansion and the CorePrep invariants] in CorePrep.
431 This means the eta-expander has to do a bit of on-the-fly
432 simplification but it's not too hard. The alernative, of relying on
433 a subsequent clean-up phase of the Simplifier to de-crapify the result,
434 means you can't really use it in CorePrep, which is painful.
437 -- | @etaExpand n us e ty@ returns an expression with
438 -- the same meaning as @e@, but with arity @n@.
442 -- > e' = etaExpand n us e ty
444 -- We should have that:
446 -- > ty = exprType e = exprType e'
447 etaExpand :: Arity -- ^ Result should have this number of value args
448 -> CoreExpr -- ^ Expression to expand
450 -- Note that SCCs are not treated specially. If we have
451 -- etaExpand 2 (\x -> scc "foo" e)
452 -- = (\xy -> (scc "foo" e) y)
453 -- So the costs of evaluating 'e' (not 'e y') are attributed to "foo"
455 -- etaExpand deals with for-alls. For example:
457 -- where E :: forall a. a -> a
459 -- (/\b. \y::a -> E b y)
461 -- It deals with coerces too, though they are now rare
462 -- so perhaps the extra code isn't worth it
464 etaExpand n orig_expr
465 | manifestArity orig_expr >= n = orig_expr -- The no-op case
469 -- Strip off existing lambdas
470 -- Note [Eta expansion and SCCs]
472 go n (Lam v body) | isTyVar v = Lam v (go n body)
473 | otherwise = Lam v (go (n-1) body)
474 go n (Note InlineMe expr) = Note InlineMe (go n expr)
475 go n (Cast expr co) = Cast (go n expr) co
476 go n expr = -- pprTrace "ee" (vcat [ppr orig_expr, ppr expr, ppr etas]) $
477 etaInfoAbs etas (etaInfoApp subst' expr etas)
479 in_scope = mkInScopeSet (exprFreeVars expr)
480 (in_scope', etas) = mkEtaWW n in_scope (exprType expr)
481 subst' = mkEmptySubst in_scope'
485 data EtaInfo = EtaVar Var -- /\a. [], [] a
487 | EtaCo Coercion -- [] |> co, [] |> (sym co)
489 instance Outputable EtaInfo where
490 ppr (EtaVar v) = ptext (sLit "EtaVar") <+> ppr v
491 ppr (EtaCo co) = ptext (sLit "EtaCo") <+> ppr co
493 pushCoercion :: Coercion -> [EtaInfo] -> [EtaInfo]
494 pushCoercion co1 (EtaCo co2 : eis)
495 | isIdentityCoercion co = eis
496 | otherwise = EtaCo co : eis
498 co = co1 `mkTransCoercion` co2
500 pushCoercion co eis = EtaCo co : eis
503 etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr
504 etaInfoAbs [] expr = expr
505 etaInfoAbs (EtaVar v : eis) expr = Lam v (etaInfoAbs eis expr)
506 etaInfoAbs (EtaCo co : eis) expr = Cast (etaInfoAbs eis expr) (mkSymCoercion co)
509 etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
510 -- (etaInfoApp s e eis) returns something equivalent to
511 -- ((substExpr s e) `appliedto` eis)
513 etaInfoApp subst (Lam v1 e) (EtaVar v2 : eis)
514 = etaInfoApp subst' e eis
516 subst' | isTyVar v1 = CoreSubst.extendTvSubst subst v1 (mkTyVarTy v2)
517 | otherwise = CoreSubst.extendIdSubst subst v1 (Var v2)
519 etaInfoApp subst (Cast e co1) eis
520 = etaInfoApp subst e (pushCoercion co' eis)
522 co' = CoreSubst.substTy subst co1
524 etaInfoApp subst (Case e b _ alts) eis
525 = Case (subst_expr subst e) b1 (coreAltsType alts') alts'
527 (subst1, b1) = substBndr subst b
528 alts' = map subst_alt alts
529 subst_alt (con, bs, rhs) = (con, bs', etaInfoApp subst2 rhs eis)
531 (subst2,bs') = substBndrs subst1 bs
533 etaInfoApp subst (Let b e) eis
534 = Let b' (etaInfoApp subst' e eis)
536 (subst', b') = subst_bind subst b
538 etaInfoApp subst (Note note e) eis
539 = Note note (etaInfoApp subst e eis)
541 etaInfoApp subst e eis
542 = go (subst_expr subst e) eis
545 go e (EtaVar v : eis) = go (App e (varToCoreExpr v)) eis
546 go e (EtaCo co : eis) = go (Cast e co) eis
549 mkEtaWW :: Arity -> InScopeSet -> Type
550 -> (InScopeSet, [EtaInfo])
551 -- EtaInfo contains fresh variables,
552 -- not free in the incoming CoreExpr
553 -- Outgoing InScopeSet includes the EtaInfo vars
554 -- and the original free vars
556 mkEtaWW n in_scope ty
557 = go n empty_subst ty []
559 empty_subst = mkTvSubst in_scope emptyTvSubstEnv
563 = (getTvInScope subst, reverse eis)
565 | Just (tv,ty') <- splitForAllTy_maybe ty
566 , let (subst', tv') = substTyVarBndr subst tv
567 -- Avoid free vars of the original expression
568 = go n subst' ty' (EtaVar tv' : eis)
570 | Just (arg_ty, res_ty) <- splitFunTy_maybe ty
571 , let (subst', eta_id') = freshEtaId n subst arg_ty
572 -- Avoid free vars of the original expression
573 = go (n-1) subst' res_ty (EtaVar eta_id' : eis)
575 | Just(ty',co) <- splitNewTypeRepCo_maybe ty
577 -- newtype T = MkT ([T] -> Int)
578 -- Consider eta-expanding this
581 -- coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
582 go n subst ty' (EtaCo (substTy subst co) : eis)
584 | otherwise -- We have an expression of arity > 0,
585 = (getTvInScope subst, reverse eis) -- but its type isn't a function.
586 -- This *can* legitmately happen:
587 -- e.g. coerce Int (\x. x) Essentially the programmer is
588 -- playing fast and loose with types (Happy does this a lot).
589 -- So we simply decline to eta-expand. Otherwise we'd end up
590 -- with an explicit lambda having a non-function type
594 -- Avoiding unnecessary substitution
596 subst_expr :: Subst -> CoreExpr -> CoreExpr
597 subst_expr s e | isEmptySubst s = e
598 | otherwise = substExpr s e
600 subst_bind :: Subst -> CoreBind -> (Subst, CoreBind)
601 subst_bind subst (NonRec b r)
602 = (subst', NonRec b' (subst_expr subst r))
604 (subst', b') = substBndr subst b
605 subst_bind subst (Rec prs)
606 = (subst', Rec (bs1 `zip` map (subst_expr subst') rhss))
608 (bs, rhss) = unzip prs
609 (subst', bs1) = substBndrs subst bs
613 freshEtaId :: Int -> TvSubst -> Type -> (TvSubst, Id)
614 -- Make a fresh Id, with specified type (after applying substitution)
615 -- It should be "fresh" in the sense that it's not in the in-scope set
616 -- of the TvSubstEnv; and it should itself then be added to the in-scope
617 -- set of the TvSubstEnv
619 -- The Int is just a reasonable starting point for generating a unique;
620 -- it does not necessarily have to be unique itself.
621 freshEtaId n subst ty
624 ty' = substTy subst ty
625 eta_id' = uniqAway (getTvInScope subst) $
626 mkSysLocal (fsLit "eta") (mkBuiltinUnique n) ty'
627 subst' = extendTvInScope subst [eta_id']