WwLib.lhs 17 KB
Newer Older
1
%
2
% (c) The GRASP/AQUA Project, Glasgow University, 1993-1998
3
4
5
6
%
\section[WwLib]{A library for the ``worker/wrapper'' back-end to the strictness analyser}

\begin{code}
7
module WwLib ( mkWwBodies, mkWWstr, mkWorkerArgs ) where
8

9
#include "HsVersions.h"
10

11
import CoreSyn
12
import CoreUtils	( exprType )
13
import Id		( Id, idType, mkSysLocal, idNewDemandInfo, setIdNewDemandInfo,
14
			  isOneShotLambda, setOneShotLambda, setIdUnfolding,
15
                          setIdInfo
16
			)
17
import IdInfo		( vanillaIdInfo )
18
import DataCon		( deepSplitProductType_maybe, deepSplitProductType )
19
import NewDemand	( Demand(..), DmdResult(..), Demands(..) ) 
20
21
import MkId		( realWorldPrimId, voidArgId, mkRuntimeErrorApp, rUNTIME_ERROR_ID,
                          mkUnpackCase, mkProductBox )
22
import TysWiredIn	( tupleCon )
23
import Type		( Type, isUnLiftedType, mkFunTys,
24
			  splitForAllTys, splitFunTys, isAlgType
25
			)
26
import Coercion         ( mkSymCoercion, splitNewTypeRepCo_maybe )
27
import BasicTypes	( Boxity(..) )
28
import Var              ( Var, isId )
29
import UniqSupply	( returnUs, thenUs, getUniquesUs, UniqSM )
sof's avatar
sof committed
30
import Util		( zipWithEqual, notNull )
sof's avatar
sof committed
31
import Outputable
32
import List		( zipWith4 )
33
34
35
36
37
38
39
40
41
42
\end{code}


%************************************************************************
%*									*
\subsection[mkWrapperAndWorker]{@mkWrapperAndWorker@}
%*									*
%************************************************************************

Here's an example.  The original function is:
43

44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
\begin{verbatim}
g :: forall a . Int -> [a] -> a

g = /\ a -> \ x ys ->
	case x of
	  0 -> head ys
	  _ -> head (tail ys)
\end{verbatim}

From this, we want to produce:
\begin{verbatim}
-- wrapper (an unfolding)
g :: forall a . Int -> [a] -> a

g = /\ a -> \ x ys ->
	case x of
60
	  I# x# -> $wg a x# ys
61
62
63
	    -- call the worker; don't forget the type args!

-- worker
64
$wg :: forall a . Int# -> [a] -> a
65

66
$wg = /\ a -> \ x# ys ->
67
68
69
70
71
72
73
74
75
	let
	    x = I# x#
	in
	    case x of		    -- note: body of g moved intact
	      0 -> head ys
	      _ -> head (tail ys)
\end{verbatim}

Something we have to be careful about:  Here's an example:
76

