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 )
16 import Id ( Id, idType, mkSysLocal, idDemandInfo, setIdDemandInfo,
17 isOneShotLambda, setOneShotLambda,
20 import IdInfo ( CprInfo(..), vanillaIdInfo )
21 import DataCon ( splitProductType )
22 import Demand ( Demand(..), wwLazy, wwPrim )
23 import PrelInfo ( realWorldPrimId, aBSENT_ERROR_ID )
24 import TysPrim ( realWorldStatePrimTy )
25 import TysWiredIn ( tupleCon )
26 import Type ( Type, isUnLiftedType, mkFunTys,
27 splitForAllTys, splitFunTys, isAlgType
29 import BasicTypes ( Arity, Boxity(..) )
30 import Var ( Var, isId )
31 import UniqSupply ( returnUs, thenUs, getUniqueUs, getUniquesUs, UniqSM )
32 import Util ( zipWithEqual )
34 import List ( zipWith4 )
38 %************************************************************************
40 \subsection[mkWrapperAndWorker]{@mkWrapperAndWorker@}
42 %************************************************************************
44 ************ WARNING ******************
45 these comments are rather out of date
46 *****************************************
48 @mkWrapperAndWorker@ is given:
51 The {\em original function} \tr{f}, of the form:
53 f = /\ tyvars -> \ args -> body
55 The original-binder \tr{f}, the \tr{tyvars}, \tr{args}, and \tr{body}
58 We use the Id \tr{f} mostly to get its type.
61 Strictness information about \tr{f}, in the form of a list of
68 @mkWrapperAndWorker@ produces (A BIT OUT-OF-DATE...):
71 Maybe @Nothing@: no worker/wrappering going on in this case. This can
72 happen (a)~if the strictness info says that there is nothing
73 interesting to do or (b)~if *any* of the argument types corresponding
74 to ``active'' arg postitions is abstract or will be to the outside
75 world (i.e., {\em this} module can see the constructors, but nobody
76 else will be able to). An ``active'' arg position is one which the
77 wrapper has to unpack. An importing module can't do this unpacking,
78 so it simply has to give up and call the wrapper only.
81 Maybe \tr{Just (wrapper_Id, wrapper_body, worker_Id, worker_body)}.
83 The @wrapper_Id@ is just the one that was passed in, with its
84 strictness IdInfo updated.
87 The \tr{body} of the original function may not be given (i.e., it's
88 BOTTOM), in which case you'd jolly well better not tug on the
91 Here's an example. The original function is:
93 g :: forall a . Int -> [a] -> a
101 From this, we want to produce:
103 -- wrapper (an unfolding)
104 g :: forall a . Int -> [a] -> a
106 g = /\ a -> \ x ys ->
108 I# x# -> g.wrk a x# ys
109 -- call the worker; don't forget the type args!
112 g.wrk :: forall a . Int# -> [a] -> a
114 g.wrk = /\ a -> \ x# ys ->
118 case x of -- note: body of g moved intact
123 Something we have to be careful about: Here's an example:
125 -- "f" strictness: U(P)U(P)
126 f (I# a) (I# b) = a +# b
128 g = f -- "g" strictness same as "f"
130 \tr{f} will get a worker all nice and friendly-like; that's good.
131 {\em But we don't want a worker for \tr{g}}, even though it has the
132 same strictness as \tr{f}. Doing so could break laziness, at best.
134 Consequently, we insist that the number of strictness-info items is
135 exactly the same as the number of lambda-bound arguments. (This is
136 probably slightly paranoid, but OK in practice.) If it isn't the
137 same, we ``revise'' the strictness info, so that we won't propagate
138 the unusable strictness-info into the interfaces.
141 %************************************************************************
143 \subsection{Functions over Demands}
145 %************************************************************************
148 mAX_WORKER_ARGS :: Int -- ToDo: set via flag
151 setUnpackStrategy :: [Demand] -> [Demand]
153 = snd (go (mAX_WORKER_ARGS - nonAbsentArgs ds) ds)
155 go :: Int -- Max number of args available for sub-components of [Demand]
157 -> (Int, [Demand]) -- Args remaining after subcomponents of [Demand] are unpacked
159 go n (WwUnpack _ cs : ds) | n' >= 0
160 = WwUnpack True cs' `cons` go n'' ds
162 = WwUnpack False cs `cons` go n ds
164 n' = n + 1 - nonAbsentArgs cs
165 -- Add one because we don't pass the top-level arg any more
166 -- Delete # of non-absent args to which we'll now be committed
169 go n (d:ds) = d `cons` go n ds
172 cons d (n,ds) = (n, d:ds)
174 nonAbsentArgs :: [Demand] -> Int
176 nonAbsentArgs (WwLazy True : ds) = nonAbsentArgs ds
177 nonAbsentArgs (d : ds) = 1 + nonAbsentArgs ds
179 worthSplitting :: [Demand]
180 -> Bool -- Result is bottom
181 -> Bool -- True <=> the wrapper would not be an identity function
182 worthSplitting ds result_bot = any worth_it ds
183 -- We used not to split if the result is bottom.
184 -- [Justification: there's no efficiency to be gained.]
185 -- But it's sometimes bad not to make a wrapper. Consider
186 -- fw = \x# -> let x = I# x# in case e of
189 -- p3 -> the real stuff
190 -- The re-boxing code won't go away unless error_fn gets a wrapper too.
193 worth_it (WwLazy True) = True -- Absent arg
194 worth_it (WwUnpack True _) = True -- Arg to unpack
195 worth_it WwStrict = False -- Don't w/w just because of strictness
196 worth_it other = False
198 allAbsent :: [Demand] -> Bool
199 allAbsent ds = all absent ds
201 absent (WwLazy is_absent) = is_absent
202 absent (WwUnpack True cs) = allAbsent cs
207 %************************************************************************
209 \subsection{The worker wrapper core}
211 %************************************************************************
213 @mkWwBodies@ is called when doing the worker/wrapper split inside a module.
216 mkWwBodies :: Type -- Type of original function
217 -> Arity -- Arity of original function
218 -> [Demand] -- Strictness of original function
219 -> Bool -- True <=> function returns bottom
220 -> [Bool] -- One-shot-ness of the function
221 -> CprInfo -- Result of CPR analysis
222 -> UniqSM ([Demand], -- Demands for worker (value) args
223 Id -> CoreExpr, -- Wrapper body, lacking only the worker Id
224 CoreExpr -> CoreExpr) -- Worker body, lacking the original function rhs
226 -- wrap_fn_args E = \x y -> E
227 -- work_fn_args E = E x y
229 -- wrap_fn_str E = case x of { (a,b) ->
230 -- case a of { (a1,a2) ->
232 -- work_fn_str E = \a2 a2 b y ->
233 -- let a = (a1,a2) in
237 mkWwBodies fun_ty arity demands res_bot one_shots cpr_info
238 = mkWWargs fun_ty arity demands' res_bot one_shots' `thenUs` \ (wrap_args, wrap_fn_args, work_fn_args, res_ty) ->
239 mkWWcpr res_ty cpr_info `thenUs` \ (wrap_fn_cpr, work_fn_cpr, cpr_res_ty) ->
240 mkWWstr cpr_res_ty wrap_args `thenUs` \ (work_dmds, wrap_fn_str, work_fn_str) ->
243 Note InlineMe . wrap_fn_args . wrap_fn_cpr . wrap_fn_str . Var,
244 work_fn_str . work_fn_cpr . work_fn_args)
245 -- We use an INLINE unconditionally, even if the wrapper turns out to be
246 -- something trivial like
248 -- f = __inline__ (coerce T fw)
249 -- The point is to propagate the coerce to f's call sites, so even though
250 -- f's RHS is now trivial (size 1) we still want the __inline__ to prevent
251 -- fw from being inlined into f's RHS
253 demands' = demands ++ repeat wwLazy
254 one_shots' = one_shots ++ repeat False
258 %************************************************************************
260 \subsection{Coercion stuff}
262 %************************************************************************
265 We really want to "look through" coerces.
266 Reason: I've seen this situation:
268 let f = coerce T (\s -> E)
274 If only we w/w'd f, we'd get
275 let f = coerce T (\s -> fw s)
279 Now we'll inline f to get
287 Now we'll see that fw has arity 1, and will arity expand
288 the \x to get what we want.
291 -- mkWWargs is driven off the function type and arity.
292 -- It chomps bites off foralls, arrows, newtypes
293 -- and keeps repeating that until it's satisfied the supplied arity
295 mkWWargs :: Type -> Arity
296 -> [Demand] -> Bool -> [Bool] -- Both these will in due course be derived
297 -- from the type. The [Bool] is True for a one-shot arg.
298 -- ** Both are infinite, extended with neutral values if necy **
299 -> UniqSM ([Var], -- Wrapper args
300 CoreExpr -> CoreExpr, -- Wrapper fn
301 CoreExpr -> CoreExpr, -- Worker fn
302 Type) -- Type of wrapper body
304 mkWWargs fun_ty arity demands res_bot one_shots
305 | (res_bot || arity > 0) && (not (null tyvars) || n_arg_tys > 0)
306 -- If the function returns bottom, we feel free to
307 -- build lots of wrapper args:
308 -- \x. let v=E in \y. bottom
309 -- = \xy. let v=E in bottom
310 = getUniquesUs `thenUs` \ wrap_uniqs ->
312 val_args = zipWith4 mk_wrap_arg wrap_uniqs arg_tys demands one_shots
313 wrap_args = tyvars ++ val_args
317 (drop n_args demands)
319 (drop n_args one_shots) `thenUs` \ (more_wrap_args, wrap_fn_args, work_fn_args, res_ty) ->
321 returnUs (wrap_args ++ more_wrap_args,
322 mkLams wrap_args . wrap_fn_args,
323 work_fn_args . applyToVars wrap_args,
326 (tyvars, tau) = splitForAllTys fun_ty
327 (arg_tys, body_ty) = splitFunTys tau
328 n_arg_tys = length arg_tys
329 n_args | res_bot = n_arg_tys
330 | otherwise = arity `min` n_arg_tys
331 new_fun_ty | n_args == n_arg_tys = body_ty
332 | otherwise = mkFunTys (drop n_args arg_tys) body_ty
334 mkWWargs fun_ty arity demands res_bot one_shots
335 = returnUs ([], id, id, fun_ty)
337 applyToVars :: [Var] -> CoreExpr -> CoreExpr
338 applyToVars vars fn = mkVarApps fn vars
340 mk_wrap_arg uniq ty dmd one_shot
341 = set_one_shot one_shot (setIdDemandInfo (mkSysLocal SLIT("w") uniq ty) dmd)
343 set_one_shot True id = setOneShotLambda id
344 set_one_shot False id = id
348 %************************************************************************
350 \subsection{Strictness stuff}
352 %************************************************************************
355 mkWWstr :: Type -- Result type
356 -> [Var] -- Wrapper args; have their demand info on them
357 -- *Includes type variables*
358 -> UniqSM ([Demand], -- Demand on worker (value) args
359 CoreExpr -> CoreExpr, -- Wrapper body, lacking the worker call
360 -- and without its lambdas
361 -- This fn adds the unboxing, and makes the
362 -- call passing the unboxed things
364 CoreExpr -> CoreExpr) -- Worker body, lacking the original body of the function,
365 -- but *with* lambdas
367 mkWWstr res_ty wrap_args
368 = mk_ww_str wrap_args `thenUs` \ (work_args, take_apart, put_together) ->
370 work_dmds = [idDemandInfo v | v <- work_args, isId v]
371 apply_to args fn = mkVarApps fn args
373 if not (null work_dmds && isUnLiftedType res_ty) then
374 returnUs ( work_dmds,
375 take_apart . apply_to work_args,
376 mkLams work_args . put_together)
378 -- Horrid special case. If the worker would have no arguments, and the
379 -- function returns a primitive type value, that would make the worker into
380 -- an unboxed value. We box it by passing a dummy void argument, thus:
382 -- f = /\abc. \xyz. fw abc void
383 -- fw = /\abc. \v. body
385 -- We use the state-token type which generates no code
386 getUniqueUs `thenUs` \ void_arg_uniq ->
388 void_arg = mk_ww_local void_arg_uniq realWorldStatePrimTy
391 take_apart . apply_to [realWorldPrimId] . apply_to work_args,
392 mkLams work_args . Lam void_arg . put_together)
397 \ wrapper_body -> wrapper_body,
398 \ worker_body -> worker_body)
403 = mk_ww_str ds `thenUs` \ (worker_args, wrap_fn, work_fn) ->
404 returnUs (arg : worker_args, wrap_fn, work_fn)
407 = case idDemandInfo arg of
411 mk_ww_str ds `thenUs` \ (worker_args, wrap_fn, work_fn) ->
412 returnUs (worker_args, wrap_fn, mk_absent_let arg . work_fn)
416 getUniquesUs `thenUs` \ uniqs ->
418 unpk_args = zipWith mk_ww_local uniqs inst_con_arg_tys
419 unpk_args_w_ds = zipWithEqual "mk_ww_str" set_worker_arg_info unpk_args cs
421 mk_ww_str (unpk_args_w_ds ++ ds) `thenUs` \ (worker_args, wrap_fn, work_fn) ->
422 returnUs (worker_args,
423 mk_unpk_case arg unpk_args data_con arg_tycon . wrap_fn,
424 work_fn . mk_pk_let arg data_con tycon_arg_tys unpk_args)
426 (arg_tycon, tycon_arg_tys, data_con, inst_con_arg_tys) = splitProductType "mk_ww_str" (idType arg)
430 mk_ww_str ds `thenUs` \ (worker_args, wrap_fn, work_fn) ->
431 returnUs (arg : worker_args, wrap_fn, work_fn)
433 -- If the wrapper argument is a one-shot lambda, then
434 -- so should (all) the corresponding worker arguments be
435 -- This bites when we do w/w on a case join point
436 set_worker_arg_info worker_arg demand = set_one_shot (setIdDemandInfo worker_arg demand)
438 set_one_shot | isOneShotLambda arg = setOneShotLambda
439 | otherwise = \x -> x
443 %************************************************************************
445 \subsection{CPR stuff}
447 %************************************************************************
450 @mkWWcpr@ takes the worker/wrapper pair produced from the strictness
451 info and adds in the CPR transformation. The worker returns an
452 unboxed tuple containing non-CPR components. The wrapper takes this
453 tuple and re-produces the correct structured output.
455 The non-CPR results appear ordered in the unboxed tuple as if by a
456 left-to-right traversal of the result structure.
460 mkWWcpr :: Type -- function body type
461 -> CprInfo -- CPR analysis results
462 -> UniqSM (CoreExpr -> CoreExpr, -- New wrapper
463 CoreExpr -> CoreExpr, -- New worker
464 Type) -- Type of worker's body
466 mkWWcpr body_ty NoCPRInfo
467 = returnUs (id, id, body_ty) -- Must be just the strictness transf.
469 mkWWcpr body_ty ReturnsCPR
470 | not (isAlgType body_ty)
471 = WARN( True, text "mkWWcpr: non-algebraic body type" <+> ppr body_ty )
472 returnUs (id, id, body_ty)
474 | n_con_args == 1 && isUnLiftedType con_arg_ty1
475 -- Special case when there is a single result of unlifted type
476 = getUniquesUs `thenUs` \ (work_uniq : arg_uniq : _) ->
478 work_wild = mk_ww_local work_uniq body_ty
479 arg = mk_ww_local arg_uniq con_arg_ty1
481 returnUs (\ wkr_call -> Case wkr_call arg [(DEFAULT, [], mkConApp data_con (map Type tycon_arg_tys ++ [Var arg]))],
482 \ body -> workerCase body work_wild [(DataAlt data_con, [arg], Var arg)],
485 | otherwise -- The general case
486 = getUniquesUs `thenUs` \ uniqs ->
488 (wrap_wild : work_wild : args) = zipWith mk_ww_local uniqs (ubx_tup_ty : body_ty : con_arg_tys)
489 arg_vars = map Var args
490 ubx_tup_con = tupleCon Unboxed n_con_args
491 ubx_tup_ty = exprType ubx_tup_app
492 ubx_tup_app = mkConApp ubx_tup_con (map Type con_arg_tys ++ arg_vars)
493 con_app = mkConApp data_con (map Type tycon_arg_tys ++ arg_vars)
495 returnUs (\ wkr_call -> Case wkr_call wrap_wild [(DataAlt ubx_tup_con, args, con_app)],
496 \ body -> workerCase body work_wild [(DataAlt data_con, args, ubx_tup_app)],
499 (_, tycon_arg_tys, data_con, con_arg_tys) = splitProductType "mkWWcpr" body_ty
500 n_con_args = length con_arg_tys
501 con_arg_ty1 = head con_arg_tys
503 -- If the original function looked like
504 -- f = \ x -> _scc_ "foo" E
506 -- then we want the CPR'd worker to look like
507 -- \ x -> _scc_ "foo" (case E of I# x -> x)
508 -- and definitely not
509 -- \ x -> case (_scc_ "foo" E) of I# x -> x)
511 -- This transform doesn't move work or allocation
512 -- from one cost centre to another
514 workerCase (Note (SCC cc) e) arg alts = Note (SCC cc) (Case e arg alts)
515 workerCase e arg alts = Case e arg alts
519 %************************************************************************
521 \subsection{Utilities}
523 %************************************************************************
527 mk_absent_let arg body
528 | not (isUnLiftedType arg_ty)
529 = Let (NonRec arg (mkTyApps (Var aBSENT_ERROR_ID) [arg_ty])) body
531 = panic "WwLib: haven't done mk_absent_let for primitives yet"
535 mk_unpk_case arg unpk_args boxing_con boxing_tycon body
538 (sanitiseCaseBndr arg)
539 [(DataAlt boxing_con, unpk_args, body)]
541 sanitiseCaseBndr :: Id -> Id
542 -- The argument we are scrutinising has the right type to be
543 -- a case binder, so it's convenient to re-use it for that purpose.
544 -- But we *must* throw away all its IdInfo. In particular, the argument
545 -- will have demand info on it, and that demand info may be incorrect for
546 -- the case binder. e.g. case ww_arg of ww_arg { I# x -> ... }
547 -- Quite likely ww_arg isn't used in '...'. The case may get discarded
548 -- if the case binder says "I'm demanded". This happened in a situation
549 -- like (x+y) `seq` ....
550 sanitiseCaseBndr id = id `setIdInfo` vanillaIdInfo
552 mk_pk_let arg boxing_con con_tys unpk_args body
553 = Let (NonRec arg (mkConApp boxing_con con_args)) body
555 con_args = map Type con_tys ++ map Var unpk_args
558 mk_ww_local uniq ty = mkSysLocal SLIT("ww") uniq ty