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"
20 import qualified CoreSubst
21 import CoreSubst ( Subst, substBndr, substBndrs, substExpr
22 , mkEmptySubst, isEmptySubst )
27 import TcType ( isDictLikeTy )
36 import GHC.Exts -- For `xori`
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 %************************************************************************
127 \subsection{Eta reduction and expansion}
129 %************************************************************************
131 exprEtaExpandArity is used when eta expanding
134 It returns 1 (or more) to:
135 case x of p -> \s -> ...
136 because for I/O ish things we really want to get that \s to the top.
137 We are prepared to evaluate x each time round the loop in order to get that
139 It's all a bit more subtle than it looks:
143 Consider one-shot lambdas
144 let x = expensive in \y z -> E
145 We want this to have arity 2 if the \y-abstraction is a 1-shot lambda
146 Hence the ArityType returned by arityType
148 2. The state-transformer hack
150 The one-shot lambda special cause is particularly important/useful for
151 IO state transformers, where we often get
152 let x = E in \ s -> ...
154 and the \s is a real-world state token abstraction. Such abstractions
155 are almost invariably 1-shot, so we want to pull the \s out, past the
156 let x=E, even if E is expensive. So we treat state-token lambdas as
157 one-shot even if they aren't really. The hack is in Id.isOneShotBndr.
159 3. Dealing with bottom
162 f = \x -> error "foo"
163 Here, arity 1 is fine. But if it is
167 then we want to get arity 2. Tecnically, this isn't quite right, because
169 should diverge, but it'll converge if we eta-expand f. Nevertheless, we
170 do so; it improves some programs significantly, and increasing convergence
171 isn't a bad thing. Hence the ABot/ATop in ArityType.
173 Actually, the situation is worse. Consider
177 Can we eta-expand here? At first the answer looks like "yes of course", but
180 This should diverge! But if we eta-expand, it won't. Again, we ignore this
181 "problem", because being scrupulous would lose an important transformation for
187 Non-recursive newtypes are transparent, and should not get in the way.
188 We do (currently) eta-expand recursive newtypes too. So if we have, say
190 newtype T = MkT ([T] -> Int)
194 where f has arity 1. Then: etaExpandArity e = 1;
195 that is, etaExpandArity looks through the coerce.
197 When we eta-expand e to arity 1: eta_expand 1 e T
198 we want to get: coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
200 HOWEVER, note that if you use coerce bogusly you can ge
202 And since negate has arity 2, you might try to eta expand. But you can't
203 decopose Int to a function type. Hence the final case in eta_expand.
207 -- ^ The Arity returned is the number of value args the
208 -- expression can be applied to without doing much work
209 exprEtaExpandArity :: DynFlags -> CoreExpr -> Arity
210 exprEtaExpandArity dflags e = arityDepth (arityType dflags e)
212 -- A limited sort of function type
213 data ArityType = AFun Bool ArityType -- True <=> one-shot
214 | ATop -- Know nothing
217 arityDepth :: ArityType -> Arity
218 arityDepth (AFun _ ty) = 1 + arityDepth ty
221 andArityType :: ArityType -> ArityType -> ArityType
222 andArityType ABot at2 = at2
223 andArityType ATop _ = ATop
224 andArityType (AFun t1 at1) (AFun t2 at2) = AFun (t1 && t2) (andArityType at1 at2)
225 andArityType at1 at2 = andArityType at2 at1
227 arityType :: DynFlags -> CoreExpr -> ArityType
228 -- (go1 e) = [b1,..,bn]
229 -- means expression can be rewritten \x_b1 -> ... \x_bn -> body
230 -- where bi is True <=> the lambda is one-shot
232 arityType dflags (Note _ e) = arityType dflags e
233 -- Not needed any more: etaExpand is cleverer
234 -- removed: | ok_note n = arityType dflags e
235 -- removed: | otherwise = ATop
237 arityType dflags (Cast e _) = arityType dflags e
240 = mk (idArity v) (arg_tys (idType v))
242 mk :: Arity -> [Type] -> ArityType
243 -- The argument types are only to steer the "state hack"
244 -- Consider case x of
246 -- False -> \(s:RealWorld) -> e
247 -- where foo has arity 1. Then we want the state hack to
248 -- apply to foo too, so we can eta expand the case.
249 mk 0 tys | isBottomingId v = ABot
250 | (ty:_) <- tys, isStateHackType ty = AFun True ATop
252 mk n (ty:tys) = AFun (isStateHackType ty) (mk (n-1) tys)
253 mk n [] = AFun False (mk (n-1) [])
255 arg_tys :: Type -> [Type] -- Ignore for-alls
257 | Just (_, ty') <- splitForAllTy_maybe ty = arg_tys ty'
258 | Just (arg,res) <- splitFunTy_maybe ty = arg : arg_tys res
261 -- Lambdas; increase arity
262 arityType dflags (Lam x e)
263 | isId x = AFun (isOneShotBndr x) (arityType dflags e)
264 | otherwise = arityType dflags e
266 -- Applications; decrease arity
267 arityType dflags (App f (Type _)) = arityType dflags f
268 arityType dflags (App f a)
269 = case arityType dflags f of
270 ABot -> ABot -- If function diverges, ignore argument
271 ATop -> ATop -- No no info about function
273 | exprIsCheap a -> xs
276 -- Case/Let; keep arity if either the expression is cheap
277 -- or it's a 1-shot lambda
278 -- The former is not really right for Haskell
279 -- f x = case x of { (a,b) -> \y. e }
281 -- f x y = case x of { (a,b) -> e }
282 -- The difference is observable using 'seq'
283 arityType dflags (Case scrut _ _ alts)
284 = case foldr1 andArityType [arityType dflags rhs | (_,_,rhs) <- alts] of
285 xs | exprIsCheap scrut -> xs
286 AFun one_shot _ | one_shot -> AFun True ATop
289 arityType dflags (Let b e)
290 = case arityType dflags e of
291 xs | cheap_bind b -> xs
292 AFun one_shot _ | one_shot -> AFun True ATop
295 cheap_bind (NonRec b e) = is_cheap (b,e)
296 cheap_bind (Rec prs) = all is_cheap prs
297 is_cheap (b,e) = (dopt Opt_DictsCheap dflags && isDictLikeTy (idType b))
299 -- If the experimental -fdicts-cheap flag is on, we eta-expand through
300 -- dictionary bindings. This improves arities. Thereby, it also
301 -- means that full laziness is less prone to floating out the
302 -- application of a function to its dictionary arguments, which
303 -- can thereby lose opportunities for fusion. Example:
304 -- foo :: Ord a => a -> ...
305 -- foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....
306 -- -- So foo has arity 1
308 -- f = \x. foo dInt $ bar x
310 -- The (foo DInt) is floated out, and makes ineffective a RULE
313 -- One could go further and make exprIsCheap reply True to any
314 -- dictionary-typed expression, but that's more work.
316 -- See Note [Dictionary-like types] in TcType.lhs for why we use
317 -- isDictLikeTy here rather than isDictTy
323 %************************************************************************
325 The main eta-expander
327 %************************************************************************
329 IMPORTANT NOTE: The eta expander is careful not to introduce "crap".
330 In particular, given a CoreExpr satisfying the 'CpeRhs' invariant (in
331 CorePrep), it returns a CoreExpr satisfying the same invariant. See
332 Note [Eta expansion and the CorePrep invariants] in CorePrep.
334 This means the eta-expander has to do a bit of on-the-fly
335 simplification but it's not too hard. The alernative, of relying on
336 a subsequent clean-up phase of the Simplifier to de-crapify the result,
337 means you can't really use it in CorePrep, which is painful.
340 -- | @etaExpand n us e ty@ returns an expression with
341 -- the same meaning as @e@, but with arity @n@.
345 -- > e' = etaExpand n us e ty
347 -- We should have that:
349 -- > ty = exprType e = exprType e'
350 etaExpand :: Arity -- ^ Result should have this number of value args
351 -> CoreExpr -- ^ Expression to expand
353 -- Note that SCCs are not treated specially. If we have
354 -- etaExpand 2 (\x -> scc "foo" e)
355 -- = (\xy -> (scc "foo" e) y)
356 -- So the costs of evaluating 'e' (not 'e y') are attributed to "foo"
358 -- etaExpand deals with for-alls. For example:
360 -- where E :: forall a. a -> a
362 -- (/\b. \y::a -> E b y)
364 -- It deals with coerces too, though they are now rare
365 -- so perhaps the extra code isn't worth it
367 etaExpand n orig_expr
368 | manifestArity orig_expr >= n = orig_expr -- The no-op case
372 -- Strip off existing lambdas
374 go n (Lam v body) | isTyVar v = Lam v (go n body)
375 | otherwise = Lam v (go (n-1) body)
376 go n (Note InlineMe expr) = Note InlineMe (go n expr)
377 -- Note [Eta expansion and SCCs]
378 go n (Cast expr co) = Cast (go n expr) co
379 go n expr = -- pprTrace "ee" (vcat [ppr orig_expr, ppr expr, ppr etas]) $
380 etaInfoAbs etas (etaInfoApp subst' expr etas)
382 in_scope = mkInScopeSet (exprFreeVars expr)
383 (in_scope', etas) = mkEtaWW n in_scope (exprType expr)
384 subst' = mkEmptySubst in_scope'
388 data EtaInfo = EtaVar Var -- /\a. [], [] a
390 | EtaCo Coercion -- [] |> co, [] |> (sym co)
392 instance Outputable EtaInfo where
393 ppr (EtaVar v) = ptext (sLit "EtaVar") <+> ppr v
394 ppr (EtaCo co) = ptext (sLit "EtaCo") <+> ppr co
396 pushCoercion :: Coercion -> [EtaInfo] -> [EtaInfo]
397 pushCoercion co1 (EtaCo co2 : eis)
398 | isIdentityCoercion co = eis
399 | otherwise = EtaCo co : eis
401 co = co1 `mkTransCoercion` co2
403 pushCoercion co eis = EtaCo co : eis
406 etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr
407 etaInfoAbs [] expr = expr
408 etaInfoAbs (EtaVar v : eis) expr = Lam v (etaInfoAbs eis expr)
409 etaInfoAbs (EtaCo co : eis) expr = Cast (etaInfoAbs eis expr) (mkSymCoercion co)
412 etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
413 -- (etaInfoApp s e eis) returns something equivalent to
414 -- ((substExpr s e) `appliedto` eis)
416 etaInfoApp subst (Lam v1 e) (EtaVar v2 : eis)
417 = etaInfoApp subst' e eis
419 subst' | isTyVar v1 = CoreSubst.extendTvSubst subst v1 (mkTyVarTy v2)
420 | otherwise = CoreSubst.extendIdSubst subst v1 (Var v2)
422 etaInfoApp subst (Cast e co1) eis
423 = etaInfoApp subst e (pushCoercion co' eis)
425 co' = CoreSubst.substTy subst co1
427 etaInfoApp subst (Case e b _ alts) eis
428 = Case (subst_expr subst e) b1 (coreAltsType alts') alts'
430 (subst1, b1) = substBndr subst b
431 alts' = map subst_alt alts
432 subst_alt (con, bs, rhs) = (con, bs', etaInfoApp subst2 rhs eis)
434 (subst2,bs') = substBndrs subst1 bs
436 etaInfoApp subst (Let b e) eis
437 = Let b' (etaInfoApp subst' e eis)
439 (subst', b') = subst_bind subst b
441 etaInfoApp subst (Note note e) eis
442 = Note note (etaInfoApp subst e eis)
444 etaInfoApp subst e eis
445 = go (subst_expr subst e) eis
448 go e (EtaVar v : eis) = go (App e (varToCoreExpr v)) eis
449 go e (EtaCo co : eis) = go (Cast e co) eis
452 mkEtaWW :: Arity -> InScopeSet -> Type
453 -> (InScopeSet, [EtaInfo])
454 -- EtaInfo contains fresh variables,
455 -- not free in the incoming CoreExpr
456 -- Outgoing InScopeSet includes the EtaInfo vars
457 -- and the original free vars
459 mkEtaWW n in_scope ty
460 = go n empty_subst ty []
462 empty_subst = mkTvSubst in_scope emptyTvSubstEnv
466 = (getTvInScope subst, reverse eis)
468 | Just (tv,ty') <- splitForAllTy_maybe ty
469 , let (subst', tv') = substTyVarBndr subst tv
470 -- Avoid free vars of the original expression
471 = go n subst' ty' (EtaVar tv' : eis)
473 | Just (arg_ty, res_ty) <- splitFunTy_maybe ty
474 , let (subst', eta_id') = freshEtaId n subst arg_ty
475 -- Avoid free vars of the original expression
476 = go (n-1) subst' res_ty (EtaVar eta_id' : eis)
478 | Just(ty',co) <- splitNewTypeRepCo_maybe ty
480 -- newtype T = MkT ([T] -> Int)
481 -- Consider eta-expanding this
484 -- coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
485 go n subst ty' (EtaCo (substTy subst co) : eis)
487 | otherwise -- We have an expression of arity > 0,
488 = (getTvInScope subst, reverse eis) -- but its type isn't a function.
489 -- This *can* legitmately happen:
490 -- e.g. coerce Int (\x. x) Essentially the programmer is
491 -- playing fast and loose with types (Happy does this a lot).
492 -- So we simply decline to eta-expand. Otherwise we'd end up
493 -- with an explicit lambda having a non-function type
497 -- Avoiding unnecessary substitution
499 subst_expr :: Subst -> CoreExpr -> CoreExpr
500 subst_expr s e | isEmptySubst s = e
501 | otherwise = substExpr s e
503 subst_bind :: Subst -> CoreBind -> (Subst, CoreBind)
504 subst_bind subst (NonRec b r)
505 = (subst', NonRec b' (subst_expr subst r))
507 (subst', b') = substBndr subst b
508 subst_bind subst (Rec prs)
509 = (subst', Rec (bs1 `zip` map (subst_expr subst') rhss))
511 (bs, rhss) = unzip prs
512 (subst', bs1) = substBndrs subst bs
516 freshEtaId :: Int -> TvSubst -> Type -> (TvSubst, Id)
517 -- Make a fresh Id, with specified type (after applying substitution)
518 -- It should be "fresh" in the sense that it's not in the in-scope set
519 -- of the TvSubstEnv; and it should itself then be added to the in-scope
520 -- set of the TvSubstEnv
522 -- The Int is just a reasonable starting point for generating a unique;
523 -- it does not necessarily have to be unique itself.
524 freshEtaId n subst ty
527 ty' = substTy subst ty
528 eta_id' = uniqAway (getTvInScope subst) $
529 mkSysLocal (fsLit "eta") (mkBuiltinUnique n) ty'
530 subst' = extendTvInScope subst [eta_id']