[project @ 2000-03-23 17:45:17 by simonpj]
[ghc-hetmet.git] / ghc / compiler / stranal / WwLib.lhs
1 %
2 % (c) The GRASP/AQUA Project, Glasgow University, 1993-1998
3 %
4 \section[WwLib]{A library for the ``worker/wrapper'' back-end to the strictness analyser}
5
6 \begin{code}
7 module WwLib (
8         mkWwBodies,
9         worthSplitting, setUnpackStrategy
10     ) where
11
12 #include "HsVersions.h"
13
14 import CoreSyn
15 import CoreUtils        ( exprType, mkInlineMe )
16 import Id               ( Id, idType, mkSysLocal, idDemandInfo, setIdDemandInfo,
17                           isOneShotLambda, setOneShotLambda,
18                           mkWildId, setIdInfo
19                         )
20 import IdInfo           ( CprInfo(..), noCprInfo, vanillaIdInfo )
21 import DataCon          ( DataCon, splitProductType )
22 import Demand           ( Demand(..), wwLazy, wwPrim )
23 import PrelInfo         ( realWorldPrimId, aBSENT_ERROR_ID )
24 import TysPrim          ( realWorldStatePrimTy )
25 import TysWiredIn       ( unboxedTupleCon, unboxedTupleTyCon )
26 import Type             ( isUnLiftedType, 
27                           splitForAllTys, splitFunTys,  isAlgType,
28                           splitAlgTyConApp_maybe, splitNewType_maybe,
29                           mkTyConApp, mkFunTys,
30                           Type
31                         )
32 import TyCon            ( isNewTyCon, isProductTyCon, TyCon )
33 import BasicTypes       ( NewOrData(..), Arity )
34 import Var              ( TyVar, Var, isId )
35 import UniqSupply       ( returnUs, thenUs, getUniqueUs, getUniquesUs, 
36                           mapUs, UniqSM )
37 import Util             ( zipWithEqual, zipEqual, lengthExceeds )
38 import Outputable
39 import List             ( zipWith4 )
40 \end{code}
41
42
43 %************************************************************************
44 %*                                                                      *
45 \subsection[mkWrapperAndWorker]{@mkWrapperAndWorker@}
46 %*                                                                      *
47 %************************************************************************
48
49         ************   WARNING  ******************
50         these comments are rather out of date
51         *****************************************
52
53 @mkWrapperAndWorker@ is given:
54 \begin{enumerate}
55 \item
56 The {\em original function} \tr{f}, of the form:
57 \begin{verbatim}
58 f = /\ tyvars -> \ args -> body
59 \end{verbatim}
60 The original-binder \tr{f}, the \tr{tyvars}, \tr{args}, and \tr{body}
61 are given separately.
62
63 We use the Id \tr{f} mostly to get its type.
64
65 \item
66 Strictness information about \tr{f}, in the form of a list of
67 @Demands@.
68
69 \item
70 A @UniqueSupply@.
71 \end{enumerate}
72
73 @mkWrapperAndWorker@ produces (A BIT OUT-OF-DATE...):
74 \begin{enumerate}
75 \item
76 Maybe @Nothing@: no worker/wrappering going on in this case. This can
77 happen (a)~if the strictness info says that there is nothing
78 interesting to do or (b)~if *any* of the argument types corresponding
79 to ``active'' arg postitions is abstract or will be to the outside
80 world (i.e., {\em this} module can see the constructors, but nobody
81 else will be able to).  An ``active'' arg position is one which the
82 wrapper has to unpack.  An importing module can't do this unpacking,
83 so it simply has to give up and call the wrapper only.
84
85 \item
86 Maybe \tr{Just (wrapper_Id, wrapper_body, worker_Id, worker_body)}.
87
88 The @wrapper_Id@ is just the one that was passed in, with its
89 strictness IdInfo updated.
90 \end{enumerate}
91
92 The \tr{body} of the original function may not be given (i.e., it's
93 BOTTOM), in which case you'd jolly well better not tug on the
94 worker-body output!
95
96 Here's an example.  The original function is:
97 \begin{verbatim}
98 g :: forall a . Int -> [a] -> a
99
100 g = /\ a -> \ x ys ->
101         case x of
102           0 -> head ys
103           _ -> head (tail ys)
104 \end{verbatim}
105
106 From this, we want to produce:
107 \begin{verbatim}
108 -- wrapper (an unfolding)
109 g :: forall a . Int -> [a] -> a
110
111 g = /\ a -> \ x ys ->
112         case x of
113           I# x# -> g.wrk a x# ys
114             -- call the worker; don't forget the type args!
115
116 -- worker
117 g.wrk :: forall a . Int# -> [a] -> a
118
119 g.wrk = /\ a -> \ x# ys ->
120         let
121             x = I# x#
122         in
123             case x of               -- note: body of g moved intact
124               0 -> head ys
125               _ -> head (tail ys)
126 \end{verbatim}
127
128 Something we have to be careful about:  Here's an example:
129 \begin{verbatim}
130 -- "f" strictness: U(P)U(P)
131 f (I# a) (I# b) = a +# b
132
133 g = f   -- "g" strictness same as "f"
134 \end{verbatim}
135 \tr{f} will get a worker all nice and friendly-like; that's good.
136 {\em But we don't want a worker for \tr{g}}, even though it has the
137 same strictness as \tr{f}.  Doing so could break laziness, at best.
138
139 Consequently, we insist that the number of strictness-info items is
140 exactly the same as the number of lambda-bound arguments.  (This is
141 probably slightly paranoid, but OK in practice.)  If it isn't the
142 same, we ``revise'' the strictness info, so that we won't propagate
143 the unusable strictness-info into the interfaces.
144
145
146 %************************************************************************
147 %*                                                                      *
148 \subsection{Functions over Demands}
149 %*                                                                      *
150 %************************************************************************
151
152 \begin{code}
153 mAX_WORKER_ARGS :: Int          -- ToDo: set via flag
154 mAX_WORKER_ARGS = 6
155
156 setUnpackStrategy :: [Demand] -> [Demand]
157 setUnpackStrategy ds
158   = snd (go (mAX_WORKER_ARGS - nonAbsentArgs ds) ds)
159   where
160     go :: Int                   -- Max number of args available for sub-components of [Demand]
161        -> [Demand]
162        -> (Int, [Demand])       -- Args remaining after subcomponents of [Demand] are unpacked
163
164     go n (WwUnpack nd _ cs : ds) | n' >= 0
165                                  = WwUnpack nd True cs' `cons` go n'' ds
166                                  | otherwise
167                                  = WwUnpack nd False cs `cons` go n ds
168                                  where
169                                    n' = n + 1 - nonAbsentArgs cs
170                                         -- Add one because we don't pass the top-level arg any more
171                                         -- Delete # of non-absent args to which we'll now be committed
172                                    (n'',cs') = go n' cs
173                                 
174     go n (d:ds) = d `cons` go n ds
175     go n []     = (n,[])
176
177     cons d (n,ds) = (n, d:ds)
178
179 nonAbsentArgs :: [Demand] -> Int
180 nonAbsentArgs []                 = 0
181 nonAbsentArgs (WwLazy True : ds) = nonAbsentArgs ds
182 nonAbsentArgs (d           : ds) = 1 + nonAbsentArgs ds
183
184 worthSplitting :: [Demand]
185                -> Bool  -- Result is bottom
186                -> Bool  -- True <=> the wrapper would not be an identity function
187 worthSplitting ds result_bot = any worth_it ds
188         -- We used not to split if the result is bottom.
189         -- [Justification:  there's no efficiency to be gained.]
190         -- But it's sometimes bad not to make a wrapper.  Consider
191         --      fw = \x# -> let x = I# x# in case e of
192         --                                      p1 -> error_fn x
193         --                                      p2 -> error_fn x
194         --                                      p3 -> the real stuff
195         -- The re-boxing code won't go away unless error_fn gets a wrapper too.
196
197   where
198     worth_it (WwLazy True)       = True         -- Absent arg
199     worth_it (WwUnpack _ True _) = True         -- Arg to unpack
200     worth_it WwStrict            = False        -- Don't w/w just because of strictness
201     worth_it other               = False
202
203 allAbsent :: [Demand] -> Bool
204 allAbsent ds = all absent ds
205   where
206     absent (WwLazy is_absent)   = is_absent
207     absent (WwUnpack _ True cs) = allAbsent cs
208     absent other                = False
209 \end{code}
210
211
212 %************************************************************************
213 %*                                                                      *
214 \subsection{The worker wrapper core}
215 %*                                                                      *
216 %************************************************************************
217
218 @mkWwBodies@ is called when doing the worker/wrapper split inside a module.
219
220 \begin{code}
221 mkWwBodies :: Type                              -- Type of original function
222            -> Arity                             -- Arity of original function
223            -> [Demand]                          -- Strictness of original function
224            -> Bool                              -- True <=> function returns bottom
225            -> [Bool]                            -- One-shot-ness of the function
226            -> CprInfo                           -- Result of CPR analysis 
227            -> UniqSM ([Demand],                 -- Demands for worker (value) args
228                       Id -> CoreExpr,           -- Wrapper body, lacking only the worker Id
229                       CoreExpr -> CoreExpr)     -- Worker body, lacking the original function rhs
230
231 mkWwBodies fun_ty arity demands res_bot one_shots cpr_info
232   = mkWWargs fun_ty arity demands' res_bot one_shots'   `thenUs` \ (wrap_args, wrap_fn_args,   work_fn_args, res_ty) ->
233     mkWWstr wrap_args                                   `thenUs` \ (work_dmds, wrap_fn_str,    work_fn_str) ->
234     mkWWcpr res_ty cpr_info                             `thenUs` \ (wrap_fn_cpr,    work_fn_cpr,  cpr_res_ty) ->
235     mkWWfixup cpr_res_ty work_dmds                      `thenUs` \ (final_work_dmds, wrap_fn_fixup,  work_fn_fixup) ->
236
237     returnUs (final_work_dmds,
238               mkInlineMe . wrap_fn_args . wrap_fn_cpr . wrap_fn_str . wrap_fn_fixup . Var,
239               work_fn_fixup . work_fn_str . work_fn_cpr . work_fn_args)
240   where
241     demands'   = demands   ++ repeat wwLazy
242     one_shots' = one_shots ++ repeat False
243 \end{code}
244
245
246 %************************************************************************
247 %*                                                                      *
248 \subsection{Coercion stuff}
249 %*                                                                      *
250 %************************************************************************
251
252
253 We really want to "look through" coerces.
254 Reason: I've seen this situation:
255
256         let f = coerce T (\s -> E)
257         in \x -> case x of
258                     p -> coerce T' f
259                     q -> \s -> E2
260                     r -> coerce T' f
261
262 If only we w/w'd f, we'd get
263         let f = coerce T (\s -> fw s)
264             fw = \s -> E
265         in ...
266
267 Now we'll inline f to get
268
269         let fw = \s -> E
270         in \x -> case x of
271                     p -> fw
272                     q -> \s -> E2
273                     r -> fw
274
275 Now we'll see that fw has arity 1, and will arity expand
276 the \x to get what we want.
277
278 \begin{code}
279 -- mkWWargs is driven off the function type and arity.
280 -- It chomps bites off foralls, arrows, newtypes
281 -- and keeps repeating that until it's satisfied the supplied arity
282
283 mkWWargs :: Type -> Arity 
284          -> [Demand] -> Bool -> [Bool]          -- Both these will in due course be derived
285                                                 -- from the type.  The [Bool] is True for a one-shot arg.
286                                                 -- ** Both are infinite, extended with neutral values if necy **
287          -> UniqSM  ([Var],             -- Wrapper args
288                      CoreExpr -> CoreExpr,      -- Wrapper fn
289                      CoreExpr -> CoreExpr,      -- Worker fn
290                      Type)                      -- Type of wrapper body
291
292 mkWWargs fun_ty arity demands res_bot one_shots
293   | (res_bot || arity > 0) && (not (null tyvars) || n_arg_tys > 0)
294         -- If the function returns bottom, we feel free to 
295         -- build lots of wrapper args:
296         --        \x. let v=E in \y. bottom
297         --      = \xy. let v=E in bottom
298   = getUniquesUs n_args         `thenUs` \ wrap_uniqs ->
299     let
300       val_args  = zipWith4 mk_wrap_arg wrap_uniqs arg_tys demands one_shots
301       wrap_args = tyvars ++ val_args
302     in
303     mkWWargs new_fun_ty
304              (arity - n_args) 
305              (drop n_args demands)
306              res_bot
307              (drop n_args one_shots)    `thenUs` \ (more_wrap_args, wrap_fn_args, work_fn_args, res_ty) ->
308
309     returnUs (wrap_args ++ more_wrap_args,
310               mkLams wrap_args . wrap_fn_args,
311               work_fn_args . applyToVars wrap_args,
312               res_ty)
313   where
314     (tyvars, tau)       = splitForAllTys fun_ty
315     (arg_tys, body_ty)  = splitFunTys tau
316     n_arg_tys           = length arg_tys
317     n_args              | res_bot   = n_arg_tys 
318                         | otherwise = arity `min` n_arg_tys
319     new_fun_ty          | n_args == n_arg_tys = body_ty
320                         | otherwise           = mkFunTys (drop n_args arg_tys) body_ty
321
322 mkWWargs fun_ty arity demands res_bot one_shots
323   = case splitNewType_maybe fun_ty of
324         Nothing     -> returnUs ([], id, id, fun_ty)
325         Just rep_ty -> mkWWargs rep_ty arity demands res_bot one_shots  `thenUs` \ (wrap_args, wrap_fn_args, work_fn_args, res_ty) ->
326                        returnUs (wrap_args,
327                                  Note (Coerce fun_ty rep_ty) . wrap_fn_args,
328                                  work_fn_args . Note (Coerce rep_ty fun_ty),
329                                  res_ty)
330
331
332 applyToVars :: [Var] -> CoreExpr -> CoreExpr
333 applyToVars vars fn = mkVarApps fn vars
334
335 mk_wrap_arg uniq ty dmd one_shot 
336   = set_one_shot one_shot (setIdDemandInfo (mkSysLocal SLIT("w") uniq ty) dmd)
337   where
338     set_one_shot True  id = setOneShotLambda id
339     set_one_shot False id = id
340 \end{code}
341
342
343 %************************************************************************
344 %*                                                                      *
345 \subsection{Fixup stuff}
346 %*                                                                      *
347 %************************************************************************
348
349 \begin{code}
350 mkWWfixup res_ty work_dmds
351   | null work_dmds && isUnLiftedType res_ty 
352         -- Horrid special case.  If the worker would have no arguments, and the
353         -- function returns a primitive type value, that would make the worker into
354         -- an unboxed value.  We box it by passing a dummy void argument, thus:
355         --
356         --      f = /\abc. \xyz. fw abc void
357         --      fw = /\abc. \v. body
358         --
359         -- We use the state-token type which generates no code
360   = getUniqueUs                 `thenUs` \ void_arg_uniq ->
361     let
362             void_arg = mk_ww_local void_arg_uniq realWorldStatePrimTy
363     in
364     returnUs ([wwPrim],         
365               \ call_to_worker -> App call_to_worker (Var realWorldPrimId),
366               \ worker_body    -> Lam void_arg worker_body)
367
368   | otherwise
369   = returnUs (work_dmds, id, id)
370 \end{code}
371
372
373 %************************************************************************
374 %*                                                                      *
375 \subsection{Strictness stuff}
376 %*                                                                      *
377 %************************************************************************
378
379 \begin{code}
380 mkWWstr :: [Var]                                -- Wrapper args; have their demand info on them
381                                                 -- *Includes type variables*
382         -> UniqSM ([Demand],                    -- Demand on worker (value) args
383                    CoreExpr -> CoreExpr,        -- Wrapper body, lacking the worker call
384                                                 -- and without its lambdas 
385                                                 -- This fn adds the unboxing, and makes the
386                                                 -- call passing the unboxed things
387                                 
388                    CoreExpr -> CoreExpr)        -- Worker body, lacking the original body of the function,
389                                                 -- but *with* lambdas
390
391 mkWWstr wrap_args
392   = mk_ww_str wrap_args         `thenUs` \ (work_args, wrap_fn, work_fn) ->
393     returnUs ( [idDemandInfo v | v <- work_args, isId v],
394                \ wrapper_body -> wrap_fn (mkVarApps wrapper_body work_args),
395                \ worker_body  -> mkLams work_args (work_fn worker_body))
396
397         -- Empty case
398 mk_ww_str []
399   = returnUs ([],
400               \ wrapper_body -> wrapper_body,
401               \ worker_body  -> worker_body)
402
403
404 mk_ww_str (arg : ds)
405   | isTyVar arg
406   = mk_ww_str ds                `thenUs` \ (worker_args, wrap_fn, work_fn) ->
407     returnUs (arg : worker_args, wrap_fn, work_fn)
408
409   | otherwise
410   = case idDemandInfo arg of
411
412         -- Absent case
413       WwLazy True ->
414         mk_ww_str ds            `thenUs` \ (worker_args, wrap_fn, work_fn) ->
415         returnUs (worker_args, wrap_fn, mk_absent_let arg . work_fn)
416
417         -- Unpack case
418       WwUnpack new_or_data True cs ->
419         getUniquesUs (length inst_con_arg_tys)          `thenUs` \ uniqs ->
420         let
421           unpk_args      = zipWith mk_ww_local uniqs inst_con_arg_tys
422           unpk_args_w_ds = zipWithEqual "mk_ww_str" set_worker_arg_info unpk_args cs
423         in
424         mk_ww_str (unpk_args_w_ds ++ ds)                `thenUs` \ (worker_args, wrap_fn, work_fn) ->
425         returnUs (worker_args,
426                   mk_unpk_case new_or_data arg unpk_args data_con arg_tycon . wrap_fn,
427                   work_fn . mk_pk_let new_or_data arg data_con tycon_arg_tys unpk_args)
428         where
429           (arg_tycon, tycon_arg_tys, data_con, inst_con_arg_tys) = splitProductType "mk_ww_str" (idType arg)
430
431         -- Other cases
432       other_demand ->
433         mk_ww_str ds            `thenUs` \ (worker_args, wrap_fn, work_fn) ->
434         returnUs (arg : worker_args, wrap_fn, work_fn)
435   where
436         -- If the wrapper argument is a one-shot lambda, then
437         -- so should (all) the corresponding worker arguments be
438         -- This bites when we do w/w on a case join point
439     set_worker_arg_info worker_arg demand = set_one_shot (setIdDemandInfo worker_arg demand)
440
441     set_one_shot | isOneShotLambda arg = setOneShotLambda
442                  | otherwise           = \x -> x
443 \end{code}
444
445
446 %************************************************************************
447 %*                                                                      *
448 \subsection{CPR stuff}
449 %*                                                                      *
450 %************************************************************************
451
452
453 @mkWWcpr@ takes the worker/wrapper pair produced from the strictness
454 info and adds in the CPR transformation.  The worker returns an
455 unboxed tuple containing non-CPR components.  The wrapper takes this
456 tuple and re-produces the correct structured output.
457
458 The non-CPR results appear ordered in the unboxed tuple as if by a
459 left-to-right traversal of the result structure.
460
461
462 \begin{code}
463 mkWWcpr :: Type                              -- function body type
464         -> CprInfo                           -- CPR analysis results
465         -> UniqSM (CoreExpr -> CoreExpr,             -- New wrapper 
466                    CoreExpr -> CoreExpr,             -- New worker
467                    Type)                        -- Type of worker's body 
468
469 mkWWcpr body_ty NoCPRInfo 
470     = returnUs (id, id, body_ty)      -- Must be just the strictness transf.
471
472 mkWWcpr body_ty ReturnsCPR
473     | not (isAlgType body_ty)
474     = WARN( True, text "mkWWcpr: non-algebraic body type" <+> ppr body_ty )
475       returnUs (id, id, body_ty)
476
477     | n_con_args == 1 && isUnLiftedType con_arg_ty1
478         -- Special case when there is a single result of unlifted type
479     = getUniquesUs 2                    `thenUs` \ [work_uniq, arg_uniq] ->
480       let
481         work_wild = mk_ww_local work_uniq body_ty
482         arg       = mk_ww_local arg_uniq  con_arg_ty1
483       in
484       returnUs (\ wkr_call -> Case wkr_call arg       [(DEFAULT, [], mkConApp data_con (map Type tycon_arg_tys ++ [Var arg]))],
485                 \ body     -> Case body     work_wild [(DataAlt data_con, [arg], Var arg)],
486                 con_arg_ty1)
487
488     | otherwise         -- The general case
489     = getUniquesUs (n_con_args + 2)     `thenUs` \ uniqs ->
490       let
491         (wrap_wild : work_wild : args) = zipWith mk_ww_local uniqs (ubx_tup_ty : body_ty : con_arg_tys)
492         arg_vars                       = map Var args
493         ubx_tup_con                    = unboxedTupleCon n_con_args
494         ubx_tup_ty                     = exprType ubx_tup_app
495         ubx_tup_app                    = mkConApp ubx_tup_con (map Type con_arg_tys   ++ arg_vars)
496         con_app                        = mkConApp data_con    (map Type tycon_arg_tys ++ arg_vars)
497       in
498       returnUs (\ wkr_call -> Case wkr_call wrap_wild [(DataAlt ubx_tup_con, args, con_app)],
499                 \ body     -> Case body     work_wild [(DataAlt data_con,    args, ubx_tup_app)],
500                 ubx_tup_ty)
501     where
502       (tycon, tycon_arg_tys, data_con, con_arg_tys) = splitProductType "mkWWcpr" body_ty
503       n_con_args  = length con_arg_tys
504       con_arg_ty1 = head con_arg_tys
505 \end{code}
506
507
508 %************************************************************************
509 %*                                                                      *
510 \subsection{Utilities}
511 %*                                                                      *
512 %************************************************************************
513
514
515 \begin{code}
516 mk_absent_let arg body
517   | not (isUnLiftedType arg_ty)
518   = Let (NonRec arg (mkTyApps (Var aBSENT_ERROR_ID) [arg_ty])) body
519   | otherwise
520   = panic "WwLib: haven't done mk_absent_let for primitives yet"
521   where
522     arg_ty = idType arg
523
524 mk_unpk_case NewType arg unpk_args boxing_con boxing_tycon body
525         -- A newtype!  Use a coercion not a case
526   = ASSERT( null other_args )
527     Case (Note (Coerce (idType unpk_arg) (idType arg)) (Var arg))
528          (sanitiseCaseBndr unpk_arg)
529          [(DEFAULT,[],body)]
530   where
531     (unpk_arg:other_args) = unpk_args
532
533 mk_unpk_case DataType arg unpk_args boxing_con boxing_tycon body
534         -- A data type
535   = Case (Var arg) 
536          (sanitiseCaseBndr arg)
537          [(DataAlt boxing_con, unpk_args, body)]
538
539 sanitiseCaseBndr :: Id -> Id
540 -- The argument we are scrutinising has the right type to be
541 -- a case binder, so it's convenient to re-use it for that purpose.
542 -- But we *must* throw away all its IdInfo.  In particular, the argument
543 -- will have demand info on it, and that demand info may be incorrect for
544 -- the case binder.  e.g.       case ww_arg of ww_arg { I# x -> ... }
545 -- Quite likely ww_arg isn't used in '...'.  The case may get discarded
546 -- if the case binder says "I'm demanded".  This happened in a situation 
547 -- like         (x+y) `seq` ....
548 sanitiseCaseBndr id = id `setIdInfo` vanillaIdInfo
549
550 mk_pk_let NewType arg boxing_con con_tys unpk_args body
551   = ASSERT( null other_args )
552     Let (NonRec arg (Note (Coerce (idType arg) (idType unpk_arg)) (Var unpk_arg))) body
553   where
554     (unpk_arg:other_args) = unpk_args
555
556 mk_pk_let DataType arg boxing_con con_tys unpk_args body
557   = Let (NonRec arg (mkConApp boxing_con con_args)) body
558   where
559     con_args = map Type con_tys ++ map Var unpk_args
560
561
562 mk_ww_local uniq ty = mkSysLocal SLIT("ww") uniq ty
563
564 \end{code}