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 )
38 import GHC.Exts -- For `xori`
41 %************************************************************************
43 manifestArity and exprArity
45 %************************************************************************
47 exprArity is a cheap-and-cheerful version of exprEtaExpandArity.
48 It tells how many things the expression can be applied to before doing
49 any work. It doesn't look inside cases, lets, etc. The idea is that
50 exprEtaExpandArity will do the hard work, leaving something that's easy
51 for exprArity to grapple with. In particular, Simplify uses exprArity to
52 compute the ArityInfo for the Id.
54 Originally I thought that it was enough just to look for top-level lambdas, but
55 it isn't. I've seen this
57 foo = PrelBase.timesInt
59 We want foo to get arity 2 even though the eta-expander will leave it
60 unchanged, in the expectation that it'll be inlined. But occasionally it
61 isn't, because foo is blacklisted (used in a rule).
63 Similarly, see the ok_note check in exprEtaExpandArity. So
64 f = __inline_me (\x -> e)
65 won't be eta-expanded.
67 And in any case it seems more robust to have exprArity be a bit more intelligent.
68 But note that (\x y z -> f x y z)
69 should have arity 3, regardless of f's arity.
71 Note [exprArity invariant]
72 ~~~~~~~~~~~~~~~~~~~~~~~~~~
73 exprArity has the following invariant:
74 (exprArity e) = n, then manifestArity (etaExpand e n) = n
76 That is, if exprArity says "the arity is n" then etaExpand really can get
77 "n" manifest lambdas to the top.
79 Why is this important? Because
80 - In TidyPgm we use exprArity to fix the *final arity* of
81 each top-level Id, and in
82 - In CorePrep we use etaExpand on each rhs, so that the visible lambdas
83 actually match that arity, which in turn means
84 that the StgRhs has the right number of lambdas
86 An alternative would be to do the eta-expansion in TidyPgm, at least
87 for top-level bindings, in which case we would not need the trim_arity
88 in exprArity. That is a less local change, so I'm going to leave it for today!
92 manifestArity :: CoreExpr -> Arity
93 -- ^ manifestArity sees how many leading value lambdas there are
94 manifestArity (Lam v e) | isId v = 1 + manifestArity e
95 | otherwise = manifestArity e
96 manifestArity (Note _ e) = manifestArity e
97 manifestArity (Cast e _) = manifestArity e
100 exprArity :: CoreExpr -> Arity
101 -- ^ An approximate, fast, version of 'exprEtaExpandArity'
104 go (Var v) = idArity v
105 go (Lam x e) | isId x = go e + 1
108 go (Cast e co) = trim_arity (go e) 0 (snd (coercionKind co))
109 go (App e (Type _)) = go e
110 go (App f a) | exprIsCheap a = (go f - 1) `max` 0
111 -- NB: exprIsCheap a!
112 -- f (fac x) does not have arity 2,
113 -- even if f has arity 3!
114 -- NB: `max 0`! (\x y -> f x) has arity 2, even if f is
115 -- unknown, hence arity 0
118 -- Note [exprArity invariant]
121 | Just (_, ty') <- splitForAllTy_maybe ty = trim_arity n a ty'
122 | Just (_, ty') <- splitFunTy_maybe ty = trim_arity n (a+1) ty'
123 | Just (ty',_) <- splitNewTypeRepCo_maybe ty = trim_arity n a ty'
127 %************************************************************************
129 \subsection{Eta reduction and expansion}
131 %************************************************************************
133 exprEtaExpandArity is used when eta expanding
136 It returns 1 (or more) to:
137 case x of p -> \s -> ...
138 because for I/O ish things we really want to get that \s to the top.
139 We are prepared to evaluate x each time round the loop in order to get that
141 It's all a bit more subtle than it looks:
145 Consider one-shot lambdas
146 let x = expensive in \y z -> E
147 We want this to have arity 2 if the \y-abstraction is a 1-shot lambda
148 Hence the ArityType returned by arityType
150 2. The state-transformer hack
152 The one-shot lambda special cause is particularly important/useful for
153 IO state transformers, where we often get
154 let x = E in \ s -> ...
156 and the \s is a real-world state token abstraction. Such abstractions
157 are almost invariably 1-shot, so we want to pull the \s out, past the
158 let x=E, even if E is expensive. So we treat state-token lambdas as
159 one-shot even if they aren't really. The hack is in Id.isOneShotBndr.
161 3. Dealing with bottom
164 f = \x -> error "foo"
165 Here, arity 1 is fine. But if it is
169 then we want to get arity 2. Tecnically, this isn't quite right, because
171 should diverge, but it'll converge if we eta-expand f. Nevertheless, we
172 do so; it improves some programs significantly, and increasing convergence
173 isn't a bad thing. Hence the ABot/ATop in ArityType.
175 Actually, the situation is worse. Consider
179 Can we eta-expand here? At first the answer looks like "yes of course", but
182 This should diverge! But if we eta-expand, it won't. Again, we ignore this
183 "problem", because being scrupulous would lose an important transformation for
189 Non-recursive newtypes are transparent, and should not get in the way.
190 We do (currently) eta-expand recursive newtypes too. So if we have, say
192 newtype T = MkT ([T] -> Int)
196 where f has arity 1. Then: etaExpandArity e = 1;
197 that is, etaExpandArity looks through the coerce.
199 When we eta-expand e to arity 1: eta_expand 1 e T
200 we want to get: coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
202 HOWEVER, note that if you use coerce bogusly you can ge
204 And since negate has arity 2, you might try to eta expand. But you can't
205 decopose Int to a function type. Hence the final case in eta_expand.
209 -- ^ The Arity returned is the number of value args the
210 -- expression can be applied to without doing much work
211 exprEtaExpandArity :: DynFlags -> CoreExpr -> Arity
212 exprEtaExpandArity dflags e = arityDepth (arityType dflags e)
214 -- A limited sort of function type
215 data ArityType = AFun Bool ArityType -- True <=> one-shot
216 | ATop -- Know nothing
219 arityDepth :: ArityType -> Arity
220 arityDepth (AFun _ ty) = 1 + arityDepth ty
223 andArityType :: ArityType -> ArityType -> ArityType
224 andArityType ABot at2 = at2
225 andArityType ATop _ = ATop
226 andArityType (AFun t1 at1) (AFun t2 at2) = AFun (t1 && t2) (andArityType at1 at2)
227 andArityType at1 at2 = andArityType at2 at1
229 arityType :: DynFlags -> CoreExpr -> ArityType
230 -- (go1 e) = [b1,..,bn]
231 -- means expression can be rewritten \x_b1 -> ... \x_bn -> body
232 -- where bi is True <=> the lambda is one-shot
234 arityType dflags (Note _ e) = arityType dflags e
235 -- Not needed any more: etaExpand is cleverer
236 -- removed: | ok_note n = arityType dflags e
237 -- removed: | otherwise = ATop
239 arityType dflags (Cast e _) = arityType dflags e
242 = mk (idArity v) (arg_tys (idType v))
244 mk :: Arity -> [Type] -> ArityType
245 -- The argument types are only to steer the "state hack"
246 -- Consider case x of
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.
251 mk 0 tys | isBottomingId v = ABot
252 | (ty:_) <- tys, isStateHackType ty = AFun True ATop
254 mk n (ty:tys) = AFun (isStateHackType ty) (mk (n-1) tys)
255 mk n [] = AFun False (mk (n-1) [])
257 arg_tys :: Type -> [Type] -- Ignore for-alls
259 | Just (_, ty') <- splitForAllTy_maybe ty = arg_tys ty'
260 | Just (arg,res) <- splitFunTy_maybe ty = arg : arg_tys res
263 -- Lambdas; increase arity
264 arityType dflags (Lam x e)
265 | isId x = AFun (isOneShotBndr x) (arityType dflags e)
266 | otherwise = arityType dflags e
268 -- Applications; decrease arity
269 arityType dflags (App f (Type _)) = arityType dflags f
270 arityType dflags (App f a)
271 = case arityType dflags f of
272 ABot -> ABot -- If function diverges, ignore argument
273 ATop -> ATop -- No no info about function
275 | exprIsCheap a -> xs
278 -- Case/Let; keep arity if either the expression is cheap
279 -- or it's a 1-shot lambda
280 -- The former is not really right for Haskell
281 -- f x = case x of { (a,b) -> \y. e }
283 -- f x y = case x of { (a,b) -> e }
284 -- The difference is observable using 'seq'
285 arityType dflags (Case scrut _ _ alts)
286 = case foldr1 andArityType [arityType dflags rhs | (_,_,rhs) <- alts] of
287 xs | exprIsCheap scrut -> xs
288 AFun one_shot _ | one_shot -> AFun True ATop
291 arityType dflags (Let b e)
292 = case arityType dflags e of
293 xs | cheap_bind b -> xs
294 AFun one_shot _ | one_shot -> AFun True ATop
297 cheap_bind (NonRec b e) = is_cheap (b,e)
298 cheap_bind (Rec prs) = all is_cheap prs
299 is_cheap (b,e) = (dopt Opt_DictsCheap dflags && isDictId b)
301 -- If the experimental -fdicts-cheap flag is on, we eta-expand through
302 -- dictionary bindings. This improves arities. Thereby, it also
303 -- means that full laziness is less prone to floating out the
304 -- application of a function to its dictionary arguments, which
305 -- can thereby lose opportunities for fusion. Example:
306 -- foo :: Ord a => a -> ...
307 -- foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....
308 -- -- So foo has arity 1
310 -- f = \x. foo dInt $ bar x
312 -- The (foo DInt) is floated out, and makes ineffective a RULE
315 -- One could go further and make exprIsCheap reply True to any
316 -- dictionary-typed expression, but that's more work.
322 %************************************************************************
324 The main eta-expander
326 %************************************************************************
328 IMPORTANT NOTE: The eta expander is careful not to introduce "crap".
329 In particular, given a CoreExpr satisfying the 'CpeRhs' invariant (in
330 CorePrep), it returns a CoreExpr satisfying the same invariant. See
331 Note [Eta expansion and the CorePrep invariants] in CorePrep.
333 This means the eta-expander has to do a bit of on-the-fly
334 simplification but it's not too hard. The alernative, of relying on
335 a subsequent clean-up phase of the Simplifier to de-crapify the result,
336 means you can't really use it in CorePrep, which is painful.
339 -- | @etaExpand n us e ty@ returns an expression with
340 -- the same meaning as @e@, but with arity @n@.
344 -- > e' = etaExpand n us e ty
346 -- We should have that:
348 -- > ty = exprType e = exprType e'
349 etaExpand :: Arity -- ^ Result should have this number of value args
350 -> CoreExpr -- ^ Expression to expand
352 -- Note that SCCs are not treated specially. If we have
353 -- etaExpand 2 (\x -> scc "foo" e)
354 -- = (\xy -> (scc "foo" e) y)
355 -- So the costs of evaluating 'e' (not 'e y') are attributed to "foo"
357 -- etaExpand deals with for-alls. For example:
359 -- where E :: forall a. a -> a
361 -- (/\b. \y::a -> E b y)
363 -- It deals with coerces too, though they are now rare
364 -- so perhaps the extra code isn't worth it
366 etaExpand n orig_expr
367 | manifestArity orig_expr >= n = orig_expr -- The no-op case
371 -- Strip off existing lambdas
373 go n (Lam v body) | isTyVar v = Lam v (go n body)
374 | otherwise = Lam v (go (n-1) body)
375 go n (Note InlineMe expr) = Note InlineMe (go n expr)
376 -- Note [Eta expansion and SCCs]
377 go n (Cast expr co) = Cast (go n expr) co
378 go n expr = -- pprTrace "ee" (vcat [ppr orig_expr, ppr expr, ppr etas]) $
379 etaInfoAbs etas (etaInfoApp subst' expr etas)
381 in_scope = mkInScopeSet (exprFreeVars expr)
382 (in_scope', etas) = mkEtaWW n in_scope (exprType expr)
383 subst' = mkEmptySubst in_scope'
387 data EtaInfo = EtaVar Var -- /\a. [], [] a
389 | EtaCo Coercion -- [] |> co, [] |> (sym co)
391 instance Outputable EtaInfo where
392 ppr (EtaVar v) = ptext (sLit "EtaVar") <+> ppr v
393 ppr (EtaCo co) = ptext (sLit "EtaCo") <+> ppr co
395 pushCoercion :: Coercion -> [EtaInfo] -> [EtaInfo]
396 pushCoercion co1 (EtaCo co2 : eis)
397 | isIdentityCoercion co = eis
398 | otherwise = EtaCo co : eis
400 co = co1 `mkTransCoercion` co2
402 pushCoercion co eis = EtaCo co : eis
405 etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr
406 etaInfoAbs [] expr = expr
407 etaInfoAbs (EtaVar v : eis) expr = Lam v (etaInfoAbs eis expr)
408 etaInfoAbs (EtaCo co : eis) expr = Cast (etaInfoAbs eis expr) (mkSymCoercion co)
411 etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
412 -- (etaInfoApp s e eis) returns something equivalent to
413 -- ((substExpr s e) `appliedto` eis)
415 etaInfoApp subst (Lam v1 e) (EtaVar v2 : eis)
416 = etaInfoApp subst' e eis
418 subst' | isTyVar v1 = CoreSubst.extendTvSubst subst v1 (mkTyVarTy v2)
419 | otherwise = CoreSubst.extendIdSubst subst v1 (Var v2)
421 etaInfoApp subst (Cast e co1) eis
422 = etaInfoApp subst e (pushCoercion co' eis)
424 co' = CoreSubst.substTy subst co1
426 etaInfoApp subst (Case e b _ alts) eis
427 = Case (subst_expr subst e) b1 (coreAltsType alts') alts'
429 (subst1, b1) = substBndr subst b
430 alts' = map subst_alt alts
431 subst_alt (con, bs, rhs) = (con, bs', etaInfoApp subst2 rhs eis)
433 (subst2,bs') = substBndrs subst1 bs
435 etaInfoApp subst (Let b e) eis
436 = Let b' (etaInfoApp subst' e eis)
438 (subst', b') = subst_bind subst b
440 etaInfoApp subst (Note note e) eis
441 = Note note (etaInfoApp subst e eis)
443 etaInfoApp subst e eis
444 = go (subst_expr subst e) eis
447 go e (EtaVar v : eis) = go (App e (varToCoreExpr v)) eis
448 go e (EtaCo co : eis) = go (Cast e co) eis
451 mkEtaWW :: Arity -> InScopeSet -> Type
452 -> (InScopeSet, [EtaInfo])
453 -- EtaInfo contains fresh variables,
454 -- not free in the incoming CoreExpr
455 -- Outgoing InScopeSet includes the EtaInfo vars
456 -- and the original free vars
458 mkEtaWW n in_scope ty
459 = go n empty_subst ty []
461 empty_subst = mkTvSubst in_scope emptyTvSubstEnv
465 = (getTvInScope subst, reverse eis)
467 | Just (tv,ty') <- splitForAllTy_maybe ty
468 , let (subst', tv') = substTyVarBndr subst tv
469 -- Avoid free vars of the original expression
470 = go n subst' ty' (EtaVar tv' : eis)
472 | Just (arg_ty, res_ty) <- splitFunTy_maybe ty
473 , let (subst', eta_id') = freshEtaId n subst arg_ty
474 -- Avoid free vars of the original expression
475 = go (n-1) subst' res_ty (EtaVar eta_id' : eis)
477 | Just(ty',co) <- splitNewTypeRepCo_maybe ty
479 -- newtype T = MkT ([T] -> Int)
480 -- Consider eta-expanding this
483 -- coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
484 go n subst ty' (EtaCo (substTy subst co) : eis)
486 | otherwise -- We have an expression of arity > 0,
487 = (getTvInScope subst, reverse eis) -- but its type isn't a function.
488 -- This *can* legitmately happen:
489 -- e.g. coerce Int (\x. x) Essentially the programmer is
490 -- playing fast and loose with types (Happy does this a lot).
491 -- So we simply decline to eta-expand. Otherwise we'd end up
492 -- with an explicit lambda having a non-function type
496 -- Avoiding unnecessary substitution
498 subst_expr :: Subst -> CoreExpr -> CoreExpr
499 subst_expr s e | isEmptySubst s = e
500 | otherwise = substExpr s e
502 subst_bind :: Subst -> CoreBind -> (Subst, CoreBind)
503 subst_bind subst (NonRec b r)
504 = (subst', NonRec b' (subst_expr subst r))
506 (subst', b') = substBndr subst b
507 subst_bind subst (Rec prs)
508 = (subst', Rec (bs1 `zip` map (subst_expr subst') rhss))
510 (bs, rhss) = unzip prs
511 (subst', bs1) = substBndrs subst bs
515 freshEtaId :: Int -> TvSubst -> Type -> (TvSubst, Id)
516 -- Make a fresh Id, with specified type (after applying substitution)
517 -- It should be "fresh" in the sense that it's not in the in-scope set
518 -- of the TvSubstEnv; and it should itself then be added to the in-scope
519 -- set of the TvSubstEnv
521 -- The Int is just a reasonable starting point for generating a unique;
522 -- it does not necessarily have to be unique itself.
523 freshEtaId n subst ty
526 ty' = substTy subst ty
527 eta_id' = uniqAway (getTvInScope subst) $
528 mkSysLocal (fsLit "eta") (mkBuiltinUnique n) ty'
529 subst' = extendTvInScope subst [eta_id']