%
% (c) The University of Glasgow 2006
% (c) The GRASP/AQUA Project, Glasgow University, 19921998
%
Arity and ete expansion
\begin{code}
module CoreArity (
manifestArity, exprArity,
exprEtaExpandArity, etaExpand
) where
#include "HsVersions.h"
import CoreSyn
import CoreFVs
import CoreUtils
import NewDemand
import TyCon ( isRecursiveTyCon )
import qualified CoreSubst
import CoreSubst ( Subst, substBndr, substBndrs, substExpr
, mkEmptySubst, isEmptySubst )
import Var
import VarEnv
import Id
import Type
import TcType ( isDictLikeTy )
import Coercion
import BasicTypes
import Unique
import Outputable
import DynFlags
import StaticFlags ( opt_NoStateHack )
import FastString
\end{code}
%************************************************************************
%* *
manifestArity and exprArity
%* *
%************************************************************************
exprArity is a cheapandcheerful 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 toplevel lambdas, but
it isn't. I've seen this
foo = PrelBase.timesInt
We want foo to get arity 2 even though the etaexpander 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 etaexpanded.
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.
Note [exprArity invariant]
~~~~~~~~~~~~~~~~~~~~~~~~~~
exprArity has the following invariant:
(exprArity e) = n, then manifestArity (etaExpand e n) = n
That is, if exprArity says "the arity is n" then etaExpand really can get
"n" manifest lambdas to the top.
Why is this important? Because
In TidyPgm we use exprArity to fix the *final arity* of
each toplevel Id, and in
In CorePrep we use etaExpand on each rhs, so that the visible lambdas
actually match that arity, which in turn means
that the StgRhs has the right number of lambdas
An alternative would be to do the etaexpansion in TidyPgm, at least
for toplevel bindings, in which case we would not need the trim_arity
in exprArity. That is a less local change, so I'm going to leave it for today!
\begin{code}
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 _ = 0
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 _ e) = go e
go (Cast e co) = trim_arity (go e) 0 (snd (coercionKind co))
go (App e (Type _)) = go e
go (App f a) | exprIsCheap a = (go f 1) `max` 0
go _ = 0
trim_arity n a ty
| n==a = a
| Just (_, ty') <- splitForAllTy_maybe ty = trim_arity n a ty'
| Just (_, ty') <- splitFunTy_maybe ty = trim_arity n (a+1) ty'
| Just (ty',_) <- splitNewTypeRepCo_maybe ty = trim_arity n a ty'
| otherwise = a
\end{code}
%************************************************************************
%* *
Eta expansion
%* *
%************************************************************************
\begin{code}
exprEtaExpandArity :: DynFlags -> CoreExpr -> Arity
exprEtaExpandArity dflags e
= applyStateHack e (arityType dicts_cheap e)
where
dicts_cheap = dopt Opt_DictsCheap dflags
\end{code}
Note [Definition of arity]
~~~~~~~~~~~~~~~~~~~~~~~~~~
The "arity" of an expression 'e' is n if
applying 'e' to *fewer* than n *value* arguments
converges rapidly
Or, to put it another way
there is no work lost in duplicating the partial
application (e x1 .. x(n1))
In the divegent case, no work is lost by duplicating because if the thing
is evaluated once, that's the end of the program.
Or, to put it another way, in any context C
C[ (\x1 .. xn. e x1 .. xn) ]
is as efficient as
C[ e ]
It's all a bit more subtle than it looks:
Note [Arity of case expressions]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
We treat the arity of
case x of p -> \s -> ...
as 1 (or more) 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.
This isn't really right in the presence of seq. Consider
f = \x -> case x of
True -> \y -> x+y
False -> \y -> xy
Can we etaexpand here? At first the answer looks like "yes of course", but
consider
(f bot) `seq` 1
This should diverge! But if we etaexpand, it won't. Again, we ignore this
"problem", because being scrupulous would lose an important transformation for
many programs.
1. Note [Oneshot lambdas]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider oneshot lambdas
let x = expensive in \y z -> E
We want this to have arity 1 if the \yabstraction is a 1shot lambda.
3. Note [Dealing with bottom]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Consider
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. Technically, this isn't quite right, because
(f True) `seq` 1
should diverge, but it'll converge if we etaexpand f. Nevertheless, we
do so; it improves some programs significantly, and increasing convergence
isn't a bad thing. Hence the ABot/ATop in ArityType.
4. Note [Newtype arity]
~~~~~~~~~~~~~~~~~~~~~~~~
Nonrecursive newtypes are transparent, and should not get in the way.
We do (currently) etaexpand 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 etaexpand 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.
Note [The statetransformer hack]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Suppose we have
f = e
where e has arity n. Then, if we know from the context that f has
a usage type like
t1 -> ... -> tn 1-> t(n+1) 1-> ... 1-> tm -> ...
then we can expand the arity to m. This usage type says that
any application (x e1 .. en) will be applied to uniquely to (mn) more args
Consider f = \x. let y = <expensive>
in 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.
Then we expect that if f is applied to one arg, it'll be applied to two
(that's the hack
See also Id.isOneShotBndr.
\begin{code}
applyStateHack :: CoreExpr -> ArityType -> Arity
applyStateHack e (AT orig_arity is_bot)
| opt_NoStateHack = orig_arity
| ABot <- is_bot = orig_arity
| otherwise = go orig_ty orig_arity
where
orig_ty = exprType e
go :: Type -> Arity -> Arity
go ty arity
| Just (_, ty') <- splitForAllTy_maybe ty = go ty' arity
| Just (tc,tys) <- splitTyConApp_maybe ty
, Just (ty', _) <- instNewTyCon_maybe tc tys
, not (isRecursiveTyCon tc) = go ty' arity
| Just (arg,res) <- splitFunTy_maybe ty
, arity > 0 || isStateHackType arg = 1 + go res (arity1)
| otherwise = WARN( arity > 0, ppr arity ) 0
\end{code}
Note [State hack and bottoming functions]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
It's a terrible idea to use the state hack on a bottoming function.
Here's what happens (Trac #2861):
f :: String -> IO T
f = \p. error "..."
Etaexpand, using the state hack:
f = \p. (\s. ((error "...") |> g1) s) |> g2
g1 :: IO T ~ (S -> (S,T))
g2 :: (S -> (S,T)) ~ IO T
Extrude the g2
f' = \p. \s. ((error "...") |> g1) s
f = f' |> (String -> g2)
Discard args for bottomming function
f' = \p. \s. ((error "...") |> g1 |> g3
g3 :: (S -> (S,T)) ~ (S,T)
Extrude g1.g3
f'' = \p. \s. (error "...")
f' = f'' |> (String -> S -> g1.g3)
And now we can repeat the whole loop. Aargh! The bug is in applying the
state hack to a function which then swallows the argument.
\begin{code}
data ArityType = AT Arity ArityRes
data ArityRes = ATop
| ABot
vanillaArityType :: ArityType
vanillaArityType = AT 0 ATop
incArity :: ArityType -> ArityType
incArity (AT a r) = AT (a+1) r
decArity :: ArityType -> ArityType
decArity (AT 0 r) = AT 0 r
decArity (AT a r) = AT (a1) r
andArityType :: ArityType -> ArityType -> ArityType
andArityType (AT a1 ATop) (AT a2 ATop) = AT (a1 `min` a2) ATop
andArityType (AT _ ABot) (AT a2 ATop) = AT a2 ATop
andArityType (AT a1 ATop) (AT _ ABot) = AT a1 ATop
andArityType (AT a1 ABot) (AT a2 ABot) = AT (a1 `max` a2) ABot
trimArity :: Bool -> ArityType -> ArityType
trimArity _cheap (AT a ABot) = AT a ABot
trimArity True (AT a ATop) = AT a ATop
trimArity False (AT _ ATop) = AT 0 ATop
arityType :: Bool -> CoreExpr -> ArityType
arityType _ (Var v)
| Just strict_sig <- idNewStrictness_maybe v
, (ds, res) <- splitStrictSig strict_sig
, isBotRes res
= AT (length ds) ABot
| otherwise
= AT (idArity v) ATop
arityType dicts_cheap (Lam x e)
| isId x = incArity (arityType dicts_cheap e)
| otherwise = arityType dicts_cheap e
arityType dicts_cheap (App fun (Type _))
= arityType dicts_cheap fun
arityType dicts_cheap (App fun arg )
= trimArity (exprIsCheap arg) (decArity (arityType dicts_cheap fun))
arityType dicts_cheap (Case scrut _ _ alts)
= trimArity (exprIsCheap scrut)
(foldr1 andArityType [arityType dicts_cheap rhs | (_,_,rhs) <- alts])
arityType dicts_cheap (Let b e)
= trimArity (cheap_bind b) (arityType dicts_cheap e)
where
cheap_bind (NonRec b e) = is_cheap (b,e)
cheap_bind (Rec prs) = all is_cheap prs
is_cheap (b,e) = (dicts_cheap && isDictLikeTy (idType b))
|| exprIsCheap e
arityType dicts_cheap (Note _ e) = arityType dicts_cheap e
arityType dicts_cheap (Cast e _) = arityType dicts_cheap e
arityType _ _ = vanillaArityType
\end{code}
%************************************************************************
%* *
The main etaexpander
%* *
%************************************************************************
IMPORTANT NOTE: The eta expander is careful not to introduce "crap".
In particular, given a CoreExpr satisfying the 'CpeRhs' invariant (in
CorePrep), it returns a CoreExpr satisfying the same invariant. See
Note [Eta expansion and the CorePrep invariants] in CorePrep.
This means the etaexpander has to do a bit of onthefly
simplification but it's not too hard. The alernative, of relying on
a subsequent cleanup phase of the Simplifier to decrapify the result,
means you can't really use it in CorePrep, which is painful.
\begin{code}
etaExpand :: Arity
-> CoreExpr
-> CoreExpr
etaExpand n orig_expr
| manifestArity orig_expr >= n = orig_expr
| otherwise
= go n orig_expr
where
go 0 expr = expr
go n (Lam v body) | isTyVar v = Lam v (go n body)
| otherwise = Lam v (go (n1) body)
go n (Note InlineMe expr) = Note InlineMe (go n expr)
go n (Cast expr co) = Cast (go n expr) co
go n expr =
etaInfoAbs etas (etaInfoApp subst' expr etas)
where
in_scope = mkInScopeSet (exprFreeVars expr)
(in_scope', etas) = mkEtaWW n in_scope (exprType expr)
subst' = mkEmptySubst in_scope'
data EtaInfo = EtaVar Var
| EtaCo Coercion
instance Outputable EtaInfo where
ppr (EtaVar v) = ptext (sLit "EtaVar") <+> ppr v
ppr (EtaCo co) = ptext (sLit "EtaCo") <+> ppr co
pushCoercion :: Coercion -> [EtaInfo] -> [EtaInfo]
pushCoercion co1 (EtaCo co2 : eis)
| isIdentityCoercion co = eis
| otherwise = EtaCo co : eis
where
co = co1 `mkTransCoercion` co2
pushCoercion co eis = EtaCo co : eis
etaInfoAbs :: [EtaInfo] -> CoreExpr -> CoreExpr
etaInfoAbs [] expr = expr
etaInfoAbs (EtaVar v : eis) expr = Lam v (etaInfoAbs eis expr)
etaInfoAbs (EtaCo co : eis) expr = Cast (etaInfoAbs eis expr) (mkSymCoercion co)
etaInfoApp :: Subst -> CoreExpr -> [EtaInfo] -> CoreExpr
etaInfoApp subst (Lam v1 e) (EtaVar v2 : eis)
= etaInfoApp subst' e eis
where
subst' | isTyVar v1 = CoreSubst.extendTvSubst subst v1 (mkTyVarTy v2)
| otherwise = CoreSubst.extendIdSubst subst v1 (Var v2)
etaInfoApp subst (Cast e co1) eis
= etaInfoApp subst e (pushCoercion co' eis)
where
co' = CoreSubst.substTy subst co1
etaInfoApp subst (Case e b _ alts) eis
= Case (subst_expr subst e) b1 (coreAltsType alts') alts'
where
(subst1, b1) = substBndr subst b
alts' = map subst_alt alts
subst_alt (con, bs, rhs) = (con, bs', etaInfoApp subst2 rhs eis)
where
(subst2,bs') = substBndrs subst1 bs
etaInfoApp subst (Let b e) eis
= Let b' (etaInfoApp subst' e eis)
where
(subst', b') = subst_bind subst b
etaInfoApp subst (Note note e) eis
= Note note (etaInfoApp subst e eis)
etaInfoApp subst e eis
= go (subst_expr subst e) eis
where
go e [] = e
go e (EtaVar v : eis) = go (App e (varToCoreExpr v)) eis
go e (EtaCo co : eis) = go (Cast e co) eis
mkEtaWW :: Arity -> InScopeSet -> Type
-> (InScopeSet, [EtaInfo])
mkEtaWW n in_scope ty
= go n empty_subst ty []
where
empty_subst = mkTvSubst in_scope emptyTvSubstEnv
go n subst ty eis
| n == 0
= (getTvInScope subst, reverse eis)
| Just (tv,ty') <- splitForAllTy_maybe ty
, let (subst', tv') = substTyVarBndr subst tv
= go n subst' ty' (EtaVar tv' : eis)
| Just (arg_ty, res_ty) <- splitFunTy_maybe ty
, let (subst', eta_id') = freshEtaId n subst arg_ty
= go (n1) subst' res_ty (EtaVar eta_id' : eis)
| Just(ty',co) <- splitNewTypeRepCo_maybe ty
=
go n subst ty' (EtaCo (substTy subst co) : eis)
| otherwise
= (getTvInScope subst, reverse eis)
subst_expr :: Subst -> CoreExpr -> CoreExpr
subst_expr s e | isEmptySubst s = e
| otherwise = substExpr s e
subst_bind :: Subst -> CoreBind -> (Subst, CoreBind)
subst_bind subst (NonRec b r)
= (subst', NonRec b' (subst_expr subst r))
where
(subst', b') = substBndr subst b
subst_bind subst (Rec prs)
= (subst', Rec (bs1 `zip` map (subst_expr subst') rhss))
where
(bs, rhss) = unzip prs
(subst', bs1) = substBndrs subst bs
freshEtaId :: Int -> TvSubst -> Type -> (TvSubst, Id)
freshEtaId n subst ty
= (subst', eta_id')
where
ty' = substTy subst ty
eta_id' = uniqAway (getTvInScope subst) $
mkSysLocal (fsLit "eta") (mkBuiltinUnique n) ty'
subst' = extendTvInScope subst [eta_id']
\end{code}