Rejig the absent-arg stuff for unlifted types
[ghc-hetmet.git] / 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 ( mkWwBodies, mkWWstr, mkWorkerArgs ) where
8
9 #include "HsVersions.h"
10
11 import CoreSyn
12 import CoreUtils        ( exprType )
13 import Id               ( Id, idType, mkSysLocal, idDemandInfo, setIdDemandInfo,
14                           isOneShotLambda, setOneShotLambda, setIdUnfolding,
15                           setIdInfo
16                         )
17 import IdInfo           ( vanillaIdInfo )
18 import DataCon
19 import Demand           ( Demand(..), DmdResult(..), Demands(..) ) 
20 import MkCore           ( mkRuntimeErrorApp, aBSENT_ERROR_ID )
21 import MkId             ( realWorldPrimId, voidArgId, 
22                           mkUnpackCase, mkProductBox )
23 import TysWiredIn       ( tupleCon )
24 import Type
25 import Coercion         ( mkSymCoercion, splitNewTypeRepCo_maybe )
26 import BasicTypes       ( Boxity(..) )
27 import Literal          ( absentLiteralOf )
28 import Var              ( Var )
29 import UniqSupply
30 import Unique
31 import Util             ( zipWithEqual )
32 import Outputable
33 import FastString
34 \end{code}
35
36
37 %************************************************************************
38 %*                                                                      *
39 \subsection[mkWrapperAndWorker]{@mkWrapperAndWorker@}
40 %*                                                                      *
41 %************************************************************************
42
43 Here's an example.  The original function is:
44
45 \begin{verbatim}
46 g :: forall a . Int -> [a] -> a
47
48 g = \/\ a -> \ x ys ->
49         case x of
50           0 -> head ys
51           _ -> head (tail ys)
52 \end{verbatim}
53
54 From this, we want to produce:
55 \begin{verbatim}
56 -- wrapper (an unfolding)
57 g :: forall a . Int -> [a] -> a
58
59 g = \/\ a -> \ x ys ->
60         case x of
61           I# x# -> $wg a x# ys
62             -- call the worker; don't forget the type args!
63
64 -- worker
65 $wg :: forall a . Int# -> [a] -> a
66
67 $wg = \/\ a -> \ x# ys ->
68         let
69             x = I# x#
70         in
71             case x of               -- note: body of g moved intact
72               0 -> head ys
73               _ -> head (tail ys)
74 \end{verbatim}
75
76 Something we have to be careful about:  Here's an example:
77
78 \begin{verbatim}
79 -- "f" strictness: U(P)U(P)
80 f (I# a) (I# b) = a +# b
81
82 g = f   -- "g" strictness same as "f"
83 \end{verbatim}
84
85 \tr{f} will get a worker all nice and friendly-like; that's good.
86 {\em But we don't want a worker for \tr{g}}, even though it has the
87 same strictness as \tr{f}.  Doing so could break laziness, at best.
88
89 Consequently, we insist that the number of strictness-info items is
90 exactly the same as the number of lambda-bound arguments.  (This is
91 probably slightly paranoid, but OK in practice.)  If it isn't the
92 same, we ``revise'' the strictness info, so that we won't propagate
93 the unusable strictness-info into the interfaces.
94
95
96 %************************************************************************
97 %*                                                                      *
98 \subsection{The worker wrapper core}
99 %*                                                                      *
100 %************************************************************************
101
102 @mkWwBodies@ is called when doing the worker\/wrapper split inside a module.
103
104 \begin{code}
105 mkWwBodies :: Type                              -- Type of original function
106            -> [Demand]                          -- Strictness of original function
107            -> DmdResult                         -- Info about function result
108            -> [Bool]                            -- One-shot-ness of the function
109            -> UniqSM ([Demand],                 -- Demands for worker (value) args
110                       Id -> CoreExpr,           -- Wrapper body, lacking only the worker Id
111                       CoreExpr -> CoreExpr)     -- Worker body, lacking the original function rhs
112
113 -- wrap_fn_args E       = \x y -> E
114 -- work_fn_args E       = E x y
115
116 -- wrap_fn_str E        = case x of { (a,b) -> 
117 --                        case a of { (a1,a2) ->
118 --                        E a1 a2 b y }}
119 -- work_fn_str E        = \a2 a2 b y ->
120 --                        let a = (a1,a2) in
121 --                        let x = (a,b) in
122 --                        E
123
124 mkWwBodies fun_ty demands res_info one_shots
125   = do  { let arg_info = demands `zip` (one_shots ++ repeat False)
126         ; (wrap_args, wrap_fn_args, work_fn_args, res_ty) <- mkWWargs emptyTvSubst fun_ty arg_info
127         ; (work_args, wrap_fn_str,  work_fn_str) <- mkWWstr wrap_args
128
129         -- Don't do CPR if the worker doesn't have any value arguments
130         -- Then the worker is just a constant, so we don't want to unbox it.
131         ; (wrap_fn_cpr, work_fn_cpr,  _cpr_res_ty)
132                <- if any isId work_args then
133                      mkWWcpr res_ty res_info
134                   else
135                      return (id, id, res_ty)
136
137         ; let (work_lam_args, work_call_args) = mkWorkerArgs work_args res_ty
138         ; return ([idDemandInfo v | v <- work_call_args, isId v],
139                   wrap_fn_args . wrap_fn_cpr . wrap_fn_str . applyToVars work_call_args . Var,
140                   mkLams work_lam_args. work_fn_str . work_fn_cpr . work_fn_args) }
141         -- We use an INLINE unconditionally, even if the wrapper turns out to be
142         -- something trivial like
143         --      fw = ...
144         --      f = __inline__ (coerce T fw)
145         -- The point is to propagate the coerce to f's call sites, so even though
146         -- f's RHS is now trivial (size 1) we still want the __inline__ to prevent
147         -- fw from being inlined into f's RHS
148 \end{code}
149
150
151 %************************************************************************
152 %*                                                                      *
153 \subsection{Making wrapper args}
154 %*                                                                      *
155 %************************************************************************
156
157 During worker-wrapper stuff we may end up with an unlifted thing
158 which we want to let-bind without losing laziness.  So we
159 add a void argument.  E.g.
160
161         f = /\a -> \x y z -> E::Int#    -- E does not mention x,y,z
162 ==>
163         fw = /\ a -> \void -> E
164         f  = /\ a -> \x y z -> fw realworld
165
166 We use the state-token type which generates no code.
167
168 \begin{code}
169 mkWorkerArgs :: [Var]
170              -> Type    -- Type of body
171              -> ([Var], -- Lambda bound args
172                  [Var]) -- Args at call site
173 mkWorkerArgs args res_ty
174     | any isId args || not (isUnLiftedType res_ty)
175     = (args, args)
176     | otherwise 
177     = (args ++ [voidArgId], args ++ [realWorldPrimId])
178 \end{code}
179
180
181 %************************************************************************
182 %*                                                                      *
183 \subsection{Coercion stuff}
184 %*                                                                      *
185 %************************************************************************
186
187 We really want to "look through" coerces.
188 Reason: I've seen this situation:
189
190         let f = coerce T (\s -> E)
191         in \x -> case x of
192                     p -> coerce T' f
193                     q -> \s -> E2
194                     r -> coerce T' f
195
196 If only we w/w'd f, we'd get
197         let f = coerce T (\s -> fw s)
198             fw = \s -> E
199         in ...
200
201 Now we'll inline f to get
202
203         let fw = \s -> E
204         in \x -> case x of
205                     p -> fw
206                     q -> \s -> E2
207                     r -> fw
208
209 Now we'll see that fw has arity 1, and will arity expand
210 the \x to get what we want.
211
212 \begin{code}
213 -- mkWWargs just does eta expansion
214 -- is driven off the function type and arity.
215 -- It chomps bites off foralls, arrows, newtypes
216 -- and keeps repeating that until it's satisfied the supplied arity
217
218 mkWWargs :: TvSubst             -- Freshening substitution to apply to the type
219                                 --   See Note [Freshen type variables]
220          -> Type                -- The type of the function
221          -> [(Demand,Bool)]     -- Demands and one-shot info for value arguments
222          -> UniqSM  ([Var],             -- Wrapper args
223                      CoreExpr -> CoreExpr,      -- Wrapper fn
224                      CoreExpr -> CoreExpr,      -- Worker fn
225                      Type)                      -- Type of wrapper body
226
227 mkWWargs subst fun_ty arg_info
228   | Just (rep_ty, co) <- splitNewTypeRepCo_maybe fun_ty
229         -- The newtype case is for when the function has
230         -- a recursive newtype after the arrow (rare)
231         -- We check for arity >= 0 to avoid looping in the case
232         -- of a function whose type is, in effect, infinite
233         -- [Arity is driven by looking at the term, not just the type.]
234         --
235         -- It's also important when we have a function returning (say) a pair
236         -- wrapped in a recursive newtype, at least if CPR analysis can look 
237         -- through such newtypes, which it probably can since they are 
238         -- simply coerces.
239         --
240         -- Note (Sept 08): This case applies even if demands is empty.
241         --                 I'm not quite sure why; perhaps it makes it
242         --                 easier for CPR
243   = do { (wrap_args, wrap_fn_args, work_fn_args, res_ty)
244             <-  mkWWargs subst rep_ty arg_info
245         ; return (wrap_args,
246                   \e -> Cast (wrap_fn_args e) (mkSymCoercion co),
247                   \e -> work_fn_args (Cast e co),
248                   res_ty) } 
249
250   | null arg_info
251   = return ([], id, id, substTy subst fun_ty)
252
253   | Just (tv, fun_ty') <- splitForAllTy_maybe fun_ty
254   = do  { let (subst', tv') = substTyVarBndr subst tv
255                 -- This substTyVarBndr clones the type variable when necy
256                 -- See Note [Freshen type variables]
257         ; (wrap_args, wrap_fn_args, work_fn_args, res_ty)
258              <- mkWWargs subst' fun_ty' arg_info
259         ; return (tv' : wrap_args,
260                   Lam tv' . wrap_fn_args,
261                   work_fn_args . (`App` Type (mkTyVarTy tv')),
262                   res_ty) }
263
264   | ((dmd,one_shot):arg_info') <- arg_info
265   , Just (arg_ty, fun_ty') <- splitFunTy_maybe fun_ty
266   = do  { uniq <- getUniqueM
267         ; let arg_ty' = substTy subst arg_ty
268               id = mk_wrap_arg uniq arg_ty' dmd one_shot
269         ; (wrap_args, wrap_fn_args, work_fn_args, res_ty)
270               <- mkWWargs subst fun_ty' arg_info'
271         ; return (id : wrap_args,
272                   Lam id . wrap_fn_args,
273                   work_fn_args . (`App` Var id),
274                   res_ty) }
275
276   | otherwise
277   = WARN( True, ppr fun_ty )                    -- Should not happen: if there is a demand
278     return ([], id, id, substTy subst fun_ty)   -- then there should be a function arrow
279
280 applyToVars :: [Var] -> CoreExpr -> CoreExpr
281 applyToVars vars fn = mkVarApps fn vars
282
283 mk_wrap_arg :: Unique -> Type -> Demand -> Bool -> Id
284 mk_wrap_arg uniq ty dmd one_shot 
285   = set_one_shot one_shot (setIdDemandInfo (mkSysLocal (fsLit "w") uniq ty) dmd)
286   where
287     set_one_shot True  id = setOneShotLambda id
288     set_one_shot False id = id
289 \end{code}
290
291 Note [Freshen type variables]
292 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
293 mkWWargs may be given a type like  (a~b) => <blah>
294 Which really means                 forall (co:a~b). <blah>
295 Because the name of the coercion variable, 'co', isn't mentioned in <blah>,
296 nested coercion foralls may all use the same variable; and sometimes do
297 see Var.mkWildCoVar.
298
299 However, when we do a worker/wrapper split, we must not use shadowed names,
300 else we'll get
301    f = /\ co /\co. fw co co
302 which is obviously wrong.  Actually, the same is true of type variables, which
303 can in principle shadow, within a type (e.g. forall a. a -> forall a. a->a).
304 But type variables *are* mentioned in <blah>, so we must substitute.
305
306 That's why we carry the TvSubst through mkWWargs
307         
308 %************************************************************************
309 %*                                                                      *
310 \subsection{Strictness stuff}
311 %*                                                                      *
312 %************************************************************************
313
314 \begin{code}
315 mkWWstr :: [Var]                                -- Wrapper args; have their demand info on them
316                                                 --  *Includes type variables*
317         -> UniqSM ([Var],                       -- Worker args
318                    CoreExpr -> CoreExpr,        -- Wrapper body, lacking the worker call
319                                                 -- and without its lambdas 
320                                                 -- This fn adds the unboxing
321                                 
322                    CoreExpr -> CoreExpr)        -- Worker body, lacking the original body of the function,
323                                                 -- and lacking its lambdas.
324                                                 -- This fn does the reboxing
325 mkWWstr []
326   = return ([], nop_fn, nop_fn)
327
328 mkWWstr (arg : args) = do
329     (args1, wrap_fn1, work_fn1) <- mkWWstr_one arg
330     (args2, wrap_fn2, work_fn2) <- mkWWstr args
331     return (args1 ++ args2, wrap_fn1 . wrap_fn2, work_fn1 . work_fn2)
332
333 ----------------------
334 -- mkWWstr_one wrap_arg = (work_args, wrap_fn, work_fn)
335 --   *  wrap_fn assumes wrap_arg is in scope,
336 --        brings into scope work_args (via cases)
337 --   * work_fn assumes work_args are in scope, a
338 --        brings into scope wrap_arg (via lets)
339 mkWWstr_one :: Var -> UniqSM ([Var], CoreExpr -> CoreExpr, CoreExpr -> CoreExpr)
340 mkWWstr_one arg
341   | isTyCoVar arg
342   = return ([arg],  nop_fn, nop_fn)
343
344   | otherwise
345   = case idDemandInfo arg of
346
347         -- Absent case.  We can't always handle absence for arbitrary
348         -- unlifted types, so we need to choose just the cases we can
349         -- (that's what mk_absent_let does)
350       Abs | Just work_fn <- mk_absent_let arg
351           -> return ([], nop_fn, work_fn)
352
353         -- Unpack case
354       Eval (Prod cs)
355         | Just (_arg_tycon, _tycon_arg_tys, data_con, inst_con_arg_tys) 
356                 <- deepSplitProductType_maybe (idType arg)
357         -> do uniqs <- getUniquesM
358               let
359                 unpk_args      = zipWith mk_ww_local uniqs inst_con_arg_tys
360                 unpk_args_w_ds = zipWithEqual "mkWWstr" set_worker_arg_info unpk_args cs
361                 unbox_fn       = mkUnpackCase (sanitiseCaseBndr arg) (Var arg) unpk_args data_con
362                 rebox_fn       = Let (NonRec arg con_app) 
363                 con_app        = mkProductBox unpk_args (idType arg)
364               (worker_args, wrap_fn, work_fn) <- mkWWstr unpk_args_w_ds
365               return (worker_args, unbox_fn . wrap_fn, work_fn . rebox_fn) 
366                            -- Don't pass the arg, rebox instead
367
368         -- `seq` demand; evaluate in wrapper in the hope
369         -- of dropping seqs in the worker
370       Eval (Poly Abs)
371         -> let
372                 arg_w_unf = arg `setIdUnfolding` evaldUnfolding
373                 -- Tell the worker arg that it's sure to be evaluated
374                 -- so that internal seqs can be dropped
375            in
376            return ([arg_w_unf], mk_seq_case arg, nop_fn)
377                 -- Pass the arg, anyway, even if it is in theory discarded
378                 -- Consider
379                 --      f x y = x `seq` y
380                 -- x gets a (Eval (Poly Abs)) demand, but if we fail to pass it to the worker
381                 -- we ABSOLUTELY MUST record that x is evaluated in the wrapper.
382                 -- Something like:
383                 --      f x y = x `seq` fw y
384                 --      fw y = let x{Evald} = error "oops" in (x `seq` y)
385                 -- If we don't pin on the "Evald" flag, the seq doesn't disappear, and
386                 -- we end up evaluating the absent thunk.
387                 -- But the Evald flag is pretty weird, and I worry that it might disappear
388                 -- during simplification, so for now I've just nuked this whole case
389                         
390         -- Other cases
391       _other_demand -> return ([arg], nop_fn, nop_fn)
392
393   where
394         -- If the wrapper argument is a one-shot lambda, then
395         -- so should (all) the corresponding worker arguments be
396         -- This bites when we do w/w on a case join point
397     set_worker_arg_info worker_arg demand = set_one_shot (setIdDemandInfo worker_arg demand)
398
399     set_one_shot | isOneShotLambda arg = setOneShotLambda
400                  | otherwise           = \x -> x
401
402 ----------------------
403 nop_fn :: CoreExpr -> CoreExpr
404 nop_fn body = body
405 \end{code}
406
407
408 %************************************************************************
409 %*                                                                      *
410 \subsection{CPR stuff}
411 %*                                                                      *
412 %************************************************************************
413
414
415 @mkWWcpr@ takes the worker/wrapper pair produced from the strictness
416 info and adds in the CPR transformation.  The worker returns an
417 unboxed tuple containing non-CPR components.  The wrapper takes this
418 tuple and re-produces the correct structured output.
419
420 The non-CPR results appear ordered in the unboxed tuple as if by a
421 left-to-right traversal of the result structure.
422
423
424 \begin{code}
425 mkWWcpr :: Type                              -- function body type
426         -> DmdResult                         -- CPR analysis results
427         -> UniqSM (CoreExpr -> CoreExpr,             -- New wrapper 
428                    CoreExpr -> CoreExpr,             -- New worker
429                    Type)                        -- Type of worker's body 
430
431 mkWWcpr body_ty RetCPR
432     | not (isClosedAlgType body_ty)
433     = WARN( True, 
434             text "mkWWcpr: non-algebraic or open body type" <+> ppr body_ty )
435       return (id, id, body_ty)
436
437     | n_con_args == 1 && isUnLiftedType con_arg_ty1 = do
438         -- Special case when there is a single result of unlifted type
439         --
440         -- Wrapper:     case (..call worker..) of x -> C x
441         -- Worker:      case (   ..body..    ) of C x -> x
442       (work_uniq : arg_uniq : _) <- getUniquesM
443       let
444         work_wild = mk_ww_local work_uniq body_ty
445         arg       = mk_ww_local arg_uniq  con_arg_ty1
446         con_app   = mkProductBox [arg] body_ty
447
448       return (\ wkr_call -> Case wkr_call (arg) (exprType con_app) [(DEFAULT, [], con_app)],
449                 \ body     -> workerCase (work_wild) body [arg] data_con (Var arg),
450                 con_arg_ty1)
451
452     | otherwise = do    -- The general case
453         -- Wrapper: case (..call worker..) of (# a, b #) -> C a b
454         -- Worker:  case (   ...body...  ) of C a b -> (# a, b #)     
455       uniqs <- getUniquesM
456       let
457         (wrap_wild : work_wild : args) = zipWith mk_ww_local uniqs (ubx_tup_ty : body_ty : con_arg_tys)
458         arg_vars                       = map Var args
459         ubx_tup_con                    = tupleCon Unboxed n_con_args
460         ubx_tup_ty                     = exprType ubx_tup_app
461         ubx_tup_app                    = mkConApp ubx_tup_con (map Type con_arg_tys   ++ arg_vars)
462         con_app                        = mkProductBox args body_ty
463
464       return (\ wkr_call -> Case wkr_call (wrap_wild) (exprType con_app)  [(DataAlt ubx_tup_con, args, con_app)],
465                 \ body     -> workerCase (work_wild) body args data_con ubx_tup_app,
466                 ubx_tup_ty)
467     where
468       (_arg_tycon, _tycon_arg_tys, data_con, con_arg_tys) = deepSplitProductType "mkWWcpr" body_ty
469       n_con_args  = length con_arg_tys
470       con_arg_ty1 = head con_arg_tys
471
472 mkWWcpr body_ty _other          -- No CPR info
473     = return (id, id, body_ty)
474
475 -- If the original function looked like
476 --      f = \ x -> _scc_ "foo" E
477 --
478 -- then we want the CPR'd worker to look like
479 --      \ x -> _scc_ "foo" (case E of I# x -> x)
480 -- and definitely not
481 --      \ x -> case (_scc_ "foo" E) of I# x -> x)
482 --
483 -- This transform doesn't move work or allocation
484 -- from one cost centre to another
485 workerCase :: Id -> CoreExpr -> [Id] -> DataCon -> CoreExpr -> CoreExpr
486 workerCase bndr (Note (SCC cc) e) args con body = Note (SCC cc) (mkUnpackCase bndr e args con body)
487 workerCase bndr e args con body = mkUnpackCase bndr e args con body
488 \end{code}
489
490
491 %************************************************************************
492 %*                                                                      *
493 \subsection{Utilities}
494 %*                                                                      *
495 %************************************************************************
496
497 Note [Absent errors]
498 ~~~~~~~~~~~~~~~~~~~~
499 We make a new binding for Ids that are marked absent, thus
500    let x = absentError "x :: Int"
501 The idea is that this binding will never be used; but if it 
502 buggily is used we'll get a runtime error message.
503
504 Coping with absence for *unlifted* types is important; see, for
505 example, Trac #4306.  For these we find a suitable literal,
506 using Literal.absentLiteralOf.  We don't have literals for
507 every primitive type, so the function is partial.
508
509     [I did try the experiment of using an error thunk for unlifted
510     things too, relying on the simplifier to drop it as dead code,
511     by making absentError 
512       (a) *not* be a bottoming Id, 
513       (b) be "ok for speculation"
514     But that relies on the simplifier finding that it really
515     is dead code, which is fragile, and indeed failed when 
516     profiling is on, which disables various optimisations.  So
517     using a literal will do.]
518
519 \begin{code}
520 mk_absent_let :: Id -> Maybe (CoreExpr -> CoreExpr)
521 mk_absent_let arg 
522   | not (isUnLiftedType arg_ty)
523   = Just (Let (NonRec arg abs_rhs))
524   | Just (tc, _) <- splitTyConApp_maybe arg_ty
525   , Just lit <- absentLiteralOf tc
526   = Just (Let (NonRec arg (Lit lit)))
527   | otherwise
528   = WARN( True, ptext (sLit "No asbent value for") <+> ppr arg_ty )
529     Nothing
530   where
531     arg_ty  = idType arg
532     abs_rhs = mkRuntimeErrorApp aBSENT_ERROR_ID arg_ty msg
533     msg     = showSDocDebug (ppr arg <+> ppr (idType arg))
534
535 mk_seq_case :: Id -> CoreExpr -> CoreExpr
536 mk_seq_case arg body = Case (Var arg) (sanitiseCaseBndr arg) (exprType body) [(DEFAULT, [], body)]
537
538 sanitiseCaseBndr :: Id -> Id
539 -- The argument we are scrutinising has the right type to be
540 -- a case binder, so it's convenient to re-use it for that purpose.
541 -- But we *must* throw away all its IdInfo.  In particular, the argument
542 -- will have demand info on it, and that demand info may be incorrect for
543 -- the case binder.  e.g.       case ww_arg of ww_arg { I# x -> ... }
544 -- Quite likely ww_arg isn't used in '...'.  The case may get discarded
545 -- if the case binder says "I'm demanded".  This happened in a situation 
546 -- like         (x+y) `seq` ....
547 sanitiseCaseBndr id = id `setIdInfo` vanillaIdInfo
548
549 mk_ww_local :: Unique -> Type -> Id
550 mk_ww_local uniq ty = mkSysLocal (fsLit "ww") uniq ty
551 \end{code}