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 mkInlineMe . 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)
241 demands' = demands ++ repeat wwLazy
242 one_shots' = one_shots ++ repeat False
246 %************************************************************************
248 \subsection{Coercion stuff}
250 %************************************************************************
253 We really want to "look through" coerces.
254 Reason: I've seen this situation:
256 let f = coerce T (\s -> E)
262 If only we w/w'd f, we'd get
263 let f = coerce T (\s -> fw s)
267 Now we'll inline f to get
275 Now we'll see that fw has arity 1, and will arity expand
276 the \x to get what we want.
279 -- mkWWargs is driven off the function type and arity.
280 -- It chomps bites off foralls, arrows, newtypes
281 -- and keeps repeating that until it's satisfied the supplied arity
283 mkWWargs :: Type -> Arity
284 -> [Demand] -> Bool -> [Bool] -- Both these will in due course be derived
285 -- from the type. The [Bool] is True for a one-shot arg.
286 -- ** Both are infinite, extended with neutral values if necy **
287 -> UniqSM ([Var], -- Wrapper args
288 CoreExpr -> CoreExpr, -- Wrapper fn
289 CoreExpr -> CoreExpr, -- Worker fn
290 Type) -- Type of wrapper body
292 mkWWargs fun_ty arity demands res_bot one_shots
293 | (res_bot || arity > 0) && (not (null tyvars) || n_arg_tys > 0)
294 -- If the function returns bottom, we feel free to
295 -- build lots of wrapper args:
296 -- \x. let v=E in \y. bottom
297 -- = \xy. let v=E in bottom
298 = getUniquesUs n_args `thenUs` \ wrap_uniqs ->
300 val_args = zipWith4 mk_wrap_arg wrap_uniqs arg_tys demands one_shots
301 wrap_args = tyvars ++ val_args
305 (drop n_args demands)
307 (drop n_args one_shots) `thenUs` \ (more_wrap_args, wrap_fn_args, work_fn_args, res_ty) ->
309 returnUs (wrap_args ++ more_wrap_args,
310 mkLams wrap_args . wrap_fn_args,
311 work_fn_args . applyToVars wrap_args,
314 (tyvars, tau) = splitForAllTys fun_ty
315 (arg_tys, body_ty) = splitFunTys tau
316 n_arg_tys = length arg_tys
317 n_args | res_bot = n_arg_tys
318 | otherwise = arity `min` n_arg_tys
319 new_fun_ty | n_args == n_arg_tys = body_ty
320 | otherwise = mkFunTys (drop n_args arg_tys) body_ty
322 mkWWargs fun_ty arity demands res_bot one_shots
323 = case splitNewType_maybe fun_ty of
324 Nothing -> returnUs ([], id, id, fun_ty)
325 Just rep_ty -> mkWWargs rep_ty arity demands res_bot one_shots `thenUs` \ (wrap_args, wrap_fn_args, work_fn_args, res_ty) ->
327 Note (Coerce fun_ty rep_ty) . wrap_fn_args,
328 work_fn_args . Note (Coerce rep_ty fun_ty),
332 applyToVars :: [Var] -> CoreExpr -> CoreExpr
333 applyToVars vars fn = mkVarApps fn vars
335 mk_wrap_arg uniq ty dmd one_shot
336 = set_one_shot one_shot (setIdDemandInfo (mkSysLocal SLIT("w") uniq ty) dmd)
338 set_one_shot True id = setOneShotLambda id
339 set_one_shot False id = id
343 %************************************************************************
345 \subsection{Fixup stuff}
347 %************************************************************************
350 mkWWfixup res_ty work_dmds
351 | null work_dmds && isUnLiftedType res_ty
352 -- Horrid special case. If the worker would have no arguments, and the
353 -- function returns a primitive type value, that would make the worker into
354 -- an unboxed value. We box it by passing a dummy void argument, thus:
356 -- f = /\abc. \xyz. fw abc void
357 -- fw = /\abc. \v. body
359 -- We use the state-token type which generates no code
360 = getUniqueUs `thenUs` \ void_arg_uniq ->
362 void_arg = mk_ww_local void_arg_uniq realWorldStatePrimTy
365 \ call_to_worker -> App call_to_worker (Var realWorldPrimId),
366 \ worker_body -> Lam void_arg worker_body)
369 = returnUs (work_dmds, id, id)
373 %************************************************************************
375 \subsection{Strictness stuff}
377 %************************************************************************
380 mkWWstr :: [Var] -- Wrapper args; have their demand info on them
381 -- *Includes type variables*
382 -> UniqSM ([Demand], -- Demand on worker (value) args
383 CoreExpr -> CoreExpr, -- Wrapper body, lacking the worker call
384 -- and without its lambdas
385 -- This fn adds the unboxing, and makes the
386 -- call passing the unboxed things
388 CoreExpr -> CoreExpr) -- Worker body, lacking the original body of the function,
389 -- but *with* lambdas
392 = mk_ww_str wrap_args `thenUs` \ (work_args, wrap_fn, work_fn) ->
393 returnUs ( [idDemandInfo v | v <- work_args, isId v],
394 \ wrapper_body -> wrap_fn (mkVarApps wrapper_body work_args),
395 \ worker_body -> mkLams work_args (work_fn worker_body))
400 \ wrapper_body -> wrapper_body,
401 \ worker_body -> worker_body)
406 = mk_ww_str ds `thenUs` \ (worker_args, wrap_fn, work_fn) ->
407 returnUs (arg : worker_args, wrap_fn, work_fn)
410 = case idDemandInfo arg of
414 mk_ww_str ds `thenUs` \ (worker_args, wrap_fn, work_fn) ->
415 returnUs (worker_args, wrap_fn, mk_absent_let arg . work_fn)
418 WwUnpack new_or_data True cs ->
419 getUniquesUs (length inst_con_arg_tys) `thenUs` \ uniqs ->
421 unpk_args = zipWith mk_ww_local uniqs inst_con_arg_tys
422 unpk_args_w_ds = zipWithEqual "mk_ww_str" set_worker_arg_info unpk_args cs
424 mk_ww_str (unpk_args_w_ds ++ ds) `thenUs` \ (worker_args, wrap_fn, work_fn) ->
425 returnUs (worker_args,
426 mk_unpk_case new_or_data arg unpk_args data_con arg_tycon . wrap_fn,
427 work_fn . mk_pk_let new_or_data arg data_con tycon_arg_tys unpk_args)
429 (arg_tycon, tycon_arg_tys, data_con, inst_con_arg_tys) = splitProductType "mk_ww_str" (idType arg)
433 mk_ww_str ds `thenUs` \ (worker_args, wrap_fn, work_fn) ->
434 returnUs (arg : worker_args, wrap_fn, work_fn)
436 -- If the wrapper argument is a one-shot lambda, then
437 -- so should (all) the corresponding worker arguments be
438 -- This bites when we do w/w on a case join point
439 set_worker_arg_info worker_arg demand = set_one_shot (setIdDemandInfo worker_arg demand)
441 set_one_shot | isOneShotLambda arg = setOneShotLambda
442 | otherwise = \x -> x
446 %************************************************************************
448 \subsection{CPR stuff}
450 %************************************************************************
453 @mkWWcpr@ takes the worker/wrapper pair produced from the strictness
454 info and adds in the CPR transformation. The worker returns an
455 unboxed tuple containing non-CPR components. The wrapper takes this
456 tuple and re-produces the correct structured output.
458 The non-CPR results appear ordered in the unboxed tuple as if by a
459 left-to-right traversal of the result structure.
463 mkWWcpr :: Type -- function body type
464 -> CprInfo -- CPR analysis results
465 -> UniqSM (CoreExpr -> CoreExpr, -- New wrapper
466 CoreExpr -> CoreExpr, -- New worker
467 Type) -- Type of worker's body
469 mkWWcpr body_ty NoCPRInfo
470 = returnUs (id, id, body_ty) -- Must be just the strictness transf.
472 mkWWcpr body_ty ReturnsCPR
473 | not (isAlgType body_ty)
474 = WARN( True, text "mkWWcpr: non-algebraic body type" <+> ppr body_ty )
475 returnUs (id, id, body_ty)
477 | n_con_args == 1 && isUnLiftedType con_arg_ty1
478 -- Special case when there is a single result of unlifted type
479 = getUniquesUs 2 `thenUs` \ [work_uniq, arg_uniq] ->
481 work_wild = mk_ww_local work_uniq body_ty
482 arg = mk_ww_local arg_uniq con_arg_ty1
484 returnUs (\ wkr_call -> Case wkr_call arg [(DEFAULT, [], mkConApp data_con (map Type tycon_arg_tys ++ [Var arg]))],
485 \ body -> Case body work_wild [(DataAlt data_con, [arg], Var arg)],
488 | otherwise -- The general case
489 = getUniquesUs (n_con_args + 2) `thenUs` \ uniqs ->
491 (wrap_wild : work_wild : args) = zipWith mk_ww_local uniqs (ubx_tup_ty : body_ty : con_arg_tys)
492 arg_vars = map Var args
493 ubx_tup_con = unboxedTupleCon n_con_args
494 ubx_tup_ty = exprType ubx_tup_app
495 ubx_tup_app = mkConApp ubx_tup_con (map Type con_arg_tys ++ arg_vars)
496 con_app = mkConApp data_con (map Type tycon_arg_tys ++ arg_vars)
498 returnUs (\ wkr_call -> Case wkr_call wrap_wild [(DataAlt ubx_tup_con, args, con_app)],
499 \ body -> Case body work_wild [(DataAlt data_con, args, ubx_tup_app)],
502 (tycon, tycon_arg_tys, data_con, con_arg_tys) = splitProductType "mkWWcpr" body_ty
503 n_con_args = length con_arg_tys
504 con_arg_ty1 = head con_arg_tys
508 %************************************************************************
510 \subsection{Utilities}
512 %************************************************************************
516 mk_absent_let arg body
517 | not (isUnLiftedType arg_ty)
518 = Let (NonRec arg (mkTyApps (Var aBSENT_ERROR_ID) [arg_ty])) body
520 = panic "WwLib: haven't done mk_absent_let for primitives yet"
524 mk_unpk_case NewType arg unpk_args boxing_con boxing_tycon body
525 -- A newtype! Use a coercion not a case
526 = ASSERT( null other_args )
527 Case (Note (Coerce (idType unpk_arg) (idType arg)) (Var arg))
528 (sanitiseCaseBndr unpk_arg)
531 (unpk_arg:other_args) = unpk_args
533 mk_unpk_case DataType arg unpk_args boxing_con boxing_tycon body
536 (sanitiseCaseBndr arg)
537 [(DataAlt boxing_con, unpk_args, body)]
539 sanitiseCaseBndr :: Id -> Id
540 -- The argument we are scrutinising has the right type to be
541 -- a case binder, so it's convenient to re-use it for that purpose.
542 -- But we *must* throw away all its IdInfo. In particular, the argument
543 -- will have demand info on it, and that demand info may be incorrect for
544 -- the case binder. e.g. case ww_arg of ww_arg { I# x -> ... }
545 -- Quite likely ww_arg isn't used in '...'. The case may get discarded
546 -- if the case binder says "I'm demanded". This happened in a situation
547 -- like (x+y) `seq` ....
548 sanitiseCaseBndr id = id `setIdInfo` vanillaIdInfo
550 mk_pk_let NewType arg boxing_con con_tys unpk_args body
551 = ASSERT( null other_args )
552 Let (NonRec arg (Note (Coerce (idType arg) (idType unpk_arg)) (Var unpk_arg))) body
554 (unpk_arg:other_args) = unpk_args
556 mk_pk_let DataType arg boxing_con con_tys unpk_args body
557 = Let (NonRec arg (mkConApp boxing_con con_args)) body
559 con_args = map Type con_tys ++ map Var unpk_args
562 mk_ww_local uniq ty = mkSysLocal SLIT("ww") uniq ty