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