add -fsimpleopt-before-flatten
[ghc-hetmet.git] / compiler / simplCore / FloatIn.lhs
1 %
2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
3 %
4 %************************************************************************
5 %*                                                                      *
6 \section[FloatIn]{Floating Inwards pass}
7 %*                                                                      *
8 %************************************************************************
9
10 The main purpose of @floatInwards@ is floating into branches of a
11 case, so that we don't allocate things, save them on the stack, and
12 then discover that they aren't needed in the chosen branch.
13
14 \begin{code}
15 module FloatIn ( floatInwards ) where
16
17 #include "HsVersions.h"
18
19 import CoreSyn
20 import CoreUtils        ( exprIsHNF, exprIsDupable )
21 import CoreFVs          ( CoreExprWithFVs, freeVars, freeVarsOf, idRuleAndUnfoldingVars )
22 import Id               ( isOneShotBndr, idType )
23 import Var
24 import Type             ( isUnLiftedType )
25 import VarSet
26 import Util             ( zipEqual, zipWithEqual, count )
27 import UniqFM
28 import Outputable
29 \end{code}
30
31 Top-level interface function, @floatInwards@.  Note that we do not
32 actually float any bindings downwards from the top-level.
33
34 \begin{code}
35 floatInwards :: [CoreBind] -> [CoreBind]
36 floatInwards = map fi_top_bind
37   where
38     fi_top_bind (NonRec binder rhs)
39       = NonRec binder (fiExpr [] (freeVars rhs))
40     fi_top_bind (Rec pairs)
41       = Rec [ (b, fiExpr [] (freeVars rhs)) | (b, rhs) <- pairs ]
42 \end{code}
43
44 %************************************************************************
45 %*                                                                      *
46 \subsection{Mail from Andr\'e [edited]}
47 %*                                                                      *
48 %************************************************************************
49
50 {\em Will wrote: What??? I thought the idea was to float as far
51 inwards as possible, no matter what.  This is dropping all bindings
52 every time it sees a lambda of any kind.  Help! }
53
54 You are assuming we DO DO full laziness AFTER floating inwards!  We
55 have to [not float inside lambdas] if we don't.
56
57 If we indeed do full laziness after the floating inwards (we could
58 check the compilation flags for that) then I agree we could be more
59 aggressive and do float inwards past lambdas.
60
61 Actually we are not doing a proper full laziness (see below), which
62 was another reason for not floating inwards past a lambda.
63
64 This can easily be fixed.  The problem is that we float lets outwards,
65 but there are a few expressions which are not let bound, like case
66 scrutinees and case alternatives.  After floating inwards the
67 simplifier could decide to inline the let and the laziness would be
68 lost, e.g.
69
70 \begin{verbatim}
71 let a = expensive             ==> \b -> case expensive of ...
72 in \ b -> case a of ...
73 \end{verbatim}
74 The fix is
75 \begin{enumerate}
76 \item
77 to let bind the algebraic case scrutinees (done, I think) and
78 the case alternatives (except the ones with an
79 unboxed type)(not done, I think). This is best done in the
80 SetLevels.lhs module, which tags things with their level numbers.
81 \item
82 do the full laziness pass (floating lets outwards).
83 \item
84 simplify. The simplifier inlines the (trivial) lets that were
85  created but were not floated outwards.
86 \end{enumerate}
87
88 With the fix I think Will's suggestion that we can gain even more from
89 strictness by floating inwards past lambdas makes sense.
90
91 We still gain even without going past lambdas, as things may be
92 strict in the (new) context of a branch (where it was floated to) or
93 of a let rhs, e.g.
94 \begin{verbatim}
95 let a = something            case x of
96 in case x of                   alt1 -> case something of a -> a + a
97      alt1 -> a + a      ==>    alt2 -> b
98      alt2 -> b
99
100 let a = something           let b = case something of a -> a + a
101 in let b = a + a        ==> in (b,b)
102 in (b,b)
103 \end{verbatim}
104 Also, even if a is not found to be strict in the new context and is
105 still left as a let, if the branch is not taken (or b is not entered)
106 the closure for a is not built.
107
108 %************************************************************************
109 %*                                                                      *
110 \subsection{Main floating-inwards code}
111 %*                                                                      *
112 %************************************************************************
113
114 \begin{code}
115 type FreeVarsSet   = IdSet
116
117 type FloatingBinds = [(CoreBind, FreeVarsSet)]
118         -- In reverse dependency order (innermost binder first)
119
120         -- The FreeVarsSet is the free variables of the binding.  In the case
121         -- of recursive bindings, the set doesn't include the bound
122         -- variables.
123
124 fiExpr :: FloatingBinds         -- Binds we're trying to drop
125                                 -- as far "inwards" as possible
126        -> CoreExprWithFVs       -- Input expr
127        -> CoreExpr              -- Result
128
129 fiExpr to_drop (_, AnnVar v) = mkCoLets' to_drop (Var v)
130
131 fiExpr to_drop (_, AnnType ty) = ASSERT( null to_drop )
132                                  Type ty
133 fiExpr to_drop (_, AnnCast expr co)
134   = Cast (fiExpr to_drop expr) co       -- Just float in past coercion
135
136 fiExpr _ (_, AnnLit lit) = Lit lit
137 \end{code}
138
139 Applications: we do float inside applications, mainly because we
140 need to get at all the arguments.  The next simplifier run will
141 pull out any silly ones.
142
143 \begin{code}
144 fiExpr to_drop (_,AnnApp fun arg)
145   = mkCoLets' drop_here (App (fiExpr fun_drop fun) (fiExpr arg_drop arg))
146   where
147     [drop_here, fun_drop, arg_drop] = sepBindsByDropPoint False [freeVarsOf fun, freeVarsOf arg] to_drop
148 \end{code}
149
150 Note [Floating in past a lambda group]
151 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
152 * We must be careful about floating inside inside a value lambda.  
153   That risks losing laziness.
154   The float-out pass might rescue us, but then again it might not.
155
156 * We must be careful about type lambdas too.  At one time we did, and
157   there is no risk of duplicating work thereby, but we do need to be
158   careful.  In particular, here is a bad case (it happened in the
159   cichelli benchmark:
160         let v = ...
161         in let f = /\t -> \a -> ...
162            ==>
163         let f = /\t -> let v = ... in \a -> ...
164   This is bad as now f is an updatable closure (update PAP)
165   and has arity 0.
166
167 * Hack alert!  We only float in through one-shot lambdas, 
168   not (as you might guess) through lone big lambdas.  
169   Reason: we float *out* past big lambdas (see the test in the Lam
170   case of FloatOut.floatExpr) and we don't want to float straight
171   back in again.
172   
173   It *is* important to float into one-shot lambdas, however;
174   see the remarks with noFloatIntoRhs.
175
176 So we treat lambda in groups, using the following rule:
177
178  Float in if (a) there is at least one Id, 
179          and (b) there are no non-one-shot Ids
180
181  Otherwise drop all the bindings outside the group.
182
183 This is what the 'go' function in the AnnLam case is doing.
184
185 Urk! if all are tyvars, and we don't float in, we may miss an 
186       opportunity to float inside a nested case branch
187
188 \begin{code}
189 fiExpr to_drop lam@(_, AnnLam _ _)
190   | go False bndrs      -- Float in
191   = mkLams bndrs (fiExpr to_drop body)
192
193   | otherwise           -- Dump it all here
194   = mkCoLets' to_drop (mkLams bndrs (fiExpr [] body))
195
196   where
197     (bndrs, body) = collectAnnBndrs lam
198
199     go seen_one_shot_id [] = seen_one_shot_id
200     go seen_one_shot_id (b:bs)
201       | isTyCoVar       b = go seen_one_shot_id bs
202       | isOneShotBndr b = go True bs
203       | otherwise       = False  -- Give up at a non-one-shot Id
204 \end{code}
205
206 We don't float lets inwards past an SCC.
207         ToDo: keep info on current cc, and when passing
208         one, if it is not the same, annotate all lets in binds with current
209         cc, change current cc to the new one and float binds into expr.
210
211 \begin{code}
212 fiExpr to_drop (_, AnnNote note@(SCC _) expr)
213   =     -- Wimp out for now
214     mkCoLets' to_drop (Note note (fiExpr [] expr))
215
216 fiExpr to_drop (_, AnnNote note@(CoreNote _) expr)
217   = Note note (fiExpr to_drop expr)
218 \end{code}
219
220 For @Lets@, the possible ``drop points'' for the \tr{to_drop}
221 bindings are: (a)~in the body, (b1)~in the RHS of a NonRec binding,
222 or~(b2), in each of the RHSs of the pairs of a @Rec@.
223
224 Note that we do {\em weird things} with this let's binding.  Consider:
225 \begin{verbatim}
226 let
227     w = ...
228 in {
229     let v = ... w ...
230     in ... v .. w ...
231 }
232 \end{verbatim}
233 Look at the inner \tr{let}.  As \tr{w} is used in both the bind and
234 body of the inner let, we could panic and leave \tr{w}'s binding where
235 it is.  But \tr{v} is floatable further into the body of the inner let, and
236 {\em then} \tr{w} will also be only in the body of that inner let.
237
238 So: rather than drop \tr{w}'s binding here, we add it onto the list of
239 things to drop in the outer let's body, and let nature take its
240 course.
241
242 Note [extra_fvs (1): avoid floating into RHS]
243 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
244 Consdider let x=\y....t... in body.  We do not necessarily want to float 
245 a binding for t into the RHS, because it'll immediately be floated out
246 again.  (It won't go inside the lambda else we risk losing work.)
247 In letrec, we need to be more careful still. We don't want to transform
248         let x# = y# +# 1#
249         in
250         letrec f = \z. ...x#...f...
251         in ...
252 into
253         letrec f = let x# = y# +# 1# in \z. ...x#...f... in ...
254 because now we can't float the let out again, because a letrec
255 can't have unboxed bindings.
256
257 So we make "extra_fvs" which is the rhs_fvs of such bindings, and
258 arrange to dump bindings that bind extra_fvs before the entire let.
259
260 Note [extra_fvs (s): free variables of rules]
261 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
262 Consider 
263   let x{rule mentioning y} = rhs in body 
264 Here y is not free in rhs or body; but we still want to dump bindings
265 that bind y outside the let.  So we augment extra_fvs with the
266 idRuleAndUnfoldingVars of x.  No need for type variables, hence not using
267 idFreeVars.
268
269
270 \begin{code}
271 fiExpr to_drop (_,AnnLet (AnnNonRec id rhs@(rhs_fvs, ann_rhs)) body)
272   = fiExpr new_to_drop body
273   where
274     body_fvs = freeVarsOf body
275
276     rule_fvs = idRuleAndUnfoldingVars id        -- See Note [extra_fvs (2): free variables of rules]
277     extra_fvs | noFloatIntoRhs ann_rhs
278               || isUnLiftedType (idType id) = rule_fvs `unionVarSet` rhs_fvs
279               | otherwise                   = rule_fvs
280         -- See Note [extra_fvs (2): avoid floating into RHS]
281         -- No point in floating in only to float straight out again
282         -- Ditto ok-for-speculation unlifted RHSs
283
284     [shared_binds, extra_binds, rhs_binds, body_binds] 
285         = sepBindsByDropPoint False [extra_fvs, rhs_fvs, body_fvs] to_drop
286
287     new_to_drop = body_binds ++                         -- the bindings used only in the body
288                   [(NonRec id rhs', rhs_fvs')] ++       -- the new binding itself
289                   extra_binds ++                        -- bindings from extra_fvs
290                   shared_binds                          -- the bindings used both in rhs and body
291
292         -- Push rhs_binds into the right hand side of the binding
293     rhs'     = fiExpr rhs_binds rhs
294     rhs_fvs' = rhs_fvs `unionVarSet` floatedBindsFVs rhs_binds `unionVarSet` rule_fvs
295                         -- Don't forget the rule_fvs; the binding mentions them!
296
297 fiExpr to_drop (_,AnnLet (AnnRec bindings) body)
298   = fiExpr new_to_drop body
299   where
300     (ids, rhss) = unzip bindings
301     rhss_fvs = map freeVarsOf rhss
302     body_fvs = freeVarsOf body
303
304         -- See Note [extra_fvs (1,2)]
305     rule_fvs = foldr (unionVarSet . idRuleAndUnfoldingVars) emptyVarSet ids
306     extra_fvs = rule_fvs `unionVarSet` 
307                 unionVarSets [ fvs | (fvs, rhs) <- rhss
308                              , noFloatIntoRhs rhs ]
309
310     (shared_binds:extra_binds:body_binds:rhss_binds) 
311         = sepBindsByDropPoint False (extra_fvs:body_fvs:rhss_fvs) to_drop
312
313     new_to_drop = body_binds ++         -- the bindings used only in the body
314                   [(Rec (fi_bind rhss_binds bindings), rhs_fvs')] ++
315                                         -- The new binding itself
316                   extra_binds ++        -- Note [extra_fvs (1,2)]
317                   shared_binds          -- Used in more than one place
318
319     rhs_fvs' = unionVarSets rhss_fvs `unionVarSet`
320                unionVarSets (map floatedBindsFVs rhss_binds) `unionVarSet`
321                rule_fvs         -- Don't forget the rule variables!
322
323     -- Push rhs_binds into the right hand side of the binding
324     fi_bind :: [FloatingBinds]      -- one per "drop pt" conjured w/ fvs_of_rhss
325             -> [(Id, CoreExprWithFVs)]
326             -> [(Id, CoreExpr)]
327
328     fi_bind to_drops pairs
329       = [ (binder, fiExpr to_drop rhs) 
330         | ((binder, rhs), to_drop) <- zipEqual "fi_bind" pairs to_drops ]
331 \end{code}
332
333 For @Case@, the possible ``drop points'' for the \tr{to_drop}
334 bindings are: (a)~inside the scrutinee, (b)~inside one of the
335 alternatives/default [default FVs always {\em first}!].
336
337 \begin{code}
338 fiExpr to_drop (_, AnnCase scrut case_bndr ty alts)
339   = mkCoLets' drop_here1 $
340     mkCoLets' drop_here2 $
341     Case (fiExpr scrut_drops scrut) case_bndr ty
342          (zipWith fi_alt alts_drops_s alts)
343   where
344         -- Float into the scrut and alts-considered-together just like App
345     [drop_here1, scrut_drops, alts_drops] = sepBindsByDropPoint False [scrut_fvs, all_alts_fvs] to_drop
346
347         -- Float into the alts with the is_case flag set
348     (drop_here2 : alts_drops_s)           = sepBindsByDropPoint True alts_fvs alts_drops
349
350     scrut_fvs    = freeVarsOf scrut
351     alts_fvs     = map alt_fvs alts
352     all_alts_fvs = unionVarSets alts_fvs
353     alt_fvs (_con, args, rhs) = foldl delVarSet (freeVarsOf rhs) (case_bndr:args)
354                                 -- Delete case_bndr and args from free vars of rhs 
355                                 -- to get free vars of alt
356
357     fi_alt to_drop (con, args, rhs) = (con, args, fiExpr to_drop rhs)
358
359 noFloatIntoRhs :: AnnExpr' Var (UniqFM Var) -> Bool
360 noFloatIntoRhs (AnnLam b _) = not (is_one_shot b)
361         -- IMPORTANT: don't say 'True' for a RHS with a one-shot lambda at the top.
362         -- This makes a big difference for things like
363         --      f x# = let x = I# x#
364         --             in let j = \() -> ...x...
365         --                in if <condition> then normal-path else j ()
366         -- If x is used only in the error case join point, j, we must float the
367         -- boxing constructor into it, else we box it every time which is very bad
368         -- news indeed.
369
370 noFloatIntoRhs rhs = exprIsHNF (deAnnotate' rhs)        -- We'd just float right back out again...
371
372 is_one_shot :: Var -> Bool
373 is_one_shot b = isId b && isOneShotBndr b
374 \end{code}
375
376
377 %************************************************************************
378 %*                                                                      *
379 \subsection{@sepBindsByDropPoint@}
380 %*                                                                      *
381 %************************************************************************
382
383 This is the crucial function.  The idea is: We have a wad of bindings
384 that we'd like to distribute inside a collection of {\em drop points};
385 insides the alternatives of a \tr{case} would be one example of some
386 drop points; the RHS and body of a non-recursive \tr{let} binding
387 would be another (2-element) collection.
388
389 So: We're given a list of sets-of-free-variables, one per drop point,
390 and a list of floating-inwards bindings.  If a binding can go into
391 only one drop point (without suddenly making something out-of-scope),
392 in it goes.  If a binding is used inside {\em multiple} drop points,
393 then it has to go in a you-must-drop-it-above-all-these-drop-points
394 point.
395
396 We have to maintain the order on these drop-point-related lists.
397
398 \begin{code}
399 sepBindsByDropPoint
400     :: Bool                 -- True <=> is case expression
401     -> [FreeVarsSet]        -- One set of FVs per drop point
402     -> FloatingBinds        -- Candidate floaters
403     -> [FloatingBinds]      -- FIRST one is bindings which must not be floated
404                             -- inside any drop point; the rest correspond
405                             -- one-to-one with the input list of FV sets
406
407 -- Every input floater is returned somewhere in the result;
408 -- none are dropped, not even ones which don't seem to be
409 -- free in *any* of the drop-point fvs.  Why?  Because, for example,
410 -- a binding (let x = E in B) might have a specialised version of
411 -- x (say x') stored inside x, but x' isn't free in E or B.
412
413 type DropBox = (FreeVarsSet, FloatingBinds)
414
415 sepBindsByDropPoint _is_case drop_pts []
416   = [] : [[] | _ <- drop_pts]   -- cut to the chase scene; it happens
417
418 sepBindsByDropPoint is_case drop_pts floaters
419   = go floaters (map (\fvs -> (fvs, [])) (emptyVarSet : drop_pts))
420   where
421     go :: FloatingBinds -> [DropBox] -> [FloatingBinds]
422         -- The *first* one in the argument list is the drop_here set
423         -- The FloatingBinds in the lists are in the reverse of
424         -- the normal FloatingBinds order; that is, they are the right way round!
425
426     go [] drop_boxes = map (reverse . snd) drop_boxes
427
428     go (bind_w_fvs@(bind, bind_fvs) : binds) drop_boxes@(here_box : fork_boxes)
429         = go binds new_boxes
430         where
431           -- "here" means the group of bindings dropped at the top of the fork
432
433           (used_here : used_in_flags) = [ any (`elemVarSet` fvs) (bindersOf bind)
434                                         | (fvs, _) <- drop_boxes]
435
436           drop_here = used_here || not can_push
437
438                 -- For case expressions we duplicate the binding if it is
439                 -- reasonably small, and if it is not used in all the RHSs
440                 -- This is good for situations like
441                 --      let x = I# y in
442                 --      case e of
443                 --        C -> error x
444                 --        D -> error x
445                 --        E -> ...not mentioning x...
446
447           n_alts      = length used_in_flags
448           n_used_alts = count id used_in_flags -- returns number of Trues in list.
449
450           can_push = n_used_alts == 1           -- Used in just one branch
451                    || (is_case &&               -- We are looking at case alternatives
452                        n_used_alts > 1 &&       -- It's used in more than one
453                        n_used_alts < n_alts &&  -- ...but not all
454                        bindIsDupable bind)      -- and we can duplicate the binding
455
456           new_boxes | drop_here = (insert here_box : fork_boxes)
457                     | otherwise = (here_box : new_fork_boxes)
458
459           new_fork_boxes = zipWithEqual "FloatIn.sepBinds" insert_maybe fork_boxes used_in_flags
460
461           insert :: DropBox -> DropBox
462           insert (fvs,drops) = (fvs `unionVarSet` bind_fvs, bind_w_fvs:drops)
463
464           insert_maybe box True  = insert box
465           insert_maybe box False = box
466
467     go _ _ = panic "sepBindsByDropPoint/go"
468
469
470 floatedBindsFVs :: FloatingBinds -> FreeVarsSet
471 floatedBindsFVs binds = unionVarSets (map snd binds)
472
473 mkCoLets' :: FloatingBinds -> CoreExpr -> CoreExpr
474 mkCoLets' to_drop e = foldl (flip (Let . fst)) e to_drop
475         -- Remember to_drop is in *reverse* dependency order
476
477 bindIsDupable :: Bind CoreBndr -> Bool
478 bindIsDupable (Rec prs)    = all (exprIsDupable . snd) prs
479 bindIsDupable (NonRec _ r) = exprIsDupable r
480 \end{code}