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 )
35 import GHC.Exts -- For `xori`
38 %************************************************************************
40 manifestArity and exprArity
42 %************************************************************************
44 exprArity is a cheap-and-cheerful version of exprEtaExpandArity.
45 It tells how many things the expression can be applied to before doing
46 any work. It doesn't look inside cases, lets, etc. The idea is that
47 exprEtaExpandArity will do the hard work, leaving something that's easy
48 for exprArity to grapple with. In particular, Simplify uses exprArity to
49 compute the ArityInfo for the Id.
51 Originally I thought that it was enough just to look for top-level lambdas, but
52 it isn't. I've seen this
54 foo = PrelBase.timesInt
56 We want foo to get arity 2 even though the eta-expander will leave it
57 unchanged, in the expectation that it'll be inlined. But occasionally it
58 isn't, because foo is blacklisted (used in a rule).
60 Similarly, see the ok_note check in exprEtaExpandArity. So
61 f = __inline_me (\x -> e)
62 won't be eta-expanded.
64 And in any case it seems more robust to have exprArity be a bit more intelligent.
65 But note that (\x y z -> f x y z)
66 should have arity 3, regardless of f's arity.
68 Note [exprArity invariant]
69 ~~~~~~~~~~~~~~~~~~~~~~~~~~
70 exprArity has the following invariant:
71 (exprArity e) = n, then manifestArity (etaExpand e n) = n
73 That is, if exprArity says "the arity is n" then etaExpand really can get
74 "n" manifest lambdas to the top.
76 Why is this important? Because
77 - In TidyPgm we use exprArity to fix the *final arity* of
78 each top-level Id, and in
79 - In CorePrep we use etaExpand on each rhs, so that the visible lambdas
80 actually match that arity, which in turn means
81 that the StgRhs has the right number of lambdas
83 An alternative would be to do the eta-expansion in TidyPgm, at least
84 for top-level bindings, in which case we would not need the trim_arity
85 in exprArity. That is a less local change, so I'm going to leave it for today!
89 manifestArity :: CoreExpr -> Arity
90 -- ^ manifestArity sees how many leading value lambdas there are
91 manifestArity (Lam v e) | isId v = 1 + manifestArity e
92 | otherwise = manifestArity e
93 manifestArity (Note _ e) = manifestArity e
94 manifestArity (Cast e _) = manifestArity e
97 exprArity :: CoreExpr -> Arity
98 -- ^ An approximate, fast, version of 'exprEtaExpandArity'
101 go (Var v) = idArity v
102 go (Lam x e) | isId x = go e + 1
105 go (Cast e co) = trim_arity (go e) 0 (snd (coercionKind co))
106 go (App e (Type _)) = go e
107 go (App f a) | exprIsCheap a = (go f - 1) `max` 0
108 -- NB: exprIsCheap a!
109 -- f (fac x) does not have arity 2,
110 -- even if f has arity 3!
111 -- NB: `max 0`! (\x y -> f x) has arity 2, even if f is
112 -- unknown, hence arity 0
115 -- Note [exprArity invariant]
118 | Just (_, ty') <- splitForAllTy_maybe ty = trim_arity n a ty'
119 | Just (_, ty') <- splitFunTy_maybe ty = trim_arity n (a+1) ty'
120 | Just (ty',_) <- splitNewTypeRepCo_maybe ty = trim_arity n a ty'
124 %************************************************************************
126 \subsection{Eta reduction and expansion}
128 %************************************************************************
130 exprEtaExpandArity is used when eta expanding
133 It returns 1 (or more) to:
134 case x of p -> \s -> ...
135 because for I/O ish things we really want to get that \s to the top.
136 We are prepared to evaluate x each time round the loop in order to get that
138 It's all a bit more subtle than it looks:
142 Consider one-shot lambdas
143 let x = expensive in \y z -> E
144 We want this to have arity 2 if the \y-abstraction is a 1-shot lambda
145 Hence the ArityType returned by arityType
147 2. The state-transformer hack
149 The one-shot lambda special cause is particularly important/useful for
150 IO state transformers, where we often get
151 let x = E in \ s -> ...
153 and the \s is a real-world state token abstraction. Such abstractions
154 are almost invariably 1-shot, so we want to pull the \s out, past the
155 let x=E, even if E is expensive. So we treat state-token lambdas as
156 one-shot even if they aren't really. The hack is in Id.isOneShotBndr.
158 3. Dealing with bottom
161 f = \x -> error "foo"
162 Here, arity 1 is fine. But if it is
166 then we want to get arity 2. Tecnically, this isn't quite right, because
168 should diverge, but it'll converge if we eta-expand f. Nevertheless, we
169 do so; it improves some programs significantly, and increasing convergence
170 isn't a bad thing. Hence the ABot/ATop in ArityType.
172 Actually, the situation is worse. Consider
176 Can we eta-expand here? At first the answer looks like "yes of course", but
179 This should diverge! But if we eta-expand, it won't. Again, we ignore this
180 "problem", because being scrupulous would lose an important transformation for
186 Non-recursive newtypes are transparent, and should not get in the way.
187 We do (currently) eta-expand recursive newtypes too. So if we have, say
189 newtype T = MkT ([T] -> Int)
193 where f has arity 1. Then: etaExpandArity e = 1;
194 that is, etaExpandArity looks through the coerce.
196 When we eta-expand e to arity 1: eta_expand 1 e T
197 we want to get: coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
199 HOWEVER, note that if you use coerce bogusly you can ge
201 And since negate has arity 2, you might try to eta expand. But you can't
202 decopose Int to a function type. Hence the final case in eta_expand.
206 -- ^ The Arity returned is the number of value args the
207 -- expression can be applied to without doing much work
208 exprEtaExpandArity :: DynFlags -> CoreExpr -> Arity
209 exprEtaExpandArity dflags e = arityDepth (arityType dflags e)
211 -- A limited sort of function type
212 data ArityType = AFun Bool ArityType -- True <=> one-shot
213 | ATop -- Know nothing
216 arityDepth :: ArityType -> Arity
217 arityDepth (AFun _ ty) = 1 + arityDepth ty
220 andArityType :: ArityType -> ArityType -> ArityType
221 andArityType ABot at2 = at2
222 andArityType ATop _ = ATop
223 andArityType (AFun t1 at1) (AFun t2 at2) = AFun (t1 && t2) (andArityType at1 at2)
224 andArityType at1 at2 = andArityType at2 at1
226 arityType :: DynFlags -> CoreExpr -> ArityType
227 -- (go1 e) = [b1,..,bn]
228 -- means expression can be rewritten \x_b1 -> ... \x_bn -> body
229 -- where bi is True <=> the lambda is one-shot
231 arityType dflags (Note _ e) = arityType dflags e
232 -- Not needed any more: etaExpand is cleverer
233 -- removed: | ok_note n = arityType dflags e
234 -- removed: | otherwise = ATop
236 arityType dflags (Cast e _) = arityType dflags e
239 = mk (idArity v) (arg_tys (idType v))
241 mk :: Arity -> [Type] -> ArityType
242 -- The argument types are only to steer the "state hack"
243 -- Consider case x of
245 -- False -> \(s:RealWorld) -> e
246 -- where foo has arity 1. Then we want the state hack to
247 -- apply to foo too, so we can eta expand the case.
248 mk 0 tys | isBottomingId v = ABot
249 | (ty:_) <- tys, isStateHackType ty = AFun True ATop
251 mk n (ty:tys) = AFun (isStateHackType ty) (mk (n-1) tys)
252 mk n [] = AFun False (mk (n-1) [])
254 arg_tys :: Type -> [Type] -- Ignore for-alls
256 | Just (_, ty') <- splitForAllTy_maybe ty = arg_tys ty'
257 | Just (arg,res) <- splitFunTy_maybe ty = arg : arg_tys res
260 -- Lambdas; increase arity
261 arityType dflags (Lam x e)
262 | isId x = AFun (isOneShotBndr x) (arityType dflags e)
263 | otherwise = arityType dflags e
265 -- Applications; decrease arity
266 arityType dflags (App f (Type _)) = arityType dflags f
267 arityType dflags (App f a)
268 = case arityType dflags f of
269 ABot -> ABot -- If function diverges, ignore argument
270 ATop -> ATop -- No no info about function
272 | exprIsCheap a -> xs
275 -- Case/Let; keep arity if either the expression is cheap
276 -- or it's a 1-shot lambda
277 -- The former is not really right for Haskell
278 -- f x = case x of { (a,b) -> \y. e }
280 -- f x y = case x of { (a,b) -> e }
281 -- The difference is observable using 'seq'
282 arityType dflags (Case scrut _ _ alts)
283 = case foldr1 andArityType [arityType dflags rhs | (_,_,rhs) <- alts] of
284 xs | exprIsCheap scrut -> xs
285 AFun one_shot _ | one_shot -> AFun True ATop
288 arityType dflags (Let b e)
289 = case arityType dflags e of
290 xs | cheap_bind b -> xs
291 AFun one_shot _ | one_shot -> AFun True ATop
294 cheap_bind (NonRec b e) = is_cheap (b,e)
295 cheap_bind (Rec prs) = all is_cheap prs
296 is_cheap (b,e) = (dopt Opt_DictsCheap dflags && isDictId b)
298 -- If the experimental -fdicts-cheap flag is on, we eta-expand through
299 -- dictionary bindings. This improves arities. Thereby, it also
300 -- means that full laziness is less prone to floating out the
301 -- application of a function to its dictionary arguments, which
302 -- can thereby lose opportunities for fusion. Example:
303 -- foo :: Ord a => a -> ...
304 -- foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....
305 -- -- So foo has arity 1
307 -- f = \x. foo dInt $ bar x
309 -- The (foo DInt) is floated out, and makes ineffective a RULE
312 -- One could go further and make exprIsCheap reply True to any
313 -- dictionary-typed expression, but that's more work.
319 %************************************************************************
321 The main eta-expander
323 %************************************************************************
325 IMPORTANT NOTE: The eta expander is careful not to introduce "crap".
326 In particular, given a CoreExpr satisfying the 'CpeRhs' invariant (in
327 CorePrep), it returns a CoreExpr satisfying the same invariant. See
328 Note [Eta expansion and the CorePrep invariants] in CorePrep.
330 This means the eta-expander has to do a bit of on-the-fly
331 simplification but it's not too hard. The alernative, of relying on
332 a subsequent clean-up phase of the Simplifier to de-crapify the result,
333 means you can't really use it in CorePrep, which is painful.
336 -- | @etaExpand n us e ty@ returns an expression with
337 -- the same meaning as @e@, but with arity @n@.
341 -- > e' = etaExpand n us e ty
343 -- We should have that:
345 -- > ty = exprType e = exprType e'
346 etaExpand :: Arity -- ^ Result should have this number of value args
347 -> CoreExpr -- ^ Expression to expand
349 -- Note that SCCs are not treated specially. If we have
350 -- etaExpand 2 (\x -> scc "foo" e)
351 -- = (\xy -> (scc "foo" e) y)
352 -- So the costs of evaluating 'e' (not 'e y') are attributed to "foo"
354 -- etaExpand deals with for-alls. For example:
356 -- where E :: forall a. a -> a
358 -- (/\b. \y::a -> E b y)
360 -- It deals with coerces too, though they are now rare
361 -- so perhaps the extra code isn't worth it
363 etaExpand n orig_expr
364 | manifestArity orig_expr >= n = orig_expr -- The no-op case
368 -- Strip off existing lambdas
370 go n (Lam v body) | isTyVar v = Lam v (go n body)
371 | otherwise = Lam v (go (n-1) body)
372 go n (Note InlineMe expr) = Note InlineMe (go n expr)
373 -- Note [Eta expansion and SCCs]
374 go n (Cast expr co) = Cast (go n expr) co
375 go n expr = -- pprTrace "ee" (vcat [ppr orig_expr, ppr expr, ppr etas]) $
376 etaInfoAbs etas (etaInfoApp subst' expr etas)
378 in_scope = mkInScopeSet (exprFreeVars expr)
379 (in_scope', etas) = mkEtaWW n in_scope (exprType expr)
380 subst' = mkEmptySubst in_scope'
384 data EtaInfo = EtaVar Var -- /\a. [], [] a
386 | EtaCo Coercion -- [] |> co, [] |> (sym co)
388 instance Outputable EtaInfo where
389 ppr (EtaVar v) = ptext (sLit "EtaVar") <+> ppr v
390 ppr (EtaCo co) = ptext (sLit "EtaCo") <+> ppr co
392 pushCoercion :: Coercion -> [EtaInfo] -> [EtaInfo]
393 pushCoercion co1 (EtaCo co2 : eis)
394 | isIdentityCoercion co = eis
395 | otherwise = EtaCo co : eis
397 co = co1 `mkTransCoercion` co2
399 pushCoercion co eis = EtaCo co : eis
402 etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr
403 etaInfoAbs [] expr = expr
404 etaInfoAbs (EtaVar v : eis) expr = Lam v (etaInfoAbs eis expr)
405 etaInfoAbs (EtaCo co : eis) expr = Cast (etaInfoAbs eis expr) (mkSymCoercion co)
408 etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
409 -- (etaInfoApp s e eis) returns something equivalent to
410 -- ((substExpr s e) `appliedto` eis)
412 etaInfoApp subst (Lam v1 e) (EtaVar v2 : eis)
413 = etaInfoApp subst' e eis
415 subst' | isTyVar v1 = CoreSubst.extendTvSubst subst v1 (mkTyVarTy v2)
416 | otherwise = CoreSubst.extendIdSubst subst v1 (Var v2)
418 etaInfoApp subst (Cast e co1) eis
419 = etaInfoApp subst e (pushCoercion co' eis)
421 co' = CoreSubst.substTy subst co1
423 etaInfoApp subst (Case e b _ alts) eis
424 = Case (subst_expr subst e) b1 (coreAltsType alts') alts'
426 (subst1, b1) = substBndr subst b
427 alts' = map subst_alt alts
428 subst_alt (con, bs, rhs) = (con, bs', etaInfoApp subst2 rhs eis)
430 (subst2,bs') = substBndrs subst1 bs
432 etaInfoApp subst (Let b e) eis
433 = Let b' (etaInfoApp subst' e eis)
435 (subst', b') = subst_bind subst b
437 etaInfoApp subst (Note note e) eis
438 = Note note (etaInfoApp subst e eis)
440 etaInfoApp subst e eis
441 = go (subst_expr subst e) eis
444 go e (EtaVar v : eis) = go (App e (varToCoreExpr v)) eis
445 go e (EtaCo co : eis) = go (Cast e co) eis
448 mkEtaWW :: Arity -> InScopeSet -> Type
449 -> (InScopeSet, [EtaInfo])
450 -- EtaInfo contains fresh variables,
451 -- not free in the incoming CoreExpr
452 -- Outgoing InScopeSet includes the EtaInfo vars
453 -- and the original free vars
455 mkEtaWW n in_scope ty
456 = go n empty_subst ty []
458 empty_subst = mkTvSubst in_scope emptyTvSubstEnv
462 = (getTvInScope subst, reverse eis)
464 | Just (tv,ty') <- splitForAllTy_maybe ty
465 , let (subst', tv') = substTyVarBndr subst tv
466 -- Avoid free vars of the original expression
467 = go n subst' ty' (EtaVar tv' : eis)
469 | Just (arg_ty, res_ty) <- splitFunTy_maybe ty
470 , let (subst', eta_id') = freshEtaId n subst arg_ty
471 -- Avoid free vars of the original expression
472 = go (n-1) subst' res_ty (EtaVar eta_id' : eis)
474 | Just(ty',co) <- splitNewTypeRepCo_maybe ty
476 -- newtype T = MkT ([T] -> Int)
477 -- Consider eta-expanding this
480 -- coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
481 go n subst ty' (EtaCo (substTy subst co) : eis)
483 | otherwise -- We have an expression of arity > 0,
484 = (getTvInScope subst, reverse eis) -- but its type isn't a function.
485 -- This *can* legitmately happen:
486 -- e.g. coerce Int (\x. x) Essentially the programmer is
487 -- playing fast and loose with types (Happy does this a lot).
488 -- So we simply decline to eta-expand. Otherwise we'd end up
489 -- with an explicit lambda having a non-function type
493 -- Avoiding unnecessary substitution
495 subst_expr :: Subst -> CoreExpr -> CoreExpr
496 subst_expr s e | isEmptySubst s = e
497 | otherwise = substExpr s e
499 subst_bind :: Subst -> CoreBind -> (Subst, CoreBind)
500 subst_bind subst (NonRec b r)
501 = (subst', NonRec b' (subst_expr subst r))
503 (subst', b') = substBndr subst b
504 subst_bind subst (Rec prs)
505 = (subst', Rec (bs1 `zip` map (subst_expr subst') rhss))
507 (bs, rhss) = unzip prs
508 (subst', bs1) = substBndrs subst bs
512 freshEtaId :: Int -> TvSubst -> Type -> (TvSubst, Id)
513 -- Make a fresh Id, with specified type (after applying substitution)
514 -- It should be "fresh" in the sense that it's not in the in-scope set
515 -- of the TvSubstEnv; and it should itself then be added to the in-scope
516 -- set of the TvSubstEnv
518 -- The Int is just a reasonable starting point for generating a unique;
519 -- it does not necessarily have to be unique itself.
520 freshEtaId n subst ty
523 ty' = substTy subst ty
524 eta_id' = uniqAway (getTvInScope subst) $
525 mkSysLocal (fsLit "eta") (mkBuiltinUnique n) ty'
526 subst' = extendTvInScope subst [eta_id']