2 % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
6 ***************************
8 ***************************
10 1. We attach binding levels to Core bindings, in preparation for floating
11 outwards (@FloatOut@).
13 2. We also let-ify many expressions (notably case scrutinees), so they
14 will have a fighting chance of being floated sensible.
16 3. We clone the binders of any floatable let-binding, so that when it is
17 floated out it will be unique. (This used to be done by the simplifier
18 but the latter now only ensures that there's no shadowing; indeed, even
19 that may not be true.)
21 NOTE: this can't be done using the uniqAway idea, because the variable
22 must be unique in the whole program, not just its current scope,
23 because two variables in different scopes may float out to the
26 NOTE: Very tiresomely, we must apply this substitution to
27 the rules stored inside a variable too.
29 We do *not* clone top-level bindings, because some of them must not change,
30 but we *do* clone bindings that are heading for the top level
33 case x of wild { p -> ...wild... }
34 we substitute x for wild in the RHS of the case alternatives:
35 case x of wild { p -> ...x... }
36 This means that a sub-expression involving x is not "trapped" inside the RHS.
37 And it's not inconvenient because we already have a substitution.
39 Note that this is EXACTLY BACKWARDS from the what the simplifier does.
40 The simplifier tries to get rid of occurrences of x, in favour of wild,
41 in the hope that there will only be one remaining occurrence of x, namely
42 the scrutinee of the case, and we can inline it.
49 LevelledBind, LevelledExpr,
51 incMinorLvl, ltMajLvl, ltLvl, isTopLvl
54 #include "HsVersions.h"
57 import CoreMonad ( FloatOutSwitches(..) )
58 import CoreUtils ( exprType, mkPiTypes )
59 import CoreArity ( exprBotStrictness_maybe )
60 import CoreFVs -- all of it
61 import CoreSubst ( Subst, emptySubst, extendInScope, extendInScopeList,
62 extendIdSubst, cloneIdBndr, cloneRecIdBndrs )
68 import Demand ( StrictSig, increaseStrictSigArity )
69 import Name ( getOccName, mkSystemVarName )
70 import OccName ( occNameString )
71 import Type ( isUnLiftedType, Type )
72 import BasicTypes ( TopLevelFlag(..), Arity )
79 %************************************************************************
81 \subsection{Level numbers}
83 %************************************************************************
86 data Level = Level Int -- Level number of enclosing lambdas
87 Int -- Number of big-lambda and/or case expressions between
88 -- here and the nearest enclosing lambda
91 The {\em level number} on a (type-)lambda-bound variable is the
92 nesting depth of the (type-)lambda which binds it. The outermost lambda
93 has level 1, so (Level 0 0) means that the variable is bound outside any lambda.
95 On an expression, it's the maximum level number of its free
96 (type-)variables. On a let(rec)-bound variable, it's the level of its
97 RHS. On a case-bound variable, it's the number of enclosing lambdas.
99 Top-level variables: level~0. Those bound on the RHS of a top-level
100 definition but ``before'' a lambda; e.g., the \tr{x} in (levels shown
101 as ``subscripts'')...
103 a_0 = let b_? = ... in
104 x_1 = ... b ... in ...
107 The main function @lvlExpr@ carries a ``context level'' (@ctxt_lvl@).
108 That's meant to be the level number of the enclosing binder in the
109 final (floated) program. If the level number of a sub-expression is
110 less than that of the context, then it might be worth let-binding the
111 sub-expression so that it will indeed float.
113 If you can float to level @Level 0 0@ worth doing so because then your
114 allocation becomes static instead of dynamic. We always start with
118 Note [FloatOut inside INLINE]
119 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
120 @InlineCtxt@ very similar to @Level 0 0@, but is used for one purpose:
121 to say "don't float anything out of here". That's exactly what we
122 want for the body of an INLINE, where we don't want to float anything
123 out at all. See notes with lvlMFE below.
127 -- At one time I tried the effect of not float anything out of an InlineMe,
128 -- but it sometimes works badly. For example, consider PrelArr.done. It
129 -- has the form __inline (\d. e)
130 -- where e doesn't mention d. If we float this to
131 -- __inline (let x = e in \d. x)
132 -- things are bad. The inliner doesn't even inline it because it doesn't look
133 -- like a head-normal form. So it seems a lesser evil to let things float.
134 -- In SetLevels we do set the context to (Level 0 0) when we get to an InlineMe
135 -- which discourages floating out.
137 So the conclusion is: don't do any floating at all inside an InlineMe.
138 (In the above example, don't float the {x=e} out of the \d.)
140 One particular case is that of workers: we don't want to float the
141 call to the worker outside the wrapper, otherwise the worker might get
142 inlined into the floated expression, and an importing module won't see
146 type LevelledExpr = TaggedExpr Level
147 type LevelledBind = TaggedBind Level
150 tOP_LEVEL = Level 0 0
152 incMajorLvl :: Level -> Level
153 incMajorLvl (Level major _) = Level (major + 1) 0
155 incMinorLvl :: Level -> Level
156 incMinorLvl (Level major minor) = Level major (minor+1)
158 maxLvl :: Level -> Level -> Level
159 maxLvl l1@(Level maj1 min1) l2@(Level maj2 min2)
160 | (maj1 > maj2) || (maj1 == maj2 && min1 > min2) = l1
163 ltLvl :: Level -> Level -> Bool
164 ltLvl (Level maj1 min1) (Level maj2 min2)
165 = (maj1 < maj2) || (maj1 == maj2 && min1 < min2)
167 ltMajLvl :: Level -> Level -> Bool
168 -- Tells if one level belongs to a difft *lambda* level to another
169 ltMajLvl (Level maj1 _) (Level maj2 _) = maj1 < maj2
171 isTopLvl :: Level -> Bool
172 isTopLvl (Level 0 0) = True
175 instance Outputable Level where
176 ppr (Level maj min) = hcat [ char '<', int maj, char ',', int min, char '>' ]
178 instance Eq Level where
179 (Level maj1 min1) == (Level maj2 min2) = maj1 == maj2 && min1 == min2
183 %************************************************************************
185 \subsection{Main level-setting code}
187 %************************************************************************
190 setLevels :: FloatOutSwitches
195 setLevels float_lams binds us
196 = initLvl us (do_them init_env binds)
198 init_env = initialEnv float_lams
200 do_them :: LevelEnv -> [CoreBind] -> LvlM [LevelledBind]
201 do_them _ [] = return []
203 = do { (lvld_bind, env') <- lvlTopBind env b
204 ; lvld_binds <- do_them env' bs
205 ; return (lvld_bind : lvld_binds) }
207 lvlTopBind :: LevelEnv -> Bind Id -> LvlM (LevelledBind, LevelEnv)
208 lvlTopBind env (NonRec binder rhs)
209 = lvlBind TopLevel tOP_LEVEL env (AnnNonRec binder (freeVars rhs))
210 -- Rhs can have no free vars!
212 lvlTopBind env (Rec pairs)
213 = lvlBind TopLevel tOP_LEVEL env (AnnRec [(b,freeVars rhs) | (b,rhs) <- pairs])
216 %************************************************************************
218 \subsection{Setting expression levels}
220 %************************************************************************
223 lvlExpr :: Level -- ctxt_lvl: Level of enclosing expression
224 -> LevelEnv -- Level of in-scope names/tyvars
225 -> CoreExprWithFVs -- input expression
226 -> LvlM LevelledExpr -- Result expression
229 The @ctxt_lvl@ is, roughly, the level of the innermost enclosing
230 binder. Here's an example
232 v = \x -> ...\y -> let r = case (..x..) of
236 When looking at the rhs of @r@, @ctxt_lvl@ will be 1 because that's
237 the level of @r@, even though it's inside a level-2 @\y@. It's
238 important that @ctxt_lvl@ is 1 and not 2 in @r@'s rhs, because we
239 don't want @lvlExpr@ to turn the scrutinee of the @case@ into an MFE
240 --- because it isn't a *maximal* free expression.
242 If there were another lambda in @r@'s rhs, it would get level-2 as well.
245 lvlExpr _ _ ( _, AnnType ty) = return (Type ty)
246 lvlExpr _ _ ( _, AnnCoercion co) = return (Coercion co)
247 lvlExpr _ env (_, AnnVar v) = return (lookupVar env v)
248 lvlExpr _ _ (_, AnnLit lit) = return (Lit lit)
250 lvlExpr ctxt_lvl env expr@(_, AnnApp _ _) = do
252 (fun, args) = collectAnnArgs expr
255 -- float out partial applications. This is very beneficial
256 -- in some cases (-7% runtime -4% alloc over nofib -O2).
257 -- In order to float a PAP, there must be a function at the
258 -- head of the application, and the application must be
259 -- over-saturated with respect to the function's arity.
260 (_, AnnVar f) | floatPAPs env &&
261 arity > 0 && arity < n_val_args ->
263 let (lapp, rargs) = left (n_val_args - arity) expr []
264 rargs' <- mapM (lvlMFE False ctxt_lvl env) rargs
265 lapp' <- lvlMFE False ctxt_lvl env lapp
266 return (foldl App lapp' rargs')
268 n_val_args = count (isValArg . deAnnotate) args
271 -- separate out the PAP that we are floating from the extra
272 -- arguments, by traversing the spine until we have collected
273 -- (n_val_args - arity) value arguments.
274 left 0 e rargs = (e, rargs)
275 left n (_, AnnApp f a) rargs
276 | isValArg (deAnnotate a) = left (n-1) f (a:rargs)
277 | otherwise = left n f (a:rargs)
278 left _ _ _ = panic "SetLevels.lvlExpr.left"
280 -- No PAPs that we can float: just carry on with the
281 -- arguments and the function.
283 args' <- mapM (lvlMFE False ctxt_lvl env) args
284 fun' <- lvlExpr ctxt_lvl env fun
285 return (foldl App fun' args')
287 lvlExpr ctxt_lvl env (_, AnnNote note expr) = do
288 expr' <- lvlExpr ctxt_lvl env expr
289 return (Note note expr')
291 lvlExpr ctxt_lvl env (_, AnnCast expr (_, co)) = do
292 expr' <- lvlExpr ctxt_lvl env expr
293 return (Cast expr' co)
295 -- We don't split adjacent lambdas. That is, given
297 -- we don't float to give
298 -- \x -> let v = x+y in \y -> (v,y)
299 -- Why not? Because partial applications are fairly rare, and splitting
300 -- lambdas makes them more expensive.
302 lvlExpr ctxt_lvl env expr@(_, AnnLam {}) = do
303 new_body <- lvlMFE True new_lvl new_env body
304 return (mkLams new_bndrs new_body)
306 (bndrs, body) = collectAnnBndrs expr
307 (new_lvl, new_bndrs) = lvlLamBndrs ctxt_lvl bndrs
308 new_env = extendLvlEnv env new_bndrs
309 -- At one time we called a special verion of collectBinders,
310 -- which ignored coercions, because we don't want to split
311 -- a lambda like this (\x -> coerce t (\s -> ...))
312 -- This used to happen quite a bit in state-transformer programs,
313 -- but not nearly so much now non-recursive newtypes are transparent.
314 -- [See SetLevels rev 1.50 for a version with this approach.]
316 lvlExpr ctxt_lvl env (_, AnnLet (AnnNonRec bndr rhs) body)
317 | isUnLiftedType (idType bndr) = do
318 -- Treat unlifted let-bindings (let x = b in e) just like (case b of x -> e)
319 -- That is, leave it exactly where it is
320 -- We used to float unlifted bindings too (e.g. to get a cheap primop
321 -- outside a lambda (to see how, look at lvlBind in rev 1.58)
322 -- but an unrelated change meant that these unlifed bindings
323 -- could get to the top level which is bad. And there's not much point;
324 -- unlifted bindings are always cheap, and so hardly worth floating.
325 rhs' <- lvlExpr ctxt_lvl env rhs
326 body' <- lvlExpr incd_lvl env' body
327 return (Let (NonRec bndr' rhs') body')
329 incd_lvl = incMinorLvl ctxt_lvl
330 bndr' = TB bndr incd_lvl
331 env' = extendLvlEnv env [bndr']
333 lvlExpr ctxt_lvl env (_, AnnLet bind body) = do
334 (bind', new_env) <- lvlBind NotTopLevel ctxt_lvl env bind
335 body' <- lvlExpr ctxt_lvl new_env body
336 return (Let bind' body')
338 lvlExpr ctxt_lvl env (_, AnnCase expr case_bndr ty alts) = do
339 expr' <- lvlMFE True ctxt_lvl env expr
340 let alts_env = extendCaseBndrLvlEnv env expr' case_bndr incd_lvl
341 alts' <- mapM (lvl_alt alts_env) alts
342 return (Case expr' (TB case_bndr incd_lvl) ty alts')
344 incd_lvl = incMinorLvl ctxt_lvl
346 lvl_alt alts_env (con, bs, rhs) = do
347 rhs' <- lvlMFE True incd_lvl new_env rhs
348 return (con, bs', rhs')
350 bs' = [ TB b incd_lvl | b <- bs ]
351 new_env = extendLvlEnv alts_env bs'
354 @lvlMFE@ is just like @lvlExpr@, except that it might let-bind
355 the expression, so that it can itself be floated.
359 We don't float unlifted MFEs, which potentially loses big opportunites.
362 where h :: Int -> Int# is expensive. We'd like to float the (h y) outside
363 the \x, but we don't because it's unboxed. Possible solution: box it.
365 Note [Bottoming floats]
366 ~~~~~~~~~~~~~~~~~~~~~~~
368 f = \x. g (error "urk")
369 we'd like to float the call to error, to get
372 Furthermore, we want to float a bottoming expression even if it has free
374 f = \x. g (let v = h x in error ("urk" ++ v))
375 Then we'd like to abstact over 'x' can float the whole arg of g:
376 lvl = \x. let v = h x in error ("urk" ++ v)
378 See Maessen's paper 1999 "Bottom extraction: factoring error handling out
379 of functional programs" (unpublished I think).
381 When we do this, we set the strictness and arity of the new bottoming
382 Id, so that it's properly exposed as such in the interface file, even if
383 this is all happening after strictness analysis.
385 Note [Bottoming floats: eta expansion] c.f Note [Bottoming floats]
386 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
387 Tiresomely, though, the simplifier has an invariant that the manifest
388 arity of the RHS should be the same as the arity; but we can't call
389 etaExpand during SetLevels because it works over a decorated form of
390 CoreExpr. So we do the eta expansion later, in FloatOut.
394 We don't float a case expression as an MFE from a strict context. Why not?
395 Because in doing so we share a tiny bit of computation (the switch) but
396 in exchange we build a thunk, which is bad. This case reduces allocation
397 by 7% in spectral/puzzle (a rather strange benchmark) and 1.2% in real/fem.
398 Doesn't change any other allocation at all.
401 lvlMFE :: Bool -- True <=> strict context [body of case or let]
402 -> Level -- Level of innermost enclosing lambda/tylam
403 -> LevelEnv -- Level of in-scope names/tyvars
404 -> CoreExprWithFVs -- input expression
405 -> LvlM LevelledExpr -- Result expression
407 lvlMFE _ _ _ (_, AnnType ty)
410 -- No point in floating out an expression wrapped in a coercion or note
411 -- If we do we'll transform lvl = e |> co
412 -- to lvl' = e; lvl = lvl' |> co
413 -- and then inline lvl. Better just to float out the payload.
414 lvlMFE strict_ctxt ctxt_lvl env (_, AnnNote n e)
415 = do { e' <- lvlMFE strict_ctxt ctxt_lvl env e
416 ; return (Note n e') }
418 lvlMFE strict_ctxt ctxt_lvl env (_, AnnCast e (_, co))
419 = do { e' <- lvlMFE strict_ctxt ctxt_lvl env e
420 ; return (Cast e' co) }
423 lvlMFE True ctxt_lvl env e@(_, AnnCase {})
424 = lvlExpr ctxt_lvl env e -- Don't share cases
426 lvlMFE strict_ctxt ctxt_lvl env ann_expr@(fvs, _)
427 | isUnLiftedType ty -- Can't let-bind it; see Note [Unlifted MFEs]
428 -- This includes coercions, which we don't
429 -- want to float anyway
430 || notWorthFloating ann_expr abs_vars
431 || not good_destination
432 = -- Don't float it out
433 lvlExpr ctxt_lvl env ann_expr
435 | otherwise -- Float it out!
436 = do expr' <- lvlFloatRhs abs_vars dest_lvl env ann_expr
437 var <- newLvlVar abs_vars ty mb_bot
438 return (Let (NonRec (TB var dest_lvl) expr')
439 (mkVarApps (Var var) abs_vars))
441 expr = deAnnotate ann_expr
443 mb_bot = exprBotStrictness_maybe expr
444 dest_lvl = destLevel env fvs (isFunction ann_expr) mb_bot
445 abs_vars = abstractVars dest_lvl env fvs
447 -- A decision to float entails let-binding this thing, and we only do
448 -- that if we'll escape a value lambda, or will go to the top level.
450 | dest_lvl `ltMajLvl` ctxt_lvl -- Escapes a value lambda
452 -- OLD CODE: not (exprIsCheap expr) || isTopLvl dest_lvl
453 -- see Note [Escaping a value lambda]
455 | otherwise -- Does not escape a value lambda
456 = isTopLvl dest_lvl -- Only float if we are going to the top level
457 && floatConsts env -- and the floatConsts flag is on
458 && not strict_ctxt -- Don't float from a strict context
459 -- We are keen to float something to the top level, even if it does not
460 -- escape a lambda, because then it needs no allocation. But it's controlled
461 -- by a flag, because doing this too early loses opportunities for RULES
462 -- which (needless to say) are important in some nofib programs
463 -- (gcd is an example).
466 -- concat = /\ a -> foldr ..a.. (++) []
467 -- was getting turned into
468 -- concat = /\ a -> lvl a
469 -- lvl = /\ a -> foldr ..a.. (++) []
470 -- which is pretty stupid. Hence the strict_ctxt test
472 annotateBotStr :: Id -> Maybe (Arity, StrictSig) -> Id
473 annotateBotStr id Nothing = id
474 annotateBotStr id (Just (arity,sig)) = id `setIdArity` arity
475 `setIdStrictness` sig
477 notWorthFloating :: CoreExprWithFVs -> [Var] -> Bool
478 -- Returns True if the expression would be replaced by
479 -- something bigger than it is now. For example:
480 -- abs_vars = tvars only: return True if e is trivial,
481 -- but False for anything bigger
482 -- abs_vars = [x] (an Id): return True for trivial, or an application (f x)
483 -- but False for (f x x)
485 -- One big goal is that floating should be idempotent. Eg if
486 -- we replace e with (lvl79 x y) and then run FloatOut again, don't want
487 -- to replace (lvl79 x y) with (lvl83 x y)!
489 notWorthFloating e abs_vars
490 = go e (count isId abs_vars)
492 go (_, AnnVar {}) n = n >= 0
493 go (_, AnnLit {}) n = n >= 0
494 go (_, AnnCast e _) n = go e n
495 go (_, AnnApp e arg) n
496 | (_, AnnType {}) <- arg = go e n
497 | (_, AnnCoercion {}) <- arg = go e n
499 | is_triv arg = go e (n-1)
503 is_triv (_, AnnLit {}) = True -- Treat all literals as trivial
504 is_triv (_, AnnVar {}) = True -- (ie not worth floating)
505 is_triv (_, AnnCast e _) = is_triv e
506 is_triv (_, AnnApp e (_, AnnType {})) = is_triv e
507 is_triv (_, AnnApp e (_, AnnCoercion {})) = is_triv e
511 Note [Escaping a value lambda]
512 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
513 We want to float even cheap expressions out of value lambdas,
514 because that saves allocation. Consider
515 f = \x. .. (\y.e) ...
516 Then we'd like to avoid allocating the (\y.e) every time we call f,
517 (assuming e does not mention x).
519 An example where this really makes a difference is simplrun009.
521 Another reason it's good is because it makes SpecContr fire on functions.
523 f = \x. ....(f (\y.e))....
524 After floating we get
526 f = \x. ....(f lvl)...
527 and that is much easier for SpecConstr to generate a robust specialisation for.
529 The OLD CODE (given where this Note is referred to) prevents floating
530 of the example above, so I just don't understand the old code. I
531 don't understand the old comment either (which appears below). I
532 measured the effect on nofib of changing OLD CODE to 'True', and got
533 zeros everywhere, but a 4% win for 'puzzle'. Very small 0.5% loss for
534 'cse'; turns out to be because our arity analysis isn't good enough
535 yet (mentioned in Simon-nofib-notes).
538 Even if it escapes a value lambda, we only
539 float if it's not cheap (unless it'll get all the
540 way to the top). I've seen cases where we
541 float dozens of tiny free expressions, which cost
542 more to allocate than to evaluate.
543 NB: exprIsCheap is also true of bottom expressions, which
544 is good; we don't want to share them
546 It's only Really Bad to float a cheap expression out of a
547 strict context, because that builds a thunk that otherwise
548 would never be built. So another alternative would be to
550 || (strict_ctxt && not (exprIsBottom expr))
551 to the condition above. We should really try this out.
554 %************************************************************************
556 \subsection{Bindings}
558 %************************************************************************
560 The binding stuff works for top level too.
563 lvlBind :: TopLevelFlag -- Used solely to decide whether to clone
564 -> Level -- Context level; might be Top even for bindings nested in the RHS
565 -- of a top level binding
568 -> LvlM (LevelledBind, LevelEnv)
570 lvlBind top_lvl ctxt_lvl env (AnnNonRec bndr rhs@(rhs_fvs,_))
571 | isTyVar bndr -- Don't do anything for TyVar binders
572 -- (simplifier gets rid of them pronto)
573 = do rhs' <- lvlExpr ctxt_lvl env rhs
574 return (NonRec (TB bndr ctxt_lvl) rhs', env)
577 = do -- No type abstraction; clone existing binder
578 rhs' <- lvlExpr dest_lvl env rhs
579 (env', bndr') <- cloneVar top_lvl env bndr ctxt_lvl dest_lvl
580 return (NonRec (TB bndr' dest_lvl) rhs', env')
583 = do -- Yes, type abstraction; create a new binder, extend substitution, etc
584 rhs' <- lvlFloatRhs abs_vars dest_lvl env rhs
585 (env', [bndr']) <- newPolyBndrs dest_lvl env abs_vars [bndr_w_str]
586 return (NonRec (TB bndr' dest_lvl) rhs', env')
589 bind_fvs = rhs_fvs `unionVarSet` idFreeVars bndr
590 abs_vars = abstractVars dest_lvl env bind_fvs
591 dest_lvl = destLevel env bind_fvs (isFunction rhs) mb_bot
592 mb_bot = exprBotStrictness_maybe (deAnnotate rhs)
593 bndr_w_str = annotateBotStr bndr mb_bot
598 lvlBind top_lvl ctxt_lvl env (AnnRec pairs)
600 = do (new_env, new_bndrs) <- cloneRecVars top_lvl env bndrs ctxt_lvl dest_lvl
601 new_rhss <- mapM (lvlExpr ctxt_lvl new_env) rhss
602 return (Rec ([TB b dest_lvl | b <- new_bndrs] `zip` new_rhss), new_env)
604 -- ToDo: when enabling the floatLambda stuff,
605 -- I think we want to stop doing this
606 | isSingleton pairs && count isId abs_vars > 1
607 = do -- Special case for self recursion where there are
608 -- several variables carried around: build a local loop:
609 -- poly_f = \abs_vars. \lam_vars . letrec f = \lam_vars. rhs in f lam_vars
610 -- This just makes the closures a bit smaller. If we don't do
611 -- this, allocation rises significantly on some programs
613 -- We could elaborate it for the case where there are several
614 -- mutually functions, but it's quite a bit more complicated
616 -- This all seems a bit ad hoc -- sigh
618 (bndr,rhs) = head pairs
619 (rhs_lvl, abs_vars_w_lvls) = lvlLamBndrs dest_lvl abs_vars
620 rhs_env = extendLvlEnv env abs_vars_w_lvls
621 (rhs_env', new_bndr) <- cloneVar NotTopLevel rhs_env bndr rhs_lvl rhs_lvl
623 (lam_bndrs, rhs_body) = collectAnnBndrs rhs
624 (body_lvl, new_lam_bndrs) = lvlLamBndrs rhs_lvl lam_bndrs
625 body_env = extendLvlEnv rhs_env' new_lam_bndrs
626 new_rhs_body <- lvlExpr body_lvl body_env rhs_body
627 (poly_env, [poly_bndr]) <- newPolyBndrs dest_lvl env abs_vars [bndr]
628 return (Rec [(TB poly_bndr dest_lvl,
629 mkLams abs_vars_w_lvls $
630 mkLams new_lam_bndrs $
631 Let (Rec [(TB new_bndr rhs_lvl, mkLams new_lam_bndrs new_rhs_body)])
632 (mkVarApps (Var new_bndr) lam_bndrs))],
635 | otherwise = do -- Non-null abs_vars
636 (new_env, new_bndrs) <- newPolyBndrs dest_lvl env abs_vars bndrs
637 new_rhss <- mapM (lvlFloatRhs abs_vars dest_lvl new_env) rhss
638 return (Rec ([TB b dest_lvl | b <- new_bndrs] `zip` new_rhss), new_env)
641 (bndrs,rhss) = unzip pairs
643 -- Finding the free vars of the binding group is annoying
644 bind_fvs = (unionVarSets [ idFreeVars bndr `unionVarSet` rhs_fvs
645 | (bndr, (rhs_fvs,_)) <- pairs])
649 dest_lvl = destLevel env bind_fvs (all isFunction rhss) Nothing
650 abs_vars = abstractVars dest_lvl env bind_fvs
652 ----------------------------------------------------
653 -- Three help functions for the type-abstraction case
655 lvlFloatRhs :: [CoreBndr] -> Level -> LevelEnv -> CoreExprWithFVs
656 -> UniqSM (Expr (TaggedBndr Level))
657 lvlFloatRhs abs_vars dest_lvl env rhs = do
658 rhs' <- lvlExpr rhs_lvl rhs_env rhs
659 return (mkLams abs_vars_w_lvls rhs')
661 (rhs_lvl, abs_vars_w_lvls) = lvlLamBndrs dest_lvl abs_vars
662 rhs_env = extendLvlEnv env abs_vars_w_lvls
666 %************************************************************************
668 \subsection{Deciding floatability}
670 %************************************************************************
673 lvlLamBndrs :: Level -> [CoreBndr] -> (Level, [TaggedBndr Level])
674 -- Compute the levels for the binders of a lambda group
675 -- The binders returned are exactly the same as the ones passed,
676 -- but they are now paired with a level
680 lvlLamBndrs lvl bndrs
681 = (new_lvl, [TB bndr new_lvl | bndr <- bndrs])
682 -- All the new binders get the same level, because
683 -- any floating binding is either going to float past
684 -- all or none. We never separate binders
686 new_lvl | any is_major bndrs = incMajorLvl lvl
687 | otherwise = incMinorLvl lvl
689 is_major bndr = isId bndr && not (isOneShotLambda bndr)
693 -- Destintion level is the max Id level of the expression
694 -- (We'll abstract the type variables, if any.)
695 destLevel :: LevelEnv -> VarSet -> Bool -> Maybe (Arity, StrictSig) -> Level
696 destLevel env fvs is_function mb_bot
697 | Just {} <- mb_bot = tOP_LEVEL -- Send bottoming bindings to the top
698 -- regardless; see Note [Bottoming floats]
699 | Just n_args <- floatLams env
700 , n_args > 0 -- n=0 case handled uniformly by the 'otherwise' case
702 , countFreeIds fvs <= n_args
703 = tOP_LEVEL -- Send functions to top level; see
704 -- the comments with isFunction
705 | otherwise = maxIdLevel env fvs
707 isFunction :: CoreExprWithFVs -> Bool
708 -- The idea here is that we want to float *functions* to
709 -- the top level. This saves no work, but
710 -- (a) it can make the host function body a lot smaller,
711 -- and hence inlinable.
712 -- (b) it can also save allocation when the function is recursive:
713 -- h = \x -> letrec f = \y -> ...f...y...x...
716 -- f = \x y -> ...(f x)...y...x...
718 -- No allocation for f now.
719 -- We may only want to do this if there are sufficiently few free
720 -- variables. We certainly only want to do it for values, and not for
721 -- constructors. So the simple thing is just to look for lambdas
722 isFunction (_, AnnLam b e) | isId b = True
723 | otherwise = isFunction e
724 isFunction (_, AnnNote _ e) = isFunction e
727 countFreeIds :: VarSet -> Int
728 countFreeIds = foldVarSet add 0
730 add :: Var -> Int -> Int
731 add v n | isId v = n+1
736 %************************************************************************
738 \subsection{Free-To-Level Monad}
740 %************************************************************************
744 = LE { le_switches :: FloatOutSwitches
745 , le_lvl_env :: VarEnv Level -- Domain is *post-cloned* TyVars and Ids
746 , le_subst :: Subst -- Domain is pre-cloned Ids; tracks the in-scope set
747 -- so that subtitution is capture-avoiding
748 , le_env :: IdEnv ([Var], LevelledExpr) -- Domain is pre-cloned Ids
750 -- We clone let-bound variables so that they are still
751 -- distinct when floated out; hence the le_subst/le_env.
752 -- (see point 3 of the module overview comment).
753 -- We also use these envs when making a variable polymorphic
754 -- because we want to float it out past a big lambda.
756 -- The le_subst and le_env always implement the same mapping, but the
757 -- le_subst maps to CoreExpr and the le_env to LevelledExpr
758 -- Since the range is always a variable or type application,
759 -- there is never any difference between the two, but sadly
760 -- the types differ. The le_subst is used when substituting in
761 -- a variable's IdInfo; the le_env when we find a Var.
763 -- In addition the le_env records a list of tyvars free in the
764 -- type application, just so we don't have to call freeVars on
765 -- the type application repeatedly.
767 -- The domain of the both envs is *pre-cloned* Ids, though
769 -- The domain of the le_lvl_env is the *post-cloned* Ids
771 initialEnv :: FloatOutSwitches -> LevelEnv
772 initialEnv float_lams
773 = LE { le_switches = float_lams, le_lvl_env = emptyVarEnv
774 , le_subst = emptySubst, le_env = emptyVarEnv }
776 floatLams :: LevelEnv -> Maybe Int
777 floatLams le = floatOutLambdas (le_switches le)
779 floatConsts :: LevelEnv -> Bool
780 floatConsts le = floatOutConstants (le_switches le)
782 floatPAPs :: LevelEnv -> Bool
783 floatPAPs le = floatOutPartialApplications (le_switches le)
785 extendLvlEnv :: LevelEnv -> [TaggedBndr Level] -> LevelEnv
786 -- Used when *not* cloning
787 extendLvlEnv le@(LE { le_lvl_env = lvl_env, le_subst = subst, le_env = id_env })
789 = le { le_lvl_env = foldl add_lvl lvl_env prs
790 , le_subst = foldl del_subst subst prs
791 , le_env = foldl del_id id_env prs }
793 add_lvl env (TB v l) = extendVarEnv env v l
794 del_subst env (TB v _) = extendInScope env v
795 del_id env (TB v _) = delVarEnv env v
796 -- We must remove any clone for this variable name in case of
797 -- shadowing. This bit me in the following case
798 -- (in nofib/real/gg/Spark.hs):
801 -- ... -> case e of wild {
802 -- ... -> ... wild ...
806 -- The inside occurrence of @wild@ was being replaced with @ds@,
807 -- incorrectly, because the SubstEnv was still lying around. Ouch!
810 extendInScopeEnv :: LevelEnv -> Var -> LevelEnv
811 extendInScopeEnv le@(LE { le_subst = subst }) v
812 = le { le_subst = extendInScope subst v }
814 extendInScopeEnvList :: LevelEnv -> [Var] -> LevelEnv
815 extendInScopeEnvList le@(LE { le_subst = subst }) vs
816 = le { le_subst = extendInScopeList subst vs }
818 -- extendCaseBndrLvlEnv adds the mapping case-bndr->scrut-var if it can
819 -- (see point 4 of the module overview comment)
820 extendCaseBndrLvlEnv :: LevelEnv -> Expr (TaggedBndr Level) -> Var -> Level
822 extendCaseBndrLvlEnv le@(LE { le_lvl_env = lvl_env, le_subst = subst, le_env = id_env })
823 (Var scrut_var) case_bndr lvl
824 = le { le_lvl_env = extendVarEnv lvl_env case_bndr lvl
825 , le_subst = extendIdSubst subst case_bndr (Var scrut_var)
826 , le_env = extendVarEnv id_env case_bndr ([scrut_var], Var scrut_var) }
828 extendCaseBndrLvlEnv env _scrut case_bndr lvl
829 = extendLvlEnv env [TB case_bndr lvl]
831 extendPolyLvlEnv :: Level -> LevelEnv -> [Var] -> [(Var, Var)] -> LevelEnv
832 extendPolyLvlEnv dest_lvl
833 le@(LE { le_lvl_env = lvl_env, le_subst = subst, le_env = id_env })
835 = le { le_lvl_env = foldl add_lvl lvl_env bndr_pairs
836 , le_subst = foldl add_subst subst bndr_pairs
837 , le_env = foldl add_id id_env bndr_pairs }
839 add_lvl env (_, v') = extendVarEnv env v' dest_lvl
840 add_subst env (v, v') = extendIdSubst env v (mkVarApps (Var v') abs_vars)
841 add_id env (v, v') = extendVarEnv env v ((v':abs_vars), mkVarApps (Var v') abs_vars)
843 extendCloneLvlEnv :: Level -> LevelEnv -> Subst -> [(Var, Var)] -> LevelEnv
844 extendCloneLvlEnv lvl le@(LE { le_lvl_env = lvl_env, le_env = id_env })
846 = le { le_lvl_env = foldl add_lvl lvl_env bndr_pairs
847 , le_subst = new_subst
848 , le_env = foldl add_id id_env bndr_pairs }
850 add_lvl env (_, v') = extendVarEnv env v' lvl
851 add_id env (v, v') = extendVarEnv env v ([v'], Var v')
853 maxIdLevel :: LevelEnv -> VarSet -> Level
854 maxIdLevel (LE { le_lvl_env = lvl_env, le_env = id_env }) var_set
855 = foldVarSet max_in tOP_LEVEL var_set
857 max_in in_var lvl = foldr max_out lvl (case lookupVarEnv id_env in_var of
858 Just (abs_vars, _) -> abs_vars
862 | isId out_var = case lookupVarEnv lvl_env out_var of
863 Just lvl' -> maxLvl lvl' lvl
865 | otherwise = lvl -- Ignore tyvars in *maxIdLevel*
867 lookupVar :: LevelEnv -> Id -> LevelledExpr
868 lookupVar le v = case lookupVarEnv (le_env le) v of
869 Just (_, expr) -> expr
872 abstractVars :: Level -> LevelEnv -> VarSet -> [Var]
873 -- Find the variables in fvs, free vars of the target expresion,
874 -- whose level is greater than the destination level
875 -- These are the ones we are going to abstract out
876 abstractVars dest_lvl (LE { le_lvl_env = lvl_env, le_env = id_env }) fvs
877 = map zap $ uniq $ sortLe le
878 [var | fv <- varSetElems fvs
879 , var <- absVarsOf id_env fv
881 -- NB: it's important to call abstract_me only on the OutIds the
882 -- come from absVarsOf (not on fv, which is an InId)
884 -- Sort the variables so the true type variables come first;
885 -- the tyvars scope over Ids and coercion vars
886 v1 `le` v2 = case (is_tv v1, is_tv v2) of
887 (True, False) -> True
888 (False, True) -> False
889 _ -> v1 <= v2 -- Same family
893 uniq :: [Var] -> [Var]
894 -- Remove adjacent duplicates; the sort will have brought them together
895 uniq (v1:v2:vs) | v1 == v2 = uniq (v2:vs)
896 | otherwise = v1 : uniq (v2:vs)
899 abstract_me v = case lookupVarEnv lvl_env v of
900 Just lvl -> dest_lvl `ltLvl` lvl
903 -- We are going to lambda-abstract, so nuke any IdInfo,
904 -- and add the tyvars of the Id (if necessary)
905 zap v | isId v = WARN( isStableUnfolding (idUnfolding v) ||
906 not (isEmptySpecInfo (idSpecialisation v)),
907 text "absVarsOf: discarding info on" <+> ppr v )
908 setIdInfo v vanillaIdInfo
911 absVarsOf :: IdEnv ([Var], LevelledExpr) -> Var -> [Var]
912 -- If f is free in the expression, and f maps to poly_f a b c in the
913 -- current substitution, then we must report a b c as candidate type
916 -- Also, if x::a is an abstracted variable, then so is a; that is,
917 -- we must look in x's type
918 -- And similarly if x is a coercion variable.
920 | isId v = [av2 | av1 <- lookup_avs v
921 , av2 <- add_tyvars av1]
922 | otherwise = ASSERT( isTyVar v ) [v]
924 lookup_avs v = case lookupVarEnv id_env v of
925 Just (abs_vars, _) -> abs_vars
928 add_tyvars v = v : varSetElems (varTypeTyVars v)
932 type LvlM result = UniqSM result
934 initLvl :: UniqSupply -> UniqSM a -> a
940 newPolyBndrs :: Level -> LevelEnv -> [Var] -> [Id] -> UniqSM (LevelEnv, [Id])
941 newPolyBndrs dest_lvl env abs_vars bndrs = do
943 let new_bndrs = zipWith mk_poly_bndr bndrs uniqs
944 return (extendPolyLvlEnv dest_lvl env abs_vars (bndrs `zip` new_bndrs), new_bndrs)
946 mk_poly_bndr bndr uniq = transferPolyIdInfo bndr abs_vars $ -- Note [transferPolyIdInfo] in Id.lhs
947 mkSysLocal (mkFastString str) uniq poly_ty
949 str = "poly_" ++ occNameString (getOccName bndr)
950 poly_ty = mkPiTypes abs_vars (idType bndr)
952 newLvlVar :: [CoreBndr] -> Type -- Abstract wrt these bndrs
953 -> Maybe (Arity, StrictSig) -- Note [Bottoming floats]
955 newLvlVar vars body_ty mb_bot
956 = do { uniq <- getUniqueM
957 ; return (mkLocalIdWithInfo (mk_name uniq) (mkPiTypes vars body_ty) info) }
959 mk_name uniq = mkSystemVarName uniq (mkFastString "lvl")
960 arity = count isId vars
961 info = case mb_bot of
962 Nothing -> vanillaIdInfo
963 Just (bot_arity, sig) -> vanillaIdInfo
964 `setArityInfo` (arity + bot_arity)
965 `setStrictnessInfo` Just (increaseStrictSigArity arity sig)
967 -- The deeply tiresome thing is that we have to apply the substitution
968 -- to the rules inside each Id. Grr. But it matters.
970 cloneVar :: TopLevelFlag -> LevelEnv -> Id -> Level -> Level -> LvlM (LevelEnv, Id)
971 cloneVar TopLevel env v _ _
972 = return (extendInScopeEnv env v, v) -- Don't clone top level things
973 -- But do extend the in-scope env, to satisfy the in-scope invariant
975 cloneVar NotTopLevel env v ctxt_lvl dest_lvl
976 = ASSERT( isId v ) do
977 us <- getUniqueSupplyM
979 (subst', v1) = cloneIdBndr (le_subst env) us v
980 v2 = zap_demand ctxt_lvl dest_lvl v1
981 env' = extendCloneLvlEnv dest_lvl env subst' [(v,v2)]
984 cloneRecVars :: TopLevelFlag -> LevelEnv -> [Id] -> Level -> Level -> LvlM (LevelEnv, [Id])
985 cloneRecVars TopLevel env vs _ _
986 = return (extendInScopeEnvList env vs, vs) -- Don't clone top level things
987 cloneRecVars NotTopLevel env vs ctxt_lvl dest_lvl
988 = ASSERT( all isId vs ) do
989 us <- getUniqueSupplyM
991 (subst', vs1) = cloneRecIdBndrs (le_subst env) us vs
992 vs2 = map (zap_demand ctxt_lvl dest_lvl) vs1
993 env' = extendCloneLvlEnv dest_lvl env subst' (vs `zip` vs2)
996 -- VERY IMPORTANT: we must zap the demand info
997 -- if the thing is going to float out past a lambda,
998 -- or if it's going to top level (where things can't be strict)
999 zap_demand :: Level -> Level -> Id -> Id
1000 zap_demand dest_lvl ctxt_lvl id
1001 | ctxt_lvl == dest_lvl,
1002 not (isTopLvl dest_lvl) = id -- Stays, and not going to top level
1003 | otherwise = zapDemandIdInfo id -- Floats out