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