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 ( 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, mkTyVarTys, mkTyVarTy, mkFunTys,
28 splitForAllTys, splitFunTysN,
29 splitAlgTyConApp_maybe, mkTyConApp,
32 import TyCon ( isNewTyCon,
34 import BasicTypes ( NewOrData(..) )
36 import UniqSupply ( returnUs, thenUs, getUniqueUs, getUniquesUs,
38 import Util ( zipWithEqual, zipEqual )
42 %************************************************************************
44 \subsection[datatype-WwLib]{@WwBinding@: a datatype for worker/wrapper-ing}
46 %************************************************************************
48 In the worker/wrapper stuff, we want to carry around @CoreBindings@ in
49 an ``intermediate form'' that can later be turned into a \tr{let} or
50 \tr{case} (depending on strictness info).
55 | WwCase (CoreExpr -> CoreExpr)
56 -- the "case" will be a "strict let" of the form:
61 -- (instead of "let <blah> = rhs in body")
63 -- The expr you pass to the function is "body" (the
64 -- expression that goes "in the corner").
67 %************************************************************************
69 \subsection[mkWrapperAndWorker]{@mkWrapperAndWorker@}
71 %************************************************************************
73 ************ WARNING ******************
74 these comments are rather out of date
75 *****************************************
77 @mkWrapperAndWorker@ is given:
80 The {\em original function} \tr{f}, of the form:
82 f = /\ tyvars -> \ args -> body
84 The original-binder \tr{f}, the \tr{tyvars}, \tr{args}, and \tr{body}
87 We use the Id \tr{f} mostly to get its type.
90 Strictness information about \tr{f}, in the form of a list of
97 @mkWrapperAndWorker@ produces (A BIT OUT-OF-DATE...):
100 Maybe @Nothing@: no worker/wrappering going on in this case. This can
101 happen (a)~if the strictness info says that there is nothing
102 interesting to do or (b)~if *any* of the argument types corresponding
103 to ``active'' arg postitions is abstract or will be to the outside
104 world (i.e., {\em this} module can see the constructors, but nobody
105 else will be able to). An ``active'' arg position is one which the
106 wrapper has to unpack. An importing module can't do this unpacking,
107 so it simply has to give up and call the wrapper only.
110 Maybe \tr{Just (wrapper_Id, wrapper_body, worker_Id, worker_body)}.
112 The @wrapper_Id@ is just the one that was passed in, with its
113 strictness IdInfo updated.
116 The \tr{body} of the original function may not be given (i.e., it's
117 BOTTOM), in which case you'd jolly well better not tug on the
120 Here's an example. The original function is:
122 g :: forall a . Int -> [a] -> a
124 g = /\ a -> \ x ys ->
130 From this, we want to produce:
132 -- wrapper (an unfolding)
133 g :: forall a . Int -> [a] -> a
135 g = /\ a -> \ x ys ->
137 I# x# -> g.wrk a x# ys
138 -- call the worker; don't forget the type args!
141 g.wrk :: forall a . Int# -> [a] -> a
143 g.wrk = /\ a -> \ x# ys ->
147 case x of -- note: body of g moved intact
152 Something we have to be careful about: Here's an example:
154 -- "f" strictness: U(P)U(P)
155 f (I# a) (I# b) = a +# b
157 g = f -- "g" strictness same as "f"
159 \tr{f} will get a worker all nice and friendly-like; that's good.
160 {\em But we don't want a worker for \tr{g}}, even though it has the
161 same strictness as \tr{f}. Doing so could break laziness, at best.
163 Consequently, we insist that the number of strictness-info items is
164 exactly the same as the number of lambda-bound arguments. (This is
165 probably slightly paranoid, but OK in practice.) If it isn't the
166 same, we ``revise'' the strictness info, so that we won't propagate
167 the unusable strictness-info into the interfaces.
170 %************************************************************************
172 \subsection{Functions over Demands}
174 %************************************************************************
177 mAX_WORKER_ARGS :: Int -- ToDo: set via flag
180 setUnpackStrategy :: [Demand] -> [Demand]
182 = snd (go (mAX_WORKER_ARGS - nonAbsentArgs ds) ds)
184 go :: Int -- Max number of args available for sub-components of [Demand]
186 -> (Int, [Demand]) -- Args remaining after subcomponents of [Demand] are unpacked
188 go n (WwUnpack nd _ cs : ds) | n' >= 0
189 = WwUnpack nd True cs' `cons` go n'' ds
191 = WwUnpack nd False cs `cons` go n ds
193 n' = n + 1 - nonAbsentArgs cs
194 -- Add one because we don't pass the top-level arg any more
195 -- Delete # of non-absent args to which we'll now be committed
198 go n (d:ds) = d `cons` go n ds
201 cons d (n,ds) = (n, d:ds)
203 nonAbsentArgs :: [Demand] -> Int
205 nonAbsentArgs (WwLazy True : ds) = nonAbsentArgs ds
206 nonAbsentArgs (d : ds) = 1 + nonAbsentArgs ds
208 worthSplitting :: [Demand]
209 -> Bool -- Result is bottom
210 -> Bool -- True <=> the wrapper would not be an identity function
211 worthSplitting ds result_bot = not result_bot && any worth_it ds
212 -- Don't split if the result is bottom; there's no efficiency to
213 -- be gained, and (worse) the wrapper body may not look like a wrapper
214 -- body to getWorkerIdAndCons
216 worth_it (WwLazy True) = True -- Absent arg
217 worth_it (WwUnpack _ True _) = True -- Arg to unpack
218 worth_it WwStrict = False -- Don't w/w just because of strictness
219 worth_it other = False
221 allAbsent :: [Demand] -> Bool
222 allAbsent ds = all absent ds
224 absent (WwLazy is_absent) = is_absent
225 absent (WwUnpack _ True cs) = allAbsent cs
230 %************************************************************************
232 \subsection{The worker wrapper core}
234 %************************************************************************
236 @mkWrapper@ is called when importing a function. We have the type of
237 the function and the name of its worker, and we want to make its body (the wrapper).
240 mkWrapper :: Type -- Wrapper type
242 -> [Demand] -- Wrapper strictness info
243 -> CprInfo -- Wrapper cpr info
244 -> UniqSM (Id -> CoreExpr) -- Wrapper body, missing worker Id
246 mkWrapper fun_ty arity demands cpr_info
247 = getUniquesUs arity `thenUs` \ wrap_uniqs ->
249 (tyvars, tau_ty) = splitForAllTys fun_ty
250 (arg_tys, body_ty) = splitFunTysN "mkWrapper" arity tau_ty
251 -- The "expanding dicts" part here is important, even for the splitForAll
252 -- The imported thing might be a dictionary, such as Functor Foo
253 -- But Functor Foo = forall a b. (a->b) -> Foo a -> Foo b
254 -- and as such might have some strictness info attached.
255 -- Then we need to have enough args to zip to the strictness info
257 wrap_args = zipWith mk_ww_local wrap_uniqs arg_tys
259 mkWwBodies tyvars wrap_args body_ty demands cpr_info `thenUs` \ (wrap_fn, _, _) ->
263 @mkWwBodies@ is called when doing the worker/wrapper split inside a module.
266 mkWwBodies :: [TyVar] -> [Id] -> Type -- Original fn args and body type
267 -> [Demand] -- Strictness info for original fn; corresp 1-1 with args
268 -> CprInfo -- Result of CPR analysis
269 -> UniqSM (Id -> CoreExpr, -- Wrapper body, lacking only the worker Id
270 CoreExpr -> CoreExpr, -- Worker body, lacking the original function body
271 [Demand]) -- Strictness info for worker
273 mkWwBodies tyvars args body_ty demands cpr_info
274 | allAbsent demands &&
275 isUnLiftedType body_ty
276 = -- Horrid special case. If the worker would have no arguments, and the
277 -- function returns a primitive type value, that would make the worker into
278 -- an unboxed value. We box it by passing a dummy void argument, thus:
280 -- f = /\abc. \xyz. fw abc void
281 -- fw = /\abc. \v. body
283 -- We use the state-token type which generates no code
284 getUniqueUs `thenUs` \ void_arg_uniq ->
286 void_arg = mk_ww_local void_arg_uniq realWorldStatePrimTy
288 returnUs (\ work_id -> Note InlineMe $ -- Inline the wrapper
289 mkLams tyvars $ mkLams args $
291 (map (Type . mkTyVarTy) tyvars ++ [Var realWorldPrimId]),
292 \ body -> mkLams (tyvars ++ [void_arg]) body,
295 mkWwBodies tyvars wrap_args body_ty demands cpr_info
298 -- demands may be longer than number of args. If we aren't doing w/w
299 -- for strictness then demands is an infinite list of 'lazy' args.
300 wrap_args_w_demands = zipWith setIdDemandInfo wrap_args demands
302 mkWW wrap_args_w_demands `thenUs` \ (wrap_fn, work_args_w_demands, work_fn) ->
304 mkWWcpr body_ty cpr_info `thenUs` \ (wrap_fn_w_cpr, work_fn_w_cpr) ->
306 returnUs (\ work_id -> Note InlineMe $
307 mkLams tyvars $ mkLams wrap_args_w_demands $
308 (wrap_fn_w_cpr . wrap_fn) (mkTyApps (Var work_id) (mkTyVarTys tyvars)),
310 \ body -> mkLams tyvars $ mkLams work_args_w_demands $
311 (work_fn_w_cpr . work_fn) body,
313 map getIdDemandInfo work_args_w_demands)
318 mkWW :: [Id] -- Wrapper args; have their demand info on them
319 -> UniqSM (CoreExpr -> CoreExpr, -- Wrapper body, lacking the inner call to the worker
320 -- and without its lambdas
321 [Id], -- Worker args; have their demand info on them
322 CoreExpr -> CoreExpr) -- Worker body, lacking the original body of the function
327 = returnUs (\ wrapper_body -> wrapper_body,
329 \ worker_body -> worker_body)
333 = case getIdDemandInfo arg of
337 mkWW ds `thenUs` \ (wrap_fn, worker_args, work_fn) ->
338 returnUs (\ wrapper_body -> wrap_fn wrapper_body,
340 \ worker_body -> mk_absent_let arg (work_fn worker_body))
344 WwUnpack new_or_data True cs ->
345 getUniquesUs (length inst_con_arg_tys) `thenUs` \ uniqs ->
347 unpk_args = zipWith mk_ww_local uniqs inst_con_arg_tys
348 unpk_args_w_ds = zipWithEqual "mkWW" setIdDemandInfo unpk_args cs
350 mkWW (unpk_args_w_ds ++ ds) `thenUs` \ (wrap_fn, worker_args, work_fn) ->
351 returnUs (\ wrapper_body -> mk_unpk_case new_or_data arg unpk_args data_con arg_tycon
352 (wrap_fn wrapper_body),
354 \ worker_body -> work_fn (mk_pk_let new_or_data arg data_con
355 tycon_arg_tys unpk_args worker_body))
357 inst_con_arg_tys = dataConArgTys data_con tycon_arg_tys
358 (arg_tycon, tycon_arg_tys, data_con)
359 = case (splitAlgTyConApp_maybe (idType arg)) of
361 Just (arg_tycon, tycon_arg_tys, [data_con]) ->
362 -- The main event: a single-constructor data type
363 (arg_tycon, tycon_arg_tys, data_con)
365 Just (_, _, data_cons) ->
366 pprPanic "mk_ww_arg_processing:"
367 (text "not one constr (interface files not consistent/up to date?)"
368 $$ (ppr arg <+> ppr (idType arg)))
371 panic "mk_ww_arg_processing: not datatype"
376 mkWW ds `thenUs` \ (wrap_fn, worker_args, work_fn) ->
377 returnUs (\ wrapper_body -> wrap_fn (App wrapper_body (Var arg)),
382 @mkWWcpr@ takes the worker/wrapper pair produced from the strictness
383 info and adds in the CPR transformation. The worker returns an
384 unboxed tuple containing non-CPR components. The wrapper takes this
385 tuple and re-produces the correct structured output.
387 The non-CPR results appear ordered in the unboxed tuple as if by a
388 left-to-right traversal of the result structure.
392 mkWWcpr :: Type -- function body type
393 -> CprInfo -- CPR analysis results
394 -> UniqSM (CoreExpr -> CoreExpr, -- New wrapper
395 CoreExpr -> CoreExpr) -- New worker
397 mkWWcpr body_ty NoCPRInfo
398 = returnUs (id, id) -- Must be just the strictness transf.
399 mkWWcpr body_ty (CPRInfo cpr_args)
400 = getUniqueUs `thenUs` \ body_arg_uniq ->
402 body_var = mk_ww_local body_arg_uniq body_ty
404 cpr_reconstruct body_ty cpr_info' `thenUs` \reconst_fn ->
405 cpr_flatten body_ty cpr_info' `thenUs` \flatten_fn ->
406 returnUs (reconst_fn, flatten_fn)
408 -- We only make use of the outer level of CprInfo, otherwise we
409 -- may lose laziness. :-( Hopefully, we will find a use for the
410 -- extra info some day (e.g. creating versions specialized to
411 -- the use made of the components of the result by the callee)
412 cpr_info' = CPRInfo (map (const NoCPRInfo) cpr_args)
416 @cpr_flatten@ takes the result type produced by the body and the info
417 from the CPR analysis and flattens the constructed product components.
418 These are returned in an unboxed tuple.
421 cpr_flatten :: Type -> CprInfo -> UniqSM (CoreExpr -> CoreExpr)
422 cpr_flatten ty cpr_info
423 = mk_cpr_case (ty, cpr_info) `thenUs` \(res_id, tup_ids, flatten_exp) ->
424 returnUs (\body -> Case body res_id
425 [(DEFAULT, [], flatten_exp (fst $ mk_unboxed_tuple tup_ids))])
429 mk_cpr_case :: (Type, CprInfo) ->
430 UniqSM (CoreBndr, -- Name of binder for this part of result
431 [(CoreExpr, Type)], -- expressions for flattened result
432 CoreExpr -> CoreExpr) -- add in code to flatten result
434 mk_cpr_case (ty, NoCPRInfo)
435 -- this component must be returned as a component of the unboxed tuple result
436 = getUniqueUs `thenUs` \id_uniq ->
437 let id_id = mk_ww_local id_uniq ty in
438 returnUs (id_id, [(Var id_id, ty)], id)
439 mk_cpr_case (ty, cpr_info@(CPRInfo ci_args))
440 | isNewTyCon tycon -- a new type: under the coercions must be a
441 -- constructed product
442 = ASSERT ( null $ tail inst_con_arg_tys )
443 mk_cpr_case (head inst_con_arg_tys, cpr_info)
444 `thenUs` \(arg, tup, exp) ->
445 getUniqueUs `thenUs` \id_uniq ->
446 let id_id = mk_ww_local id_uniq ty
447 new_exp_case = \var -> Case (Note (Coerce (idType arg) ty) (Var id_id))
449 [(DEFAULT,[], exp var)]
451 returnUs (id_id, tup, new_exp_case)
453 | otherwise -- a data type
454 -- flatten components
455 = mapUs mk_cpr_case (zip inst_con_arg_tys ci_args)
456 `thenUs` \sub_builds ->
457 getUniqueUs `thenUs` \id_uniq ->
458 let id_id = mk_ww_local id_uniq ty
459 (args, tup, exp) = unzip3 sub_builds
460 con_app = mkConApp data_con (map Var args)
462 new_exp_case = \var -> Case (Var id_id) (mkWildId ty)
463 [(DataCon data_con, args,
464 foldl (\e f -> f e) var exp)]
466 returnUs (id_id, new_tup, new_exp_case)
468 (data_con, tycon, tycon_arg_tys, inst_con_arg_tys) = splitType "mk_cpr_case" ty
472 @cpr_reconstruct@ does the opposite of @cpr_flatten@. It takes the unboxed
473 tuple produced by the worker and reconstructs the structured result.
476 cpr_reconstruct :: Type -> CprInfo -> UniqSM (CoreExpr -> CoreExpr)
477 cpr_reconstruct ty cpr_info
478 = mk_cpr_let (ty,cpr_info) `thenUs` \(res_id, tup_ids, reconstruct_exp) ->
479 returnUs (\worker -> Case worker (mkWildId $ worker_type tup_ids)
480 [(DataCon $ unboxedTupleCon $ length tup_ids,
481 tup_ids, reconstruct_exp $ Var res_id)])
484 worker_type ids = mkTyConApp (unboxedTupleTyCon (length ids)) (map idType ids)
487 mk_cpr_let :: (Type, CprInfo) ->
488 UniqSM (CoreBndr, -- Binder for this component of result
489 [CoreBndr], -- Binders which will appear in worker's result
490 CoreExpr -> CoreExpr) -- Code to produce structured result.
491 mk_cpr_let (ty, NoCPRInfo)
492 -- this component will appear explicitly in the unboxed tuple.
493 = getUniqueUs `thenUs` \id_uniq ->
495 id_id = mk_ww_local id_uniq ty
497 returnUs (id_id, [id_id], id)
499 mk_cpr_let (ty, cpr_info@(CPRInfo ci_args))
500 | isNewTyCon tycon -- a new type: must coerce the argument to this type
501 = ASSERT ( null $ tail inst_con_arg_tys )
502 mk_cpr_let (head inst_con_arg_tys, cpr_info)
503 `thenUs` \(arg, tup, exp) ->
504 getUniqueUs `thenUs` \id_uniq ->
505 let id_id = mk_ww_local id_uniq ty
506 new_exp = \var -> exp (Let (NonRec id_id (Note (Coerce ty (idType arg)) (Var arg))) var)
508 returnUs (id_id, tup, new_exp)
510 | otherwise -- a data type
511 -- reconstruct components then apply data con
512 = mapUs mk_cpr_let (zip inst_con_arg_tys ci_args)
513 `thenUs` \sub_builds ->
514 getUniqueUs `thenUs` \id_uniq ->
515 let id_id = mk_ww_local id_uniq ty
516 (args, tup, exp) = unzip3 sub_builds
517 con_app = mkConApp data_con $ (map Type tycon_arg_tys) ++ (map Var args)
519 new_exp = \var -> foldl (\e f -> f e) (Let (NonRec id_id con_app) var) exp
521 returnUs (id_id, new_tup, new_exp)
523 (data_con, tycon, tycon_arg_tys, inst_con_arg_tys) = splitType "mk_cpr_let" ty
525 splitType :: String -> Type -> (DataCon, TyCon, [Type], [Type])
526 splitType fname ty = (data_con, tycon, tycon_arg_tys, dataConArgTys data_con tycon_arg_tys)
528 (data_con, tycon, tycon_arg_tys)
529 = case (splitAlgTyConApp_maybe ty) of
530 Just (arg_tycon, tycon_arg_tys, [data_con]) ->
531 -- The main event: a single-constructor data type
532 (data_con, arg_tycon, tycon_arg_tys)
534 Just (_, _, data_cons) ->
535 pprPanic (fname ++ ":")
536 (text "not one constr (interface files not consistent/up to date?)"
540 pprPanic (fname ++ ":")
541 (text "not a datatype" $$ ppr ty)
545 %************************************************************************
547 \subsection{Utilities}
549 %************************************************************************
553 mk_absent_let arg body
554 | not (isUnLiftedType arg_ty)
555 = Let (NonRec arg (mkTyApps (Var aBSENT_ERROR_ID) [arg_ty])) body
557 = panic "WwLib: haven't done mk_absent_let for primitives yet"
561 mk_unpk_case NewType arg unpk_args boxing_con boxing_tycon body
562 -- A newtype! Use a coercion not a case
563 = ASSERT( null other_args )
564 Case (Note (Coerce (idType unpk_arg) (idType arg)) (Var arg))
565 (sanitiseCaseBndr unpk_arg)
568 (unpk_arg:other_args) = unpk_args
570 mk_unpk_case DataType arg unpk_args boxing_con boxing_tycon body
573 (sanitiseCaseBndr arg)
574 [(DataCon boxing_con, unpk_args, body)]
576 sanitiseCaseBndr :: Id -> Id
577 -- The argument we are scrutinising has the right type to be
578 -- a case binder, so it's convenient to re-use it for that purpose.
579 -- But we *must* throw away all its IdInfo. In particular, the argument
580 -- will have demand info on it, and that demand info may be incorrect for
581 -- the case binder. e.g. case ww_arg of ww_arg { I# x -> ... }
582 -- Quite likely ww_arg isn't used in '...'. The case may get discarded
583 -- if the case binder says "I'm demanded". This happened in a situation
584 -- like (x+y) `seq` ....
585 sanitiseCaseBndr id = id `setIdInfo` vanillaIdInfo
587 mk_pk_let NewType arg boxing_con con_tys unpk_args body
588 = ASSERT( null other_args )
589 Let (NonRec arg (Note (Coerce (idType arg) (idType unpk_arg)) (Var unpk_arg))) body
591 (unpk_arg:other_args) = unpk_args
593 mk_pk_let DataType arg boxing_con con_tys unpk_args body
594 = Let (NonRec arg (Con (DataCon boxing_con) con_args)) body
596 con_args = map Type con_tys ++ map Var unpk_args
599 mk_ww_local uniq ty = mkSysLocal SLIT("ww") uniq ty
602 mk_unboxed_tuple :: [(CoreExpr, Type)] -> (CoreExpr, Type)
603 mk_unboxed_tuple contents
604 = (mkConApp (unboxedTupleCon (length contents))
605 (map (Type . snd) contents ++
607 mkTyConApp (unboxedTupleTyCon (length contents))