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