%
% (c) The GRASP/AQUA Project, Glasgow University, 1993-1998
%
\section[WwLib]{A library for the ``worker\/wrapper'' back-end to the strictness analyser}
\begin{code}
module WwLib ( mkWwBodies, mkWWstr, mkWorkerArgs ) where
#include "HsVersions.h"
import CoreSyn
import CoreUtils ( exprType )
import Id ( Id, idType, mkSysLocal, idNewDemandInfo, setIdNewDemandInfo,
isOneShotLambda, setOneShotLambda, setIdUnfolding,
setIdInfo
)
import IdInfo ( vanillaIdInfo )
import DataCon
import NewDemand ( Demand(..), DmdResult(..), Demands(..) )
import MkId ( realWorldPrimId, voidArgId, mkRuntimeErrorApp, rUNTIME_ERROR_ID,
mkUnpackCase, mkProductBox )
import TysWiredIn ( tupleCon )
import Type
import Coercion ( mkSymCoercion, splitNewTypeRepCo_maybe )
import BasicTypes ( Boxity(..) )
import Var ( Var )
import UniqSupply
import Unique
import Util ( zipWithEqual )
import Outputable
import FastString
\end{code}
%************************************************************************
%* *
\subsection[mkWrapperAndWorker]{@mkWrapperAndWorker@}
%* *
%************************************************************************
Here's an example. The original function is:
\begin{verbatim}
g :: forall a . Int -> [a] -> a
g = \/\ a -> \ x ys ->
case x of
0 -> head ys
_ -> head (tail ys)
\end{verbatim}
From this, we want to produce:
\begin{verbatim}
-- wrapper (an unfolding)
g :: forall a . Int -> [a] -> a
g = \/\ a -> \ x ys ->
case x of
I# x# -> $wg a x# ys
-- call the worker; don't forget the type args!
-- worker
$wg :: forall a . Int# -> [a] -> a
$wg = \/\ a -> \ x# ys ->
let
x = I# x#
in
case x of -- note: body of g moved intact
0 -> head ys
_ -> head (tail ys)
\end{verbatim}
Something we have to be careful about: Here's an example:
\begin{verbatim}
-- "f" strictness: U(P)U(P)
f (I# a) (I# b) = a +# b
g = f -- "g" strictness same as "f"
\end{verbatim}
\tr{f} will get a worker all nice and friendly-like; that's good.
{\em But we don't want a worker for \tr{g}}, even though it has the
same strictness as \tr{f}. Doing so could break laziness, at best.
Consequently, we insist that the number of strictness-info items is
exactly the same as the number of lambda-bound arguments. (This is
probably slightly paranoid, but OK in practice.) If it isn't the
same, we ``revise'' the strictness info, so that we won't propagate
the unusable strictness-info into the interfaces.
%************************************************************************
%* *
\subsection{The worker wrapper core}
%* *
%************************************************************************
@mkWwBodies@ is called when doing the worker\/wrapper split inside a module.
\begin{code}
mkWwBodies :: Type
-> [Demand]
-> DmdResult
-> [Bool]
-> UniqSM ([Demand],
Id -> CoreExpr,
CoreExpr -> CoreExpr)
mkWwBodies fun_ty demands res_info one_shots
= do { let arg_info = demands `zip` (one_shots ++ repeat False)
; (wrap_args, wrap_fn_args, work_fn_args, res_ty) <- mkWWargs emptyTvSubst fun_ty arg_info
; (work_args, wrap_fn_str, work_fn_str) <- mkWWstr wrap_args
; (wrap_fn_cpr, work_fn_cpr, _cpr_res_ty)
<- if any isId work_args then
mkWWcpr res_ty res_info
else
return (id, id, res_ty)
; let (work_lam_args, work_call_args) = mkWorkerArgs work_args res_ty
; return ([idNewDemandInfo v | v <- work_call_args, isId v],
Note InlineMe . wrap_fn_args . wrap_fn_cpr . wrap_fn_str . applyToVars work_call_args . Var,
mkLams work_lam_args. work_fn_str . work_fn_cpr . work_fn_args) }
\end{code}
%************************************************************************
%* *
\subsection{Making wrapper args}
%* *
%************************************************************************
During worker-wrapper stuff we may end up with an unlifted thing
which we want to let-bind without losing laziness. So we
add a void argument. E.g.
f = /\a -> \x y z -> E::Int# -- E does not mention x,y,z
==>
fw = /\ a -> \void -> E
f = /\ a -> \x y z -> fw realworld
We use the state-token type which generates no code.
\begin{code}
mkWorkerArgs :: [Var]
-> Type
-> ([Var],
[Var])
mkWorkerArgs args res_ty
| any isId args || not (isUnLiftedType res_ty)
= (args, args)
| otherwise
= (args ++ [voidArgId], args ++ [realWorldPrimId])
\end{code}
%************************************************************************
%* *
\subsection{Coercion stuff}
%* *
%************************************************************************
We really want to "look through" coerces.
Reason: I've seen this situation:
let f = coerce T (\s -> E)
in \x -> case x of
p -> coerce T' f
q -> \s -> E2
r -> coerce T' f
If only we w/w'd f, we'd get
let f = coerce T (\s -> fw s)
fw = \s -> E
in ...
Now we'll inline f to get
let fw = \s -> E
in \x -> case x of
p -> fw
q -> \s -> E2
r -> fw
Now we'll see that fw has arity 1, and will arity expand
the \x to get what we want.
\begin{code}
mkWWargs :: TvSubst
-> Type
-> [(Demand,Bool)]
-> UniqSM ([Var],
CoreExpr -> CoreExpr,
CoreExpr -> CoreExpr,
Type)
mkWWargs subst fun_ty arg_info
| Just (rep_ty, co) <- splitNewTypeRepCo_maybe fun_ty
= do { (wrap_args, wrap_fn_args, work_fn_args, res_ty)
<- mkWWargs subst rep_ty arg_info
; return (wrap_args,
\e -> Cast (wrap_fn_args e) (mkSymCoercion co),
\e -> work_fn_args (Cast e co),
res_ty) }
| null arg_info
= return ([], id, id, substTy subst fun_ty)
| Just (tv, fun_ty') <- splitForAllTy_maybe fun_ty
= do { let (subst', tv') = substTyVarBndr subst tv
; (wrap_args, wrap_fn_args, work_fn_args, res_ty)
<- mkWWargs subst' fun_ty' arg_info
; return (tv' : wrap_args,
Lam tv' . wrap_fn_args,
work_fn_args . (`App` Type (mkTyVarTy tv')),
res_ty) }
| ((dmd,one_shot):arg_info') <- arg_info
, Just (arg_ty, fun_ty') <- splitFunTy_maybe fun_ty
= do { uniq <- getUniqueM
; let arg_ty' = substTy subst arg_ty
id = mk_wrap_arg uniq arg_ty' dmd one_shot
; (wrap_args, wrap_fn_args, work_fn_args, res_ty)
<- mkWWargs subst fun_ty' arg_info'
; return (id : wrap_args,
Lam id . wrap_fn_args,
work_fn_args . (`App` Var id),
res_ty) }
| otherwise
= WARN( True, ppr fun_ty )
return ([], id, id, substTy subst fun_ty)
applyToVars :: [Var] -> CoreExpr -> CoreExpr
applyToVars vars fn = mkVarApps fn vars
mk_wrap_arg :: Unique -> Type -> NewDemand.Demand -> Bool -> Id
mk_wrap_arg uniq ty dmd one_shot
= set_one_shot one_shot (setIdNewDemandInfo (mkSysLocal (fsLit "w") uniq ty) dmd)
where
set_one_shot True id = setOneShotLambda id
set_one_shot False id = id
\end{code}
Note [Freshen type variables]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~
mkWWargs may be given a type like (a~b) =>
Which really means forall (co:a~b).
Because the name of the coercion variable, 'co', isn't mentioned in ,
nested coercion foralls may all use the same variable; and sometimes do
see Var.mkWildCoVar.
However, when we do a worker/wrapper split, we must not use shadowed names,
else we'll get
f = /\ co /\co. fw co co
which is obviously wrong. Actually, the same is true of type variables, which
can in principle shadow, within a type (e.g. forall a. a -> forall a. a->a).
But type variables *are* mentioned in , so we must substitute.
That's why we carry the TvSubst through mkWWargs
%************************************************************************
%* *
\subsection{Strictness stuff}
%* *
%************************************************************************
\begin{code}
mkWWstr :: [Var]
-> UniqSM ([Var],
CoreExpr -> CoreExpr,
CoreExpr -> CoreExpr)
mkWWstr []
= return ([], nop_fn, nop_fn)
mkWWstr (arg : args) = do
(args1, wrap_fn1, work_fn1) <- mkWWstr_one arg
(args2, wrap_fn2, work_fn2) <- mkWWstr args
return (args1 ++ args2, wrap_fn1 . wrap_fn2, work_fn1 . work_fn2)
mkWWstr_one :: Var -> UniqSM ([Var], CoreExpr -> CoreExpr, CoreExpr -> CoreExpr)
mkWWstr_one arg
| isTyVar arg
= return ([arg], nop_fn, nop_fn)
| otherwise
= case idNewDemandInfo arg of
Abs | not (isUnLiftedType (idType arg)) ->
return ([], nop_fn, mk_absent_let arg)
Eval (Prod cs)
| Just (_arg_tycon, _tycon_arg_tys, data_con, inst_con_arg_tys)
<- deepSplitProductType_maybe (idType arg)
-> do uniqs <- getUniquesM
let
unpk_args = zipWith mk_ww_local uniqs inst_con_arg_tys
unpk_args_w_ds = zipWithEqual "mkWWstr" set_worker_arg_info unpk_args cs
unbox_fn = mkUnpackCase (sanitiseCaseBndr arg) (Var arg) unpk_args data_con
rebox_fn = Let (NonRec arg con_app)
con_app = mkProductBox unpk_args (idType arg)
(worker_args, wrap_fn, work_fn) <- mkWWstr unpk_args_w_ds
return (worker_args, unbox_fn . wrap_fn, work_fn . rebox_fn)
Eval (Poly Abs)
-> let
arg_w_unf = arg `setIdUnfolding` evaldUnfolding
in
return ([arg_w_unf], mk_seq_case arg, nop_fn)
_other_demand -> return ([arg], nop_fn, nop_fn)
where
set_worker_arg_info worker_arg demand = set_one_shot (setIdNewDemandInfo worker_arg demand)
set_one_shot | isOneShotLambda arg = setOneShotLambda
| otherwise = \x -> x
nop_fn :: CoreExpr -> CoreExpr
nop_fn body = body
\end{code}
%************************************************************************
%* *
\subsection{CPR stuff}
%* *
%************************************************************************
@mkWWcpr@ takes the worker/wrapper pair produced from the strictness
info and adds in the CPR transformation. The worker returns an
unboxed tuple containing non-CPR components. The wrapper takes this
tuple and re-produces the correct structured output.
The non-CPR results appear ordered in the unboxed tuple as if by a
left-to-right traversal of the result structure.
\begin{code}
mkWWcpr :: Type
-> DmdResult
-> UniqSM (CoreExpr -> CoreExpr,
CoreExpr -> CoreExpr,
Type)
mkWWcpr body_ty RetCPR
| not (isClosedAlgType body_ty)
= WARN( True,
text "mkWWcpr: non-algebraic or open body type" <+> ppr body_ty )
return (id, id, body_ty)
| n_con_args == 1 && isUnLiftedType con_arg_ty1 = do
(work_uniq : arg_uniq : _) <- getUniquesM
let
work_wild = mk_ww_local work_uniq body_ty
arg = mk_ww_local arg_uniq con_arg_ty1
con_app = mkProductBox [arg] body_ty
return (\ wkr_call -> Case wkr_call (arg) (exprType con_app) [(DEFAULT, [], con_app)],
\ body -> workerCase (work_wild) body [arg] data_con (Var arg),
con_arg_ty1)
| otherwise = do
uniqs <- getUniquesM
let
(wrap_wild : work_wild : args) = zipWith mk_ww_local uniqs (ubx_tup_ty : body_ty : con_arg_tys)
arg_vars = map Var args
ubx_tup_con = tupleCon Unboxed n_con_args
ubx_tup_ty = exprType ubx_tup_app
ubx_tup_app = mkConApp ubx_tup_con (map Type con_arg_tys ++ arg_vars)
con_app = mkProductBox args body_ty
return (\ wkr_call -> Case wkr_call (wrap_wild) (exprType con_app) [(DataAlt ubx_tup_con, args, con_app)],
\ body -> workerCase (work_wild) body args data_con ubx_tup_app,
ubx_tup_ty)
where
(_arg_tycon, _tycon_arg_tys, data_con, con_arg_tys) = deepSplitProductType "mkWWcpr" body_ty
n_con_args = length con_arg_tys
con_arg_ty1 = head con_arg_tys
mkWWcpr body_ty _other
= return (id, id, body_ty)
workerCase :: Id -> CoreExpr -> [Id] -> DataCon -> CoreExpr -> CoreExpr
workerCase bndr (Note (SCC cc) e) args con body = Note (SCC cc) (mkUnpackCase bndr e args con body)
workerCase bndr e args con body = mkUnpackCase bndr e args con body
\end{code}
%************************************************************************
%* *
\subsection{Utilities}
%* *
%************************************************************************
\begin{code}
mk_absent_let :: Id -> CoreExpr -> CoreExpr
mk_absent_let arg body
| not (isUnLiftedType arg_ty)
= Let (NonRec arg abs_rhs) body
| otherwise
= panic "WwLib: haven't done mk_absent_let for primitives yet"
where
arg_ty = idType arg
abs_rhs = mkRuntimeErrorApp rUNTIME_ERROR_ID arg_ty msg
msg = "Oops! Entered absent arg " ++ showSDocDebug (ppr arg <+> ppr (idType arg))
mk_seq_case :: Id -> CoreExpr -> CoreExpr
mk_seq_case arg body = Case (Var arg) (sanitiseCaseBndr arg) (exprType body) [(DEFAULT, [], body)]
sanitiseCaseBndr :: Id -> Id
sanitiseCaseBndr id = id `setIdInfo` vanillaIdInfo
mk_ww_local :: Unique -> Type -> Id
mk_ww_local uniq ty = mkSysLocal (fsLit "ww") uniq ty
\end{code}