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,
27 splitForAllTys, splitFunTys, isAlgType,
28 splitNewType_maybe, mkFunTys
30 import BasicTypes ( NewOrData(..), Arity, Boxity(..) )
31 import Var ( Var, isId )
32 import UniqSupply ( returnUs, thenUs, getUniqueUs, getUniquesUs, UniqSM )
33 import Util ( zipWithEqual )
35 import List ( zipWith4 )
39 %************************************************************************
41 \subsection[mkWrapperAndWorker]{@mkWrapperAndWorker@}
43 %************************************************************************
45 ************ WARNING ******************
46 these comments are rather out of date
47 *****************************************
49 @mkWrapperAndWorker@ is given:
52 The {\em original function} \tr{f}, of the form:
54 f = /\ tyvars -> \ args -> body
56 The original-binder \tr{f}, the \tr{tyvars}, \tr{args}, and \tr{body}
59 We use the Id \tr{f} mostly to get its type.
62 Strictness information about \tr{f}, in the form of a list of
69 @mkWrapperAndWorker@ produces (A BIT OUT-OF-DATE...):
72 Maybe @Nothing@: no worker/wrappering going on in this case. This can
73 happen (a)~if the strictness info says that there is nothing
74 interesting to do or (b)~if *any* of the argument types corresponding
75 to ``active'' arg postitions is abstract or will be to the outside
76 world (i.e., {\em this} module can see the constructors, but nobody
77 else will be able to). An ``active'' arg position is one which the
78 wrapper has to unpack. An importing module can't do this unpacking,
79 so it simply has to give up and call the wrapper only.
82 Maybe \tr{Just (wrapper_Id, wrapper_body, worker_Id, worker_body)}.
84 The @wrapper_Id@ is just the one that was passed in, with its
85 strictness IdInfo updated.
88 The \tr{body} of the original function may not be given (i.e., it's
89 BOTTOM), in which case you'd jolly well better not tug on the
92 Here's an example. The original function is:
94 g :: forall a . Int -> [a] -> a
102 From this, we want to produce:
104 -- wrapper (an unfolding)
105 g :: forall a . Int -> [a] -> a
107 g = /\ a -> \ x ys ->
109 I# x# -> g.wrk a x# ys
110 -- call the worker; don't forget the type args!
113 g.wrk :: forall a . Int# -> [a] -> a
115 g.wrk = /\ a -> \ x# ys ->
119 case x of -- note: body of g moved intact
124 Something we have to be careful about: Here's an example:
126 -- "f" strictness: U(P)U(P)
127 f (I# a) (I# b) = a +# b
129 g = f -- "g" strictness same as "f"
131 \tr{f} will get a worker all nice and friendly-like; that's good.
132 {\em But we don't want a worker for \tr{g}}, even though it has the
133 same strictness as \tr{f}. Doing so could break laziness, at best.
135 Consequently, we insist that the number of strictness-info items is
136 exactly the same as the number of lambda-bound arguments. (This is
137 probably slightly paranoid, but OK in practice.) If it isn't the
138 same, we ``revise'' the strictness info, so that we won't propagate
139 the unusable strictness-info into the interfaces.
142 %************************************************************************
144 \subsection{Functions over Demands}
146 %************************************************************************
149 mAX_WORKER_ARGS :: Int -- ToDo: set via flag
152 setUnpackStrategy :: [Demand] -> [Demand]
154 = snd (go (mAX_WORKER_ARGS - nonAbsentArgs ds) ds)
156 go :: Int -- Max number of args available for sub-components of [Demand]
158 -> (Int, [Demand]) -- Args remaining after subcomponents of [Demand] are unpacked
160 go n (WwUnpack nd _ cs : ds) | n' >= 0
161 = WwUnpack nd True cs' `cons` go n'' ds
163 = WwUnpack nd False cs `cons` go n ds
165 n' = n + 1 - nonAbsentArgs cs
166 -- Add one because we don't pass the top-level arg any more
167 -- Delete # of non-absent args to which we'll now be committed
170 go n (d:ds) = d `cons` go n ds
173 cons d (n,ds) = (n, d:ds)
175 nonAbsentArgs :: [Demand] -> Int
177 nonAbsentArgs (WwLazy True : ds) = nonAbsentArgs ds
178 nonAbsentArgs (d : ds) = 1 + nonAbsentArgs ds
180 worthSplitting :: [Demand]
181 -> Bool -- Result is bottom
182 -> Bool -- True <=> the wrapper would not be an identity function
183 worthSplitting ds result_bot = any worth_it ds
184 -- We used not to split if the result is bottom.
185 -- [Justification: there's no efficiency to be gained.]
186 -- But it's sometimes bad not to make a wrapper. Consider
187 -- fw = \x# -> let x = I# x# in case e of
190 -- p3 -> the real stuff
191 -- The re-boxing code won't go away unless error_fn gets a wrapper too.
194 worth_it (WwLazy True) = True -- Absent arg
195 worth_it (WwUnpack _ True _) = True -- Arg to unpack
196 worth_it WwStrict = False -- Don't w/w just because of strictness
197 worth_it other = False
199 allAbsent :: [Demand] -> Bool
200 allAbsent ds = all absent ds
202 absent (WwLazy is_absent) = is_absent
203 absent (WwUnpack _ True cs) = allAbsent cs
208 %************************************************************************
210 \subsection{The worker wrapper core}
212 %************************************************************************
214 @mkWwBodies@ is called when doing the worker/wrapper split inside a module.
217 mkWwBodies :: Type -- Type of original function
218 -> Arity -- Arity of original function
219 -> [Demand] -- Strictness of original function
220 -> Bool -- True <=> function returns bottom
221 -> [Bool] -- One-shot-ness of the function
222 -> CprInfo -- Result of CPR analysis
223 -> UniqSM ([Demand], -- Demands for worker (value) args
224 Id -> CoreExpr, -- Wrapper body, lacking only the worker Id
225 CoreExpr -> CoreExpr) -- Worker body, lacking the original function rhs
227 mkWwBodies fun_ty arity demands res_bot one_shots cpr_info
228 = mkWWargs fun_ty arity demands' res_bot one_shots' `thenUs` \ (wrap_args, wrap_fn_args, work_fn_args, res_ty) ->
229 mkWWstr wrap_args `thenUs` \ (work_dmds, wrap_fn_str, work_fn_str) ->
230 mkWWcpr res_ty cpr_info `thenUs` \ (wrap_fn_cpr, work_fn_cpr, cpr_res_ty) ->
231 mkWWfixup cpr_res_ty work_dmds `thenUs` \ (final_work_dmds, wrap_fn_fixup, work_fn_fixup) ->
233 returnUs (final_work_dmds,
234 Note InlineMe . wrap_fn_args . wrap_fn_cpr . wrap_fn_str . wrap_fn_fixup . Var,
235 work_fn_fixup . work_fn_str . work_fn_cpr . work_fn_args)
236 -- We use an INLINE unconditionally, even if the wrapper turns out to be
237 -- something trivial like
239 -- f = __inline__ (coerce T fw)
240 -- The point is to propagate the coerce to f's call sites, so even though
241 -- f's RHS is now trivial (size 1) we still want the __inline__ to prevent
242 -- fw from being inlined into f's RHS
244 demands' = demands ++ repeat wwLazy
245 one_shots' = one_shots ++ repeat False
249 %************************************************************************
251 \subsection{Coercion stuff}
253 %************************************************************************
256 We really want to "look through" coerces.
257 Reason: I've seen this situation:
259 let f = coerce T (\s -> E)
265 If only we w/w'd f, we'd get
266 let f = coerce T (\s -> fw s)
270 Now we'll inline f to get
278 Now we'll see that fw has arity 1, and will arity expand
279 the \x to get what we want.
282 -- mkWWargs is driven off the function type and arity.
283 -- It chomps bites off foralls, arrows, newtypes
284 -- and keeps repeating that until it's satisfied the supplied arity
286 mkWWargs :: Type -> Arity
287 -> [Demand] -> Bool -> [Bool] -- Both these will in due course be derived
288 -- from the type. The [Bool] is True for a one-shot arg.
289 -- ** Both are infinite, extended with neutral values if necy **
290 -> UniqSM ([Var], -- Wrapper args
291 CoreExpr -> CoreExpr, -- Wrapper fn
292 CoreExpr -> CoreExpr, -- Worker fn
293 Type) -- Type of wrapper body
295 mkWWargs fun_ty arity demands res_bot one_shots
296 | (res_bot || arity > 0) && (not (null tyvars) || n_arg_tys > 0)
297 -- If the function returns bottom, we feel free to
298 -- build lots of wrapper args:
299 -- \x. let v=E in \y. bottom
300 -- = \xy. let v=E in bottom
301 = getUniquesUs n_args `thenUs` \ wrap_uniqs ->
303 val_args = zipWith4 mk_wrap_arg wrap_uniqs arg_tys demands one_shots
304 wrap_args = tyvars ++ val_args
308 (drop n_args demands)
310 (drop n_args one_shots) `thenUs` \ (more_wrap_args, wrap_fn_args, work_fn_args, res_ty) ->
312 returnUs (wrap_args ++ more_wrap_args,
313 mkLams wrap_args . wrap_fn_args,
314 work_fn_args . applyToVars wrap_args,
317 (tyvars, tau) = splitForAllTys fun_ty
318 (arg_tys, body_ty) = splitFunTys tau
319 n_arg_tys = length arg_tys
320 n_args | res_bot = n_arg_tys
321 | otherwise = arity `min` n_arg_tys
322 new_fun_ty | n_args == n_arg_tys = body_ty
323 | otherwise = mkFunTys (drop n_args arg_tys) body_ty
325 mkWWargs fun_ty arity demands res_bot one_shots
326 = case splitNewType_maybe fun_ty of
327 Nothing -> returnUs ([], id, id, fun_ty)
328 Just rep_ty -> mkWWargs rep_ty arity demands res_bot one_shots `thenUs` \ (wrap_args, wrap_fn_args, work_fn_args, res_ty) ->
330 Note (Coerce fun_ty rep_ty) . wrap_fn_args,
331 work_fn_args . Note (Coerce rep_ty fun_ty),
335 applyToVars :: [Var] -> CoreExpr -> CoreExpr
336 applyToVars vars fn = mkVarApps fn vars
338 mk_wrap_arg uniq ty dmd one_shot
339 = set_one_shot one_shot (setIdDemandInfo (mkSysLocal SLIT("w") uniq ty) dmd)
341 set_one_shot True id = setOneShotLambda id
342 set_one_shot False id = id
346 %************************************************************************
348 \subsection{Fixup stuff}
350 %************************************************************************
353 mkWWfixup res_ty work_dmds
354 | null work_dmds && isUnLiftedType res_ty
355 -- Horrid special case. If the worker would have no arguments, and the
356 -- function returns a primitive type value, that would make the worker into
357 -- an unboxed value. We box it by passing a dummy void argument, thus:
359 -- f = /\abc. \xyz. fw abc void
360 -- fw = /\abc. \v. body
362 -- We use the state-token type which generates no code
363 = getUniqueUs `thenUs` \ void_arg_uniq ->
365 void_arg = mk_ww_local void_arg_uniq realWorldStatePrimTy
368 \ call_to_worker -> App call_to_worker (Var realWorldPrimId),
369 \ worker_body -> Lam void_arg worker_body)
372 = returnUs (work_dmds, id, id)
376 %************************************************************************
378 \subsection{Strictness stuff}
380 %************************************************************************
383 mkWWstr :: [Var] -- Wrapper args; have their demand info on them
384 -- *Includes type variables*
385 -> UniqSM ([Demand], -- Demand on worker (value) args
386 CoreExpr -> CoreExpr, -- Wrapper body, lacking the worker call
387 -- and without its lambdas
388 -- This fn adds the unboxing, and makes the
389 -- call passing the unboxed things
391 CoreExpr -> CoreExpr) -- Worker body, lacking the original body of the function,
392 -- but *with* lambdas
395 = mk_ww_str wrap_args `thenUs` \ (work_args, wrap_fn, work_fn) ->
396 returnUs ( [idDemandInfo v | v <- work_args, isId v],
397 \ wrapper_body -> wrap_fn (mkVarApps wrapper_body work_args),
398 \ worker_body -> mkLams work_args (work_fn worker_body))
403 \ wrapper_body -> wrapper_body,
404 \ worker_body -> worker_body)
409 = mk_ww_str ds `thenUs` \ (worker_args, wrap_fn, work_fn) ->
410 returnUs (arg : worker_args, wrap_fn, work_fn)
413 = case idDemandInfo arg of
417 mk_ww_str ds `thenUs` \ (worker_args, wrap_fn, work_fn) ->
418 returnUs (worker_args, wrap_fn, mk_absent_let arg . work_fn)
421 WwUnpack new_or_data True cs ->
422 getUniquesUs (length inst_con_arg_tys) `thenUs` \ uniqs ->
424 unpk_args = zipWith mk_ww_local uniqs inst_con_arg_tys
425 unpk_args_w_ds = zipWithEqual "mk_ww_str" set_worker_arg_info unpk_args cs
427 mk_ww_str (unpk_args_w_ds ++ ds) `thenUs` \ (worker_args, wrap_fn, work_fn) ->
428 returnUs (worker_args,
429 mk_unpk_case new_or_data arg unpk_args data_con arg_tycon . wrap_fn,
430 work_fn . mk_pk_let new_or_data arg data_con tycon_arg_tys unpk_args)
432 (arg_tycon, tycon_arg_tys, data_con, inst_con_arg_tys) = splitProductType "mk_ww_str" (idType arg)
436 mk_ww_str ds `thenUs` \ (worker_args, wrap_fn, work_fn) ->
437 returnUs (arg : worker_args, wrap_fn, work_fn)
439 -- If the wrapper argument is a one-shot lambda, then
440 -- so should (all) the corresponding worker arguments be
441 -- This bites when we do w/w on a case join point
442 set_worker_arg_info worker_arg demand = set_one_shot (setIdDemandInfo worker_arg demand)
444 set_one_shot | isOneShotLambda arg = setOneShotLambda
445 | otherwise = \x -> x
449 %************************************************************************
451 \subsection{CPR stuff}
453 %************************************************************************
456 @mkWWcpr@ takes the worker/wrapper pair produced from the strictness
457 info and adds in the CPR transformation. The worker returns an
458 unboxed tuple containing non-CPR components. The wrapper takes this
459 tuple and re-produces the correct structured output.
461 The non-CPR results appear ordered in the unboxed tuple as if by a
462 left-to-right traversal of the result structure.
466 mkWWcpr :: Type -- function body type
467 -> CprInfo -- CPR analysis results
468 -> UniqSM (CoreExpr -> CoreExpr, -- New wrapper
469 CoreExpr -> CoreExpr, -- New worker
470 Type) -- Type of worker's body
472 mkWWcpr body_ty NoCPRInfo
473 = returnUs (id, id, body_ty) -- Must be just the strictness transf.
475 mkWWcpr body_ty ReturnsCPR
476 | not (isAlgType body_ty)
477 = WARN( True, text "mkWWcpr: non-algebraic body type" <+> ppr body_ty )
478 returnUs (id, id, body_ty)
480 | n_con_args == 1 && isUnLiftedType con_arg_ty1
481 -- Special case when there is a single result of unlifted type
482 = getUniquesUs 2 `thenUs` \ [work_uniq, arg_uniq] ->
484 work_wild = mk_ww_local work_uniq body_ty
485 arg = mk_ww_local arg_uniq con_arg_ty1
487 returnUs (\ wkr_call -> Case wkr_call arg [(DEFAULT, [], mkConApp data_con (map Type tycon_arg_tys ++ [Var arg]))],
488 \ body -> Case body work_wild [(DataAlt data_con, [arg], Var arg)],
491 | otherwise -- The general case
492 = getUniquesUs (n_con_args + 2) `thenUs` \ uniqs ->
494 (wrap_wild : work_wild : args) = zipWith mk_ww_local uniqs (ubx_tup_ty : body_ty : con_arg_tys)
495 arg_vars = map Var args
496 ubx_tup_con = tupleCon Unboxed n_con_args
497 ubx_tup_ty = exprType ubx_tup_app
498 ubx_tup_app = mkConApp ubx_tup_con (map Type con_arg_tys ++ arg_vars)
499 con_app = mkConApp data_con (map Type tycon_arg_tys ++ arg_vars)
501 returnUs (\ wkr_call -> Case wkr_call wrap_wild [(DataAlt ubx_tup_con, args, con_app)],
502 \ body -> Case body work_wild [(DataAlt data_con, args, ubx_tup_app)],
505 (tycon, tycon_arg_tys, data_con, con_arg_tys) = splitProductType "mkWWcpr" body_ty
506 n_con_args = length con_arg_tys
507 con_arg_ty1 = head con_arg_tys
511 %************************************************************************
513 \subsection{Utilities}
515 %************************************************************************
519 mk_absent_let arg body
520 | not (isUnLiftedType arg_ty)
521 = Let (NonRec arg (mkTyApps (Var aBSENT_ERROR_ID) [arg_ty])) body
523 = panic "WwLib: haven't done mk_absent_let for primitives yet"
527 mk_unpk_case NewType arg unpk_args boxing_con boxing_tycon body
528 -- A newtype! Use a coercion not a case
529 = ASSERT( null other_args )
530 Case (Note (Coerce (idType unpk_arg) (idType arg)) (Var arg))
531 (sanitiseCaseBndr unpk_arg)
534 (unpk_arg:other_args) = unpk_args
536 mk_unpk_case DataType arg unpk_args boxing_con boxing_tycon body
539 (sanitiseCaseBndr arg)
540 [(DataAlt boxing_con, unpk_args, body)]
542 sanitiseCaseBndr :: Id -> Id
543 -- The argument we are scrutinising has the right type to be
544 -- a case binder, so it's convenient to re-use it for that purpose.
545 -- But we *must* throw away all its IdInfo. In particular, the argument
546 -- will have demand info on it, and that demand info may be incorrect for
547 -- the case binder. e.g. case ww_arg of ww_arg { I# x -> ... }
548 -- Quite likely ww_arg isn't used in '...'. The case may get discarded
549 -- if the case binder says "I'm demanded". This happened in a situation
550 -- like (x+y) `seq` ....
551 sanitiseCaseBndr id = id `setIdInfo` vanillaIdInfo
553 mk_pk_let NewType arg boxing_con con_tys unpk_args body
554 = ASSERT( null other_args )
555 Let (NonRec arg (Note (Coerce (idType arg) (idType unpk_arg)) (Var unpk_arg))) body
557 (unpk_arg:other_args) = unpk_args
559 mk_pk_let DataType arg boxing_con con_tys unpk_args body
560 = Let (NonRec arg (mkConApp boxing_con con_args)) body
562 con_args = map Type con_tys ++ map Var unpk_args
565 mk_ww_local uniq ty = mkSysLocal SLIT("ww") uniq ty