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 ( coreExprType )
16 import Id ( Id, idType, mkSysLocal, getIdDemandInfo, setIdDemandInfo,
17 isOneShotLambda, setOneShotLambda,
20 import IdInfo ( CprInfo(..), noCprInfo, vanillaIdInfo )
21 import Const ( Con(..), DataCon )
22 import DataCon ( isExistentialDataCon, dataConArgTys )
23 import Demand ( Demand(..) )
24 import PrelInfo ( realWorldPrimId, aBSENT_ERROR_ID )
25 import TysPrim ( realWorldStatePrimTy )
26 import TysWiredIn ( unboxedTupleCon, unboxedTupleTyCon )
27 import Type ( isUnLiftedType,
28 splitForAllTys, splitFunTys,
29 splitAlgTyConApp_maybe, splitNewType_maybe,
33 import TyCon ( isNewTyCon, isProductTyCon, TyCon )
34 import BasicTypes ( NewOrData(..), Arity )
35 import Var ( TyVar, IdOrTyVar )
36 import UniqSupply ( returnUs, thenUs, getUniqueUs, getUniquesUs,
38 import Util ( zipWithEqual, zipEqual, lengthExceeds )
40 import List ( zipWith4 )
44 %************************************************************************
46 \subsection[mkWrapperAndWorker]{@mkWrapperAndWorker@}
48 %************************************************************************
50 ************ WARNING ******************
51 these comments are rather out of date
52 *****************************************
54 @mkWrapperAndWorker@ is given:
57 The {\em original function} \tr{f}, of the form:
59 f = /\ tyvars -> \ args -> body
61 The original-binder \tr{f}, the \tr{tyvars}, \tr{args}, and \tr{body}
64 We use the Id \tr{f} mostly to get its type.
67 Strictness information about \tr{f}, in the form of a list of
74 @mkWrapperAndWorker@ produces (A BIT OUT-OF-DATE...):
77 Maybe @Nothing@: no worker/wrappering going on in this case. This can
78 happen (a)~if the strictness info says that there is nothing
79 interesting to do or (b)~if *any* of the argument types corresponding
80 to ``active'' arg postitions is abstract or will be to the outside
81 world (i.e., {\em this} module can see the constructors, but nobody
82 else will be able to). An ``active'' arg position is one which the
83 wrapper has to unpack. An importing module can't do this unpacking,
84 so it simply has to give up and call the wrapper only.
87 Maybe \tr{Just (wrapper_Id, wrapper_body, worker_Id, worker_body)}.
89 The @wrapper_Id@ is just the one that was passed in, with its
90 strictness IdInfo updated.
93 The \tr{body} of the original function may not be given (i.e., it's
94 BOTTOM), in which case you'd jolly well better not tug on the
97 Here's an example. The original function is:
99 g :: forall a . Int -> [a] -> a
101 g = /\ a -> \ x ys ->
107 From this, we want to produce:
109 -- wrapper (an unfolding)
110 g :: forall a . Int -> [a] -> a
112 g = /\ a -> \ x ys ->
114 I# x# -> g.wrk a x# ys
115 -- call the worker; don't forget the type args!
118 g.wrk :: forall a . Int# -> [a] -> a
120 g.wrk = /\ a -> \ x# ys ->
124 case x of -- note: body of g moved intact
129 Something we have to be careful about: Here's an example:
131 -- "f" strictness: U(P)U(P)
132 f (I# a) (I# b) = a +# b
134 g = f -- "g" strictness same as "f"
136 \tr{f} will get a worker all nice and friendly-like; that's good.
137 {\em But we don't want a worker for \tr{g}}, even though it has the
138 same strictness as \tr{f}. Doing so could break laziness, at best.
140 Consequently, we insist that the number of strictness-info items is
141 exactly the same as the number of lambda-bound arguments. (This is
142 probably slightly paranoid, but OK in practice.) If it isn't the
143 same, we ``revise'' the strictness info, so that we won't propagate
144 the unusable strictness-info into the interfaces.
147 %************************************************************************
149 \subsection{Functions over Demands}
151 %************************************************************************
154 mAX_WORKER_ARGS :: Int -- ToDo: set via flag
157 setUnpackStrategy :: [Demand] -> [Demand]
159 = snd (go (mAX_WORKER_ARGS - nonAbsentArgs ds) ds)
161 go :: Int -- Max number of args available for sub-components of [Demand]
163 -> (Int, [Demand]) -- Args remaining after subcomponents of [Demand] are unpacked
165 go n (WwUnpack nd _ cs : ds) | n' >= 0
166 = WwUnpack nd True cs' `cons` go n'' ds
168 = WwUnpack nd False cs `cons` go n ds
170 n' = n + 1 - nonAbsentArgs cs
171 -- Add one because we don't pass the top-level arg any more
172 -- Delete # of non-absent args to which we'll now be committed
175 go n (d:ds) = d `cons` go n ds
178 cons d (n,ds) = (n, d:ds)
180 nonAbsentArgs :: [Demand] -> Int
182 nonAbsentArgs (WwLazy True : ds) = nonAbsentArgs ds
183 nonAbsentArgs (d : ds) = 1 + nonAbsentArgs ds
185 worthSplitting :: [Demand]
186 -> Bool -- Result is bottom
187 -> Bool -- True <=> the wrapper would not be an identity function
188 worthSplitting ds result_bot = any worth_it ds
189 -- We used not to split if the result is bottom.
190 -- [Justification: there's no efficiency to be gained,
191 -- and (worse) the wrapper body may not look like a wrapper
192 -- body to getWorkerIdAndCons]
193 -- But now (a) we don't have getWorkerIdAndCons, and
194 -- (b) it's sometimes bad not to make a wrapper. Consider
195 -- fw = \x# -> let x = I# x# in case e of
198 -- p3 -> the real stuff
199 -- The re-boxing code won't go away unless error_fn gets a wrapper too.
202 worth_it (WwLazy True) = True -- Absent arg
203 worth_it (WwUnpack _ True _) = True -- Arg to unpack
204 worth_it WwStrict = False -- Don't w/w just because of strictness
205 worth_it other = False
207 allAbsent :: [Demand] -> Bool
208 allAbsent ds = all absent ds
210 absent (WwLazy is_absent) = is_absent
211 absent (WwUnpack _ True cs) = allAbsent cs
216 %************************************************************************
218 \subsection{The worker wrapper core}
220 %************************************************************************
222 @mkWwBodies@ is called when doing the worker/wrapper split inside a module.
225 mkWwBodies :: Type -- Type of original function
226 -> Arity -- Arity of original function
227 -> [Demand] -- Strictness of original function
228 -> [Bool] -- One-shot-ness of the function
229 -> CprInfo -- Result of CPR analysis
230 -> UniqSM ([IdOrTyVar], -- Worker args
231 Id -> CoreExpr, -- Wrapper body, lacking only the worker Id
232 CoreExpr -> CoreExpr) -- Worker body, lacking the original function rhs
234 mkWwBodies fun_ty arity demands one_shots cpr_info
235 = WARN( not (lengthExceeds demands (arity-1))
236 || not (lengthExceeds one_shots (arity-1)),
237 text "mkWrapper" <+> ppr fun_ty <+> ppr arity <+> ppr (take arity demands) <+> ppr (take arity one_shots) )
238 mkWWargs fun_ty arity demands one_shots `thenUs` \ (wrap_args, wrap_fn_args, work_fn_args, res_ty) ->
239 mkWWstr wrap_args `thenUs` \ (work_args, wrap_fn_str, work_fn_str) ->
240 mkWWcpr res_ty cpr_info `thenUs` \ (wrap_fn_cpr, work_fn_cpr, cpr_res_ty) ->
241 mkWWfixup cpr_res_ty work_args `thenUs` \ (wrap_fn_fixup, work_fn_fixup) ->
244 Note InlineMe . wrap_fn_args . wrap_fn_cpr . wrap_fn_str . wrap_fn_fixup . Var,
245 work_fn_fixup . work_fn_str . work_fn_cpr . work_fn_args)
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.
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] -- Both these will in due course be derived
288 -- from the type. The [Bool] is True for a one-shot arg.
289 -> UniqSM ([IdOrTyVar], -- Wrapper args
290 CoreExpr -> CoreExpr, -- Wrapper fn
291 CoreExpr -> CoreExpr, -- Worker fn
292 Type) -- Type of wrapper body
294 mkWWargs fun_ty arity demands one_shots
296 = returnUs ([], id, id, fun_ty)
299 = getUniquesUs n_args `thenUs` \ wrap_uniqs ->
301 val_args = zipWith4 mk_wrap_arg wrap_uniqs arg_tys demands one_shots
302 wrap_args = tyvars ++ val_args
306 (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_coerce_fn . wrap_fn_args,
311 work_fn_args . work_coerce_fn . 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 = arity `min` n_arg_tys
318 (wrap_coerce_fn, work_coerce_fn, body_rep_ty)
319 | n_arg_tys == n_args -- All arg_tys used up
320 = case splitNewType_maybe body_ty of
321 Just rep_ty -> (Note (Coerce body_ty rep_ty), Note (Coerce rep_ty body_ty), rep_ty)
322 Nothing -> ASSERT2( n_args /= 0, text "mkWWargs" <+> ppr arity <+> ppr fun_ty )
324 | otherwise -- Leftover arg-tys
325 = (id, id, mkFunTys (drop n_args arg_tys) body_ty)
327 applyToVars :: [IdOrTyVar] -> CoreExpr -> CoreExpr
328 applyToVars vars fn = mkVarApps fn vars
330 mk_wrap_arg uniq ty dmd one_shot
331 = set_one_shot one_shot (setIdDemandInfo (mkSysLocal SLIT("w") uniq ty) dmd)
333 set_one_shot True id = setOneShotLambda id
334 set_one_shot False id = id
338 %************************************************************************
340 \subsection{Fixup stuff}
342 %************************************************************************
345 mkWWfixup res_ty work_args
346 | null work_args && isUnLiftedType res_ty
347 -- Horrid special case. If the worker would have no arguments, and the
348 -- function returns a primitive type value, that would make the worker into
349 -- an unboxed value. We box it by passing a dummy void argument, thus:
351 -- f = /\abc. \xyz. fw abc void
352 -- fw = /\abc. \v. body
354 -- We use the state-token type which generates no code
355 = getUniqueUs `thenUs` \ void_arg_uniq ->
357 void_arg = mk_ww_local void_arg_uniq realWorldStatePrimTy
359 returnUs (\ call_to_worker -> App call_to_worker (Var realWorldPrimId),
360 \ worker_body -> Lam void_arg worker_body)
367 %************************************************************************
369 \subsection{Strictness stuff}
371 %************************************************************************
374 mkWWstr :: [IdOrTyVar] -- Wrapper args; have their demand info on them
375 -- *Includes type variables*
376 -> UniqSM ([IdOrTyVar], -- Worker args
377 CoreExpr -> CoreExpr, -- Wrapper body, lacking the worker call
378 -- and without its lambdas
379 -- This fn adds the unboxing, and makes the
380 -- call passing the unboxed things
382 CoreExpr -> CoreExpr) -- Worker body, lacking the original body of the function,
383 -- but *with* lambdas
386 = mk_ww_str wrap_args `thenUs` \ (work_args, wrap_fn, work_fn) ->
387 returnUs ( work_args,
388 \ wrapper_body -> wrap_fn (mkVarApps wrapper_body work_args),
389 \ worker_body -> mkLams work_args (work_fn worker_body))
394 \ wrapper_body -> wrapper_body,
395 \ worker_body -> worker_body)
400 = mk_ww_str ds `thenUs` \ (worker_args, wrap_fn, work_fn) ->
401 returnUs (arg : worker_args, wrap_fn, work_fn)
404 = case getIdDemandInfo arg of
408 mk_ww_str ds `thenUs` \ (worker_args, wrap_fn, work_fn) ->
409 returnUs (worker_args, wrap_fn, mk_absent_let arg . work_fn)
412 WwUnpack new_or_data True cs ->
413 getUniquesUs (length inst_con_arg_tys) `thenUs` \ uniqs ->
415 unpk_args = zipWith mk_ww_local uniqs inst_con_arg_tys
416 unpk_args_w_ds = zipWithEqual "mk_ww_str" set_worker_arg_info unpk_args cs
418 mk_ww_str (unpk_args_w_ds ++ ds) `thenUs` \ (worker_args, wrap_fn, work_fn) ->
419 returnUs (worker_args,
420 mk_unpk_case new_or_data arg unpk_args data_con arg_tycon . wrap_fn,
421 work_fn . mk_pk_let new_or_data arg data_con tycon_arg_tys unpk_args)
423 (arg_tycon, tycon_arg_tys, data_con, inst_con_arg_tys) = splitProductType "mk_ww_str" (idType arg)
427 mk_ww_str ds `thenUs` \ (worker_args, wrap_fn, work_fn) ->
428 returnUs (arg : worker_args, wrap_fn, work_fn)
430 -- If the wrapper argument is a one-shot lambda, then
431 -- so should (all) the corresponding worker arguments be
432 -- This bites when we do w/w on a case join point
433 set_worker_arg_info worker_arg demand = set_one_shot (setIdDemandInfo worker_arg demand)
435 set_one_shot | isOneShotLambda arg = setOneShotLambda
436 | otherwise = \x -> x
440 %************************************************************************
442 \subsection{CPR stuff}
444 %************************************************************************
447 @mkWWcpr@ takes the worker/wrapper pair produced from the strictness
448 info and adds in the CPR transformation. The worker returns an
449 unboxed tuple containing non-CPR components. The wrapper takes this
450 tuple and re-produces the correct structured output.
452 The non-CPR results appear ordered in the unboxed tuple as if by a
453 left-to-right traversal of the result structure.
457 mkWWcpr :: Type -- function body type
458 -> CprInfo -- CPR analysis results
459 -> UniqSM (CoreExpr -> CoreExpr, -- New wrapper
460 CoreExpr -> CoreExpr, -- New worker
461 Type) -- Type of worker's body
463 mkWWcpr body_ty NoCPRInfo
464 = returnUs (id, id, body_ty) -- Must be just the strictness transf.
466 mkWWcpr body_ty (CPRInfo cpr_args)
467 | n_con_args == 1 && isUnLiftedType con_arg_ty1
468 -- Special case when there is a single result of unlifted type
469 = getUniquesUs 2 `thenUs` \ [work_uniq, arg_uniq] ->
471 work_wild = mk_ww_local work_uniq body_ty
472 arg = mk_ww_local arg_uniq con_arg_ty1
474 returnUs (\ wkr_call -> Case wkr_call arg [(DEFAULT, [], mkConApp data_con (map Type tycon_arg_tys ++ [Var arg]))],
475 \ body -> Case body work_wild [(DataCon data_con, [arg], Var arg)],
478 | otherwise -- The general case
479 = getUniquesUs (n_con_args + 2) `thenUs` \ uniqs ->
481 (wrap_wild : work_wild : args) = zipWith mk_ww_local uniqs (ubx_tup_ty : body_ty : con_arg_tys)
482 arg_vars = map Var args
483 ubx_tup_con = unboxedTupleCon n_con_args
484 ubx_tup_ty = coreExprType ubx_tup_app
485 ubx_tup_app = mkConApp ubx_tup_con (map Type con_arg_tys ++ arg_vars)
486 con_app = mkConApp data_con (map Type tycon_arg_tys ++ arg_vars)
488 returnUs (\ wkr_call -> Case wkr_call wrap_wild [(DataCon ubx_tup_con, args, con_app)],
489 \ body -> Case body work_wild [(DataCon data_con, args, ubx_tup_app)],
492 (tycon, tycon_arg_tys, data_con, con_arg_tys) = splitProductType "mkWWcpr" body_ty
493 n_con_args = length con_arg_tys
494 con_arg_ty1 = head con_arg_tys
497 splitProductType :: String -> Type -> (TyCon, [Type], DataCon, [Type])
498 -- For a tiresome reason, the type might not look like a product type
499 -- This happens when compiling the compiler! The module Name
500 -- imports {-# SOURCE #-} TyCon and Id
501 -- data Name = Name NameSort Unique OccName Provenance
502 -- data NameSort = WiredInId Module Id | ...
503 -- So Name does not look recursive (because Id is imported via a hi-boot file,
504 -- which says nothing about Id's rep) but actually it is, because Ids have Names.
505 -- Modules that *import* Name have a more complete view, see that Name is recursive,
506 -- and therefore that it isn't a ProductType. This conflicts with the CPR info
507 -- in exports from Name that say "do CPR".
509 -- Arguably we should regard Name as a product anyway because it isn't recursive
510 -- via products all the way... but we don't have that info to hand, and even if
511 -- we did this case might *still* arise.
514 -- So we hack our way out for now, by trusting the pragma that said "do CPR"
515 -- that means we can't use splitProductType_maybe
517 splitProductType fname ty
518 = case splitAlgTyConApp_maybe ty of
519 Just (tycon, tycon_args, (con:other_cons))
520 | null other_cons && not (isExistentialDataCon con)
521 -> WARN( not (isProductTyCon tycon),
522 text "splitProductType hack: I happened!" <+> ppr ty )
523 (tycon, tycon_args, con, dataConArgTys con tycon_args)
525 other -> pprPanic (fname ++ ": not a product") (ppr ty)
529 %************************************************************************
531 \subsection{Utilities}
533 %************************************************************************
537 mk_absent_let arg body
538 | not (isUnLiftedType arg_ty)
539 = Let (NonRec arg (mkTyApps (Var aBSENT_ERROR_ID) [arg_ty])) body
541 = panic "WwLib: haven't done mk_absent_let for primitives yet"
545 mk_unpk_case NewType arg unpk_args boxing_con boxing_tycon body
546 -- A newtype! Use a coercion not a case
547 = ASSERT( null other_args )
548 Case (Note (Coerce (idType unpk_arg) (idType arg)) (Var arg))
549 (sanitiseCaseBndr unpk_arg)
552 (unpk_arg:other_args) = unpk_args
554 mk_unpk_case DataType arg unpk_args boxing_con boxing_tycon body
557 (sanitiseCaseBndr arg)
558 [(DataCon boxing_con, unpk_args, body)]
560 sanitiseCaseBndr :: Id -> Id
561 -- The argument we are scrutinising has the right type to be
562 -- a case binder, so it's convenient to re-use it for that purpose.
563 -- But we *must* throw away all its IdInfo. In particular, the argument
564 -- will have demand info on it, and that demand info may be incorrect for
565 -- the case binder. e.g. case ww_arg of ww_arg { I# x -> ... }
566 -- Quite likely ww_arg isn't used in '...'. The case may get discarded
567 -- if the case binder says "I'm demanded". This happened in a situation
568 -- like (x+y) `seq` ....
569 sanitiseCaseBndr id = id `setIdInfo` vanillaIdInfo
571 mk_pk_let NewType arg boxing_con con_tys unpk_args body
572 = ASSERT( null other_args )
573 Let (NonRec arg (Note (Coerce (idType arg) (idType unpk_arg)) (Var unpk_arg))) body
575 (unpk_arg:other_args) = unpk_args
577 mk_pk_let DataType arg boxing_con con_tys unpk_args body
578 = Let (NonRec arg (Con (DataCon boxing_con) con_args)) body
580 con_args = map Type con_tys ++ map Var unpk_args
583 mk_ww_local uniq ty = mkSysLocal SLIT("ww") uniq ty