2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
4 %************************************************************************
6 \section[FloatIn]{Floating Inwards pass}
8 %************************************************************************
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.
15 module FloatIn ( floatInwards ) where
17 #include "HsVersions.h"
20 import CoreUtils ( exprIsHNF, exprIsDupable )
21 import CoreFVs ( CoreExprWithFVs, freeVars, freeVarsOf, idRuleVars )
22 import Id ( isOneShotBndr, idType )
24 import Type ( isUnLiftedType )
26 import Util ( zipEqual, zipWithEqual, count )
31 Top-level interface function, @floatInwards@. Note that we do not
32 actually float any bindings downwards from the top-level.
35 floatInwards :: [CoreBind] -> [CoreBind]
36 floatInwards = map fi_top_bind
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 ]
44 %************************************************************************
46 \subsection{Mail from Andr\'e [edited]}
48 %************************************************************************
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! }
54 You are assuming we DO DO full laziness AFTER floating inwards! We
55 have to [not float inside lambdas] if we don't.
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.
61 Actually we are not doing a proper full laziness (see below), which
62 was another reason for not floating inwards past a lambda.
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
71 let a = expensive ==> \b -> case expensive of ...
72 in \ b -> case a of ...
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.
82 do the full laziness pass (floating lets outwards).
84 simplify. The simplifier inlines the (trivial) lets that were
85 created but were not floated outwards.
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.
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
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
100 let a = something let b = case something of a -> a + a
101 in let b = a + a ==> in (b,b)
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.
108 %************************************************************************
110 \subsection{Main floating-inwards code}
112 %************************************************************************
115 type FreeVarsSet = IdSet
117 type FloatingBinds = [(CoreBind, FreeVarsSet)]
118 -- In reverse dependency order (innermost binder first)
120 -- The FreeVarsSet is the free variables of the binding. In the case
121 -- of recursive bindings, the set doesn't include the bound
124 fiExpr :: FloatingBinds -- Binds we're trying to drop
125 -- as far "inwards" as possible
126 -> CoreExprWithFVs -- Input expr
127 -> CoreExpr -- Result
129 fiExpr to_drop (_, AnnVar v) = mkCoLets' to_drop (Var v)
131 fiExpr to_drop (_, AnnType ty) = ASSERT( null to_drop )
133 fiExpr to_drop (_, AnnCast expr co)
134 = Cast (fiExpr to_drop expr) co -- Just float in past coercion
136 fiExpr _ (_, AnnLit lit) = Lit lit
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.
144 fiExpr to_drop (_,AnnApp fun arg)
145 = mkCoLets' drop_here (App (fiExpr fun_drop fun) (fiExpr arg_drop arg))
147 [drop_here, fun_drop, arg_drop] = sepBindsByDropPoint False [freeVarsOf fun, freeVarsOf arg] to_drop
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.
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
161 in let f = /\t -> \a -> ...
163 let f = /\t -> let v = ... in \a -> ...
164 This is bad as now f is an updatable closure (update PAP)
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
173 It *is* important to float into one-shot lambdas, however;
174 see the remarks with noFloatIntoRhs.
176 So we treat lambda in groups, using the following rule:
178 Float in if (a) there is at least one Id,
179 and (b) there are no non-one-shot Ids
181 Otherwise drop all the bindings outside the group.
183 This is what the 'go' function in the AnnLam case is doing.
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
189 fiExpr to_drop lam@(_, AnnLam _ _)
190 | go False bndrs -- Float in
191 = mkLams bndrs (fiExpr to_drop body)
193 | otherwise -- Dump it all here
194 = mkCoLets' to_drop (mkLams bndrs (fiExpr [] body))
197 (bndrs, body) = collectAnnBndrs lam
199 go seen_one_shot_id [] = seen_one_shot_id
200 go seen_one_shot_id (b:bs)
201 | isTyVar b = go seen_one_shot_id bs
202 | isOneShotBndr b = go True bs
203 | otherwise = False -- Give up at a non-one-shot Id
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.
212 fiExpr to_drop (_, AnnNote note@(SCC _) expr)
213 = -- Wimp out for now
214 mkCoLets' to_drop (Note note (fiExpr [] expr))
216 fiExpr to_drop (_, AnnNote InlineMe expr)
217 = -- Ditto... don't float anything into an INLINE expression
218 mkCoLets' to_drop (Note InlineMe (fiExpr [] expr))
220 fiExpr to_drop (_, AnnNote note@(CoreNote _) expr)
221 = Note note (fiExpr to_drop expr)
224 For @Lets@, the possible ``drop points'' for the \tr{to_drop}
225 bindings are: (a)~in the body, (b1)~in the RHS of a NonRec binding,
226 or~(b2), in each of the RHSs of the pairs of a @Rec@.
228 Note that we do {\em weird things} with this let's binding. Consider:
237 Look at the inner \tr{let}. As \tr{w} is used in both the bind and
238 body of the inner let, we could panic and leave \tr{w}'s binding where
239 it is. But \tr{v} is floatable further into the body of the inner let, and
240 {\em then} \tr{w} will also be only in the body of that inner let.
242 So: rather than drop \tr{w}'s binding here, we add it onto the list of
243 things to drop in the outer let's body, and let nature take its
246 Note [extra_fvs (1): avoid floating into RHS]
247 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
248 Consdider let x=\y....t... in body. We do not necessarily want to float
249 a binding for t into the RHS, because it'll immediately be floated out
250 again. (It won't go inside the lambda else we risk losing work.)
251 In letrec, we need to be more careful still. We don't want to transform
254 letrec f = \z. ...x#...f...
257 letrec f = let x# = y# +# 1# in \z. ...x#...f... in ...
258 because now we can't float the let out again, because a letrec
259 can't have unboxed bindings.
261 So we make "extra_fvs" which is the rhs_fvs of such bindings, and
262 arrange to dump bindings that bind extra_fvs before the entire let.
264 Note [extra_fvs (s): free variables of rules]
265 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
266 Consider let x{rule mentioning y} = rhs in body
267 Here y is not free in rhs or body; but we still want to dump bindings
268 that bind y outside the let. So we augment extra_fvs with the
273 fiExpr to_drop (_,AnnLet (AnnNonRec id rhs@(rhs_fvs, ann_rhs)) body)
274 = fiExpr new_to_drop body
276 body_fvs = freeVarsOf body
278 rule_fvs = idRuleVars id -- See Note [extra_fvs (2): free variables of rules]
279 extra_fvs | noFloatIntoRhs ann_rhs
280 || isUnLiftedType (idType id) = rule_fvs `unionVarSet` rhs_fvs
281 | otherwise = rule_fvs
282 -- See Note [extra_fvs (2): avoid floating into RHS]
283 -- No point in floating in only to float straight out again
284 -- Ditto ok-for-speculation unlifted RHSs
286 [shared_binds, extra_binds, rhs_binds, body_binds]
287 = sepBindsByDropPoint False [extra_fvs, rhs_fvs, body_fvs] to_drop
289 new_to_drop = body_binds ++ -- the bindings used only in the body
290 [(NonRec id rhs', rhs_fvs')] ++ -- the new binding itself
291 extra_binds ++ -- bindings from extra_fvs
292 shared_binds -- the bindings used both in rhs and body
294 -- Push rhs_binds into the right hand side of the binding
295 rhs' = fiExpr rhs_binds rhs
296 rhs_fvs' = rhs_fvs `unionVarSet` floatedBindsFVs rhs_binds `unionVarSet` rule_fvs
297 -- Don't forget the rule_fvs; the binding mentions them!
299 fiExpr to_drop (_,AnnLet (AnnRec bindings) body)
300 = fiExpr new_to_drop body
302 (ids, rhss) = unzip bindings
303 rhss_fvs = map freeVarsOf rhss
304 body_fvs = freeVarsOf body
306 -- See Note [extra_fvs (1,2)]
307 rule_fvs = foldr (unionVarSet . idRuleVars) emptyVarSet ids
308 extra_fvs = rule_fvs `unionVarSet`
309 unionVarSets [ fvs | (fvs, rhs) <- rhss
310 , noFloatIntoRhs rhs ]
312 (shared_binds:extra_binds:body_binds:rhss_binds)
313 = sepBindsByDropPoint False (extra_fvs:body_fvs:rhss_fvs) to_drop
315 new_to_drop = body_binds ++ -- the bindings used only in the body
316 [(Rec (fi_bind rhss_binds bindings), rhs_fvs')] ++
317 -- The new binding itself
318 extra_binds ++ -- Note [extra_fvs (1,2)]
319 shared_binds -- Used in more than one place
321 rhs_fvs' = unionVarSets rhss_fvs `unionVarSet`
322 unionVarSets (map floatedBindsFVs rhss_binds) `unionVarSet`
323 rule_fvs -- Don't forget the rule variables!
325 -- Push rhs_binds into the right hand side of the binding
326 fi_bind :: [FloatingBinds] -- one per "drop pt" conjured w/ fvs_of_rhss
327 -> [(Id, CoreExprWithFVs)]
330 fi_bind to_drops pairs
331 = [ (binder, fiExpr to_drop rhs)
332 | ((binder, rhs), to_drop) <- zipEqual "fi_bind" pairs to_drops ]
335 For @Case@, the possible ``drop points'' for the \tr{to_drop}
336 bindings are: (a)~inside the scrutinee, (b)~inside one of the
337 alternatives/default [default FVs always {\em first}!].
340 fiExpr to_drop (_, AnnCase scrut case_bndr ty alts)
341 = mkCoLets' drop_here1 $
342 mkCoLets' drop_here2 $
343 Case (fiExpr scrut_drops scrut) case_bndr ty
344 (zipWith fi_alt alts_drops_s alts)
346 -- Float into the scrut and alts-considered-together just like App
347 [drop_here1, scrut_drops, alts_drops] = sepBindsByDropPoint False [scrut_fvs, all_alts_fvs] to_drop
349 -- Float into the alts with the is_case flag set
350 (drop_here2 : alts_drops_s) = sepBindsByDropPoint True alts_fvs alts_drops
352 scrut_fvs = freeVarsOf scrut
353 alts_fvs = map alt_fvs alts
354 all_alts_fvs = unionVarSets alts_fvs
355 alt_fvs (_con, args, rhs) = foldl delVarSet (freeVarsOf rhs) (case_bndr:args)
356 -- Delete case_bndr and args from free vars of rhs
357 -- to get free vars of alt
359 fi_alt to_drop (con, args, rhs) = (con, args, fiExpr to_drop rhs)
361 noFloatIntoRhs :: AnnExpr' Var (UniqFM Var) -> Bool
362 noFloatIntoRhs (AnnNote InlineMe _) = True
363 noFloatIntoRhs (AnnLam b _) = not (is_one_shot b)
364 -- IMPORTANT: don't say 'True' for a RHS with a one-shot lambda at the top.
365 -- This makes a big difference for things like
366 -- f x# = let x = I# x#
367 -- in let j = \() -> ...x...
368 -- in if <condition> then normal-path else j ()
369 -- If x is used only in the error case join point, j, we must float the
370 -- boxing constructor into it, else we box it every time which is very bad
373 noFloatIntoRhs rhs = exprIsHNF (deAnnotate' rhs) -- We'd just float right back out again...
375 is_one_shot :: Var -> Bool
376 is_one_shot b = isIdVar b && isOneShotBndr b
380 %************************************************************************
382 \subsection{@sepBindsByDropPoint@}
384 %************************************************************************
386 This is the crucial function. The idea is: We have a wad of bindings
387 that we'd like to distribute inside a collection of {\em drop points};
388 insides the alternatives of a \tr{case} would be one example of some
389 drop points; the RHS and body of a non-recursive \tr{let} binding
390 would be another (2-element) collection.
392 So: We're given a list of sets-of-free-variables, one per drop point,
393 and a list of floating-inwards bindings. If a binding can go into
394 only one drop point (without suddenly making something out-of-scope),
395 in it goes. If a binding is used inside {\em multiple} drop points,
396 then it has to go in a you-must-drop-it-above-all-these-drop-points
399 We have to maintain the order on these drop-point-related lists.
403 :: Bool -- True <=> is case expression
404 -> [FreeVarsSet] -- One set of FVs per drop point
405 -> FloatingBinds -- Candidate floaters
406 -> [FloatingBinds] -- FIRST one is bindings which must not be floated
407 -- inside any drop point; the rest correspond
408 -- one-to-one with the input list of FV sets
410 -- Every input floater is returned somewhere in the result;
411 -- none are dropped, not even ones which don't seem to be
412 -- free in *any* of the drop-point fvs. Why? Because, for example,
413 -- a binding (let x = E in B) might have a specialised version of
414 -- x (say x') stored inside x, but x' isn't free in E or B.
416 type DropBox = (FreeVarsSet, FloatingBinds)
418 sepBindsByDropPoint _is_case drop_pts []
419 = [] : [[] | _ <- drop_pts] -- cut to the chase scene; it happens
421 sepBindsByDropPoint is_case drop_pts floaters
422 = go floaters (map (\fvs -> (fvs, [])) (emptyVarSet : drop_pts))
424 go :: FloatingBinds -> [DropBox] -> [FloatingBinds]
425 -- The *first* one in the argument list is the drop_here set
426 -- The FloatingBinds in the lists are in the reverse of
427 -- the normal FloatingBinds order; that is, they are the right way round!
429 go [] drop_boxes = map (reverse . snd) drop_boxes
431 go (bind_w_fvs@(bind, bind_fvs) : binds) drop_boxes@(here_box : fork_boxes)
434 -- "here" means the group of bindings dropped at the top of the fork
436 (used_here : used_in_flags) = [ any (`elemVarSet` fvs) (bindersOf bind)
437 | (fvs, _) <- drop_boxes]
439 drop_here = used_here || not can_push
441 -- For case expressions we duplicate the binding if it is
442 -- reasonably small, and if it is not used in all the RHSs
443 -- This is good for situations like
448 -- E -> ...not mentioning x...
450 n_alts = length used_in_flags
451 n_used_alts = count id used_in_flags -- returns number of Trues in list.
453 can_push = n_used_alts == 1 -- Used in just one branch
454 || (is_case && -- We are looking at case alternatives
455 n_used_alts > 1 && -- It's used in more than one
456 n_used_alts < n_alts && -- ...but not all
457 bindIsDupable bind) -- and we can duplicate the binding
459 new_boxes | drop_here = (insert here_box : fork_boxes)
460 | otherwise = (here_box : new_fork_boxes)
462 new_fork_boxes = zipWithEqual "FloatIn.sepBinds" insert_maybe fork_boxes used_in_flags
464 insert :: DropBox -> DropBox
465 insert (fvs,drops) = (fvs `unionVarSet` bind_fvs, bind_w_fvs:drops)
467 insert_maybe box True = insert box
468 insert_maybe box False = box
470 go _ _ = panic "sepBindsByDropPoint/go"
473 floatedBindsFVs :: FloatingBinds -> FreeVarsSet
474 floatedBindsFVs binds = unionVarSets (map snd binds)
476 mkCoLets' :: FloatingBinds -> CoreExpr -> CoreExpr
477 mkCoLets' to_drop e = foldl (flip (Let . fst)) e to_drop
478 -- Remember to_drop is in *reverse* dependency order
480 bindIsDupable :: Bind CoreBndr -> Bool
481 bindIsDupable (Rec prs) = all (exprIsDupable . snd) prs
482 bindIsDupable (NonRec _ r) = exprIsDupable r