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 _ env (_, AnnVar v) = return (lookupVar env v)
247 lvlExpr _ _ (_, AnnLit lit) = return (Lit lit)
249 lvlExpr ctxt_lvl env expr@(_, AnnApp _ _) = do
251 (fun, args) = collectAnnArgs expr
254 -- float out partial applications. This is very beneficial
255 -- in some cases (-7% runtime -4% alloc over nofib -O2).
256 -- In order to float a PAP, there must be a function at the
257 -- head of the application, and the application must be
258 -- over-saturated with respect to the function's arity.
259 (_, AnnVar f) | floatPAPs env &&
260 arity > 0 && arity < n_val_args ->
262 let (lapp, rargs) = left (n_val_args - arity) expr []
263 rargs' <- mapM (lvlMFE False ctxt_lvl env) rargs
264 lapp' <- lvlMFE False ctxt_lvl env lapp
265 return (foldl App lapp' rargs')
267 n_val_args = count (isValArg . deAnnotate) args
270 -- separate out the PAP that we are floating from the extra
271 -- arguments, by traversing the spine until we have collected
272 -- (n_val_args - arity) value arguments.
273 left 0 e rargs = (e, rargs)
274 left n (_, AnnApp f a) rargs
275 | isValArg (deAnnotate a) = left (n-1) f (a:rargs)
276 | otherwise = left n f (a:rargs)
277 left _ _ _ = panic "SetLevels.lvlExpr.left"
279 -- No PAPs that we can float: just carry on with the
280 -- arguments and the function.
282 args' <- mapM (lvlMFE False ctxt_lvl env) args
283 fun' <- lvlExpr ctxt_lvl env fun
284 return (foldl App fun' args')
286 lvlExpr ctxt_lvl env (_, AnnNote note expr) = do
287 expr' <- lvlExpr ctxt_lvl env expr
288 return (Note note expr')
290 lvlExpr ctxt_lvl env (_, AnnCast expr co) = do
291 expr' <- lvlExpr ctxt_lvl env expr
292 return (Cast expr' co)
294 -- We don't split adjacent lambdas. That is, given
296 -- we don't float to give
297 -- \x -> let v = x+y in \y -> (v,y)
298 -- Why not? Because partial applications are fairly rare, and splitting
299 -- lambdas makes them more expensive.
301 lvlExpr ctxt_lvl env expr@(_, AnnLam {}) = do
302 new_body <- lvlMFE True new_lvl new_env body
303 return (mkLams new_bndrs new_body)
305 (bndrs, body) = collectAnnBndrs expr
306 (new_lvl, new_bndrs) = lvlLamBndrs ctxt_lvl bndrs
307 new_env = extendLvlEnv env new_bndrs
308 -- At one time we called a special verion of collectBinders,
309 -- which ignored coercions, because we don't want to split
310 -- a lambda like this (\x -> coerce t (\s -> ...))
311 -- This used to happen quite a bit in state-transformer programs,
312 -- but not nearly so much now non-recursive newtypes are transparent.
313 -- [See SetLevels rev 1.50 for a version with this approach.]
315 lvlExpr ctxt_lvl env (_, AnnLet (AnnNonRec bndr rhs) body)
316 | isUnLiftedType (idType bndr) = do
317 -- Treat unlifted let-bindings (let x = b in e) just like (case b of x -> e)
318 -- That is, leave it exactly where it is
319 -- We used to float unlifted bindings too (e.g. to get a cheap primop
320 -- outside a lambda (to see how, look at lvlBind in rev 1.58)
321 -- but an unrelated change meant that these unlifed bindings
322 -- could get to the top level which is bad. And there's not much point;
323 -- unlifted bindings are always cheap, and so hardly worth floating.
324 rhs' <- lvlExpr ctxt_lvl env rhs
325 body' <- lvlExpr incd_lvl env' body
326 return (Let (NonRec bndr' rhs') body')
328 incd_lvl = incMinorLvl ctxt_lvl
329 bndr' = TB bndr incd_lvl
330 env' = extendLvlEnv env [bndr']
332 lvlExpr ctxt_lvl env (_, AnnLet bind body) = do
333 (bind', new_env) <- lvlBind NotTopLevel ctxt_lvl env bind
334 body' <- lvlExpr ctxt_lvl new_env body
335 return (Let bind' body')
337 lvlExpr ctxt_lvl env (_, AnnCase expr case_bndr ty alts) = do
338 expr' <- lvlMFE True ctxt_lvl env expr
339 let alts_env = extendCaseBndrLvlEnv env expr' case_bndr incd_lvl
340 alts' <- mapM (lvl_alt alts_env) alts
341 return (Case expr' (TB case_bndr incd_lvl) ty alts')
343 incd_lvl = incMinorLvl ctxt_lvl
345 lvl_alt alts_env (con, bs, rhs) = do
346 rhs' <- lvlMFE True incd_lvl new_env rhs
347 return (con, bs', rhs')
349 bs' = [ TB b incd_lvl | b <- bs ]
350 new_env = extendLvlEnv alts_env bs'
353 @lvlMFE@ is just like @lvlExpr@, except that it might let-bind
354 the expression, so that it can itself be floated.
358 We don't float unlifted MFEs, which potentially loses big opportunites.
361 where h :: Int -> Int# is expensive. We'd like to float the (h y) outside
362 the \x, but we don't because it's unboxed. Possible solution: box it.
364 Note [Bottoming floats]
365 ~~~~~~~~~~~~~~~~~~~~~~~
367 f = \x. g (error "urk")
368 we'd like to float the call to error, to get
371 Furthermore, we want to float a bottoming expression even if it has free
373 f = \x. g (let v = h x in error ("urk" ++ v))
374 Then we'd like to abstact over 'x' can float the whole arg of g:
375 lvl = \x. let v = h x in error ("urk" ++ v)
377 See Maessen's paper 1999 "Bottom extraction: factoring error handling out
378 of functional programs" (unpublished I think).
380 When we do this, we set the strictness and arity of the new bottoming
381 Id, so that it's properly exposed as such in the interface file, even if
382 this is all happening after strictness analysis.
384 Note [Bottoming floats: eta expansion] c.f Note [Bottoming floats]
385 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
386 Tiresomely, though, the simplifier has an invariant that the manifest
387 arity of the RHS should be the same as the arity; but we can't call
388 etaExpand during SetLevels because it works over a decorated form of
389 CoreExpr. So we do the eta expansion later, in FloatOut.
393 We don't float a case expression as an MFE from a strict context. Why not?
394 Because in doing so we share a tiny bit of computation (the switch) but
395 in exchange we build a thunk, which is bad. This case reduces allocation
396 by 7% in spectral/puzzle (a rather strange benchmark) and 1.2% in real/fem.
397 Doesn't change any other allocation at all.
400 lvlMFE :: Bool -- True <=> strict context [body of case or let]
401 -> Level -- Level of innermost enclosing lambda/tylam
402 -> LevelEnv -- Level of in-scope names/tyvars
403 -> CoreExprWithFVs -- input expression
404 -> LvlM LevelledExpr -- Result expression
406 lvlMFE _ _ _ (_, AnnType ty)
409 -- No point in floating out an expression wrapped in a coercion or note
410 -- If we do we'll transform lvl = e |> co
411 -- to lvl' = e; lvl = lvl' |> co
412 -- and then inline lvl. Better just to float out the payload.
413 lvlMFE strict_ctxt ctxt_lvl env (_, AnnNote n e)
414 = do { e' <- lvlMFE strict_ctxt ctxt_lvl env e
415 ; return (Note n e') }
417 lvlMFE strict_ctxt ctxt_lvl env (_, AnnCast e co)
418 = do { e' <- lvlMFE strict_ctxt ctxt_lvl env e
419 ; return (Cast e' co) }
422 lvlMFE True ctxt_lvl env e@(_, AnnCase {})
423 = lvlExpr ctxt_lvl env e -- Don't share cases
425 lvlMFE strict_ctxt ctxt_lvl env ann_expr@(fvs, _)
426 | isUnLiftedType ty -- Can't let-bind it; see Note [Unlifted MFEs]
427 || notWorthFloating ann_expr abs_vars
428 || not good_destination
429 = -- Don't float it out
430 lvlExpr ctxt_lvl env ann_expr
432 | otherwise -- Float it out!
433 = do expr' <- lvlFloatRhs abs_vars dest_lvl env ann_expr
434 var <- newLvlVar abs_vars ty mb_bot
435 return (Let (NonRec (TB var dest_lvl) expr')
436 (mkVarApps (Var var) abs_vars))
438 expr = deAnnotate ann_expr
440 mb_bot = exprBotStrictness_maybe expr
441 dest_lvl = destLevel env fvs (isFunction ann_expr) mb_bot
442 abs_vars = abstractVars dest_lvl env fvs
444 -- A decision to float entails let-binding this thing, and we only do
445 -- that if we'll escape a value lambda, or will go to the top level.
447 | dest_lvl `ltMajLvl` ctxt_lvl -- Escapes a value lambda
449 -- OLD CODE: not (exprIsCheap expr) || isTopLvl dest_lvl
450 -- see Note [Escaping a value lambda]
452 | otherwise -- Does not escape a value lambda
453 = isTopLvl dest_lvl -- Only float if we are going to the top level
454 && floatConsts env -- and the floatConsts flag is on
455 && not strict_ctxt -- Don't float from a strict context
456 -- We are keen to float something to the top level, even if it does not
457 -- escape a lambda, because then it needs no allocation. But it's controlled
458 -- by a flag, because doing this too early loses opportunities for RULES
459 -- which (needless to say) are important in some nofib programs
460 -- (gcd is an example).
463 -- concat = /\ a -> foldr ..a.. (++) []
464 -- was getting turned into
465 -- concat = /\ a -> lvl a
466 -- lvl = /\ a -> foldr ..a.. (++) []
467 -- which is pretty stupid. Hence the strict_ctxt test
469 annotateBotStr :: Id -> Maybe (Arity, StrictSig) -> Id
470 annotateBotStr id Nothing = id
471 annotateBotStr id (Just (arity,sig)) = id `setIdArity` arity
472 `setIdStrictness` sig
474 notWorthFloating :: CoreExprWithFVs -> [Var] -> Bool
475 -- Returns True if the expression would be replaced by
476 -- something bigger than it is now. For example:
477 -- abs_vars = tvars only: return True if e is trivial,
478 -- but False for anything bigger
479 -- abs_vars = [x] (an Id): return True for trivial, or an application (f x)
480 -- but False for (f x x)
482 -- One big goal is that floating should be idempotent. Eg if
483 -- we replace e with (lvl79 x y) and then run FloatOut again, don't want
484 -- to replace (lvl79 x y) with (lvl83 x y)!
486 notWorthFloating e abs_vars
487 = go e (count isId abs_vars)
489 go (_, AnnVar {}) n = n >= 0
490 go (_, AnnLit {}) n = n >= 0
491 go (_, AnnCast e _) n = go e n
492 go (_, AnnApp e arg) n
493 | (_, AnnType {}) <- arg = go e n
495 | is_triv arg = go e (n-1)
499 is_triv (_, AnnLit {}) = True -- Treat all literals as trivial
500 is_triv (_, AnnVar {}) = True -- (ie not worth floating)
501 is_triv (_, AnnCast e _) = is_triv e
502 is_triv (_, AnnApp e (_, AnnType {})) = is_triv e
506 Note [Escaping a value lambda]
507 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
508 We want to float even cheap expressions out of value lambdas,
509 because that saves allocation. Consider
510 f = \x. .. (\y.e) ...
511 Then we'd like to avoid allocating the (\y.e) every time we call f,
512 (assuming e does not mention x).
514 An example where this really makes a difference is simplrun009.
516 Another reason it's good is because it makes SpecContr fire on functions.
518 f = \x. ....(f (\y.e))....
519 After floating we get
521 f = \x. ....(f lvl)...
522 and that is much easier for SpecConstr to generate a robust specialisation for.
524 The OLD CODE (given where this Note is referred to) prevents floating
525 of the example above, so I just don't understand the old code. I
526 don't understand the old comment either (which appears below). I
527 measured the effect on nofib of changing OLD CODE to 'True', and got
528 zeros everywhere, but a 4% win for 'puzzle'. Very small 0.5% loss for
529 'cse'; turns out to be because our arity analysis isn't good enough
530 yet (mentioned in Simon-nofib-notes).
533 Even if it escapes a value lambda, we only
534 float if it's not cheap (unless it'll get all the
535 way to the top). I've seen cases where we
536 float dozens of tiny free expressions, which cost
537 more to allocate than to evaluate.
538 NB: exprIsCheap is also true of bottom expressions, which
539 is good; we don't want to share them
541 It's only Really Bad to float a cheap expression out of a
542 strict context, because that builds a thunk that otherwise
543 would never be built. So another alternative would be to
545 || (strict_ctxt && not (exprIsBottom expr))
546 to the condition above. We should really try this out.
549 %************************************************************************
551 \subsection{Bindings}
553 %************************************************************************
555 The binding stuff works for top level too.
558 lvlBind :: TopLevelFlag -- Used solely to decide whether to clone
559 -> Level -- Context level; might be Top even for bindings nested in the RHS
560 -- of a top level binding
563 -> LvlM (LevelledBind, LevelEnv)
565 lvlBind top_lvl ctxt_lvl env (AnnNonRec bndr rhs@(rhs_fvs,_))
566 | isTyCoVar bndr -- Don't do anything for TyVar binders
567 -- (simplifier gets rid of them pronto)
568 = do rhs' <- lvlExpr ctxt_lvl env rhs
569 return (NonRec (TB bndr ctxt_lvl) rhs', env)
572 = do -- No type abstraction; clone existing binder
573 rhs' <- lvlExpr dest_lvl env rhs
574 (env', bndr') <- cloneVar top_lvl env bndr ctxt_lvl dest_lvl
575 return (NonRec (TB bndr' dest_lvl) rhs', env')
578 = do -- Yes, type abstraction; create a new binder, extend substitution, etc
579 rhs' <- lvlFloatRhs abs_vars dest_lvl env rhs
580 (env', [bndr']) <- newPolyBndrs dest_lvl env abs_vars [bndr_w_str]
581 return (NonRec (TB bndr' dest_lvl) rhs', env')
584 bind_fvs = rhs_fvs `unionVarSet` idFreeVars bndr
585 abs_vars = abstractVars dest_lvl env bind_fvs
586 dest_lvl = destLevel env bind_fvs (isFunction rhs) mb_bot
587 mb_bot = exprBotStrictness_maybe (deAnnotate rhs)
588 bndr_w_str = annotateBotStr bndr mb_bot
593 lvlBind top_lvl ctxt_lvl env (AnnRec pairs)
595 = do (new_env, new_bndrs) <- cloneRecVars top_lvl env bndrs ctxt_lvl dest_lvl
596 new_rhss <- mapM (lvlExpr ctxt_lvl new_env) rhss
597 return (Rec ([TB b dest_lvl | b <- new_bndrs] `zip` new_rhss), new_env)
599 -- ToDo: when enabling the floatLambda stuff,
600 -- I think we want to stop doing this
601 | isSingleton pairs && count isId abs_vars > 1
602 = do -- Special case for self recursion where there are
603 -- several variables carried around: build a local loop:
604 -- poly_f = \abs_vars. \lam_vars . letrec f = \lam_vars. rhs in f lam_vars
605 -- This just makes the closures a bit smaller. If we don't do
606 -- this, allocation rises significantly on some programs
608 -- We could elaborate it for the case where there are several
609 -- mutually functions, but it's quite a bit more complicated
611 -- This all seems a bit ad hoc -- sigh
613 (bndr,rhs) = head pairs
614 (rhs_lvl, abs_vars_w_lvls) = lvlLamBndrs dest_lvl abs_vars
615 rhs_env = extendLvlEnv env abs_vars_w_lvls
616 (rhs_env', new_bndr) <- cloneVar NotTopLevel rhs_env bndr rhs_lvl rhs_lvl
618 (lam_bndrs, rhs_body) = collectAnnBndrs rhs
619 (body_lvl, new_lam_bndrs) = lvlLamBndrs rhs_lvl lam_bndrs
620 body_env = extendLvlEnv rhs_env' new_lam_bndrs
621 new_rhs_body <- lvlExpr body_lvl body_env rhs_body
622 (poly_env, [poly_bndr]) <- newPolyBndrs dest_lvl env abs_vars [bndr]
623 return (Rec [(TB poly_bndr dest_lvl,
624 mkLams abs_vars_w_lvls $
625 mkLams new_lam_bndrs $
626 Let (Rec [(TB new_bndr rhs_lvl, mkLams new_lam_bndrs new_rhs_body)])
627 (mkVarApps (Var new_bndr) lam_bndrs))],
630 | otherwise = do -- Non-null abs_vars
631 (new_env, new_bndrs) <- newPolyBndrs dest_lvl env abs_vars bndrs
632 new_rhss <- mapM (lvlFloatRhs abs_vars dest_lvl new_env) rhss
633 return (Rec ([TB b dest_lvl | b <- new_bndrs] `zip` new_rhss), new_env)
636 (bndrs,rhss) = unzip pairs
638 -- Finding the free vars of the binding group is annoying
639 bind_fvs = (unionVarSets [ idFreeVars bndr `unionVarSet` rhs_fvs
640 | (bndr, (rhs_fvs,_)) <- pairs])
644 dest_lvl = destLevel env bind_fvs (all isFunction rhss) Nothing
645 abs_vars = abstractVars dest_lvl env bind_fvs
647 ----------------------------------------------------
648 -- Three help functions for the type-abstraction case
650 lvlFloatRhs :: [CoreBndr] -> Level -> LevelEnv -> CoreExprWithFVs
651 -> UniqSM (Expr (TaggedBndr Level))
652 lvlFloatRhs abs_vars dest_lvl env rhs = do
653 rhs' <- lvlExpr rhs_lvl rhs_env rhs
654 return (mkLams abs_vars_w_lvls rhs')
656 (rhs_lvl, abs_vars_w_lvls) = lvlLamBndrs dest_lvl abs_vars
657 rhs_env = extendLvlEnv env abs_vars_w_lvls
661 %************************************************************************
663 \subsection{Deciding floatability}
665 %************************************************************************
668 lvlLamBndrs :: Level -> [CoreBndr] -> (Level, [TaggedBndr Level])
669 -- Compute the levels for the binders of a lambda group
670 -- The binders returned are exactly the same as the ones passed,
671 -- but they are now paired with a level
675 lvlLamBndrs lvl bndrs
676 = (new_lvl, [TB bndr new_lvl | bndr <- bndrs])
677 -- All the new binders get the same level, because
678 -- any floating binding is either going to float past
679 -- all or none. We never separate binders
681 new_lvl | any is_major bndrs = incMajorLvl lvl
682 | otherwise = incMinorLvl lvl
684 is_major bndr = isId bndr && not (isOneShotLambda bndr)
688 -- Destintion level is the max Id level of the expression
689 -- (We'll abstract the type variables, if any.)
690 destLevel :: LevelEnv -> VarSet -> Bool -> Maybe (Arity, StrictSig) -> Level
691 destLevel env fvs is_function mb_bot
692 | Just {} <- mb_bot = tOP_LEVEL -- Send bottoming bindings to the top
693 -- regardless; see Note [Bottoming floats]
694 | Just n_args <- floatLams env
695 , n_args > 0 -- n=0 case handled uniformly by the 'otherwise' case
697 , countFreeIds fvs <= n_args
698 = tOP_LEVEL -- Send functions to top level; see
699 -- the comments with isFunction
700 | otherwise = maxIdLevel env fvs
702 isFunction :: CoreExprWithFVs -> Bool
703 -- The idea here is that we want to float *functions* to
704 -- the top level. This saves no work, but
705 -- (a) it can make the host function body a lot smaller,
706 -- and hence inlinable.
707 -- (b) it can also save allocation when the function is recursive:
708 -- h = \x -> letrec f = \y -> ...f...y...x...
711 -- f = \x y -> ...(f x)...y...x...
713 -- No allocation for f now.
714 -- We may only want to do this if there are sufficiently few free
715 -- variables. We certainly only want to do it for values, and not for
716 -- constructors. So the simple thing is just to look for lambdas
717 isFunction (_, AnnLam b e) | isId b = True
718 | otherwise = isFunction e
719 isFunction (_, AnnNote _ e) = isFunction e
722 countFreeIds :: VarSet -> Int
723 countFreeIds = foldVarSet add 0
725 add :: Var -> Int -> Int
726 add v n | isId v = n+1
731 %************************************************************************
733 \subsection{Free-To-Level Monad}
735 %************************************************************************
739 = LE { le_switches :: FloatOutSwitches
740 , le_lvl_env :: VarEnv Level -- Domain is *post-cloned* TyVars and Ids
741 , le_subst :: Subst -- Domain is pre-cloned Ids; tracks the in-scope set
742 -- so that subtitution is capture-avoiding
743 , le_env :: IdEnv ([Var], LevelledExpr) -- Domain is pre-cloned Ids
745 -- We clone let-bound variables so that they are still
746 -- distinct when floated out; hence the le_subst/le_env.
747 -- (see point 3 of the module overview comment).
748 -- We also use these envs when making a variable polymorphic
749 -- because we want to float it out past a big lambda.
751 -- The le_subst and le_env always implement the same mapping, but the
752 -- le_subst maps to CoreExpr and the le_env to LevelledExpr
753 -- Since the range is always a variable or type application,
754 -- there is never any difference between the two, but sadly
755 -- the types differ. The le_subst is used when substituting in
756 -- a variable's IdInfo; the le_env when we find a Var.
758 -- In addition the le_env records a list of tyvars free in the
759 -- type application, just so we don't have to call freeVars on
760 -- the type application repeatedly.
762 -- The domain of the both envs is *pre-cloned* Ids, though
764 -- The domain of the le_lvl_env is the *post-cloned* Ids
766 initialEnv :: FloatOutSwitches -> LevelEnv
767 initialEnv float_lams
768 = LE { le_switches = float_lams, le_lvl_env = emptyVarEnv
769 , le_subst = emptySubst, le_env = emptyVarEnv }
771 floatLams :: LevelEnv -> Maybe Int
772 floatLams le = floatOutLambdas (le_switches le)
774 floatConsts :: LevelEnv -> Bool
775 floatConsts le = floatOutConstants (le_switches le)
777 floatPAPs :: LevelEnv -> Bool
778 floatPAPs le = floatOutPartialApplications (le_switches le)
780 extendLvlEnv :: LevelEnv -> [TaggedBndr Level] -> LevelEnv
781 -- Used when *not* cloning
782 extendLvlEnv le@(LE { le_lvl_env = lvl_env, le_subst = subst, le_env = id_env })
784 = le { le_lvl_env = foldl add_lvl lvl_env prs
785 , le_subst = foldl del_subst subst prs
786 , le_env = foldl del_id id_env prs }
788 add_lvl env (TB v l) = extendVarEnv env v l
789 del_subst env (TB v _) = extendInScope env v
790 del_id env (TB v _) = delVarEnv env v
791 -- We must remove any clone for this variable name in case of
792 -- shadowing. This bit me in the following case
793 -- (in nofib/real/gg/Spark.hs):
796 -- ... -> case e of wild {
797 -- ... -> ... wild ...
801 -- The inside occurrence of @wild@ was being replaced with @ds@,
802 -- incorrectly, because the SubstEnv was still lying around. Ouch!
805 extendInScopeEnv :: LevelEnv -> Var -> LevelEnv
806 extendInScopeEnv le@(LE { le_subst = subst }) v
807 = le { le_subst = extendInScope subst v }
809 extendInScopeEnvList :: LevelEnv -> [Var] -> LevelEnv
810 extendInScopeEnvList le@(LE { le_subst = subst }) vs
811 = le { le_subst = extendInScopeList subst vs }
813 -- extendCaseBndrLvlEnv adds the mapping case-bndr->scrut-var if it can
814 -- (see point 4 of the module overview comment)
815 extendCaseBndrLvlEnv :: LevelEnv -> Expr (TaggedBndr Level) -> Var -> Level
817 extendCaseBndrLvlEnv le@(LE { le_lvl_env = lvl_env, le_subst = subst, le_env = id_env })
818 (Var scrut_var) case_bndr lvl
819 = le { le_lvl_env = extendVarEnv lvl_env case_bndr lvl
820 , le_subst = extendIdSubst subst case_bndr (Var scrut_var)
821 , le_env = extendVarEnv id_env case_bndr ([scrut_var], Var scrut_var) }
823 extendCaseBndrLvlEnv env _scrut case_bndr lvl
824 = extendLvlEnv env [TB case_bndr lvl]
826 extendPolyLvlEnv :: Level -> LevelEnv -> [Var] -> [(Var, Var)] -> LevelEnv
827 extendPolyLvlEnv dest_lvl
828 le@(LE { le_lvl_env = lvl_env, le_subst = subst, le_env = id_env })
830 = le { le_lvl_env = foldl add_lvl lvl_env bndr_pairs
831 , le_subst = foldl add_subst subst bndr_pairs
832 , le_env = foldl add_id id_env bndr_pairs }
834 add_lvl env (_, v') = extendVarEnv env v' dest_lvl
835 add_subst env (v, v') = extendIdSubst env v (mkVarApps (Var v') abs_vars)
836 add_id env (v, v') = extendVarEnv env v ((v':abs_vars), mkVarApps (Var v') abs_vars)
838 extendCloneLvlEnv :: Level -> LevelEnv -> Subst -> [(Var, Var)] -> LevelEnv
839 extendCloneLvlEnv lvl le@(LE { le_lvl_env = lvl_env, le_env = id_env })
841 = le { le_lvl_env = foldl add_lvl lvl_env bndr_pairs
842 , le_subst = new_subst
843 , le_env = foldl add_id id_env bndr_pairs }
845 add_lvl env (_, v') = extendVarEnv env v' lvl
846 add_id env (v, v') = extendVarEnv env v ([v'], Var v')
848 maxIdLevel :: LevelEnv -> VarSet -> Level
849 maxIdLevel (LE { le_lvl_env = lvl_env, le_env = id_env }) var_set
850 = foldVarSet max_in tOP_LEVEL var_set
852 max_in in_var lvl = foldr max_out lvl (case lookupVarEnv id_env in_var of
853 Just (abs_vars, _) -> abs_vars
857 | isId out_var = case lookupVarEnv lvl_env out_var of
858 Just lvl' -> maxLvl lvl' lvl
860 | otherwise = lvl -- Ignore tyvars in *maxIdLevel*
862 lookupVar :: LevelEnv -> Id -> LevelledExpr
863 lookupVar le v = case lookupVarEnv (le_env le) v of
864 Just (_, expr) -> expr
867 abstractVars :: Level -> LevelEnv -> VarSet -> [Var]
868 -- Find the variables in fvs, free vars of the target expresion,
869 -- whose level is greater than the destination level
870 -- These are the ones we are going to abstract out
871 abstractVars dest_lvl (LE { le_lvl_env = lvl_env, le_env = id_env }) fvs
872 = map zap $ uniq $ sortLe le
873 [var | fv <- varSetElems fvs
874 , var <- absVarsOf id_env fv
876 -- NB: it's important to call abstract_me only on the OutIds the
877 -- come from absVarsOf (not on fv, which is an InId)
879 -- Sort the variables so the true type variables come first;
880 -- the tyvars scope over Ids and coercion vars
881 v1 `le` v2 = case (is_tv v1, is_tv v2) of
882 (True, False) -> True
883 (False, True) -> False
884 _ -> v1 <= v2 -- Same family
886 is_tv v = isTyCoVar v && not (isCoVar v)
888 uniq :: [Var] -> [Var]
889 -- Remove adjacent duplicates; the sort will have brought them together
890 uniq (v1:v2:vs) | v1 == v2 = uniq (v2:vs)
891 | otherwise = v1 : uniq (v2:vs)
894 abstract_me v = case lookupVarEnv lvl_env v of
895 Just lvl -> dest_lvl `ltLvl` lvl
898 -- We are going to lambda-abstract, so nuke any IdInfo,
899 -- and add the tyvars of the Id (if necessary)
900 zap v | isId v = WARN( isStableUnfolding (idUnfolding v) ||
901 not (isEmptySpecInfo (idSpecialisation v)),
902 text "absVarsOf: discarding info on" <+> ppr v )
903 setIdInfo v vanillaIdInfo
906 absVarsOf :: IdEnv ([Var], LevelledExpr) -> Var -> [Var]
907 -- If f is free in the expression, and f maps to poly_f a b c in the
908 -- current substitution, then we must report a b c as candidate type
911 -- Also, if x::a is an abstracted variable, then so is a; that is,
912 -- we must look in x's type
913 -- And similarly if x is a coercion variable.
915 | isId v = [av2 | av1 <- lookup_avs v
916 , av2 <- add_tyvars av1]
917 | isCoVar v = add_tyvars v
921 lookup_avs v = case lookupVarEnv id_env v of
922 Just (abs_vars, _) -> abs_vars
925 add_tyvars v = v : varSetElems (varTypeTyVars v)
929 type LvlM result = UniqSM result
931 initLvl :: UniqSupply -> UniqSM a -> a
937 newPolyBndrs :: Level -> LevelEnv -> [Var] -> [Id] -> UniqSM (LevelEnv, [Id])
938 newPolyBndrs dest_lvl env abs_vars bndrs = do
940 let new_bndrs = zipWith mk_poly_bndr bndrs uniqs
941 return (extendPolyLvlEnv dest_lvl env abs_vars (bndrs `zip` new_bndrs), new_bndrs)
943 mk_poly_bndr bndr uniq = transferPolyIdInfo bndr abs_vars $ -- Note [transferPolyIdInfo] in Id.lhs
944 mkSysLocal (mkFastString str) uniq poly_ty
946 str = "poly_" ++ occNameString (getOccName bndr)
947 poly_ty = mkPiTypes abs_vars (idType bndr)
949 newLvlVar :: [CoreBndr] -> Type -- Abstract wrt these bndrs
950 -> Maybe (Arity, StrictSig) -- Note [Bottoming floats]
952 newLvlVar vars body_ty mb_bot
953 = do { uniq <- getUniqueM
954 ; return (mkLocalIdWithInfo (mk_name uniq) (mkPiTypes vars body_ty) info) }
956 mk_name uniq = mkSystemVarName uniq (mkFastString "lvl")
957 arity = count isId vars
958 info = case mb_bot of
959 Nothing -> vanillaIdInfo
960 Just (bot_arity, sig) -> vanillaIdInfo
961 `setArityInfo` (arity + bot_arity)
962 `setStrictnessInfo` Just (increaseStrictSigArity arity sig)
964 -- The deeply tiresome thing is that we have to apply the substitution
965 -- to the rules inside each Id. Grr. But it matters.
967 cloneVar :: TopLevelFlag -> LevelEnv -> Id -> Level -> Level -> LvlM (LevelEnv, Id)
968 cloneVar TopLevel env v _ _
969 = return (extendInScopeEnv env v, v) -- Don't clone top level things
970 -- But do extend the in-scope env, to satisfy the in-scope invariant
972 cloneVar NotTopLevel env v ctxt_lvl dest_lvl
973 = ASSERT( isId v ) do
974 us <- getUniqueSupplyM
976 (subst', v1) = cloneIdBndr (le_subst env) us v
977 v2 = zap_demand ctxt_lvl dest_lvl v1
978 env' = extendCloneLvlEnv dest_lvl env subst' [(v,v2)]
981 cloneRecVars :: TopLevelFlag -> LevelEnv -> [Id] -> Level -> Level -> LvlM (LevelEnv, [Id])
982 cloneRecVars TopLevel env vs _ _
983 = return (extendInScopeEnvList env vs, vs) -- Don't clone top level things
984 cloneRecVars NotTopLevel env vs ctxt_lvl dest_lvl
985 = ASSERT( all isId vs ) do
986 us <- getUniqueSupplyM
988 (subst', vs1) = cloneRecIdBndrs (le_subst env) us vs
989 vs2 = map (zap_demand ctxt_lvl dest_lvl) vs1
990 env' = extendCloneLvlEnv dest_lvl env subst' (vs `zip` vs2)
993 -- VERY IMPORTANT: we must zap the demand info
994 -- if the thing is going to float out past a lambda,
995 -- or if it's going to top level (where things can't be strict)
996 zap_demand :: Level -> Level -> Id -> Id
997 zap_demand dest_lvl ctxt_lvl id
998 | ctxt_lvl == dest_lvl,
999 not (isTopLvl dest_lvl) = id -- Stays, and not going to top level
1000 | otherwise = zapDemandIdInfo id -- Floats out