77
78
79
80
81
82
\begin{verbatim}
-- "f" strictness: U(P)U(P)
f (I# a) (I# b) = a +# b

g = f	-- "g" strictness same as "f"
\end{verbatim}
83

84
85
86
87
88
89
90
91
92
93
94
\tr{f} will get a worker all nice and friendly-like; that's good.
{\em But we don't want a worker for \tr{g}}, even though it has the
same strictness as \tr{f}.  Doing so could break laziness, at best.

Consequently, we insist that the number of strictness-info items is
exactly the same as the number of lambda-bound arguments.  (This is
probably slightly paranoid, but OK in practice.)  If it isn't the
same, we ``revise'' the strictness info, so that we won't propagate
the unusable strictness-info into the interfaces.


95
96
97
98
99
%************************************************************************
%*									*
\subsection{The worker wrapper core}
%*									*
%************************************************************************
100

101
@mkWwBodies@ is called when doing the worker/wrapper split inside a module.
102
103

\begin{code}
104
105
mkWwBodies :: Type				-- Type of original function
	   -> [Demand]				-- Strictness of original function
106
	   -> DmdResult				-- Info about function result
107
	   -> [Bool]				-- One-shot-ness of the function
108
	   -> UniqSM ([Demand],			-- Demands for worker (value) args
109
110
111
		      Id -> CoreExpr,		-- Wrapper body, lacking only the worker Id
		      CoreExpr -> CoreExpr)	-- Worker body, lacking the original function rhs

112
113
114
115
116
117
118
119
120
121
122
-- wrap_fn_args E	= \x y -> E
-- work_fn_args E	= E x y

-- wrap_fn_str E 	= case x of { (a,b) -> 
--			  case a of { (a1,a2) ->
--			  E a1 a2 b y }}
-- work_fn_str E	= \a2 a2 b y ->
--			  let a = (a1,a2) in
--			  let x = (a,b) in
--			  E

123
124
mkWwBodies fun_ty demands res_info one_shots
  = mkWWargs fun_ty demands one_shots'	`thenUs` \ (wrap_args,   wrap_fn_args, work_fn_args, res_ty) ->
125
    mkWWstr wrap_args			`thenUs` \ (work_args,   wrap_fn_str,  work_fn_str) ->
126
    let
127
	(work_lam_args, work_call_args) = mkWorkerArgs work_args res_ty
128
    in
129
130
131
132
133
134
	-- Don't do CPR if the worker doesn't have any value arguments
	-- Then the worker is just a constant, so we don't want to unbox it.
    (if any isId work_args then
	mkWWcpr res_ty res_info
     else
	returnUs (id, id, res_ty)
135
    )					`thenUs` \ (wrap_fn_cpr, work_fn_cpr,  _cpr_res_ty) ->
136

137
    returnUs ([idNewDemandInfo v | v <- work_call_args, isId v],
138
	      Note InlineMe . wrap_fn_args . wrap_fn_cpr . wrap_fn_str . applyToVars work_call_args . Var,
139
	      mkLams work_lam_args. work_fn_str . work_fn_cpr . work_fn_args)
140
141
142
143
144
145
146
	-- We use an INLINE unconditionally, even if the wrapper turns out to be
	-- something trivial like
	--	fw = ...
	--	f = __inline__ (coerce T fw)
	-- The point is to propagate the coerce to f's call sites, so even though
	-- f's RHS is now trivial (size 1) we still want the __inline__ to prevent
	-- fw from being inlined into f's RHS
147
148
  where
    one_shots' = one_shots ++ repeat False
149
\end{code}
150

151
152
153
154
155
156
157
158
159
160
161

%************************************************************************
%*									*
\subsection{Making wrapper args}
%*									*
%************************************************************************

During worker-wrapper stuff we may end up with an unlifted thing
which we want to let-bind without losing laziness.  So we
add a void argument.  E.g.

162
	f = /\a -> \x y z -> E::Int#	-- E does not mention x,y,z
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
==>
	fw = /\ a -> \void -> E
	f  = /\ a -> \x y z -> fw realworld

We use the state-token type which generates no code.

\begin{code}
mkWorkerArgs :: [Var]
	     -> Type	-- Type of body
	     -> ([Var],	-- Lambda bound args
		 [Var])	-- Args at call site
mkWorkerArgs args res_ty
    | any isId args || not (isUnLiftedType res_ty)
    = (args, args)
    | otherwise	
    = (args ++ [voidArgId], args ++ [realWorldPrimId])
179
180
181
182
183
184
185
186
187
188
\end{code}


%************************************************************************
%*									*
\subsection{Coercion stuff}
%*									*
%************************************************************************


189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
We really want to "look through" coerces.
Reason: I've seen this situation:

	let f = coerce T (\s -> E)
	in \x -> case x of
	   	    p -> coerce T' f
		    q -> \s -> E2
	   	    r -> coerce T' f

If only we w/w'd f, we'd get
	let f = coerce T (\s -> fw s)
	    fw = \s -> E
	in ...

Now we'll inline f to get

	let fw = \s -> E
	in \x -> case x of
	   	    p -> fw
		    q -> \s -> E2
	   	    r -> fw

Now we'll see that fw has arity 1, and will arity expand
the \x to get what we want.
213
214

\begin{code}
215
-- mkWWargs is driven off the function type and arity.
216
217
-- It chomps bites off foralls, arrows, newtypes
-- and keeps repeating that until it's satisfied the supplied arity
218

219
220
221
mkWWargs :: Type
	 -> [Demand]
	 -> [Bool]			-- True for a one-shot arg; ** may be infinite **
222
	 -> UniqSM  ([Var],		-- Wrapper args
223
224
225
		     CoreExpr -> CoreExpr,	-- Wrapper fn
		     CoreExpr -> CoreExpr,	-- Worker fn
		     Type)			-- Type of wrapper body
226

227
mkWWargs fun_ty demands one_shots
228
  | Just (rep_ty, co) <- splitNewTypeRepCo_maybe fun_ty
229
230
231
232
233
234
235
236
237
238
   	-- The newtype case is for when the function has
	-- a recursive newtype after the arrow (rare)
	-- We check for arity >= 0 to avoid looping in the case
	-- of a function whose type is, in effect, infinite
	-- [Arity is driven by looking at the term, not just the type.]
	--
	-- It's also important when we have a function returning (say) a pair
	-- wrapped in a recursive newtype, at least if CPR analysis can look 
	-- through such newtypes, which it probably can since they are 
	-- simply coerces.
239
  = mkWWargs rep_ty demands one_shots	`thenUs` \ (wrap_args, wrap_fn_args, work_fn_args, res_ty) ->
240
    returnUs (wrap_args,
241
242
	      \ e -> Cast (wrap_fn_args e) (mkSymCoercion co),
	      \ e -> work_fn_args (Cast e co),
243
	      res_ty)
sof's avatar
sof committed
244
  | notNull demands
245
246
  = getUniquesUs 		`thenUs` \ wrap_uniqs ->
    let
247
248
      (tyvars, tau)      = splitForAllTys fun_ty
      (arg_tys, body_ty) = splitFunTys tau
249
250
251
252
253
254
255
256
257
258
259
260

      n_demands	= length demands
      n_arg_tys	= length arg_tys
      n_args    = n_demands `min` n_arg_tys

      new_fun_ty    = mkFunTys (drop n_demands arg_tys) body_ty
      new_demands   = drop n_arg_tys demands
      new_one_shots = drop n_args one_shots

      val_args	= zipWith4 mk_wrap_arg wrap_uniqs arg_tys demands one_shots
      wrap_args = tyvars ++ val_args
    in
sof's avatar
sof committed
261
{-     ASSERT( notNull tyvars || notNull arg_tys ) -}
262
263
264
265
266
    if (null tyvars) && (null arg_tys) then
	pprTrace "mkWWargs" (ppr fun_ty $$ ppr demands) 
		returnUs ([], id, id, fun_ty)
	else

267
268
269
270
271
272
273
274
275
    mkWWargs new_fun_ty
	     new_demands
	     new_one_shots	`thenUs` \ (more_wrap_args, wrap_fn_args, work_fn_args, res_ty) ->

    returnUs (wrap_args ++ more_wrap_args,
	      mkLams wrap_args . wrap_fn_args,
	      work_fn_args . applyToVars wrap_args,
	      res_ty)

276
277
278
  | otherwise
  = returnUs ([], id, id, fun_ty)

279
280

applyToVars :: [Var] -> CoreExpr -> CoreExpr
281
282
applyToVars vars fn = mkVarApps fn vars

283
mk_wrap_arg uniq ty dmd one_shot 
284
  = set_one_shot one_shot (setIdNewDemandInfo (mkSysLocal FSLIT("w") uniq ty) dmd)
285
286
287
  where
    set_one_shot True  id = setOneShotLambda id
    set_one_shot False id = id
288
\end{code}
289

290

291
292
293
294
295
296
297
%************************************************************************
%*									*
\subsection{Strictness stuff}
%*									*
%************************************************************************

\begin{code}
298
mkWWstr :: [Var]				-- Wrapper args; have their demand info on them
299
						--  *Includes type variables*
300
        -> UniqSM ([Var],			-- Worker args
301
		   CoreExpr -> CoreExpr,	-- Wrapper body, lacking the worker call
302
						-- and without its lambdas 
303
						-- This fn adds the unboxing
304
				
305
		   CoreExpr -> CoreExpr)	-- Worker body, lacking the original body of the function,
306
307
						-- and lacking its lambdas.
						-- This fn does the reboxing
308

309
310
----------------------
nop_fn body = body
311

312
----------------------
313
mkWWstr []
314
  = returnUs ([], nop_fn, nop_fn)
315

316
317
318
mkWWstr (arg : args)
  = mkWWstr_one arg		`thenUs` \ (args1, wrap_fn1, work_fn1) ->
    mkWWstr args		`thenUs` \ (args2, wrap_fn2, work_fn2) ->
319
320
321
322
    returnUs (args1 ++ args2, wrap_fn1 . wrap_fn2, work_fn1 . work_fn2)


----------------------
323
324
325
326
327
328
-- mkWWstr_one wrap_arg = (work_args, wrap_fn, work_fn)
--   *  wrap_fn assumes wrap_arg is in scope,
--	  brings into scope work_args (via cases)
--   * work_fn assumes work_args are in scope, a
--	  brings into scope wrap_arg (via lets)

329
mkWWstr_one arg
330
  | isTyVar arg
331
  = returnUs ([arg],  nop_fn, nop_fn)
332
333

  | otherwise
334
  = case idNewDemandInfo arg of
335

336
337
338
339
	-- Absent case.  We don't deal with absence for unlifted types,
	-- though, because it's not so easy to manufacture a placeholder
	-- We'll see if this turns out to be a problem
      Abs | not (isUnLiftedType (idType arg)) ->
340
	returnUs ([], nop_fn, mk_absent_let arg) 
341

342
343
	-- Unpack case
      Eval (Prod cs)
344
	| Just (_arg_tycon, _tycon_arg_tys, data_con, inst_con_arg_tys) 
345
		<- deepSplitProductType_maybe (idType arg)
346
347
348
349
	-> getUniquesUs 		`thenUs` \ uniqs ->
	   let
	     unpk_args	    = zipWith mk_ww_local uniqs inst_con_arg_tys
	     unpk_args_w_ds = zipWithEqual "mkWWstr" set_worker_arg_info unpk_args cs
350
	     unbox_fn       = mkUnpackCase (sanitiseCaseBndr arg) (Var arg) unpk_args data_con
351
	     rebox_fn	    = Let (NonRec arg con_app) 
352
	     con_app	    = mkProductBox unpk_args (idType arg)
353
354
	   in
	   mkWWstr unpk_args_w_ds		`thenUs` \ (worker_args, wrap_fn, work_fn) ->
355
	   returnUs (worker_args, unbox_fn . wrap_fn, work_fn . rebox_fn) 
356
357
358
359
360
	  		   -- Don't pass the arg, rebox instead

	-- `seq` demand; evaluate in wrapper in the hope
	-- of dropping seqs in the worker
      Eval (Poly Abs)
361
	-> let
362
		arg_w_unf = arg `setIdUnfolding` evaldUnfolding
363
364
365
		-- Tell the worker arg that it's sure to be evaluated
		-- so that internal seqs can be dropped
	   in
366
	   returnUs ([arg_w_unf], mk_seq_case arg, nop_fn)
367
368
369
	  	-- Pass the arg, anyway, even if it is in theory discarded
		-- Consider
		--	f x y = x `seq` y
370
		-- x gets a (Eval (Poly Abs)) demand, but if we fail to pass it to the worker
371
372
373
374
375
376
		-- we ABSOLUTELY MUST record that x is evaluated in the wrapper.
		-- Something like:
		--	f x y = x `seq` fw y
		--	fw y = let x{Evald} = error "oops" in (x `seq` y)
		-- If we don't pin on the "Evald" flag, the seq doesn't disappear, and
		-- we end up evaluating the absent thunk.
377
		-- But the Evald flag is pretty weird, and I worry that it might disappear
378
379
		-- during simplification, so for now I've just nuked this whole case
			
380
	-- Other cases
381
382
      other_demand -> returnUs ([arg], nop_fn, nop_fn)

383
384
385
386
  where
	-- If the wrapper argument is a one-shot lambda, then
	-- so should (all) the corresponding worker arguments be
	-- This bites when we do w/w on a case join point
387
    set_worker_arg_info worker_arg demand = set_one_shot (setIdNewDemandInfo worker_arg demand)
388
389
390

    set_one_shot | isOneShotLambda arg = setOneShotLambda
		 | otherwise	       = \x -> x
391
\end{code}
392

393
394
395
396
397
398
399
400

%************************************************************************
%*									*
\subsection{CPR stuff}
%*									*
%************************************************************************


401
402
403
404
405
406
407
408
409
410
@mkWWcpr@ takes the worker/wrapper pair produced from the strictness
info and adds in the CPR transformation.  The worker returns an
unboxed tuple containing non-CPR components.  The wrapper takes this
tuple and re-produces the correct structured output.

The non-CPR results appear ordered in the unboxed tuple as if by a
left-to-right traversal of the result structure.


\begin{code}
411
mkWWcpr :: Type                              -- function body type
412
        -> DmdResult                         -- CPR analysis results
413
        -> UniqSM (CoreExpr -> CoreExpr,             -- New wrapper 
414
415
                   CoreExpr -> CoreExpr,	     -- New worker
		   Type)			-- Type of worker's body 
416

417
mkWWcpr body_ty RetCPR
418
419
420
421
    | not (isAlgType body_ty)
    = WARN( True, text "mkWWcpr: non-algebraic body type" <+> ppr body_ty )
      returnUs (id, id, body_ty)

422
423
    | n_con_args == 1 && isUnLiftedType con_arg_ty1
	-- Special case when there is a single result of unlifted type
424
425
426
	--
	-- Wrapper:	case (..call worker..) of x -> C x
	-- Worker:	case (   ..body..    ) of C x -> x
427
    = getUniquesUs 			`thenUs` \ (work_uniq : arg_uniq : _) ->
428
      let
429
430
	work_wild = mk_ww_local work_uniq body_ty
	arg	  = mk_ww_local arg_uniq  con_arg_ty1
431
	con_app   = mkProductBox [arg] body_ty
432
      in
433
      returnUs (\ wkr_call -> Case wkr_call (arg) (exprType con_app) [(DEFAULT, [], con_app)],
434
		\ body     -> workerCase (work_wild) body [arg] data_con (Var arg),
435
		con_arg_ty1)
436

437
    | otherwise		-- The general case
438
439
	-- Wrapper: case (..call worker..) of (# a, b #) -> C a b
	-- Worker:  case (   ...body...  ) of C a b -> (# a, b #)     
440
    = getUniquesUs 	  	`thenUs` \ uniqs ->
441
      let
442
443
        (wrap_wild : work_wild : args) = zipWith mk_ww_local uniqs (ubx_tup_ty : body_ty : con_arg_tys)
	arg_vars		       = map Var args
444
	ubx_tup_con		       = tupleCon Unboxed n_con_args
445
	ubx_tup_ty		       = exprType ubx_tup_app
446
	ubx_tup_app		       = mkConApp ubx_tup_con (map Type con_arg_tys   ++ arg_vars)
447
        con_app			       = mkProductBox args body_ty
448
      in
449
      returnUs (\ wkr_call -> Case wkr_call (wrap_wild) (exprType con_app)  [(DataAlt ubx_tup_con, args, con_app)],
450
		\ body     -> workerCase (work_wild) body args data_con ubx_tup_app,
451
		ubx_tup_ty)
452
    where
453
      (_arg_tycon, _tycon_arg_tys, data_con, con_arg_tys) = deepSplitProductType "mkWWcpr" body_ty
454
455
      n_con_args  = length con_arg_tys
      con_arg_ty1 = head con_arg_tys
456

457
458
459
mkWWcpr body_ty other		-- No CPR info
    = returnUs (id, id, body_ty)

460
461
462
463
464
465
466
467
468
469
470
-- If the original function looked like
--	f = \ x -> _scc_ "foo" E
--
-- then we want the CPR'd worker to look like
--	\ x -> _scc_ "foo" (case E of I# x -> x)
-- and definitely not
--	\ x -> case (_scc_ "foo" E) of I# x -> x)
--
-- This transform doesn't move work or allocation
-- from one cost centre to another

471
472
workerCase bndr (Note (SCC cc) e) args con body = Note (SCC cc) (mkUnpackCase bndr e args con body)
workerCase bndr e args con body = mkUnpackCase bndr e args con body
473
\end{code}
474

475

476
477
478
479
480
481
482
483
484
%************************************************************************
%*									*
\subsection{Utilities}
%*									*
%************************************************************************


\begin{code}
mk_absent_let arg body
485
  | not (isUnLiftedType arg_ty)
486
  = Let (NonRec arg abs_rhs) body
487
488
489
490
  | otherwise
  = panic "WwLib: haven't done mk_absent_let for primitives yet"
  where
    arg_ty = idType arg
491
    abs_rhs = mkRuntimeErrorApp rUNTIME_ERROR_ID arg_ty msg
492
    msg     = "Oops!  Entered absent arg " ++ showSDocDebug (ppr arg <+> ppr (idType arg))
493

494
mk_seq_case arg body = Case (Var arg) (sanitiseCaseBndr arg) (exprType body) [(DEFAULT, [], body)]
495

496
497
498
499
500
501
502
503
504
505
sanitiseCaseBndr :: Id -> Id
-- The argument we are scrutinising has the right type to be
-- a case binder, so it's convenient to re-use it for that purpose.
-- But we *must* throw away all its IdInfo.  In particular, the argument
-- will have demand info on it, and that demand info may be incorrect for
-- the case binder.  e.g.  	case ww_arg of ww_arg { I# x -> ... }
-- Quite likely ww_arg isn't used in '...'.  The case may get discarded
-- if the case binder says "I'm demanded".  This happened in a situation 
-- like		(x+y) `seq` ....
sanitiseCaseBndr id = id `setIdInfo` vanillaIdInfo
506

507
mk_ww_local uniq ty = mkSysLocal FSLIT("ww") uniq ty
508
\end{code}