-%************************************************************************
-%* *
-\subsection{Eta reduction and expansion}
-%* *
-%************************************************************************
-
-\begin{code}
-exprEtaExpandArity :: DynFlags -> CoreExpr -> Arity
-{- The Arity returned is the number of value args the
- thing can be applied to without doing much work
-
-exprEtaExpandArity is used when eta expanding
- e ==> \xy -> e x y
-
-It returns 1 (or more) to:
- case x of p -> \s -> ...
-because for I/O ish things we really want to get that \s to the top.
-We are prepared to evaluate x each time round the loop in order to get that
-
-It's all a bit more subtle than it looks:
-
-1. One-shot lambdas
-
-Consider one-shot lambdas
- let x = expensive in \y z -> E
-We want this to have arity 2 if the \y-abstraction is a 1-shot lambda
-Hence the ArityType returned by arityType
-
-2. The state-transformer hack
-
-The one-shot lambda special cause is particularly important/useful for
-IO state transformers, where we often get
- let x = E in \ s -> ...
-
-and the \s is a real-world state token abstraction. Such abstractions
-are almost invariably 1-shot, so we want to pull the \s out, past the
-let x=E, even if E is expensive. So we treat state-token lambdas as
-one-shot even if they aren't really. The hack is in Id.isOneShotBndr.
-
-3. Dealing with bottom
-
-Consider also
- f = \x -> error "foo"
-Here, arity 1 is fine. But if it is
- f = \x -> case x of
- True -> error "foo"
- False -> \y -> x+y
-then we want to get arity 2. Tecnically, this isn't quite right, because
- (f True) `seq` 1
-should diverge, but it'll converge if we eta-expand f. Nevertheless, we
-do so; it improves some programs significantly, and increasing convergence
-isn't a bad thing. Hence the ABot/ATop in ArityType.
-
-Actually, the situation is worse. Consider
- f = \x -> case x of
- True -> \y -> x+y
- False -> \y -> x-y
-Can we eta-expand here? At first the answer looks like "yes of course", but
-consider
- (f bot) `seq` 1
-This should diverge! But if we eta-expand, it won't. Again, we ignore this
-"problem", because being scrupulous would lose an important transformation for
-many programs.
-
-
-4. Newtypes
-
-Non-recursive newtypes are transparent, and should not get in the way.
-We do (currently) eta-expand recursive newtypes too. So if we have, say
-
- newtype T = MkT ([T] -> Int)
-
-Suppose we have
- e = coerce T f
-where f has arity 1. Then: etaExpandArity e = 1;
-that is, etaExpandArity looks through the coerce.
-
-When we eta-expand e to arity 1: eta_expand 1 e T
-we want to get: coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
-
-HOWEVER, note that if you use coerce bogusly you can ge
- coerce Int negate
-And since negate has arity 2, you might try to eta expand. But you can't
-decopose Int to a function type. Hence the final case in eta_expand.
--}
-
-
-exprEtaExpandArity dflags e = arityDepth (arityType dflags e)
-
--- A limited sort of function type
-data ArityType = AFun Bool ArityType -- True <=> one-shot
- | ATop -- Know nothing
- | ABot -- Diverges
-
-arityDepth :: ArityType -> Arity
-arityDepth (AFun _ ty) = 1 + arityDepth ty
-arityDepth ty = 0
-
-andArityType ABot at2 = at2
-andArityType ATop at2 = ATop
-andArityType (AFun t1 at1) (AFun t2 at2) = AFun (t1 && t2) (andArityType at1 at2)
-andArityType at1 at2 = andArityType at2 at1
-
-arityType :: DynFlags -> CoreExpr -> ArityType
- -- (go1 e) = [b1,..,bn]
- -- means expression can be rewritten \x_b1 -> ... \x_bn -> body
- -- where bi is True <=> the lambda is one-shot
-
-arityType dflags (Note n e) = arityType dflags e
--- Not needed any more: etaExpand is cleverer
--- | ok_note n = arityType dflags e
--- | otherwise = ATop
-
-arityType dflags (Cast e co) = arityType dflags e
-
-arityType dflags (Var v)
- = mk (idArity v) (arg_tys (idType v))
- where
- mk :: Arity -> [Type] -> ArityType
- -- The argument types are only to steer the "state hack"
- -- Consider case x of
- -- True -> foo
- -- False -> \(s:RealWorld) -> e
- -- where foo has arity 1. Then we want the state hack to
- -- apply to foo too, so we can eta expand the case.
- mk 0 tys | isBottomingId v = ABot
- | (ty:tys) <- tys, isStateHackType ty = AFun True ATop
- | otherwise = ATop
- mk n (ty:tys) = AFun (isStateHackType ty) (mk (n-1) tys)
- mk n [] = AFun False (mk (n-1) [])
-
- arg_tys :: Type -> [Type] -- Ignore for-alls
- arg_tys ty
- | Just (_, ty') <- splitForAllTy_maybe ty = arg_tys ty'
- | Just (arg,res) <- splitFunTy_maybe ty = arg : arg_tys res
- | otherwise = []
-
- -- Lambdas; increase arity
-arityType dflags (Lam x e)
- | isId x = AFun (isOneShotBndr x) (arityType dflags e)
- | otherwise = arityType dflags e
-
- -- Applications; decrease arity
-arityType dflags (App f (Type _)) = arityType dflags f
-arityType dflags (App f a) = case arityType dflags f of
- AFun one_shot xs | exprIsCheap a -> xs
- other -> ATop
-
- -- Case/Let; keep arity if either the expression is cheap
- -- or it's a 1-shot lambda
- -- The former is not really right for Haskell
- -- f x = case x of { (a,b) -> \y. e }
- -- ===>
- -- f x y = case x of { (a,b) -> e }
- -- The difference is observable using 'seq'
-arityType dflags (Case scrut _ _ alts)
- = case foldr1 andArityType [arityType dflags rhs | (_,_,rhs) <- alts] of
- xs | exprIsCheap scrut -> xs
- xs@(AFun one_shot _) | one_shot -> AFun True ATop
- other -> ATop
-
-arityType dflags (Let b e)
- = case arityType dflags e of
- xs | cheap_bind b -> xs
- xs@(AFun one_shot _) | one_shot -> AFun True ATop
- other -> ATop
- where
- cheap_bind (NonRec b e) = is_cheap (b,e)
- cheap_bind (Rec prs) = all is_cheap prs
- is_cheap (b,e) = (dopt Opt_DictsCheap dflags && isDictId b)
- || exprIsCheap e
- -- If the experimental -fdicts-cheap flag is on, we eta-expand through
- -- dictionary bindings. This improves arities. Thereby, it also
- -- means that full laziness is less prone to floating out the
- -- application of a function to its dictionary arguments, which
- -- can thereby lose opportunities for fusion. Example:
- -- foo :: Ord a => a -> ...
- -- foo = /\a \(d:Ord a). let d' = ...d... in \(x:a). ....
- -- -- So foo has arity 1
- --
- -- f = \x. foo dInt $ bar x
- --
- -- The (foo DInt) is floated out, and makes ineffective a RULE
- -- foo (bar x) = ...
- --
- -- One could go further and make exprIsCheap reply True to any
- -- dictionary-typed expression, but that's more work.
-
-arityType dflags other = ATop
-
-{- NOT NEEDED ANY MORE: etaExpand is cleverer
-ok_note InlineMe = False
-ok_note other = True
- -- Notice that we do not look through __inline_me__
- -- This may seem surprising, but consider
- -- f = _inline_me (\x -> e)
- -- We DO NOT want to eta expand this to
- -- f = \x -> (_inline_me (\x -> e)) x
- -- because the _inline_me gets dropped now it is applied,
- -- giving just
- -- f = \x -> e
- -- A Bad Idea
--}
-\end{code}
-
-
-\begin{code}
-etaExpand :: Arity -- Result should have this number of value args
- -> [Unique]
- -> CoreExpr -> Type -- Expression and its type
- -> CoreExpr
--- (etaExpand n us e ty) returns an expression with
--- the same meaning as 'e', but with arity 'n'.
---
--- Given e' = etaExpand n us e ty
--- We should have
--- ty = exprType e = exprType e'
---
--- Note that SCCs are not treated specially. If we have
--- etaExpand 2 (\x -> scc "foo" e)
--- = (\xy -> (scc "foo" e) y)
--- So the costs of evaluating 'e' (not 'e y') are attributed to "foo"
-
-etaExpand n us expr ty
- | manifestArity expr >= n = expr -- The no-op case
- | otherwise
- = eta_expand n us expr ty
- where
-
--- manifestArity sees how many leading value lambdas there are
-manifestArity :: CoreExpr -> Arity
-manifestArity (Lam v e) | isId v = 1 + manifestArity e
- | otherwise = manifestArity e
-manifestArity (Note _ e) = manifestArity e
-manifestArity (Cast e _) = manifestArity e
-manifestArity e = 0
-
--- etaExpand deals with for-alls. For example:
--- etaExpand 1 E
--- where E :: forall a. a -> a
--- would return
--- (/\b. \y::a -> E b y)
---
--- It deals with coerces too, though they are now rare
--- so perhaps the extra code isn't worth it
-
-eta_expand n us expr ty
- | n == 0 &&
- -- The ILX code generator requires eta expansion for type arguments
- -- too, but alas the 'n' doesn't tell us how many of them there
- -- may be. So we eagerly eta expand any big lambdas, and just
- -- cross our fingers about possible loss of sharing in the ILX case.
- -- The Right Thing is probably to make 'arity' include
- -- type variables throughout the compiler. (ToDo.)
- not (isForAllTy ty)
- -- Saturated, so nothing to do
- = expr
-
- -- Short cut for the case where there already
- -- is a lambda; no point in gratuitously adding more
-eta_expand n us (Lam v body) ty
- | isTyVar v
- = Lam v (eta_expand n us body (applyTy ty (mkTyVarTy v)))
-
- | otherwise
- = Lam v (eta_expand (n-1) us body (funResultTy ty))
-
--- We used to have a special case that stepped inside Coerces here,
--- thus: eta_expand n us (Note note@(Coerce _ ty) e) _
--- = Note note (eta_expand n us e ty)
--- BUT this led to an infinite loop
--- Example: newtype T = MkT (Int -> Int)
--- eta_expand 1 (coerce (Int->Int) e)
--- --> coerce (Int->Int) (eta_expand 1 T e)
--- by the bogus eqn
--- --> coerce (Int->Int) (coerce T
--- (\x::Int -> eta_expand 1 (coerce (Int->Int) e)))
--- by the splitNewType_maybe case below
--- and round we go
-
-eta_expand n us expr ty
- = ASSERT2 (exprType expr `coreEqType` ty, ppr (exprType expr) $$ ppr ty)
- case splitForAllTy_maybe ty of {
- Just (tv,ty') ->
-
- Lam lam_tv (eta_expand n us2 (App expr (Type (mkTyVarTy lam_tv))) (substTyWith [tv] [mkTyVarTy lam_tv] ty'))
- where
- lam_tv = mkTyVar (mkSysTvName uniq FSLIT("etaT")) (tyVarKind tv)
- (uniq:us2) = us
- ; Nothing ->
-
- case splitFunTy_maybe ty of {
- Just (arg_ty, res_ty) -> Lam arg1 (eta_expand (n-1) us2 (App expr (Var arg1)) res_ty)
- where
- arg1 = mkSysLocal FSLIT("eta") uniq arg_ty
- (uniq:us2) = us
-
- ; Nothing ->
-
- -- Given this:
- -- newtype T = MkT ([T] -> Int)
- -- Consider eta-expanding this
- -- eta_expand 1 e T
- -- We want to get
- -- coerce T (\x::[T] -> (coerce ([T]->Int) e) x)
-
- case splitNewTypeRepCo_maybe ty of {
- Just(ty1,co) ->
- mkCoerce co (eta_expand n us (mkCoerce (mkSymCoercion co) expr) ty1) ;
- Nothing ->
-
- -- We have an expression of arity > 0, but its type isn't a function
- -- This *can* legitmately happen: e.g. coerce Int (\x. x)
- -- Essentially the programmer is playing fast and loose with types
- -- (Happy does this a lot). So we simply decline to eta-expand.
- expr
- }}}
-\end{code}
-
-exprArity is a cheap-and-cheerful version of exprEtaExpandArity.
-It tells how many things the expression can be applied to before doing
-any work. It doesn't look inside cases, lets, etc. The idea is that
-exprEtaExpandArity will do the hard work, leaving something that's easy
-for exprArity to grapple with. In particular, Simplify uses exprArity to
-compute the ArityInfo for the Id.
-
-Originally I thought that it was enough just to look for top-level lambdas, but
-it isn't. I've seen this
-
- foo = PrelBase.timesInt
-
-We want foo to get arity 2 even though the eta-expander will leave it
-unchanged, in the expectation that it'll be inlined. But occasionally it
-isn't, because foo is blacklisted (used in a rule).
-
-Similarly, see the ok_note check in exprEtaExpandArity. So
- f = __inline_me (\x -> e)
-won't be eta-expanded.
-
-And in any case it seems more robust to have exprArity be a bit more intelligent.
-But note that (\x y z -> f x y z)
-should have arity 3, regardless of f's arity.
-
-\begin{code}
-exprArity :: CoreExpr -> Arity
-exprArity e = go e
- where
- go (Var v) = idArity v
- go (Lam x e) | isId x = go e + 1
- | otherwise = go e
- go (Note n e) = go e
- go (Cast e _) = go e
- go (App e (Type t)) = go e
- go (App f a) | exprIsCheap a = (go f - 1) `max` 0
- -- NB: exprIsCheap a!
- -- f (fac x) does not have arity 2,
- -- even if f has arity 3!
- -- NB: `max 0`! (\x y -> f x) has arity 2, even if f is
- -- unknown, hence arity 0
- go _ = 0