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}
10 worthSplitting, setUnpackStrategy,
14 #include "HsVersions.h"
17 import Id ( Id, idType, mkSysLocal, getIdDemandInfo, setIdDemandInfo,
20 import IdInfo ( CprInfo(..), noCprInfo, vanillaIdInfo )
21 import Const ( Con(..), DataCon )
22 import DataCon ( splitProductType_maybe )
23 import Demand ( Demand(..) )
24 import PrelInfo ( realWorldPrimId, aBSENT_ERROR_ID )
25 import TysPrim ( realWorldStatePrimTy )
26 import TysWiredIn ( unboxedTupleCon, unboxedTupleTyCon )
27 import Type ( isUnLiftedType, mkTyVarTys, mkTyVarTy, mkFunTys,
28 splitForAllTys, splitFunTys, splitFunTysN,
29 splitAlgTyConApp_maybe, splitAlgTyConApp,
30 mkTyConApp, splitNewType_maybe,
33 import TyCon ( isNewTyCon,
35 import BasicTypes ( NewOrData(..) )
37 import UniqSupply ( returnUs, thenUs, getUniqueUs, getUniquesUs,
39 import Util ( zipWithEqual, zipEqual )
43 %************************************************************************
45 \subsection[datatype-WwLib]{@WwBinding@: a datatype for worker/wrapper-ing}
47 %************************************************************************
49 In the worker/wrapper stuff, we want to carry around @CoreBindings@ in
50 an ``intermediate form'' that can later be turned into a \tr{let} or
51 \tr{case} (depending on strictness info).
56 | WwCase (CoreExpr -> CoreExpr)
57 -- the "case" will be a "strict let" of the form:
62 -- (instead of "let <blah> = rhs in body")
64 -- The expr you pass to the function is "body" (the
65 -- expression that goes "in the corner").
68 %************************************************************************
70 \subsection[mkWrapperAndWorker]{@mkWrapperAndWorker@}
72 %************************************************************************
74 ************ WARNING ******************
75 these comments are rather out of date
76 *****************************************
78 @mkWrapperAndWorker@ is given:
81 The {\em original function} \tr{f}, of the form:
83 f = /\ tyvars -> \ args -> body
85 The original-binder \tr{f}, the \tr{tyvars}, \tr{args}, and \tr{body}
88 We use the Id \tr{f} mostly to get its type.
91 Strictness information about \tr{f}, in the form of a list of
98 @mkWrapperAndWorker@ produces (A BIT OUT-OF-DATE...):
101 Maybe @Nothing@: no worker/wrappering going on in this case. This can
102 happen (a)~if the strictness info says that there is nothing
103 interesting to do or (b)~if *any* of the argument types corresponding
104 to ``active'' arg postitions is abstract or will be to the outside
105 world (i.e., {\em this} module can see the constructors, but nobody
106 else will be able to). An ``active'' arg position is one which the
107 wrapper has to unpack. An importing module can't do this unpacking,
108 so it simply has to give up and call the wrapper only.
111 Maybe \tr{Just (wrapper_Id, wrapper_body, worker_Id, worker_body)}.
113 The @wrapper_Id@ is just the one that was passed in, with its
114 strictness IdInfo updated.
117 The \tr{body} of the original function may not be given (i.e., it's
118 BOTTOM), in which case you'd jolly well better not tug on the
121 Here's an example. The original function is:
123 g :: forall a . Int -> [a] -> a
125 g = /\ a -> \ x ys ->
131 From this, we want to produce:
133 -- wrapper (an unfolding)
134 g :: forall a . Int -> [a] -> a
136 g = /\ a -> \ x ys ->
138 I# x# -> g.wrk a x# ys
139 -- call the worker; don't forget the type args!
142 g.wrk :: forall a . Int# -> [a] -> a
144 g.wrk = /\ a -> \ x# ys ->
148 case x of -- note: body of g moved intact
153 Something we have to be careful about: Here's an example:
155 -- "f" strictness: U(P)U(P)
156 f (I# a) (I# b) = a +# b
158 g = f -- "g" strictness same as "f"
160 \tr{f} will get a worker all nice and friendly-like; that's good.
161 {\em But we don't want a worker for \tr{g}}, even though it has the
162 same strictness as \tr{f}. Doing so could break laziness, at best.
164 Consequently, we insist that the number of strictness-info items is
165 exactly the same as the number of lambda-bound arguments. (This is
166 probably slightly paranoid, but OK in practice.) If it isn't the
167 same, we ``revise'' the strictness info, so that we won't propagate
168 the unusable strictness-info into the interfaces.
171 %************************************************************************
173 \subsection{Functions over Demands}
175 %************************************************************************
178 mAX_WORKER_ARGS :: Int -- ToDo: set via flag
181 setUnpackStrategy :: [Demand] -> [Demand]
183 = snd (go (mAX_WORKER_ARGS - nonAbsentArgs ds) ds)
185 go :: Int -- Max number of args available for sub-components of [Demand]
187 -> (Int, [Demand]) -- Args remaining after subcomponents of [Demand] are unpacked
189 go n (WwUnpack nd _ cs : ds) | n' >= 0
190 = WwUnpack nd True cs' `cons` go n'' ds
192 = WwUnpack nd False cs `cons` go n ds
194 n' = n + 1 - nonAbsentArgs cs
195 -- Add one because we don't pass the top-level arg any more
196 -- Delete # of non-absent args to which we'll now be committed
199 go n (d:ds) = d `cons` go n ds
202 cons d (n,ds) = (n, d:ds)
204 nonAbsentArgs :: [Demand] -> Int
206 nonAbsentArgs (WwLazy True : ds) = nonAbsentArgs ds
207 nonAbsentArgs (d : ds) = 1 + nonAbsentArgs ds
209 worthSplitting :: [Demand]
210 -> Bool -- Result is bottom
211 -> Bool -- True <=> the wrapper would not be an identity function
212 worthSplitting ds result_bot = not result_bot && any worth_it ds
213 -- Don't split if the result is bottom; there's no efficiency to
214 -- be gained, and (worse) the wrapper body may not look like a wrapper
215 -- body to getWorkerIdAndCons
217 worth_it (WwLazy True) = True -- Absent arg
218 worth_it (WwUnpack _ True _) = True -- Arg to unpack
219 worth_it WwStrict = False -- Don't w/w just because of strictness
220 worth_it other = False
222 allAbsent :: [Demand] -> Bool
223 allAbsent ds = all absent ds
225 absent (WwLazy is_absent) = is_absent
226 absent (WwUnpack _ True cs) = allAbsent cs
231 %************************************************************************
233 \subsection{The worker wrapper core}
235 %************************************************************************
237 @mkWrapper@ is called when importing a function. We have the type of
238 the function and the name of its worker, and we want to make its body (the wrapper).
241 mkWrapper :: Type -- Wrapper type
243 -> [Demand] -- Wrapper strictness info
244 -> CprInfo -- Wrapper cpr info
245 -> UniqSM (Id -> CoreExpr) -- Wrapper body, missing worker Id
247 mkWrapper fun_ty arity demands cpr_info
248 = getUniquesUs arity `thenUs` \ wrap_uniqs ->
250 (tyvars, tau_ty) = splitForAllTys fun_ty
251 (arg_tys, body_ty) = splitFunTysN "mkWrapper" arity tau_ty
252 -- The "expanding dicts" part here is important, even for the splitForAll
253 -- The imported thing might be a dictionary, such as Functor Foo
254 -- But Functor Foo = forall a b. (a->b) -> Foo a -> Foo b
255 -- and as such might have some strictness info attached.
256 -- Then we need to have enough args to zip to the strictness info
258 wrap_args = zipWith mk_ww_local wrap_uniqs arg_tys
260 mkWwBodies tyvars wrap_args body_ty demands cpr_info `thenUs` \ (wrap_fn, _, _) ->
264 @mkWwBodies@ is called when doing the worker/wrapper split inside a module.
267 mkWwBodies :: [TyVar] -> [Id] -> Type -- Original fn args and body type
268 -> [Demand] -- Strictness info for original fn; corresp 1-1 with args
269 -> CprInfo -- Result of CPR analysis
270 -> UniqSM (Id -> CoreExpr, -- Wrapper body, lacking only the worker Id
271 CoreExpr -> CoreExpr, -- Worker body, lacking the original function body
272 [Demand]) -- Strictness info for worker
274 mkWwBodies tyvars wrap_args body_ty demands cpr_info
276 -- demands may be longer than number of args. If we aren't doing w/w
277 -- for strictness then demands is an infinite list of 'lazy' args.
278 wrap_args_w_demands = zipWith setIdDemandInfo wrap_args demands
279 (wrap_fn_coerce, work_fn_coerce) = mkWWcoerce body_ty
281 mkWWstr body_ty wrap_args_w_demands `thenUs` \ (work_args_w_demands, wrap_fn_str, work_fn_str) ->
283 mkWWcpr body_ty cpr_info `thenUs` \ (wrap_fn_cpr, work_fn_cpr) ->
285 returnUs (\ work_id -> Note InlineMe $
286 mkLams tyvars $ mkLams wrap_args_w_demands $
287 (wrap_fn_coerce . wrap_fn_str . wrap_fn_cpr) $
288 mkVarApps (Var work_id) (tyvars ++ work_args_w_demands),
290 \ work_body -> mkLams tyvars $ mkLams work_args_w_demands $
291 (work_fn_coerce . work_fn_str . work_fn_cpr)
294 map getIdDemandInfo work_args_w_demands)
298 %************************************************************************
300 \subsection{Coercion stuff}
302 %************************************************************************
304 The "coerce" transformation is
308 f = \xy -> coerce R R' (fw x y)
309 fw = \xy -> coerce R' R e
311 where R' is the representation type for R.
315 = case splitNewType_maybe body_ty of
317 Just rep_ty -> (mkNote (Coerce body_ty rep_ty),
318 mkNote (Coerce rep_ty body_ty))
323 %************************************************************************
325 \subsection{Strictness stuff}
327 %************************************************************************
331 mkWWstr :: Type -- Body type
332 -> [Id] -- Wrapper args; have their demand info on them
333 -> UniqSM ([Id], -- Worker args; have their demand info on them
335 CoreExpr -> CoreExpr, -- Wrapper body, lacking the inner call to the worker
336 -- and without its lambdas
337 -- At the call site, the worker args are bound
339 CoreExpr -> CoreExpr) -- Worker body, lacking the original body of the function,
340 -- and without its lambdas
342 mkWWstr body_ty wrap_args
343 = mk_ww wrap_args `thenUs` \ (work_args, wrap_fn, work_fn) ->
345 if null work_args && isUnLiftedType body_ty then
346 -- Horrid special case. If the worker would have no arguments, and the
347 -- function returns a primitive type value, that would make the worker into
348 -- an unboxed value. We box it by passing a dummy void argument, thus:
350 -- f = /\abc. \xyz. fw abc void
351 -- fw = /\abc. \v. body
353 -- We use the state-token type which generates no code
354 getUniqueUs `thenUs` \ void_arg_uniq ->
356 void_arg = mk_ww_local void_arg_uniq realWorldStatePrimTy
358 returnUs ([void_arg],
359 wrap_fn . Let (NonRec void_arg (Var realWorldPrimId)),
362 returnUs (work_args, wrap_fn, work_fn)
369 \ wrapper_body -> wrapper_body,
370 \ worker_body -> worker_body)
374 = case getIdDemandInfo arg of
378 mk_ww ds `thenUs` \ (worker_args, wrap_fn, work_fn) ->
379 returnUs (worker_args, wrap_fn, mk_absent_let arg . work_fn)
382 WwUnpack new_or_data True cs ->
383 getUniquesUs (length inst_con_arg_tys) `thenUs` \ uniqs ->
385 unpk_args = zipWith mk_ww_local uniqs inst_con_arg_tys
386 unpk_args_w_ds = zipWithEqual "mk_ww" setIdDemandInfo unpk_args cs
388 mk_ww (unpk_args_w_ds ++ ds) `thenUs` \ (worker_args, wrap_fn, work_fn) ->
389 returnUs (worker_args,
390 mk_unpk_case new_or_data arg unpk_args data_con arg_tycon . wrap_fn,
391 work_fn . mk_pk_let new_or_data arg data_con tycon_arg_tys unpk_args)
393 (arg_tycon, tycon_arg_tys, data_con, inst_con_arg_tys) = splitProductType "mk_ww" (idType arg)
397 mk_ww ds `thenUs` \ (worker_args, wrap_fn, work_fn) ->
398 returnUs (arg : worker_args, wrap_fn, work_fn)
402 %************************************************************************
404 \subsection{CPR stuff}
406 %************************************************************************
409 @mkWWcpr@ takes the worker/wrapper pair produced from the strictness
410 info and adds in the CPR transformation. The worker returns an
411 unboxed tuple containing non-CPR components. The wrapper takes this
412 tuple and re-produces the correct structured output.
414 The non-CPR results appear ordered in the unboxed tuple as if by a
415 left-to-right traversal of the result structure.
419 mkWWcpr :: Type -- function body type
420 -> CprInfo -- CPR analysis results
421 -> UniqSM (CoreExpr -> CoreExpr, -- New wrapper
422 CoreExpr -> CoreExpr) -- New worker
424 mkWWcpr body_ty NoCPRInfo
425 = returnUs (id, id) -- Must be just the strictness transf.
426 mkWWcpr body_ty (CPRInfo cpr_args)
427 = getUniqueUs `thenUs` \ body_arg_uniq ->
429 body_var = mk_ww_local body_arg_uniq body_ty
431 cpr_reconstruct body_ty cpr_info' `thenUs` \reconst_fn ->
432 cpr_flatten body_ty cpr_info' `thenUs` \flatten_fn ->
433 returnUs (reconst_fn, flatten_fn)
435 -- We only make use of the outer level of CprInfo, otherwise we
436 -- may lose laziness. :-( Hopefully, we will find a use for the
437 -- extra info some day (e.g. creating versions specialized to
438 -- the use made of the components of the result by the callee)
439 cpr_info' = CPRInfo (map (const NoCPRInfo) cpr_args)
443 @cpr_flatten@ takes the result type produced by the body and the info
444 from the CPR analysis and flattens the constructed product components.
445 These are returned in an unboxed tuple.
448 cpr_flatten :: Type -> CprInfo -> UniqSM (CoreExpr -> CoreExpr)
449 cpr_flatten ty cpr_info
450 = mk_cpr_case (ty, cpr_info) `thenUs` \(res_id, tup_ids, flatten_exp) ->
451 returnUs (\body -> Case body res_id
452 [(DEFAULT, [], flatten_exp (fst $ mk_unboxed_tuple tup_ids))])
456 mk_cpr_case :: (Type, CprInfo) ->
457 UniqSM (CoreBndr, -- Name of binder for this part of result
458 [(CoreExpr, Type)], -- expressions for flattened result
459 CoreExpr -> CoreExpr) -- add in code to flatten result
461 mk_cpr_case (ty, NoCPRInfo)
462 -- this component must be returned as a component of the unboxed tuple result
463 = getUniqueUs `thenUs` \id_uniq ->
464 let id_id = mk_ww_local id_uniq ty in
465 returnUs (id_id, [(Var id_id, ty)], id)
466 mk_cpr_case (ty, cpr_info@(CPRInfo ci_args))
467 | isNewTyCon tycon -- a new type: under the coercions must be a
468 -- constructed product
469 = ASSERT ( null $ tail inst_con_arg_tys )
470 mk_cpr_case (target_of_from_type, cpr_info)
471 `thenUs` \(arg, tup, exp) ->
472 getUniqueUs `thenUs` \id_uniq ->
473 let id_id = mk_ww_local id_uniq ty
474 new_exp_case = \var -> Case (Note (Coerce (idType arg) ty) (Var id_id))
476 [(DEFAULT,[], exp var)]
478 returnUs (id_id, tup, new_exp_case)
480 | otherwise -- a data type
481 -- flatten components
482 = mapUs mk_cpr_case (zip inst_con_arg_tys ci_args)
483 `thenUs` \sub_builds ->
484 getUniqueUs `thenUs` \id_uniq ->
485 let id_id = mk_ww_local id_uniq ty
486 (args, tup, exp) = unzip3 sub_builds
487 con_app = mkConApp data_con (map Var args)
489 new_exp_case = \var -> Case (Var id_id) (mkWildId ty)
490 [(DataCon data_con, args,
491 foldl (\e f -> f e) var exp)]
493 returnUs (id_id, new_tup, new_exp_case)
495 (tycon, tycon_arg_tys, data_con, inst_con_arg_tys) = splitProductType "mk_cpr_case" ty
496 from_type = head inst_con_arg_tys
497 -- if coerced from a function 'look through' to find result type
498 target_of_from_type = (snd.splitFunTys.snd.splitForAllTys) from_type
502 @cpr_reconstruct@ does the opposite of @cpr_flatten@. It takes the unboxed
503 tuple produced by the worker and reconstructs the structured result.
506 cpr_reconstruct :: Type -> CprInfo -> UniqSM (CoreExpr -> CoreExpr)
507 cpr_reconstruct ty cpr_info
508 = mk_cpr_let (ty,cpr_info) `thenUs` \(res_id, tup_ids, reconstruct_exp) ->
509 returnUs (\worker -> Case worker (mkWildId $ worker_type tup_ids)
510 [(DataCon $ unboxedTupleCon $ length tup_ids,
511 tup_ids, reconstruct_exp $ Var res_id)])
514 worker_type ids = mkTyConApp (unboxedTupleTyCon (length ids)) (map idType ids)
517 mk_cpr_let :: (Type, CprInfo) ->
518 UniqSM (CoreBndr, -- Binder for this component of result
519 [CoreBndr], -- Binders which will appear in worker's result
520 CoreExpr -> CoreExpr) -- Code to produce structured result.
521 mk_cpr_let (ty, NoCPRInfo)
522 -- this component will appear explicitly in the unboxed tuple.
523 = getUniqueUs `thenUs` \id_uniq ->
525 id_id = mk_ww_local id_uniq ty
527 returnUs (id_id, [id_id], id)
529 mk_cpr_let (ty, cpr_info@(CPRInfo ci_args))
530 | isNewTyCon tycon -- a new type: must coerce the argument to this type
531 = ASSERT ( null $ tail inst_con_arg_tys )
532 mk_cpr_let (target_of_from_type, cpr_info)
533 `thenUs` \(arg, tup, exp) ->
534 getUniqueUs `thenUs` \id_uniq ->
535 let id_id = mk_ww_local id_uniq ty
536 new_exp = \var -> exp (Let (NonRec id_id (Note (Coerce ty (idType arg)) (Var arg))) var)
538 returnUs (id_id, tup, new_exp)
540 | otherwise -- a data type
541 -- reconstruct components then apply data con
542 = mapUs mk_cpr_let (zip inst_con_arg_tys ci_args)
543 `thenUs` \sub_builds ->
544 getUniqueUs `thenUs` \id_uniq ->
545 let id_id = mk_ww_local id_uniq ty
546 (args, tup, exp) = unzip3 sub_builds
547 con_app = mkConApp data_con $ (map Type tycon_arg_tys) ++ (map Var args)
549 new_exp = \var -> foldl (\e f -> f e) (Let (NonRec id_id con_app) var) exp
551 returnUs (id_id, new_tup, new_exp)
553 (tycon, tycon_arg_tys, data_con, inst_con_arg_tys) = splitProductType "mk_cpr_let" ty
554 from_type = head inst_con_arg_tys
555 -- if coerced from a function 'look through' to find result type
556 target_of_from_type = (snd.splitFunTys.snd.splitForAllTys) from_type
559 splitProductType :: String -> Type -> (TyCon, [Type], DataCon, [Type])
560 splitProductType fname ty = case splitProductType_maybe ty of
562 Nothing -> pprPanic (fname ++ ": not a product") (ppr ty)
566 %************************************************************************
568 \subsection{Utilities}
570 %************************************************************************
574 mk_absent_let arg body
575 | not (isUnLiftedType arg_ty)
576 = Let (NonRec arg (mkTyApps (Var aBSENT_ERROR_ID) [arg_ty])) body
578 = panic "WwLib: haven't done mk_absent_let for primitives yet"
582 mk_unpk_case NewType arg unpk_args boxing_con boxing_tycon body
583 -- A newtype! Use a coercion not a case
584 = ASSERT( null other_args )
585 Case (Note (Coerce (idType unpk_arg) (idType arg)) (Var arg))
586 (sanitiseCaseBndr unpk_arg)
589 (unpk_arg:other_args) = unpk_args
591 mk_unpk_case DataType arg unpk_args boxing_con boxing_tycon body
594 (sanitiseCaseBndr arg)
595 [(DataCon boxing_con, unpk_args, body)]
597 sanitiseCaseBndr :: Id -> Id
598 -- The argument we are scrutinising has the right type to be
599 -- a case binder, so it's convenient to re-use it for that purpose.
600 -- But we *must* throw away all its IdInfo. In particular, the argument
601 -- will have demand info on it, and that demand info may be incorrect for
602 -- the case binder. e.g. case ww_arg of ww_arg { I# x -> ... }
603 -- Quite likely ww_arg isn't used in '...'. The case may get discarded
604 -- if the case binder says "I'm demanded". This happened in a situation
605 -- like (x+y) `seq` ....
606 sanitiseCaseBndr id = id `setIdInfo` vanillaIdInfo
608 mk_pk_let NewType arg boxing_con con_tys unpk_args body
609 = ASSERT( null other_args )
610 Let (NonRec arg (Note (Coerce (idType arg) (idType unpk_arg)) (Var unpk_arg))) body
612 (unpk_arg:other_args) = unpk_args
614 mk_pk_let DataType arg boxing_con con_tys unpk_args body
615 = Let (NonRec arg (Con (DataCon boxing_con) con_args)) body
617 con_args = map Type con_tys ++ map Var unpk_args
620 mk_ww_local uniq ty = mkSysLocal SLIT("ww") uniq ty
623 mk_unboxed_tuple :: [(CoreExpr, Type)] -> (CoreExpr, Type)
624 mk_unboxed_tuple contents
625 = (mkConApp (unboxedTupleCon (length contents))
626 (map (Type . snd) contents ++
628 mkTyConApp (unboxedTupleTyCon (length contents))