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}
7 module WwLib ( mkWwBodies, mkWWstr, mkWorkerArgs ) where
9 #include "HsVersions.h"
12 import CoreUtils ( exprType )
13 import Id ( Id, idType, mkSysLocal, idDemandInfo, setIdDemandInfo,
14 isOneShotLambda, setOneShotLambda, setIdUnfolding,
17 import IdInfo ( vanillaIdInfo )
19 import Demand ( Demand(..), DmdResult(..), Demands(..) )
20 import MkCore ( mkRuntimeErrorApp, aBSENT_ERROR_ID )
21 import MkId ( realWorldPrimId, voidArgId,
22 mkUnpackCase, mkProductBox )
23 import TysPrim ( realWorldStatePrimTy )
24 import TysWiredIn ( tupleCon )
26 import Coercion ( mkSymCo, splitNewTypeRepCo_maybe )
27 import BasicTypes ( Boxity(..) )
28 import Literal ( absentLiteralOf )
31 import Util ( zipWithEqual )
37 %************************************************************************
39 \subsection[mkWrapperAndWorker]{@mkWrapperAndWorker@}
41 %************************************************************************
43 Here's an example. The original function is:
46 g :: forall a . Int -> [a] -> a
48 g = \/\ a -> \ x ys ->
54 From this, we want to produce:
56 -- wrapper (an unfolding)
57 g :: forall a . Int -> [a] -> a
59 g = \/\ a -> \ x ys ->
62 -- call the worker; don't forget the type args!
65 $wg :: forall a . Int# -> [a] -> a
67 $wg = \/\ a -> \ x# ys ->
71 case x of -- note: body of g moved intact
76 Something we have to be careful about: Here's an example:
79 -- "f" strictness: U(P)U(P)
80 f (I# a) (I# b) = a +# b
82 g = f -- "g" strictness same as "f"
85 \tr{f} will get a worker all nice and friendly-like; that's good.
86 {\em But we don't want a worker for \tr{g}}, even though it has the
87 same strictness as \tr{f}. Doing so could break laziness, at best.
89 Consequently, we insist that the number of strictness-info items is
90 exactly the same as the number of lambda-bound arguments. (This is
91 probably slightly paranoid, but OK in practice.) If it isn't the
92 same, we ``revise'' the strictness info, so that we won't propagate
93 the unusable strictness-info into the interfaces.
96 %************************************************************************
98 \subsection{The worker wrapper core}
100 %************************************************************************
102 @mkWwBodies@ is called when doing the worker\/wrapper split inside a module.
105 mkWwBodies :: Type -- Type of original function
106 -> [Demand] -- Strictness of original function
107 -> DmdResult -- Info about function result
108 -> [Bool] -- One-shot-ness of the function
109 -> UniqSM ([Demand], -- Demands for worker (value) args
110 Id -> CoreExpr, -- Wrapper body, lacking only the worker Id
111 CoreExpr -> CoreExpr) -- Worker body, lacking the original function rhs
113 -- wrap_fn_args E = \x y -> E
114 -- work_fn_args E = E x y
116 -- wrap_fn_str E = case x of { (a,b) ->
117 -- case a of { (a1,a2) ->
119 -- work_fn_str E = \a2 a2 b y ->
120 -- let a = (a1,a2) in
124 mkWwBodies fun_ty demands res_info one_shots
125 = do { let arg_info = demands `zip` (one_shots ++ repeat False)
126 ; (wrap_args, wrap_fn_args, work_fn_args, res_ty) <- mkWWargs emptyTvSubst fun_ty arg_info
127 ; (work_args, wrap_fn_str, work_fn_str) <- mkWWstr wrap_args
129 -- Don't do CPR if the worker doesn't have any value arguments
130 -- Then the worker is just a constant, so we don't want to unbox it.
131 ; (wrap_fn_cpr, work_fn_cpr, _cpr_res_ty)
132 <- if any isId work_args then
133 mkWWcpr res_ty res_info
135 return (id, id, res_ty)
137 ; let (work_lam_args, work_call_args) = mkWorkerArgs work_args res_ty
138 ; return ([idDemandInfo v | v <- work_call_args, isId v],
139 wrap_fn_args . wrap_fn_cpr . wrap_fn_str . applyToVars work_call_args . Var,
140 mkLams work_lam_args. work_fn_str . work_fn_cpr . work_fn_args) }
141 -- We use an INLINE unconditionally, even if the wrapper turns out to be
142 -- something trivial like
144 -- f = __inline__ (coerce T fw)
145 -- The point is to propagate the coerce to f's call sites, so even though
146 -- f's RHS is now trivial (size 1) we still want the __inline__ to prevent
147 -- fw from being inlined into f's RHS
151 %************************************************************************
153 \subsection{Making wrapper args}
155 %************************************************************************
157 During worker-wrapper stuff we may end up with an unlifted thing
158 which we want to let-bind without losing laziness. So we
159 add a void argument. E.g.
161 f = /\a -> \x y z -> E::Int# -- E does not mention x,y,z
163 fw = /\ a -> \void -> E
164 f = /\ a -> \x y z -> fw realworld
166 We use the state-token type which generates no code.
169 mkWorkerArgs :: [Var]
170 -> Type -- Type of body
171 -> ([Var], -- Lambda bound args
172 [Var]) -- Args at call site
173 mkWorkerArgs args res_ty
174 | any isId args || not (isUnLiftedType res_ty)
177 = (args ++ [voidArgId], args ++ [realWorldPrimId])
181 %************************************************************************
183 \subsection{Coercion stuff}
185 %************************************************************************
187 We really want to "look through" coerces.
188 Reason: I've seen this situation:
190 let f = coerce T (\s -> E)
196 If only we w/w'd f, we'd get
197 let f = coerce T (\s -> fw s)
201 Now we'll inline f to get
209 Now we'll see that fw has arity 1, and will arity expand
210 the \x to get what we want.
213 -- mkWWargs just does eta expansion
214 -- is driven off the function type and arity.
215 -- It chomps bites off foralls, arrows, newtypes
216 -- and keeps repeating that until it's satisfied the supplied arity
218 mkWWargs :: TvSubst -- Freshening substitution to apply to the type
219 -- See Note [Freshen type variables]
220 -> Type -- The type of the function
221 -> [(Demand,Bool)] -- Demands and one-shot info for value arguments
222 -> UniqSM ([Var], -- Wrapper args
223 CoreExpr -> CoreExpr, -- Wrapper fn
224 CoreExpr -> CoreExpr, -- Worker fn
225 Type) -- Type of wrapper body
227 mkWWargs subst fun_ty arg_info
228 | Just (rep_ty, co) <- splitNewTypeRepCo_maybe fun_ty
229 -- The newtype case is for when the function has
230 -- a recursive newtype after the arrow (rare)
231 -- We check for arity >= 0 to avoid looping in the case
232 -- of a function whose type is, in effect, infinite
233 -- [Arity is driven by looking at the term, not just the type.]
235 -- It's also important when we have a function returning (say) a pair
236 -- wrapped in a recursive newtype, at least if CPR analysis can look
237 -- through such newtypes, which it probably can since they are
240 -- Note (Sept 08): This case applies even if demands is empty.
241 -- I'm not quite sure why; perhaps it makes it
243 = do { (wrap_args, wrap_fn_args, work_fn_args, res_ty)
244 <- mkWWargs subst rep_ty arg_info
246 \e -> Cast (wrap_fn_args e) (mkSymCo co),
247 \e -> work_fn_args (Cast e co),
251 = return ([], id, id, substTy subst fun_ty)
253 | Just (tv, fun_ty') <- splitForAllTy_maybe fun_ty
254 = do { let (subst', tv') = substTyVarBndr subst tv
255 -- This substTyVarBndr clones the type variable when necy
256 -- See Note [Freshen type variables]
257 ; (wrap_args, wrap_fn_args, work_fn_args, res_ty)
258 <- mkWWargs subst' fun_ty' arg_info
259 ; return (tv' : wrap_args,
260 Lam tv' . wrap_fn_args,
261 work_fn_args . (`App` Type (mkTyVarTy tv')),
264 | ((dmd,one_shot):arg_info') <- arg_info
265 , Just (arg_ty, fun_ty') <- splitFunTy_maybe fun_ty
266 = do { uniq <- getUniqueM
267 ; let arg_ty' = substTy subst arg_ty
268 id = mk_wrap_arg uniq arg_ty' dmd one_shot
269 ; (wrap_args, wrap_fn_args, work_fn_args, res_ty)
270 <- mkWWargs subst fun_ty' arg_info'
271 ; return (id : wrap_args,
272 Lam id . wrap_fn_args,
273 work_fn_args . (`App` varToCoreExpr id),
277 = WARN( True, ppr fun_ty ) -- Should not happen: if there is a demand
278 return ([], id, id, substTy subst fun_ty) -- then there should be a function arrow
280 applyToVars :: [Var] -> CoreExpr -> CoreExpr
281 applyToVars vars fn = mkVarApps fn vars
283 mk_wrap_arg :: Unique -> Type -> Demand -> Bool -> Id
284 mk_wrap_arg uniq ty dmd one_shot
285 = set_one_shot one_shot (setIdDemandInfo (mkSysLocal (fsLit "w") uniq ty) dmd)
287 set_one_shot True id = setOneShotLambda id
288 set_one_shot False id = id
291 Note [Freshen type variables]
292 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
293 Wen we do a worker/wrapper split, we must not use shadowed names,
296 which is obviously wrong. Type variables can can in principle shadow,
297 within a type (e.g. forall a. a -> forall a. a->a). But type
298 variables *are* mentioned in <blah>, so we must substitute.
300 That's why we carry the TvSubst through mkWWargs
302 %************************************************************************
304 \subsection{Strictness stuff}
306 %************************************************************************
309 mkWWstr :: [Var] -- Wrapper args; have their demand info on them
310 -- *Includes type variables*
311 -> UniqSM ([Var], -- Worker args
312 CoreExpr -> CoreExpr, -- Wrapper body, lacking the worker call
313 -- and without its lambdas
314 -- This fn adds the unboxing
316 CoreExpr -> CoreExpr) -- Worker body, lacking the original body of the function,
317 -- and lacking its lambdas.
318 -- This fn does the reboxing
320 = return ([], nop_fn, nop_fn)
322 mkWWstr (arg : args) = do
323 (args1, wrap_fn1, work_fn1) <- mkWWstr_one arg
324 (args2, wrap_fn2, work_fn2) <- mkWWstr args
325 return (args1 ++ args2, wrap_fn1 . wrap_fn2, work_fn1 . work_fn2)
327 ----------------------
328 -- mkWWstr_one wrap_arg = (work_args, wrap_fn, work_fn)
329 -- * wrap_fn assumes wrap_arg is in scope,
330 -- brings into scope work_args (via cases)
331 -- * work_fn assumes work_args are in scope, a
332 -- brings into scope wrap_arg (via lets)
333 mkWWstr_one :: Var -> UniqSM ([Var], CoreExpr -> CoreExpr, CoreExpr -> CoreExpr)
336 = return ([arg], nop_fn, nop_fn)
339 = case idDemandInfo arg of
341 -- Absent case. We can't always handle absence for arbitrary
342 -- unlifted types, so we need to choose just the cases we can
343 -- (that's what mk_absent_let does)
344 Abs | Just work_fn <- mk_absent_let arg
345 -> return ([], nop_fn, work_fn)
349 | Just (_arg_tycon, _tycon_arg_tys, data_con, inst_con_arg_tys)
350 <- deepSplitProductType_maybe (idType arg)
351 -> do uniqs <- getUniquesM
353 unpk_args = zipWith mk_ww_local uniqs inst_con_arg_tys
354 unpk_args_w_ds = zipWithEqual "mkWWstr" set_worker_arg_info unpk_args cs
355 unbox_fn = mkUnpackCase (sanitiseCaseBndr arg) (Var arg) unpk_args data_con
356 rebox_fn = Let (NonRec arg con_app)
357 con_app = mkProductBox unpk_args (idType arg)
358 (worker_args, wrap_fn, work_fn) <- mkWWstr unpk_args_w_ds
359 return (worker_args, unbox_fn . wrap_fn, work_fn . rebox_fn)
360 -- Don't pass the arg, rebox instead
362 -- `seq` demand; evaluate in wrapper in the hope
363 -- of dropping seqs in the worker
366 arg_w_unf = arg `setIdUnfolding` evaldUnfolding
367 -- Tell the worker arg that it's sure to be evaluated
368 -- so that internal seqs can be dropped
370 return ([arg_w_unf], mk_seq_case arg, nop_fn)
371 -- Pass the arg, anyway, even if it is in theory discarded
374 -- x gets a (Eval (Poly Abs)) demand, but if we fail to pass it to the worker
375 -- we ABSOLUTELY MUST record that x is evaluated in the wrapper.
377 -- f x y = x `seq` fw y
378 -- fw y = let x{Evald} = error "oops" in (x `seq` y)
379 -- If we don't pin on the "Evald" flag, the seq doesn't disappear, and
380 -- we end up evaluating the absent thunk.
381 -- But the Evald flag is pretty weird, and I worry that it might disappear
382 -- during simplification, so for now I've just nuked this whole case
385 _other_demand -> return ([arg], nop_fn, nop_fn)
388 -- If the wrapper argument is a one-shot lambda, then
389 -- so should (all) the corresponding worker arguments be
390 -- This bites when we do w/w on a case join point
391 set_worker_arg_info worker_arg demand = set_one_shot (setIdDemandInfo worker_arg demand)
393 set_one_shot | isOneShotLambda arg = setOneShotLambda
394 | otherwise = \x -> x
396 ----------------------
397 nop_fn :: CoreExpr -> CoreExpr
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 -> DmdResult -- CPR analysis results
421 -> UniqSM (CoreExpr -> CoreExpr, -- New wrapper
422 CoreExpr -> CoreExpr, -- New worker
423 Type) -- Type of worker's body
425 mkWWcpr body_ty RetCPR
426 | not (isClosedAlgType body_ty)
428 text "mkWWcpr: non-algebraic or open body type" <+> ppr body_ty )
429 return (id, id, body_ty)
431 | n_con_args == 1 && isUnLiftedType con_arg_ty1 = do
432 -- Special case when there is a single result of unlifted type
434 -- Wrapper: case (..call worker..) of x -> C x
435 -- Worker: case ( ..body.. ) of C x -> x
436 (work_uniq : arg_uniq : _) <- getUniquesM
438 work_wild = mk_ww_local work_uniq body_ty
439 arg = mk_ww_local arg_uniq con_arg_ty1
440 con_app = mkProductBox [arg] body_ty
442 return (\ wkr_call -> Case wkr_call (arg) (exprType con_app) [(DEFAULT, [], con_app)],
443 \ body -> workerCase (work_wild) body [arg] data_con (Var arg),
446 | otherwise = do -- The general case
447 -- Wrapper: case (..call worker..) of (# a, b #) -> C a b
448 -- Worker: case ( ...body... ) of C a b -> (# a, b #)
451 (wrap_wild : work_wild : args) = zipWith mk_ww_local uniqs (ubx_tup_ty : body_ty : con_arg_tys)
452 arg_vars = map Var args
453 ubx_tup_con = tupleCon Unboxed n_con_args
454 ubx_tup_ty = exprType ubx_tup_app
455 ubx_tup_app = mkConApp ubx_tup_con (map Type con_arg_tys ++ arg_vars)
456 con_app = mkProductBox args body_ty
458 return (\ wkr_call -> Case wkr_call (wrap_wild) (exprType con_app) [(DataAlt ubx_tup_con, args, con_app)],
459 \ body -> workerCase (work_wild) body args data_con ubx_tup_app,
462 (_arg_tycon, _tycon_arg_tys, data_con, con_arg_tys) = deepSplitProductType "mkWWcpr" body_ty
463 n_con_args = length con_arg_tys
464 con_arg_ty1 = head con_arg_tys
466 mkWWcpr body_ty _other -- No CPR info
467 = return (id, id, body_ty)
469 -- If the original function looked like
470 -- f = \ x -> _scc_ "foo" E
472 -- then we want the CPR'd worker to look like
473 -- \ x -> _scc_ "foo" (case E of I# x -> x)
474 -- and definitely not
475 -- \ x -> case (_scc_ "foo" E) of I# x -> x)
477 -- This transform doesn't move work or allocation
478 -- from one cost centre to another
479 workerCase :: Id -> CoreExpr -> [Id] -> DataCon -> CoreExpr -> CoreExpr
480 workerCase bndr (Note (SCC cc) e) args con body = Note (SCC cc) (mkUnpackCase bndr e args con body)
481 workerCase bndr e args con body = mkUnpackCase bndr e args con body
485 %************************************************************************
487 \subsection{Utilities}
489 %************************************************************************
493 We make a new binding for Ids that are marked absent, thus
494 let x = absentError "x :: Int"
495 The idea is that this binding will never be used; but if it
496 buggily is used we'll get a runtime error message.
498 Coping with absence for *unlifted* types is important; see, for
499 example, Trac #4306. For these we find a suitable literal,
500 using Literal.absentLiteralOf. We don't have literals for
501 every primitive type, so the function is partial.
503 [I did try the experiment of using an error thunk for unlifted
504 things too, relying on the simplifier to drop it as dead code,
505 by making absentError
506 (a) *not* be a bottoming Id,
507 (b) be "ok for speculation"
508 But that relies on the simplifier finding that it really
509 is dead code, which is fragile, and indeed failed when
510 profiling is on, which disables various optimisations. So
511 using a literal will do.]
514 mk_absent_let :: Id -> Maybe (CoreExpr -> CoreExpr)
516 | not (isUnLiftedType arg_ty)
517 = Just (Let (NonRec arg abs_rhs))
518 | Just (tc, _) <- splitTyConApp_maybe arg_ty
519 , Just lit <- absentLiteralOf tc
520 = Just (Let (NonRec arg (Lit lit)))
521 | arg_ty `eqType` realWorldStatePrimTy
522 = Just (Let (NonRec arg (Var realWorldPrimId)))
524 = WARN( True, ptext (sLit "No absent value for") <+> ppr arg_ty )
528 abs_rhs = mkRuntimeErrorApp aBSENT_ERROR_ID arg_ty msg
529 msg = showSDocDebug (ppr arg <+> ppr (idType arg))
531 mk_seq_case :: Id -> CoreExpr -> CoreExpr
532 mk_seq_case arg body = Case (Var arg) (sanitiseCaseBndr arg) (exprType body) [(DEFAULT, [], body)]
534 sanitiseCaseBndr :: Id -> Id
535 -- The argument we are scrutinising has the right type to be
536 -- a case binder, so it's convenient to re-use it for that purpose.
537 -- But we *must* throw away all its IdInfo. In particular, the argument
538 -- will have demand info on it, and that demand info may be incorrect for
539 -- the case binder. e.g. case ww_arg of ww_arg { I# x -> ... }
540 -- Quite likely ww_arg isn't used in '...'. The case may get discarded
541 -- if the case binder says "I'm demanded". This happened in a situation
542 -- like (x+y) `seq` ....
543 sanitiseCaseBndr id = id `setIdInfo` vanillaIdInfo
545 mk_ww_local :: Unique -> Type -> Id
546 mk_ww_local uniq ty = mkSysLocal (fsLit "ww") uniq ty