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,
19 import IdInfo ( CprInfo(..), noCprInfo )
20 import Const ( Con(..), DataCon )
21 import DataCon ( dataConArgTys )
22 import Demand ( Demand(..) )
23 import PrelVals ( aBSENT_ERROR_ID )
24 import TysWiredIn ( unitTy, unitDataCon,
25 unboxedTupleCon, unboxedTupleTyCon )
26 import Type ( isUnLiftedType, mkTyVarTys, mkTyVarTy, mkFunTys,
27 splitForAllTys, splitFunTys,
28 splitAlgTyConApp_maybe, mkTyConApp,
31 import TyCon ( isNewTyCon,
33 import BasicTypes ( NewOrData(..) )
35 import UniqSupply ( returnUs, thenUs, getUniqueUs, getUniquesUs,
37 import Util ( zipWithEqual, zipEqual )
41 %************************************************************************
43 \subsection[datatype-WwLib]{@WwBinding@: a datatype for worker/wrapper-ing}
45 %************************************************************************
47 In the worker/wrapper stuff, we want to carry around @CoreBindings@ in
48 an ``intermediate form'' that can later be turned into a \tr{let} or
49 \tr{case} (depending on strictness info).
54 | WwCase (CoreExpr -> CoreExpr)
55 -- the "case" will be a "strict let" of the form:
60 -- (instead of "let <blah> = rhs in body")
62 -- The expr you pass to the function is "body" (the
63 -- expression that goes "in the corner").
66 %************************************************************************
68 \subsection[mkWrapperAndWorker]{@mkWrapperAndWorker@}
70 %************************************************************************
72 ************ WARNING ******************
73 these comments are rather out of date
74 *****************************************
76 @mkWrapperAndWorker@ is given:
79 The {\em original function} \tr{f}, of the form:
81 f = /\ tyvars -> \ args -> body
83 The original-binder \tr{f}, the \tr{tyvars}, \tr{args}, and \tr{body}
86 We use the Id \tr{f} mostly to get its type.
89 Strictness information about \tr{f}, in the form of a list of
96 @mkWrapperAndWorker@ produces (A BIT OUT-OF-DATE...):
99 Maybe @Nothing@: no worker/wrappering going on in this case. This can
100 happen (a)~if the strictness info says that there is nothing
101 interesting to do or (b)~if *any* of the argument types corresponding
102 to ``active'' arg postitions is abstract or will be to the outside
103 world (i.e., {\em this} module can see the constructors, but nobody
104 else will be able to). An ``active'' arg position is one which the
105 wrapper has to unpack. An importing module can't do this unpacking,
106 so it simply has to give up and call the wrapper only.
109 Maybe \tr{Just (wrapper_Id, wrapper_body, worker_Id, worker_body)}.
111 The @wrapper_Id@ is just the one that was passed in, with its
112 strictness IdInfo updated.
115 The \tr{body} of the original function may not be given (i.e., it's
116 BOTTOM), in which case you'd jolly well better not tug on the
119 Here's an example. The original function is:
121 g :: forall a . Int -> [a] -> a
123 g = /\ a -> \ x ys ->
129 From this, we want to produce:
131 -- wrapper (an unfolding)
132 g :: forall a . Int -> [a] -> a
134 g = /\ a -> \ x ys ->
136 I# x# -> g.wrk a x# ys
137 -- call the worker; don't forget the type args!
140 g.wrk :: forall a . Int# -> [a] -> a
142 g.wrk = /\ a -> \ x# ys ->
146 case x of -- note: body of g moved intact
151 Something we have to be careful about: Here's an example:
153 -- "f" strictness: U(P)U(P)
154 f (I# a) (I# b) = a +# b
156 g = f -- "g" strictness same as "f"
158 \tr{f} will get a worker all nice and friendly-like; that's good.
159 {\em But we don't want a worker for \tr{g}}, even though it has the
160 same strictness as \tr{f}. Doing so could break laziness, at best.
162 Consequently, we insist that the number of strictness-info items is
163 exactly the same as the number of lambda-bound arguments. (This is
164 probably slightly paranoid, but OK in practice.) If it isn't the
165 same, we ``revise'' the strictness info, so that we won't propagate
166 the unusable strictness-info into the interfaces.
169 %************************************************************************
171 \subsection{Functions over Demands}
173 %************************************************************************
176 mAX_WORKER_ARGS :: Int -- ToDo: set via flag
179 setUnpackStrategy :: [Demand] -> [Demand]
181 = snd (go (mAX_WORKER_ARGS - nonAbsentArgs ds) ds)
183 go :: Int -- Max number of args available for sub-components of [Demand]
185 -> (Int, [Demand]) -- Args remaining after subcomponents of [Demand] are unpacked
187 go n (WwUnpack nd _ cs : ds) | n' >= 0
188 = WwUnpack nd True cs' `cons` go n'' ds
190 = WwUnpack nd False cs `cons` go n ds
192 n' = n + 1 - nonAbsentArgs cs
193 -- Add one because we don't pass the top-level arg any more
194 -- Delete # of non-absent args to which we'll now be committed
197 go n (d:ds) = d `cons` go n ds
200 cons d (n,ds) = (n, d:ds)
202 nonAbsentArgs :: [Demand] -> Int
204 nonAbsentArgs (WwLazy True : ds) = nonAbsentArgs ds
205 nonAbsentArgs (d : ds) = 1 + nonAbsentArgs ds
207 worthSplitting :: [Demand] -> Bool -- True <=> the wrapper would not be an identity function
208 worthSplitting ds = any worth_it ds
210 worth_it (WwLazy True) = True -- Absent arg
211 worth_it (WwUnpack _ True _) = True -- Arg to unpack
212 worth_it WwStrict = False -- Don't w/w just because of strictness
213 worth_it other = False
215 allAbsent :: [Demand] -> Bool
216 allAbsent ds = all absent ds
218 absent (WwLazy is_absent) = is_absent
219 absent (WwUnpack _ True cs) = allAbsent cs
224 %************************************************************************
226 \subsection{The worker wrapper core}
228 %************************************************************************
230 @mkWrapper@ is called when importing a function. We have the type of
231 the function and the name of its worker, and we want to make its body (the wrapper).
234 mkWrapper :: Type -- Wrapper type
235 -> [Demand] -- Wrapper strictness info
236 -> CprInfo -- Wrapper cpr info
237 -> UniqSM (Id -> CoreExpr) -- Wrapper body, missing worker Id
239 mkWrapper fun_ty demands cpr_info
241 n_wrap_args = length demands
243 getUniquesUs n_wrap_args `thenUs` \ wrap_uniqs ->
245 (tyvars, tau_ty) = splitForAllTys fun_ty
246 (arg_tys, body_ty) = splitFunTys tau_ty
247 -- The "expanding dicts" part here is important, even for the splitForAll
248 -- The imported thing might be a dictionary, such as Functor Foo
249 -- But Functor Foo = forall a b. (a->b) -> Foo a -> Foo b
250 -- and as such might have some strictness info attached.
251 -- Then we need to have enough args to zip to the strictness info
253 wrap_args = ASSERT( n_wrap_args <= length arg_tys )
254 zipWith mk_ww_local wrap_uniqs arg_tys
256 leftover_arg_tys = drop n_wrap_args arg_tys
257 final_body_ty = mkFunTys leftover_arg_tys body_ty
259 mkWwBodies tyvars wrap_args final_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 getUniqueUs `thenUs` \ void_arg_uniq ->
285 void_arg = mk_ww_local void_arg_uniq unitTy
287 returnUs (\ work_id -> mkLams tyvars $ mkLams args $
289 (map (Type . mkTyVarTy) tyvars ++ [mkConApp unitDataCon []]),
290 \ body -> mkLams (tyvars ++ [void_arg]) body,
293 mkWwBodies tyvars wrap_args body_ty demands cpr_info
296 -- demands may be longer than number of args. If we aren't doing w/w
297 -- for strictness then demands is an infinite list of 'lazy' args.
298 wrap_args_w_demands = zipWith setIdDemandInfo wrap_args demands
300 mkWW wrap_args_w_demands `thenUs` \ (wrap_fn, work_args_w_demands, work_fn) ->
301 mkWWcpr body_ty cpr_info
302 `thenUs` \ (wrap_fn_w_cpr, work_fn_w_cpr) ->
303 returnUs (\ work_id -> mkLams tyvars $ mkLams wrap_args_w_demands $
304 (wrap_fn_w_cpr . wrap_fn) (mkTyApps (Var work_id) (mkTyVarTys tyvars)),
306 \ body -> mkLams tyvars $ mkLams work_args_w_demands $
307 (work_fn_w_cpr . work_fn) body,
309 map getIdDemandInfo work_args_w_demands)
314 mkWW :: [Id] -- Wrapper args; have their demand info on them
315 -> UniqSM (CoreExpr -> CoreExpr, -- Wrapper body, lacking the inner call to the worker
316 -- and without its lambdas
317 [Id], -- Worker args; have their demand info on them
318 CoreExpr -> CoreExpr) -- Worker body, lacking the original body of the function
323 = returnUs (\ wrapper_body -> wrapper_body,
325 \ worker_body -> worker_body)
329 = case getIdDemandInfo arg of
333 mkWW ds `thenUs` \ (wrap_fn, worker_args, work_fn) ->
334 returnUs (\ wrapper_body -> wrap_fn wrapper_body,
336 \ worker_body -> mk_absent_let arg (work_fn worker_body))
340 WwUnpack new_or_data True cs ->
341 getUniquesUs (length inst_con_arg_tys) `thenUs` \ uniqs ->
343 unpk_args = zipWith mk_ww_local uniqs inst_con_arg_tys
344 unpk_args_w_ds = zipWithEqual "mkWW" setIdDemandInfo unpk_args cs
346 mkWW (unpk_args_w_ds ++ ds) `thenUs` \ (wrap_fn, worker_args, work_fn) ->
347 returnUs (\ wrapper_body -> mk_unpk_case new_or_data arg unpk_args data_con arg_tycon
348 (wrap_fn wrapper_body),
350 \ worker_body -> work_fn (mk_pk_let new_or_data arg data_con
351 tycon_arg_tys unpk_args worker_body))
353 inst_con_arg_tys = dataConArgTys data_con tycon_arg_tys
354 (arg_tycon, tycon_arg_tys, data_con)
355 = case (splitAlgTyConApp_maybe (idType arg)) of
357 Just (arg_tycon, tycon_arg_tys, [data_con]) ->
358 -- The main event: a single-constructor data type
359 (arg_tycon, tycon_arg_tys, data_con)
361 Just (_, _, data_cons) ->
362 pprPanic "mk_ww_arg_processing:"
363 (text "not one constr (interface files not consistent/up to date?)"
364 $$ (ppr arg <+> ppr (idType arg)))
367 panic "mk_ww_arg_processing: not datatype"
372 mkWW ds `thenUs` \ (wrap_fn, worker_args, work_fn) ->
373 returnUs (\ wrapper_body -> wrap_fn (App wrapper_body (Var arg)),
378 @mkWWcpr@ takes the worker/wrapper pair produced from the strictness
379 info and adds in the CPR transformation. The worker returns an
380 unboxed tuple containing non-CPR components. The wrapper takes this
381 tuple and re-produces the correct structured output.
383 The non-CPR results appear ordered in the unboxed tuple as if by a
384 left-to-right traversal of the result structure.
389 mkWWcpr :: Type -- function body type
390 -> CprInfo -- CPR analysis results
391 -> UniqSM (CoreExpr -> CoreExpr, -- New wrapper
392 CoreExpr -> CoreExpr) -- New worker
394 mkWWcpr body_ty NoCPRInfo
395 = returnUs (id, id) -- Must be just the strictness transf.
396 mkWWcpr body_ty (CPRInfo cpr_args)
397 = getUniqueUs `thenUs` \ body_arg_uniq ->
399 body_var = mk_ww_local body_arg_uniq body_ty
401 cpr_reconstruct body_ty cpr_info' `thenUs` \reconst_fn ->
402 cpr_flatten body_ty cpr_info' `thenUs` \flatten_fn ->
403 returnUs (reconst_fn, flatten_fn)
404 -- We only make use of the outer level of CprInfo, otherwise we
405 -- may lose laziness. :-( Hopefully, we will find a use for the
406 -- extra info some day (e.g. creating versions specialized to
407 -- the use made of the components of the result by the callee)
408 where cpr_info' = CPRInfo (map (const NoCPRInfo) cpr_args)
412 @cpr_flatten@ takes the result type produced by the body and the info
413 from the CPR analysis and flattens the constructed product components.
414 These are returned in an unboxed tuple.
418 cpr_flatten :: Type -> CprInfo -> UniqSM (CoreExpr -> CoreExpr)
419 cpr_flatten ty cpr_info
420 = mk_cpr_case (ty, cpr_info) `thenUs` \(res_id, tup_ids, flatten_exp) ->
421 returnUs (\body -> Case body res_id
422 [(DEFAULT, [], flatten_exp (fst $ mk_unboxed_tuple tup_ids))])
426 mk_cpr_case :: (Type, CprInfo) ->
427 UniqSM (CoreBndr, -- Name of binder for this part of result
428 [(CoreExpr, Type)], -- expressions for flattened result
429 CoreExpr -> CoreExpr) -- add in code to flatten result
431 mk_cpr_case (ty, NoCPRInfo)
432 -- this component must be returned as a component of the unboxed tuple result
433 = getUniqueUs `thenUs` \id_uniq ->
434 let id_id = mk_ww_local id_uniq ty in
435 returnUs (id_id, [(Var id_id, ty)], id)
436 mk_cpr_case (ty, cpr_info@(CPRInfo ci_args))
437 | isNewTyCon tycon -- a new type: under the coercions must be a
438 -- constructed product
439 = ASSERT ( null $ tail inst_con_arg_tys )
440 mk_cpr_case (head inst_con_arg_tys, cpr_info)
441 `thenUs` \(arg, tup, exp) ->
442 getUniqueUs `thenUs` \id_uniq ->
443 let id_id = mk_ww_local id_uniq ty
444 new_exp_case = \var -> Case (Note (Coerce (idType arg) ty) (Var id_id))
446 [(DEFAULT,[], exp var)]
448 returnUs (id_id, tup, new_exp_case)
450 | otherwise -- a data type
451 -- flatten components
452 = mapUs mk_cpr_case (zip inst_con_arg_tys ci_args)
453 `thenUs` \sub_builds ->
454 getUniqueUs `thenUs` \id_uniq ->
455 let id_id = mk_ww_local id_uniq ty
456 (args, tup, exp) = unzip3 sub_builds
457 con_app = mkConApp data_con (map Var args)
459 new_exp_case = \var -> Case (Var id_id) (mkWildId ty)
460 [(DataCon data_con, args,
461 foldl (\e f -> f e) var exp)]
463 returnUs (id_id, new_tup, new_exp_case)
465 (data_con, tycon, tycon_arg_tys, inst_con_arg_tys) = splitType "mk_cpr_case" ty
469 @cpr_reconstruct@ does the opposite of @cpr_flatten@. It takes the unboxed
470 tuple produced by the worker and reconstructs the structured result.
473 cpr_reconstruct :: Type -> CprInfo -> UniqSM (CoreExpr -> CoreExpr)
474 cpr_reconstruct ty cpr_info
475 = mk_cpr_let (ty,cpr_info) `thenUs` \(res_id, tup_ids, reconstruct_exp) ->
476 returnUs (\worker -> Case worker (mkWildId $ worker_type tup_ids)
477 [(DataCon $ unboxedTupleCon $ length tup_ids,
478 tup_ids, reconstruct_exp $ Var res_id)])
481 worker_type ids = mkTyConApp (unboxedTupleTyCon (length ids)) (map idType ids)
484 mk_cpr_let :: (Type, CprInfo) ->
485 UniqSM (CoreBndr, -- Binder for this component of result
486 [CoreBndr], -- Binders which will appear in worker's result
487 CoreExpr -> CoreExpr) -- Code to produce structured result.
488 mk_cpr_let (ty, NoCPRInfo)
489 -- this component will appear explicitly in the unboxed tuple.
490 = getUniqueUs `thenUs` \id_uniq ->
491 let id_id = mk_ww_local id_uniq ty in
492 returnUs (id_id, [id_id], id)
493 mk_cpr_let (ty, cpr_info@(CPRInfo ci_args))
494 | isNewTyCon tycon -- a new type: must coerce the argument to this type
495 = ASSERT ( null $ tail inst_con_arg_tys )
496 mk_cpr_let (head inst_con_arg_tys, cpr_info)
497 `thenUs` \(arg, tup, exp) ->
498 getUniqueUs `thenUs` \id_uniq ->
499 let id_id = mk_ww_local id_uniq ty
500 new_exp = \var -> exp (Let (NonRec id_id (Note (Coerce ty (idType arg)) (Var arg))) var)
502 returnUs (id_id, tup, new_exp)
504 | otherwise -- a data type
505 -- reconstruct components then apply data con
506 = mapUs mk_cpr_let (zip inst_con_arg_tys ci_args)
507 `thenUs` \sub_builds ->
508 getUniqueUs `thenUs` \id_uniq ->
509 let id_id = mk_ww_local id_uniq ty
510 (args, tup, exp) = unzip3 sub_builds
511 con_app = mkConApp data_con $ (map Type tycon_arg_tys) ++ (map Var args)
513 new_exp = \var -> foldl (\e f -> f e) (Let (NonRec id_id con_app) var) exp
515 returnUs (id_id, new_tup, new_exp)
517 (data_con, tycon, tycon_arg_tys, inst_con_arg_tys) = splitType "mk_cpr_let" ty
519 splitType :: String -> Type -> (DataCon, TyCon, [Type], [Type])
520 splitType fname ty = (data_con, tycon, tycon_arg_tys, dataConArgTys data_con tycon_arg_tys)
522 (data_con, tycon, tycon_arg_tys)
523 = case (splitAlgTyConApp_maybe ty) of
524 Just (arg_tycon, tycon_arg_tys, [data_con]) ->
525 -- The main event: a single-constructor data type
526 (data_con, arg_tycon, tycon_arg_tys)
528 Just (_, _, data_cons) ->
529 pprPanic (fname ++ ":")
530 (text "not one constr (interface files not consistent/up to date?)"
534 pprPanic (fname ++ ":")
535 (text "not a datatype" $$ ppr ty)
540 %************************************************************************
542 \subsection{Utilities}
544 %************************************************************************
548 mk_absent_let arg body
549 | not (isUnLiftedType arg_ty)
550 = Let (NonRec arg (mkTyApps (Var aBSENT_ERROR_ID) [arg_ty])) body
552 = panic "WwLib: haven't done mk_absent_let for primitives yet"
556 mk_unpk_case NewType arg unpk_args boxing_con boxing_tycon body
557 -- A newtype! Use a coercion not a case
558 = ASSERT( null other_args )
559 Case (Note (Coerce (idType unpk_arg) (idType arg)) (Var arg))
563 (unpk_arg:other_args) = unpk_args
565 mk_unpk_case DataType arg unpk_args boxing_con boxing_tycon body
567 = Case (Var arg) arg [(DataCon boxing_con, unpk_args, body)]
569 mk_pk_let NewType arg boxing_con con_tys unpk_args body
570 = ASSERT( null other_args )
571 Let (NonRec arg (Note (Coerce (idType arg) (idType unpk_arg)) (Var unpk_arg))) body
573 (unpk_arg:other_args) = unpk_args
575 mk_pk_let DataType arg boxing_con con_tys unpk_args body
576 = Let (NonRec arg (Con (DataCon boxing_con) con_args)) body
578 con_args = map Type con_tys ++ map Var unpk_args
581 mk_ww_local uniq ty = mkSysLocal SLIT("ww") uniq ty
584 mk_unboxed_tuple :: [(CoreExpr, Type)] -> (CoreExpr, Type)
585 mk_unboxed_tuple contents
586 = (mkConApp (unboxedTupleCon (length contents))
587 (map (Type . snd) contents ++
589 mkTyConApp (unboxedTupleTyCon (length contents))