2 % (c) The GRASP/AQUA Project, Glasgow University, 1993-1998
4 \section[WwLib]{A library for the ``worker/wrapper'' back-end to the strictness analyser}
9 worthSplitting, setUnpackStrategy
12 #include "HsVersions.h"
15 import CoreUtils ( exprType, mkInlineMe )
16 import Id ( Id, idType, mkSysLocal, idDemandInfo, setIdDemandInfo,
17 isOneShotLambda, setOneShotLambda,
20 import IdInfo ( CprInfo(..), noCprInfo, vanillaIdInfo )
21 import DataCon ( DataCon, splitProductType )
22 import Demand ( Demand(..), wwLazy, wwPrim )
23 import PrelInfo ( realWorldPrimId, aBSENT_ERROR_ID )
24 import TysPrim ( realWorldStatePrimTy )
25 import TysWiredIn ( unboxedTupleCon, unboxedTupleTyCon )
26 import Type ( isUnLiftedType,
27 splitForAllTys, splitFunTys, isAlgType,
28 splitAlgTyConApp_maybe, splitNewType_maybe,
32 import TyCon ( isNewTyCon, isProductTyCon, TyCon )
33 import BasicTypes ( NewOrData(..), Arity )
34 import Var ( TyVar, Var, isId )
35 import UniqSupply ( returnUs, thenUs, getUniqueUs, getUniquesUs,
37 import Util ( zipWithEqual, zipEqual, lengthExceeds )
39 import List ( zipWith4 )
43 %************************************************************************
45 \subsection[mkWrapperAndWorker]{@mkWrapperAndWorker@}
47 %************************************************************************
49 ************ WARNING ******************
50 these comments are rather out of date
51 *****************************************
53 @mkWrapperAndWorker@ is given:
56 The {\em original function} \tr{f}, of the form:
58 f = /\ tyvars -> \ args -> body
60 The original-binder \tr{f}, the \tr{tyvars}, \tr{args}, and \tr{body}
63 We use the Id \tr{f} mostly to get its type.
66 Strictness information about \tr{f}, in the form of a list of
73 @mkWrapperAndWorker@ produces (A BIT OUT-OF-DATE...):
76 Maybe @Nothing@: no worker/wrappering going on in this case. This can
77 happen (a)~if the strictness info says that there is nothing
78 interesting to do or (b)~if *any* of the argument types corresponding
79 to ``active'' arg postitions is abstract or will be to the outside
80 world (i.e., {\em this} module can see the constructors, but nobody
81 else will be able to). An ``active'' arg position is one which the
82 wrapper has to unpack. An importing module can't do this unpacking,
83 so it simply has to give up and call the wrapper only.
86 Maybe \tr{Just (wrapper_Id, wrapper_body, worker_Id, worker_body)}.
88 The @wrapper_Id@ is just the one that was passed in, with its
89 strictness IdInfo updated.
92 The \tr{body} of the original function may not be given (i.e., it's
93 BOTTOM), in which case you'd jolly well better not tug on the
96 Here's an example. The original function is:
98 g :: forall a . Int -> [a] -> a
100 g = /\ a -> \ x ys ->
106 From this, we want to produce:
108 -- wrapper (an unfolding)
109 g :: forall a . Int -> [a] -> a
111 g = /\ a -> \ x ys ->
113 I# x# -> g.wrk a x# ys
114 -- call the worker; don't forget the type args!
117 g.wrk :: forall a . Int# -> [a] -> a
119 g.wrk = /\ a -> \ x# ys ->
123 case x of -- note: body of g moved intact
128 Something we have to be careful about: Here's an example:
130 -- "f" strictness: U(P)U(P)
131 f (I# a) (I# b) = a +# b
133 g = f -- "g" strictness same as "f"
135 \tr{f} will get a worker all nice and friendly-like; that's good.
136 {\em But we don't want a worker for \tr{g}}, even though it has the
137 same strictness as \tr{f}. Doing so could break laziness, at best.
139 Consequently, we insist that the number of strictness-info items is
140 exactly the same as the number of lambda-bound arguments. (This is
141 probably slightly paranoid, but OK in practice.) If it isn't the
142 same, we ``revise'' the strictness info, so that we won't propagate
143 the unusable strictness-info into the interfaces.
146 %************************************************************************
148 \subsection{Functions over Demands}
150 %************************************************************************
153 mAX_WORKER_ARGS :: Int -- ToDo: set via flag
156 setUnpackStrategy :: [Demand] -> [Demand]
158 = snd (go (mAX_WORKER_ARGS - nonAbsentArgs ds) ds)
160 go :: Int -- Max number of args available for sub-components of [Demand]
162 -> (Int, [Demand]) -- Args remaining after subcomponents of [Demand] are unpacked
164 go n (WwUnpack nd _ cs : ds) | n' >= 0
165 = WwUnpack nd True cs' `cons` go n'' ds
167 = WwUnpack nd False cs `cons` go n ds
169 n' = n + 1 - nonAbsentArgs cs
170 -- Add one because we don't pass the top-level arg any more
171 -- Delete # of non-absent args to which we'll now be committed
174 go n (d:ds) = d `cons` go n ds
177 cons d (n,ds) = (n, d:ds)
179 nonAbsentArgs :: [Demand] -> Int
181 nonAbsentArgs (WwLazy True : ds) = nonAbsentArgs ds
182 nonAbsentArgs (d : ds) = 1 + nonAbsentArgs ds
184 worthSplitting :: [Demand]
185 -> Bool -- Result is bottom
186 -> Bool -- True <=> the wrapper would not be an identity function
187 worthSplitting ds result_bot = any worth_it ds
188 -- We used not to split if the result is bottom.
189 -- [Justification: there's no efficiency to be gained.]
190 -- But it's sometimes bad not to make a wrapper. Consider
191 -- fw = \x# -> let x = I# x# in case e of
194 -- p3 -> the real stuff
195 -- The re-boxing code won't go away unless error_fn gets a wrapper too.
198 worth_it (WwLazy True) = True -- Absent arg
199 worth_it (WwUnpack _ True _) = True -- Arg to unpack
200 worth_it WwStrict = False -- Don't w/w just because of strictness
201 worth_it other = False
203 allAbsent :: [Demand] -> Bool
204 allAbsent ds = all absent ds
206 absent (WwLazy is_absent) = is_absent
207 absent (WwUnpack _ True cs) = allAbsent cs
212 %************************************************************************
214 \subsection{The worker wrapper core}
216 %************************************************************************
218 @mkWwBodies@ is called when doing the worker/wrapper split inside a module.
221 mkWwBodies :: Type -- Type of original function
222 -> Arity -- Arity of original function
223 -> [Demand] -- Strictness of original function
224 -> Bool -- True <=> function returns bottom
225 -> [Bool] -- One-shot-ness of the function
226 -> CprInfo -- Result of CPR analysis
227 -> UniqSM ([Demand], -- Demands for worker (value) args
228 Id -> CoreExpr, -- Wrapper body, lacking only the worker Id
229 CoreExpr -> CoreExpr) -- Worker body, lacking the original function rhs
231 mkWwBodies fun_ty arity demands res_bot one_shots cpr_info
232 = mkWWargs fun_ty arity demands' res_bot one_shots' `thenUs` \ (wrap_args, wrap_fn_args, work_fn_args, res_ty) ->
233 mkWWstr wrap_args `thenUs` \ (work_dmds, wrap_fn_str, work_fn_str) ->
234 mkWWcpr res_ty cpr_info `thenUs` \ (wrap_fn_cpr, work_fn_cpr, cpr_res_ty) ->
235 mkWWfixup cpr_res_ty work_dmds `thenUs` \ (final_work_dmds, wrap_fn_fixup, work_fn_fixup) ->
237 returnUs (final_work_dmds,
238 Note InlineMe . wrap_fn_args . wrap_fn_cpr . wrap_fn_str . wrap_fn_fixup . Var,
239 work_fn_fixup . work_fn_str . work_fn_cpr . work_fn_args)
240 -- We use an INLINE unconditionally, even if the wrapper turns out to be
241 -- something trivial like
243 -- f = __inline__ (coerce T fw)
244 -- The point is to propagate the coerce to f's call sites, so even though
245 -- f's RHS is now trivial (size 1) we still want the __inline__ to prevent
246 -- fw from being inlined into f's RHS
248 demands' = demands ++ repeat wwLazy
249 one_shots' = one_shots ++ repeat False
253 %************************************************************************
255 \subsection{Coercion stuff}
257 %************************************************************************
260 We really want to "look through" coerces.
261 Reason: I've seen this situation:
263 let f = coerce T (\s -> E)
269 If only we w/w'd f, we'd get
270 let f = coerce T (\s -> fw s)
274 Now we'll inline f to get
282 Now we'll see that fw has arity 1, and will arity expand
283 the \x to get what we want.
286 -- mkWWargs is driven off the function type and arity.
287 -- It chomps bites off foralls, arrows, newtypes
288 -- and keeps repeating that until it's satisfied the supplied arity
290 mkWWargs :: Type -> Arity
291 -> [Demand] -> Bool -> [Bool] -- Both these will in due course be derived
292 -- from the type. The [Bool] is True for a one-shot arg.
293 -- ** Both are infinite, extended with neutral values if necy **
294 -> UniqSM ([Var], -- Wrapper args
295 CoreExpr -> CoreExpr, -- Wrapper fn
296 CoreExpr -> CoreExpr, -- Worker fn
297 Type) -- Type of wrapper body
299 mkWWargs fun_ty arity demands res_bot one_shots
300 | (res_bot || arity > 0) && (not (null tyvars) || n_arg_tys > 0)
301 -- If the function returns bottom, we feel free to
302 -- build lots of wrapper args:
303 -- \x. let v=E in \y. bottom
304 -- = \xy. let v=E in bottom
305 = getUniquesUs n_args `thenUs` \ wrap_uniqs ->
307 val_args = zipWith4 mk_wrap_arg wrap_uniqs arg_tys demands one_shots
308 wrap_args = tyvars ++ val_args
312 (drop n_args demands)
314 (drop n_args one_shots) `thenUs` \ (more_wrap_args, wrap_fn_args, work_fn_args, res_ty) ->
316 returnUs (wrap_args ++ more_wrap_args,
317 mkLams wrap_args . wrap_fn_args,
318 work_fn_args . applyToVars wrap_args,
321 (tyvars, tau) = splitForAllTys fun_ty
322 (arg_tys, body_ty) = splitFunTys tau
323 n_arg_tys = length arg_tys
324 n_args | res_bot = n_arg_tys
325 | otherwise = arity `min` n_arg_tys
326 new_fun_ty | n_args == n_arg_tys = body_ty
327 | otherwise = mkFunTys (drop n_args arg_tys) body_ty
329 mkWWargs fun_ty arity demands res_bot one_shots
330 = case splitNewType_maybe fun_ty of
331 Nothing -> returnUs ([], id, id, fun_ty)
332 Just rep_ty -> mkWWargs rep_ty arity demands res_bot one_shots `thenUs` \ (wrap_args, wrap_fn_args, work_fn_args, res_ty) ->
334 Note (Coerce fun_ty rep_ty) . wrap_fn_args,
335 work_fn_args . Note (Coerce rep_ty fun_ty),
339 applyToVars :: [Var] -> CoreExpr -> CoreExpr
340 applyToVars vars fn = mkVarApps fn vars
342 mk_wrap_arg uniq ty dmd one_shot
343 = set_one_shot one_shot (setIdDemandInfo (mkSysLocal SLIT("w") uniq ty) dmd)
345 set_one_shot True id = setOneShotLambda id
346 set_one_shot False id = id
350 %************************************************************************
352 \subsection{Fixup stuff}
354 %************************************************************************
357 mkWWfixup res_ty work_dmds
358 | null work_dmds && isUnLiftedType res_ty
359 -- Horrid special case. If the worker would have no arguments, and the
360 -- function returns a primitive type value, that would make the worker into
361 -- an unboxed value. We box it by passing a dummy void argument, thus:
363 -- f = /\abc. \xyz. fw abc void
364 -- fw = /\abc. \v. body
366 -- We use the state-token type which generates no code
367 = getUniqueUs `thenUs` \ void_arg_uniq ->
369 void_arg = mk_ww_local void_arg_uniq realWorldStatePrimTy
372 \ call_to_worker -> App call_to_worker (Var realWorldPrimId),
373 \ worker_body -> Lam void_arg worker_body)
376 = returnUs (work_dmds, id, id)
380 %************************************************************************
382 \subsection{Strictness stuff}
384 %************************************************************************
387 mkWWstr :: [Var] -- Wrapper args; have their demand info on them
388 -- *Includes type variables*
389 -> UniqSM ([Demand], -- Demand on worker (value) args
390 CoreExpr -> CoreExpr, -- Wrapper body, lacking the worker call
391 -- and without its lambdas
392 -- This fn adds the unboxing, and makes the
393 -- call passing the unboxed things
395 CoreExpr -> CoreExpr) -- Worker body, lacking the original body of the function,
396 -- but *with* lambdas
399 = mk_ww_str wrap_args `thenUs` \ (work_args, wrap_fn, work_fn) ->
400 returnUs ( [idDemandInfo v | v <- work_args, isId v],
401 \ wrapper_body -> wrap_fn (mkVarApps wrapper_body work_args),
402 \ worker_body -> mkLams work_args (work_fn worker_body))
407 \ wrapper_body -> wrapper_body,
408 \ worker_body -> worker_body)
413 = mk_ww_str ds `thenUs` \ (worker_args, wrap_fn, work_fn) ->
414 returnUs (arg : worker_args, wrap_fn, work_fn)
417 = case idDemandInfo arg of
421 mk_ww_str ds `thenUs` \ (worker_args, wrap_fn, work_fn) ->
422 returnUs (worker_args, wrap_fn, mk_absent_let arg . work_fn)
425 WwUnpack new_or_data True cs ->
426 getUniquesUs (length inst_con_arg_tys) `thenUs` \ uniqs ->
428 unpk_args = zipWith mk_ww_local uniqs inst_con_arg_tys
429 unpk_args_w_ds = zipWithEqual "mk_ww_str" set_worker_arg_info unpk_args cs
431 mk_ww_str (unpk_args_w_ds ++ ds) `thenUs` \ (worker_args, wrap_fn, work_fn) ->
432 returnUs (worker_args,
433 mk_unpk_case new_or_data arg unpk_args data_con arg_tycon . wrap_fn,
434 work_fn . mk_pk_let new_or_data arg data_con tycon_arg_tys unpk_args)
436 (arg_tycon, tycon_arg_tys, data_con, inst_con_arg_tys) = splitProductType "mk_ww_str" (idType arg)
440 mk_ww_str ds `thenUs` \ (worker_args, wrap_fn, work_fn) ->
441 returnUs (arg : worker_args, wrap_fn, work_fn)
443 -- If the wrapper argument is a one-shot lambda, then
444 -- so should (all) the corresponding worker arguments be
445 -- This bites when we do w/w on a case join point
446 set_worker_arg_info worker_arg demand = set_one_shot (setIdDemandInfo worker_arg demand)
448 set_one_shot | isOneShotLambda arg = setOneShotLambda
449 | otherwise = \x -> x
453 %************************************************************************
455 \subsection{CPR stuff}
457 %************************************************************************
460 @mkWWcpr@ takes the worker/wrapper pair produced from the strictness
461 info and adds in the CPR transformation. The worker returns an
462 unboxed tuple containing non-CPR components. The wrapper takes this
463 tuple and re-produces the correct structured output.
465 The non-CPR results appear ordered in the unboxed tuple as if by a
466 left-to-right traversal of the result structure.
470 mkWWcpr :: Type -- function body type
471 -> CprInfo -- CPR analysis results
472 -> UniqSM (CoreExpr -> CoreExpr, -- New wrapper
473 CoreExpr -> CoreExpr, -- New worker
474 Type) -- Type of worker's body
476 mkWWcpr body_ty NoCPRInfo
477 = returnUs (id, id, body_ty) -- Must be just the strictness transf.
479 mkWWcpr body_ty ReturnsCPR
480 | not (isAlgType body_ty)
481 = WARN( True, text "mkWWcpr: non-algebraic body type" <+> ppr body_ty )
482 returnUs (id, id, body_ty)
484 | n_con_args == 1 && isUnLiftedType con_arg_ty1
485 -- Special case when there is a single result of unlifted type
486 = getUniquesUs 2 `thenUs` \ [work_uniq, arg_uniq] ->
488 work_wild = mk_ww_local work_uniq body_ty
489 arg = mk_ww_local arg_uniq con_arg_ty1
491 returnUs (\ wkr_call -> Case wkr_call arg [(DEFAULT, [], mkConApp data_con (map Type tycon_arg_tys ++ [Var arg]))],
492 \ body -> Case body work_wild [(DataAlt data_con, [arg], Var arg)],
495 | otherwise -- The general case
496 = getUniquesUs (n_con_args + 2) `thenUs` \ uniqs ->
498 (wrap_wild : work_wild : args) = zipWith mk_ww_local uniqs (ubx_tup_ty : body_ty : con_arg_tys)
499 arg_vars = map Var args
500 ubx_tup_con = unboxedTupleCon n_con_args
501 ubx_tup_ty = exprType ubx_tup_app
502 ubx_tup_app = mkConApp ubx_tup_con (map Type con_arg_tys ++ arg_vars)
503 con_app = mkConApp data_con (map Type tycon_arg_tys ++ arg_vars)
505 returnUs (\ wkr_call -> Case wkr_call wrap_wild [(DataAlt ubx_tup_con, args, con_app)],
506 \ body -> Case body work_wild [(DataAlt data_con, args, ubx_tup_app)],
509 (tycon, tycon_arg_tys, data_con, con_arg_tys) = splitProductType "mkWWcpr" body_ty
510 n_con_args = length con_arg_tys
511 con_arg_ty1 = head con_arg_tys
515 %************************************************************************
517 \subsection{Utilities}
519 %************************************************************************
523 mk_absent_let arg body
524 | not (isUnLiftedType arg_ty)
525 = Let (NonRec arg (mkTyApps (Var aBSENT_ERROR_ID) [arg_ty])) body
527 = panic "WwLib: haven't done mk_absent_let for primitives yet"
531 mk_unpk_case NewType arg unpk_args boxing_con boxing_tycon body
532 -- A newtype! Use a coercion not a case
533 = ASSERT( null other_args )
534 Case (Note (Coerce (idType unpk_arg) (idType arg)) (Var arg))
535 (sanitiseCaseBndr unpk_arg)
538 (unpk_arg:other_args) = unpk_args
540 mk_unpk_case DataType arg unpk_args boxing_con boxing_tycon body
543 (sanitiseCaseBndr arg)
544 [(DataAlt boxing_con, unpk_args, body)]
546 sanitiseCaseBndr :: Id -> Id
547 -- The argument we are scrutinising has the right type to be
548 -- a case binder, so it's convenient to re-use it for that purpose.
549 -- But we *must* throw away all its IdInfo. In particular, the argument
550 -- will have demand info on it, and that demand info may be incorrect for
551 -- the case binder. e.g. case ww_arg of ww_arg { I# x -> ... }
552 -- Quite likely ww_arg isn't used in '...'. The case may get discarded
553 -- if the case binder says "I'm demanded". This happened in a situation
554 -- like (x+y) `seq` ....
555 sanitiseCaseBndr id = id `setIdInfo` vanillaIdInfo
557 mk_pk_let NewType arg boxing_con con_tys unpk_args body
558 = ASSERT( null other_args )
559 Let (NonRec arg (Note (Coerce (idType arg) (idType unpk_arg)) (Var unpk_arg))) body
561 (unpk_arg:other_args) = unpk_args
563 mk_pk_let DataType arg boxing_con con_tys unpk_args body
564 = Let (NonRec arg (mkConApp boxing_con con_args)) body
566 con_args = map Type con_tys ++ map Var unpk_args
569 mk_ww_local uniq ty = mkSysLocal SLIT("ww") uniq ty