-- Same rhs as foo
If CSE produces
foo = bar
then foo will never be inlined (when it should be); but if it produces
bar = foo
bar will be inlined (when it should not be). Even if we remove INLINE foo,
we'd still like foo to be inlined if rhs is small. This won't happen
with foo = bar.
Not CSE-ing inside INLINE also solves an annoying bug in CSE. Consider
a worker/wrapper, in which the worker has turned into a single variable:
$wf = h
f = \x -> ...$wf...
Now CSE may transform to
f = \x -> ...h...
But the WorkerInfo for f still says $wf, which is now dead! This won't
happen now that we don't look inside INLINEs (which wrappers are).
%************************************************************************
%* *
\section{Common subexpression}
%* *
%************************************************************************
\begin{code}
cseProgram :: [CoreBind] -> [CoreBind]
cseProgram binds = cseBinds emptyCSEnv binds
cseBinds :: CSEnv -> [CoreBind] -> [CoreBind]
cseBinds _ [] = []
cseBinds env (b:bs) = (b':bs')
where
(env1, b') = cseBind env b
bs' = cseBinds env1 bs
cseBind :: CSEnv -> CoreBind -> (CSEnv, CoreBind)
cseBind env (NonRec b e)
= (env2, NonRec b' e')
where
(env1, b') = addBinder env b
(env2, e') = cseRhs env1 (b',e)
cseBind env (Rec pairs)
= (env2, Rec (bs' `zip` es'))
where
(bs,es) = unzip pairs
(env1, bs') = addRecBinders env bs
(env2, es') = mapAccumL cseRhs env1 (bs' `zip` es)
cseRhs :: CSEnv -> (OutBndr, InExpr) -> (CSEnv, OutExpr)
cseRhs env (id',rhs)
= case lookupCSEnv env rhs' of
Just other_expr -> (env, other_expr)
Nothing -> (addCSEnvItem env rhs' (Var id'), rhs')
where
rhs' | isAlwaysActive (idInlineActivation id') = cseExpr env rhs
| otherwise = rhs
tryForCSE :: CSEnv -> InExpr -> OutExpr
tryForCSE _ (Type t) = Type t
tryForCSE _ (Coercion c) = Coercion c
tryForCSE env expr = case lookupCSEnv env expr' of
Just smaller_expr -> smaller_expr
Nothing -> expr'
where
expr' = cseExpr env expr
cseExpr :: CSEnv -> InExpr -> OutExpr
cseExpr _ (Type t) = Type t
cseExpr _ (Coercion co) = Coercion co
cseExpr _ (Lit lit) = Lit lit
cseExpr env (Var v) = lookupSubst env v
cseExpr env (App f a) = App (cseExpr env f) (tryForCSE env a)
cseExpr env (Note n e) = Note n (cseExpr env e)
cseExpr env (Cast e co) = Cast (cseExpr env e) co
cseExpr env (Lam b e) = let (env', b') = addBinder env b
in Lam b' (cseExpr env' e)
cseExpr env (Let bind e) = let (env', bind') = cseBind env bind
in Let bind' (cseExpr env' e)
cseExpr env (Case scrut bndr ty alts) = Case scrut' bndr'' ty alts'
where
alts' = cseAlts env' scrut' bndr bndr'' alts
scrut' = tryForCSE env scrut
(env', bndr') = addBinder env bndr
bndr'' = zapIdOccInfo bndr'
cseAlts :: CSEnv -> OutExpr -> InBndr -> InBndr -> [InAlt] -> [OutAlt]
cseAlts env scrut' bndr _bndr' [(DataAlt con, args, rhs)]
| isUnboxedTupleCon con
= [(DataAlt con, args'', tryForCSE new_env rhs)]
where
(env', args') = addBinders env args
args'' = map zapIdOccInfo args'
new_env | exprIsCheap scrut' = env'
| otherwise = extendCSEnv env' scrut' tup_value
tup_value = mkAltExpr (DataAlt con) args'' (tyConAppArgs (idType bndr))
cseAlts env scrut' bndr bndr' alts
= map cse_alt alts
where
(con_target, alt_env)
= case scrut' of
Var v' -> (v', extendCSSubst env bndr v')
_ -> (bndr', extendCSEnv env scrut' (Var bndr'))
arg_tys = tyConAppArgs (idType bndr)
cse_alt (DataAlt con, args, rhs)
| not (null args)
= (DataAlt con, args', tryForCSE new_env rhs)
where
(env', args') = addBinders alt_env args
new_env = extendCSEnv env' (mkAltExpr (DataAlt con) args' arg_tys)
(Var con_target)
cse_alt (con, args, rhs)
= (con, args', tryForCSE env' rhs)
where
(env', args') = addBinders alt_env args
\end{code}
%************************************************************************
%* *
\section{The CSE envt}
%* *
%************************************************************************
\begin{code}
type InExpr = CoreExpr
type InBndr = CoreBndr
type InAlt = CoreAlt
type OutExpr = CoreExpr
type OutBndr = CoreBndr
type OutAlt = CoreAlt
data CSEnv = CS CSEMap Subst
type CSEMap = UniqFM [(OutExpr, OutExpr)]
emptyCSEnv :: CSEnv
emptyCSEnv = CS emptyUFM emptySubst
lookupCSEnv :: CSEnv -> OutExpr -> Maybe OutExpr
lookupCSEnv (CS cs sub) expr
= case lookupUFM cs (hashExpr expr) of
Nothing -> Nothing
Just pairs -> lookup_list pairs
where
lookup_list :: [(OutExpr,OutExpr)] -> Maybe OutExpr
lookup_list ((e,e'):es)
| eqExpr (substInScope sub) e expr = Just e'
| otherwise = lookup_list es
lookup_list [] = Nothing
addCSEnvItem :: CSEnv -> OutExpr -> OutExpr -> CSEnv
addCSEnvItem env expr expr' | exprIsBig expr = env
| otherwise = extendCSEnv env expr expr'
extendCSEnv :: CSEnv -> OutExpr -> OutExpr -> CSEnv
extendCSEnv (CS cs sub) expr expr'
= CS (addToUFM_C combine cs hash [(expr, expr')]) sub
where
hash = hashExpr expr
combine old new
= WARN( result `lengthExceeds` 4, short_msg $$ nest 2 long_msg ) result
where
result = new ++ old
short_msg = ptext (sLit "extendCSEnv: long list, length") <+> int (length result)
long_msg | opt_PprStyle_Debug = (text "hash code" <+> text (show hash)) $$ ppr result
| otherwise = empty
lookupSubst :: CSEnv -> Id -> OutExpr
lookupSubst (CS _ sub) x = lookupIdSubst (text "CSE.lookupSubst") sub x
extendCSSubst :: CSEnv -> Id -> Id -> CSEnv
extendCSSubst (CS cs sub) x y = CS cs (extendIdSubst sub x (Var y))
addBinder :: CSEnv -> Var -> (CSEnv, Var)
addBinder (CS cs sub) v = (CS cs sub', v')
where
(sub', v') = substBndr sub v
addBinders :: CSEnv -> [Var] -> (CSEnv, [Var])
addBinders (CS cs sub) vs = (CS cs sub', vs')
where
(sub', vs') = substBndrs sub vs
addRecBinders :: CSEnv -> [Id] -> (CSEnv, [Id])
addRecBinders (CS cs sub) vs = (CS cs sub', vs')
where
(sub', vs') = substRecBndrs sub vs
\end{code}