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 )
63 import Id ( idType, mkLocalIdWithInfo, mkSysLocal, isOneShotLambda,
64 zapDemandIdInfo, transferPolyIdInfo,
65 idSpecialisation, idUnfolding, setIdInfo,
66 setIdStrictness, setIdArity
72 import Demand ( StrictSig, increaseStrictSigArity )
73 import Name ( getOccName, mkSystemVarName )
74 import OccName ( occNameString )
75 import Type ( isUnLiftedType, Type )
76 import BasicTypes ( TopLevelFlag(..), Arity )
78 import Util ( sortLe, isSingleton, count )
83 %************************************************************************
85 \subsection{Level numbers}
87 %************************************************************************
90 data Level = Level Int -- Level number of enclosing lambdas
91 Int -- Number of big-lambda and/or case expressions between
92 -- here and the nearest enclosing lambda
95 The {\em level number} on a (type-)lambda-bound variable is the
96 nesting depth of the (type-)lambda which binds it. The outermost lambda
97 has level 1, so (Level 0 0) means that the variable is bound outside any lambda.
99 On an expression, it's the maximum level number of its free
100 (type-)variables. On a let(rec)-bound variable, it's the level of its
101 RHS. On a case-bound variable, it's the number of enclosing lambdas.
103 Top-level variables: level~0. Those bound on the RHS of a top-level
104 definition but ``before'' a lambda; e.g., the \tr{x} in (levels shown
105 as ``subscripts'')...
107 a_0 = let b_? = ... in
108 x_1 = ... b ... in ...
111 The main function @lvlExpr@ carries a ``context level'' (@ctxt_lvl@).
112 That's meant to be the level number of the enclosing binder in the
113 final (floated) program. If the level number of a sub-expression is
114 less than that of the context, then it might be worth let-binding the
115 sub-expression so that it will indeed float.
117 If you can float to level @Level 0 0@ worth doing so because then your
118 allocation becomes static instead of dynamic. We always start with
122 Note [FloatOut inside INLINE]
123 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
124 @InlineCtxt@ very similar to @Level 0 0@, but is used for one purpose:
125 to say "don't float anything out of here". That's exactly what we
126 want for the body of an INLINE, where we don't want to float anything
127 out at all. See notes with lvlMFE below.
131 -- At one time I tried the effect of not float anything out of an InlineMe,
132 -- but it sometimes works badly. For example, consider PrelArr.done. It
133 -- has the form __inline (\d. e)
134 -- where e doesn't mention d. If we float this to
135 -- __inline (let x = e in \d. x)
136 -- things are bad. The inliner doesn't even inline it because it doesn't look
137 -- like a head-normal form. So it seems a lesser evil to let things float.
138 -- In SetLevels we do set the context to (Level 0 0) when we get to an InlineMe
139 -- which discourages floating out.
141 So the conclusion is: don't do any floating at all inside an InlineMe.
142 (In the above example, don't float the {x=e} out of the \d.)
144 One particular case is that of workers: we don't want to float the
145 call to the worker outside the wrapper, otherwise the worker might get
146 inlined into the floated expression, and an importing module won't see
150 type LevelledExpr = TaggedExpr Level
151 type LevelledBind = TaggedBind Level
154 tOP_LEVEL = Level 0 0
156 incMajorLvl :: Level -> Level
157 incMajorLvl (Level major _) = Level (major + 1) 0
159 incMinorLvl :: Level -> Level
160 incMinorLvl (Level major minor) = Level major (minor+1)
162 maxLvl :: Level -> Level -> Level
163 maxLvl l1@(Level maj1 min1) l2@(Level maj2 min2)
164 | (maj1 > maj2) || (maj1 == maj2 && min1 > min2) = l1
167 ltLvl :: Level -> Level -> Bool
168 ltLvl (Level maj1 min1) (Level maj2 min2)
169 = (maj1 < maj2) || (maj1 == maj2 && min1 < min2)
171 ltMajLvl :: Level -> Level -> Bool
172 -- Tells if one level belongs to a difft *lambda* level to another
173 ltMajLvl (Level maj1 _) (Level maj2 _) = maj1 < maj2
175 isTopLvl :: Level -> Bool
176 isTopLvl (Level 0 0) = True
179 instance Outputable Level where
180 ppr (Level maj min) = hcat [ char '<', int maj, char ',', int min, char '>' ]
182 instance Eq Level where
183 (Level maj1 min1) == (Level maj2 min2) = maj1 == maj2 && min1 == min2
187 %************************************************************************
189 \subsection{Main level-setting code}
191 %************************************************************************
194 setLevels :: FloatOutSwitches
199 setLevels float_lams binds us
200 = initLvl us (do_them init_env binds)
202 init_env = initialEnv float_lams
204 do_them :: LevelEnv -> [CoreBind] -> LvlM [LevelledBind]
205 do_them _ [] = return []
207 = do { (lvld_bind, env') <- lvlTopBind env b
208 ; lvld_binds <- do_them env' bs
209 ; return (lvld_bind : lvld_binds) }
211 lvlTopBind :: LevelEnv -> Bind Id -> LvlM (LevelledBind, LevelEnv)
212 lvlTopBind env (NonRec binder rhs)
213 = lvlBind TopLevel tOP_LEVEL env (AnnNonRec binder (freeVars rhs))
214 -- Rhs can have no free vars!
216 lvlTopBind env (Rec pairs)
217 = lvlBind TopLevel tOP_LEVEL env (AnnRec [(b,freeVars rhs) | (b,rhs) <- pairs])
220 %************************************************************************
222 \subsection{Setting expression levels}
224 %************************************************************************
227 lvlExpr :: Level -- ctxt_lvl: Level of enclosing expression
228 -> LevelEnv -- Level of in-scope names/tyvars
229 -> CoreExprWithFVs -- input expression
230 -> LvlM LevelledExpr -- Result expression
233 The @ctxt_lvl@ is, roughly, the level of the innermost enclosing
234 binder. Here's an example
236 v = \x -> ...\y -> let r = case (..x..) of
240 When looking at the rhs of @r@, @ctxt_lvl@ will be 1 because that's
241 the level of @r@, even though it's inside a level-2 @\y@. It's
242 important that @ctxt_lvl@ is 1 and not 2 in @r@'s rhs, because we
243 don't want @lvlExpr@ to turn the scrutinee of the @case@ into an MFE
244 --- because it isn't a *maximal* free expression.
246 If there were another lambda in @r@'s rhs, it would get level-2 as well.
249 lvlExpr _ _ ( _, AnnType ty) = return (Type ty)
250 lvlExpr _ env (_, AnnVar v) = return (lookupVar env v)
251 lvlExpr _ _ (_, AnnLit lit) = return (Lit lit)
253 lvlExpr ctxt_lvl env (_, AnnApp fun arg) = do
254 fun' <- lvlExpr ctxt_lvl env fun -- We don't do MFE on partial applications
255 arg' <- lvlMFE False ctxt_lvl env arg
256 return (App fun' arg')
258 lvlExpr ctxt_lvl env (_, AnnNote note expr) = do
259 expr' <- lvlExpr ctxt_lvl env expr
260 return (Note note expr')
262 lvlExpr ctxt_lvl env (_, AnnCast expr co) = do
263 expr' <- lvlExpr ctxt_lvl env expr
264 return (Cast expr' co)
266 -- We don't split adjacent lambdas. That is, given
268 -- we don't float to give
269 -- \x -> let v = x+y in \y -> (v,y)
270 -- Why not? Because partial applications are fairly rare, and splitting
271 -- lambdas makes them more expensive.
273 lvlExpr ctxt_lvl env expr@(_, AnnLam {}) = do
274 new_body <- lvlMFE True new_lvl new_env body
275 return (mkLams new_bndrs new_body)
277 (bndrs, body) = collectAnnBndrs expr
278 (new_lvl, new_bndrs) = lvlLamBndrs ctxt_lvl bndrs
279 new_env = extendLvlEnv env new_bndrs
280 -- At one time we called a special verion of collectBinders,
281 -- which ignored coercions, because we don't want to split
282 -- a lambda like this (\x -> coerce t (\s -> ...))
283 -- This used to happen quite a bit in state-transformer programs,
284 -- but not nearly so much now non-recursive newtypes are transparent.
285 -- [See SetLevels rev 1.50 for a version with this approach.]
287 lvlExpr ctxt_lvl env (_, AnnLet (AnnNonRec bndr rhs) body)
288 | isUnLiftedType (idType bndr) = do
289 -- Treat unlifted let-bindings (let x = b in e) just like (case b of x -> e)
290 -- That is, leave it exactly where it is
291 -- We used to float unlifted bindings too (e.g. to get a cheap primop
292 -- outside a lambda (to see how, look at lvlBind in rev 1.58)
293 -- but an unrelated change meant that these unlifed bindings
294 -- could get to the top level which is bad. And there's not much point;
295 -- unlifted bindings are always cheap, and so hardly worth floating.
296 rhs' <- lvlExpr ctxt_lvl env rhs
297 body' <- lvlExpr incd_lvl env' body
298 return (Let (NonRec bndr' rhs') body')
300 incd_lvl = incMinorLvl ctxt_lvl
301 bndr' = TB bndr incd_lvl
302 env' = extendLvlEnv env [bndr']
304 lvlExpr ctxt_lvl env (_, AnnLet bind body) = do
305 (bind', new_env) <- lvlBind NotTopLevel ctxt_lvl env bind
306 body' <- lvlExpr ctxt_lvl new_env body
307 return (Let bind' body')
309 lvlExpr ctxt_lvl env (_, AnnCase expr case_bndr ty alts) = do
310 expr' <- lvlMFE True ctxt_lvl env expr
311 let alts_env = extendCaseBndrLvlEnv env expr' case_bndr incd_lvl
312 alts' <- mapM (lvl_alt alts_env) alts
313 return (Case expr' (TB case_bndr incd_lvl) ty alts')
315 incd_lvl = incMinorLvl ctxt_lvl
317 lvl_alt alts_env (con, bs, rhs) = do
318 rhs' <- lvlMFE True incd_lvl new_env rhs
319 return (con, bs', rhs')
321 bs' = [ TB b incd_lvl | b <- bs ]
322 new_env = extendLvlEnv alts_env bs'
325 @lvlMFE@ is just like @lvlExpr@, except that it might let-bind
326 the expression, so that it can itself be floated.
330 We don't float unlifted MFEs, which potentially loses big opportunites.
333 where h :: Int -> Int# is expensive. We'd like to float the (h y) outside
334 the \x, but we don't because it's unboxed. Possible solution: box it.
336 Note [Bottoming floats]
337 ~~~~~~~~~~~~~~~~~~~~~~~
339 f = \x. g (error "urk")
340 we'd like to float the call to error, to get
343 Furthermore, we want to float a bottoming expression even if it has free
345 f = \x. g (let v = h x in error ("urk" ++ v))
346 Then we'd like to abstact over 'x' can float the whole arg of g:
347 lvl = \x. let v = h x in error ("urk" ++ v)
349 See Maessen's paper 1999 "Bottom extraction: factoring error handling out
350 of functional programs" (unpublished I think).
352 When we do this, we set the strictness and arity of the new bottoming
353 Id, so that it's properly exposed as such in the interface file, even if
354 this is all happening after strictness analysis.
356 Note [Bottoming floats: eta expansion] c.f Note [Bottoming floats]
357 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
358 Tiresomely, though, the simplifier has an invariant that the manifest
359 arity of the RHS should be the same as the arity; but we can't call
360 etaExpand during SetLevels because it works over a decorated form of
361 CoreExpr. So we do the eta expansion later, in FloatOut.
365 We don't float a case expression as an MFE from a strict context. Why not?
366 Because in doing so we share a tiny bit of computation (the switch) but
367 in exchange we build a thunk, which is bad. This case reduces allocation
368 by 7% in spectral/puzzle (a rather strange benchmark) and 1.2% in real/fem.
369 Doesn't change any other allocation at all.
372 lvlMFE :: Bool -- True <=> strict context [body of case or let]
373 -> Level -- Level of innermost enclosing lambda/tylam
374 -> LevelEnv -- Level of in-scope names/tyvars
375 -> CoreExprWithFVs -- input expression
376 -> LvlM LevelledExpr -- Result expression
378 lvlMFE _ _ _ (_, AnnType ty)
381 -- No point in floating out an expression wrapped in a coercion or note
382 -- If we do we'll transform lvl = e |> co
383 -- to lvl' = e; lvl = lvl' |> co
384 -- and then inline lvl. Better just to float out the payload.
385 lvlMFE strict_ctxt ctxt_lvl env (_, AnnNote n e)
386 = do { e' <- lvlMFE strict_ctxt ctxt_lvl env e
387 ; return (Note n e') }
389 lvlMFE strict_ctxt ctxt_lvl env (_, AnnCast e co)
390 = do { e' <- lvlMFE strict_ctxt ctxt_lvl env e
391 ; return (Cast e' co) }
394 lvlMFE True ctxt_lvl env e@(_, AnnCase {})
395 = lvlExpr ctxt_lvl env e -- Don't share cases
397 lvlMFE strict_ctxt ctxt_lvl env ann_expr@(fvs, _)
398 | isUnLiftedType ty -- Can't let-bind it; see Note [Unlifted MFEs]
399 || notWorthFloating ann_expr abs_vars
400 || not good_destination
401 = -- Don't float it out
402 lvlExpr ctxt_lvl env ann_expr
404 | otherwise -- Float it out!
405 = do expr' <- lvlFloatRhs abs_vars dest_lvl env ann_expr
406 var <- newLvlVar abs_vars ty mb_bot
407 return (Let (NonRec (TB var dest_lvl) expr')
408 (mkVarApps (Var var) abs_vars))
410 expr = deAnnotate ann_expr
412 mb_bot = exprBotStrictness_maybe expr
413 dest_lvl = destLevel env fvs (isFunction ann_expr) mb_bot
414 abs_vars = abstractVars dest_lvl env fvs
416 -- A decision to float entails let-binding this thing, and we only do
417 -- that if we'll escape a value lambda, or will go to the top level.
419 | dest_lvl `ltMajLvl` ctxt_lvl -- Escapes a value lambda
421 -- OLD CODE: not (exprIsCheap expr) || isTopLvl dest_lvl
422 -- see Note [Escaping a value lambda]
424 | otherwise -- Does not escape a value lambda
425 = isTopLvl dest_lvl -- Only float if we are going to the top level
426 && floatConsts env -- and the floatConsts flag is on
427 && not strict_ctxt -- Don't float from a strict context
428 -- We are keen to float something to the top level, even if it does not
429 -- escape a lambda, because then it needs no allocation. But it's controlled
430 -- by a flag, because doing this too early loses opportunities for RULES
431 -- which (needless to say) are important in some nofib programs
432 -- (gcd is an example).
435 -- concat = /\ a -> foldr ..a.. (++) []
436 -- was getting turned into
437 -- concat = /\ a -> lvl a
438 -- lvl = /\ a -> foldr ..a.. (++) []
439 -- which is pretty stupid. Hence the strict_ctxt test
441 annotateBotStr :: Id -> Maybe (Arity, StrictSig) -> Id
442 annotateBotStr id Nothing = id
443 annotateBotStr id (Just (arity,sig)) = id `setIdArity` arity
444 `setIdStrictness` sig
446 notWorthFloating :: CoreExprWithFVs -> [Var] -> Bool
447 -- Returns True if the expression would be replaced by
448 -- something bigger than it is now. For example:
449 -- abs_vars = tvars only: return True if e is trivial,
450 -- but False for anything bigger
451 -- abs_vars = [x] (an Id): return True for trivial, or an application (f x)
452 -- but False for (f x x)
454 -- One big goal is that floating should be idempotent. Eg if
455 -- we replace e with (lvl79 x y) and then run FloatOut again, don't want
456 -- to replace (lvl79 x y) with (lvl83 x y)!
458 notWorthFloating e abs_vars
459 = go e (count isId abs_vars)
461 go (_, AnnVar {}) n = n == 0
462 go (_, AnnLit {}) n = n == 0
463 go (_, AnnCast e _) n = go e n
464 go (_, AnnApp e arg) n
465 | (_, AnnType {}) <- arg = go e n
467 | is_triv arg = go e (n-1)
471 is_triv (_, AnnLit {}) = True -- Treat all literals as trivial
472 is_triv (_, AnnVar {}) = True -- (ie not worth floating)
473 is_triv (_, AnnCast e _) = is_triv e
474 is_triv (_, AnnApp e (_, AnnType {})) = is_triv e
478 Note [Escaping a value lambda]
479 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
480 We want to float even cheap expressions out of value lambdas,
481 because that saves allocation. Consider
482 f = \x. .. (\y.e) ...
483 Then we'd like to avoid allocating the (\y.e) every time we call f,
484 (assuming e does not mention x).
486 An example where this really makes a difference is simplrun009.
488 Another reason it's good is because it makes SpecContr fire on functions.
490 f = \x. ....(f (\y.e))....
491 After floating we get
493 f = \x. ....(f lvl)...
494 and that is much easier for SpecConstr to generate a robust specialisation for.
496 The OLD CODE (given where this Note is referred to) prevents floating
497 of the example above, so I just don't understand the old code. I
498 don't understand the old comment either (which appears below). I
499 measured the effect on nofib of changing OLD CODE to 'True', and got
500 zeros everywhere, but a 4% win for 'puzzle'. Very small 0.5% loss for
501 'cse'; turns out to be because our arity analysis isn't good enough
502 yet (mentioned in Simon-nofib-notes).
505 Even if it escapes a value lambda, we only
506 float if it's not cheap (unless it'll get all the
507 way to the top). I've seen cases where we
508 float dozens of tiny free expressions, which cost
509 more to allocate than to evaluate.
510 NB: exprIsCheap is also true of bottom expressions, which
511 is good; we don't want to share them
513 It's only Really Bad to float a cheap expression out of a
514 strict context, because that builds a thunk that otherwise
515 would never be built. So another alternative would be to
517 || (strict_ctxt && not (exprIsBottom expr))
518 to the condition above. We should really try this out.
521 %************************************************************************
523 \subsection{Bindings}
525 %************************************************************************
527 The binding stuff works for top level too.
530 lvlBind :: TopLevelFlag -- Used solely to decide whether to clone
531 -> Level -- Context level; might be Top even for bindings nested in the RHS
532 -- of a top level binding
535 -> LvlM (LevelledBind, LevelEnv)
537 lvlBind top_lvl ctxt_lvl env (AnnNonRec bndr rhs@(rhs_fvs,_))
538 | isTyVar bndr -- Don't do anything for TyVar binders
539 -- (simplifier gets rid of them pronto)
540 = do rhs' <- lvlExpr ctxt_lvl env rhs
541 return (NonRec (TB bndr ctxt_lvl) rhs', env)
544 = do -- No type abstraction; clone existing binder
545 rhs' <- lvlExpr dest_lvl env rhs
546 (env', bndr') <- cloneVar top_lvl env bndr ctxt_lvl dest_lvl
547 return (NonRec (TB bndr' dest_lvl) rhs', env')
550 = do -- Yes, type abstraction; create a new binder, extend substitution, etc
551 rhs' <- lvlFloatRhs abs_vars dest_lvl env rhs
552 (env', [bndr']) <- newPolyBndrs dest_lvl env abs_vars [bndr_w_str]
553 return (NonRec (TB bndr' dest_lvl) rhs', env')
556 bind_fvs = rhs_fvs `unionVarSet` idFreeVars bndr
557 abs_vars = abstractVars dest_lvl env bind_fvs
558 dest_lvl = destLevel env bind_fvs (isFunction rhs) mb_bot
559 mb_bot = exprBotStrictness_maybe (deAnnotate rhs)
560 bndr_w_str = annotateBotStr bndr mb_bot
565 lvlBind top_lvl ctxt_lvl env (AnnRec pairs)
567 = do (new_env, new_bndrs) <- cloneRecVars top_lvl env bndrs ctxt_lvl dest_lvl
568 new_rhss <- mapM (lvlExpr ctxt_lvl new_env) rhss
569 return (Rec ([TB b dest_lvl | b <- new_bndrs] `zip` new_rhss), new_env)
571 | isSingleton pairs && count isId abs_vars > 1
572 = do -- Special case for self recursion where there are
573 -- several variables carried around: build a local loop:
574 -- poly_f = \abs_vars. \lam_vars . letrec f = \lam_vars. rhs in f lam_vars
575 -- This just makes the closures a bit smaller. If we don't do
576 -- this, allocation rises significantly on some programs
578 -- We could elaborate it for the case where there are several
579 -- mutually functions, but it's quite a bit more complicated
581 -- This all seems a bit ad hoc -- sigh
583 (bndr,rhs) = head pairs
584 (rhs_lvl, abs_vars_w_lvls) = lvlLamBndrs dest_lvl abs_vars
585 rhs_env = extendLvlEnv env abs_vars_w_lvls
586 (rhs_env', new_bndr) <- cloneVar NotTopLevel rhs_env bndr rhs_lvl rhs_lvl
588 (lam_bndrs, rhs_body) = collectAnnBndrs rhs
589 (body_lvl, new_lam_bndrs) = lvlLamBndrs rhs_lvl lam_bndrs
590 body_env = extendLvlEnv rhs_env' new_lam_bndrs
591 new_rhs_body <- lvlExpr body_lvl body_env rhs_body
592 (poly_env, [poly_bndr]) <- newPolyBndrs dest_lvl env abs_vars [bndr]
593 return (Rec [(TB poly_bndr dest_lvl,
594 mkLams abs_vars_w_lvls $
595 mkLams new_lam_bndrs $
596 Let (Rec [(TB new_bndr rhs_lvl, mkLams new_lam_bndrs new_rhs_body)])
597 (mkVarApps (Var new_bndr) lam_bndrs))],
600 | otherwise = do -- Non-null abs_vars
601 (new_env, new_bndrs) <- newPolyBndrs dest_lvl env abs_vars bndrs
602 new_rhss <- mapM (lvlFloatRhs abs_vars dest_lvl new_env) rhss
603 return (Rec ([TB b dest_lvl | b <- new_bndrs] `zip` new_rhss), new_env)
606 (bndrs,rhss) = unzip pairs
608 -- Finding the free vars of the binding group is annoying
609 bind_fvs = (unionVarSets [ idFreeVars bndr `unionVarSet` rhs_fvs
610 | (bndr, (rhs_fvs,_)) <- pairs])
614 dest_lvl = destLevel env bind_fvs (all isFunction rhss) Nothing
615 abs_vars = abstractVars dest_lvl env bind_fvs
617 ----------------------------------------------------
618 -- Three help functons for the type-abstraction case
620 lvlFloatRhs :: [CoreBndr] -> Level -> LevelEnv -> CoreExprWithFVs
621 -> UniqSM (Expr (TaggedBndr Level))
622 lvlFloatRhs abs_vars dest_lvl env rhs = do
623 rhs' <- lvlExpr rhs_lvl rhs_env rhs
624 return (mkLams abs_vars_w_lvls rhs')
626 (rhs_lvl, abs_vars_w_lvls) = lvlLamBndrs dest_lvl abs_vars
627 rhs_env = extendLvlEnv env abs_vars_w_lvls
631 %************************************************************************
633 \subsection{Deciding floatability}
635 %************************************************************************
638 lvlLamBndrs :: Level -> [CoreBndr] -> (Level, [TaggedBndr Level])
639 -- Compute the levels for the binders of a lambda group
640 -- The binders returned are exactly the same as the ones passed,
641 -- but they are now paired with a level
645 lvlLamBndrs lvl bndrs
646 = go (incMinorLvl lvl)
647 False -- Havn't bumped major level in this group
650 go old_lvl bumped_major rev_lvld_bndrs (bndr:bndrs)
651 | isId bndr && -- Go to the next major level if this is a value binder,
652 not bumped_major && -- and we havn't already gone to the next level (one jump per group)
653 not (isOneShotLambda bndr) -- and it isn't a one-shot lambda
654 = go new_lvl True (TB bndr new_lvl : rev_lvld_bndrs) bndrs
657 = go old_lvl bumped_major (TB bndr old_lvl : rev_lvld_bndrs) bndrs
660 new_lvl = incMajorLvl old_lvl
662 go old_lvl _ rev_lvld_bndrs []
663 = (old_lvl, reverse rev_lvld_bndrs)
664 -- a lambda like this (\x -> coerce t (\s -> ...))
665 -- This happens quite a bit in state-transformer programs
669 -- Destintion level is the max Id level of the expression
670 -- (We'll abstract the type variables, if any.)
671 destLevel :: LevelEnv -> VarSet -> Bool -> Maybe (Arity, StrictSig) -> Level
672 destLevel env fvs is_function mb_bot
673 | Just {} <- mb_bot = tOP_LEVEL -- Send bottoming bindings to the top
674 -- regardless; see Note [Bottoming floats]
676 && is_function = tOP_LEVEL -- Send functions to top level; see
677 -- the comments with isFunction
678 | otherwise = maxIdLevel env fvs
680 isFunction :: CoreExprWithFVs -> Bool
681 -- The idea here is that we want to float *functions* to
682 -- the top level. This saves no work, but
683 -- (a) it can make the host function body a lot smaller,
684 -- and hence inlinable.
685 -- (b) it can also save allocation when the function is recursive:
686 -- h = \x -> letrec f = \y -> ...f...y...x...
689 -- f = \x y -> ...(f x)...y...x...
691 -- No allocation for f now.
692 -- We may only want to do this if there are sufficiently few free
693 -- variables. We certainly only want to do it for values, and not for
694 -- constructors. So the simple thing is just to look for lambdas
695 isFunction (_, AnnLam b e) | isId b = True
696 | otherwise = isFunction e
697 isFunction (_, AnnNote _ e) = isFunction e
702 %************************************************************************
704 \subsection{Free-To-Level Monad}
706 %************************************************************************
709 type LevelEnv = (FloatOutSwitches,
710 VarEnv Level, -- Domain is *post-cloned* TyVars and Ids
711 Subst, -- Domain is pre-cloned Ids; tracks the in-scope set
712 -- so that subtitution is capture-avoiding
713 IdEnv ([Var], LevelledExpr)) -- Domain is pre-cloned Ids
714 -- We clone let-bound variables so that they are still
715 -- distinct when floated out; hence the SubstEnv/IdEnv.
716 -- (see point 3 of the module overview comment).
717 -- We also use these envs when making a variable polymorphic
718 -- because we want to float it out past a big lambda.
720 -- The Subst and IdEnv always implement the same mapping, but the
721 -- Subst maps to CoreExpr and the IdEnv to LevelledExpr
722 -- Since the range is always a variable or type application,
723 -- there is never any difference between the two, but sadly
724 -- the types differ. The SubstEnv is used when substituting in
725 -- a variable's IdInfo; the IdEnv when we find a Var.
727 -- In addition the IdEnv records a list of tyvars free in the
728 -- type application, just so we don't have to call freeVars on
729 -- the type application repeatedly.
731 -- The domain of the both envs is *pre-cloned* Ids, though
733 -- The domain of the VarEnv Level is the *post-cloned* Ids
735 initialEnv :: FloatOutSwitches -> LevelEnv
736 initialEnv float_lams = (float_lams, emptyVarEnv, emptySubst, emptyVarEnv)
738 floatLams :: LevelEnv -> Bool
739 floatLams (fos, _, _, _) = floatOutLambdas fos
741 floatConsts :: LevelEnv -> Bool
742 floatConsts (fos, _, _, _) = floatOutConstants fos
744 extendLvlEnv :: LevelEnv -> [TaggedBndr Level] -> LevelEnv
745 -- Used when *not* cloning
746 extendLvlEnv (float_lams, lvl_env, subst, id_env) prs
748 foldl add_lvl lvl_env prs,
749 foldl del_subst subst prs,
750 foldl del_id id_env prs)
752 add_lvl env (TB v l) = extendVarEnv env v l
753 del_subst env (TB v _) = extendInScope env v
754 del_id env (TB v _) = delVarEnv env v
755 -- We must remove any clone for this variable name in case of
756 -- shadowing. This bit me in the following case
757 -- (in nofib/real/gg/Spark.hs):
760 -- ... -> case e of wild {
761 -- ... -> ... wild ...
765 -- The inside occurrence of @wild@ was being replaced with @ds@,
766 -- incorrectly, because the SubstEnv was still lying around. Ouch!
769 extendInScopeEnv :: LevelEnv -> Var -> LevelEnv
770 extendInScopeEnv (fl, le, subst, ids) v = (fl, le, extendInScope subst v, ids)
772 extendInScopeEnvList :: LevelEnv -> [Var] -> LevelEnv
773 extendInScopeEnvList (fl, le, subst, ids) vs = (fl, le, extendInScopeList subst vs, ids)
775 -- extendCaseBndrLvlEnv adds the mapping case-bndr->scrut-var if it can
776 -- (see point 4 of the module overview comment)
777 extendCaseBndrLvlEnv :: LevelEnv -> Expr (TaggedBndr Level) -> Var -> Level
779 extendCaseBndrLvlEnv (float_lams, lvl_env, subst, id_env) (Var scrut_var) case_bndr lvl
781 extendVarEnv lvl_env case_bndr lvl,
782 extendIdSubst subst case_bndr (Var scrut_var),
783 extendVarEnv id_env case_bndr ([scrut_var], Var scrut_var))
785 extendCaseBndrLvlEnv env _scrut case_bndr lvl
786 = extendLvlEnv env [TB case_bndr lvl]
788 extendPolyLvlEnv :: Level -> LevelEnv -> [Var] -> [(Var, Var)] -> LevelEnv
789 extendPolyLvlEnv dest_lvl (float_lams, lvl_env, subst, id_env) abs_vars bndr_pairs
791 foldl add_lvl lvl_env bndr_pairs,
792 foldl add_subst subst bndr_pairs,
793 foldl add_id id_env bndr_pairs)
795 add_lvl env (_, v') = extendVarEnv env v' dest_lvl
796 add_subst env (v, v') = extendIdSubst env v (mkVarApps (Var v') abs_vars)
797 add_id env (v, v') = extendVarEnv env v ((v':abs_vars), mkVarApps (Var v') abs_vars)
799 extendCloneLvlEnv :: Level -> LevelEnv -> Subst -> [(Var, Var)] -> LevelEnv
800 extendCloneLvlEnv lvl (float_lams, lvl_env, _, id_env) new_subst bndr_pairs
802 foldl add_lvl lvl_env bndr_pairs,
804 foldl add_id id_env bndr_pairs)
806 add_lvl env (_, v') = extendVarEnv env v' lvl
807 add_id env (v, v') = extendVarEnv env v ([v'], Var v')
810 maxIdLevel :: LevelEnv -> VarSet -> Level
811 maxIdLevel (_, lvl_env,_,id_env) var_set
812 = foldVarSet max_in tOP_LEVEL var_set
814 max_in in_var lvl = foldr max_out lvl (case lookupVarEnv id_env in_var of
815 Just (abs_vars, _) -> abs_vars
819 | isId out_var = case lookupVarEnv lvl_env out_var of
820 Just lvl' -> maxLvl lvl' lvl
822 | otherwise = lvl -- Ignore tyvars in *maxIdLevel*
824 lookupVar :: LevelEnv -> Id -> LevelledExpr
825 lookupVar (_, _, _, id_env) v = case lookupVarEnv id_env v of
826 Just (_, expr) -> expr
829 abstractVars :: Level -> LevelEnv -> VarSet -> [Var]
830 -- Find the variables in fvs, free vars of the target expresion,
831 -- whose level is greater than the destination level
832 -- These are the ones we are going to abstract out
833 abstractVars dest_lvl (_, lvl_env, _, id_env) fvs
834 = map zap $ uniq $ sortLe le
835 [var | fv <- varSetElems fvs
836 , var <- absVarsOf id_env fv
838 -- NB: it's important to call abstract_me only on the OutIds the
839 -- come from absVarsOf (not on fv, which is an InId)
841 -- Sort the variables so the true type variables come first;
842 -- the tyvars scope over Ids and coercion vars
843 v1 `le` v2 = case (is_tv v1, is_tv v2) of
844 (True, False) -> True
845 (False, True) -> False
846 _ -> v1 <= v2 -- Same family
848 is_tv v = isTyVar v && not (isCoVar v)
850 uniq :: [Var] -> [Var]
851 -- Remove adjacent duplicates; the sort will have brought them together
852 uniq (v1:v2:vs) | v1 == v2 = uniq (v2:vs)
853 | otherwise = v1 : uniq (v2:vs)
856 abstract_me v = case lookupVarEnv lvl_env v of
857 Just lvl -> dest_lvl `ltLvl` lvl
860 -- We are going to lambda-abstract, so nuke any IdInfo,
861 -- and add the tyvars of the Id (if necessary)
862 zap v | isId v = WARN( isInlineRule (idUnfolding v) ||
863 not (isEmptySpecInfo (idSpecialisation v)),
864 text "absVarsOf: discarding info on" <+> ppr v )
865 setIdInfo v vanillaIdInfo
868 absVarsOf :: IdEnv ([Var], LevelledExpr) -> Var -> [Var]
869 -- If f is free in the expression, and f maps to poly_f a b c in the
870 -- current substitution, then we must report a b c as candidate type
873 -- Also, if x::a is an abstracted variable, then so is a; that is,
874 -- we must look in x's type
875 -- And similarly if x is a coercion variable.
877 | isId v = [av2 | av1 <- lookup_avs v
878 , av2 <- add_tyvars av1]
879 | isCoVar v = add_tyvars v
883 lookup_avs v = case lookupVarEnv id_env v of
884 Just (abs_vars, _) -> abs_vars
887 add_tyvars v = v : varSetElems (varTypeTyVars v)
891 type LvlM result = UniqSM result
893 initLvl :: UniqSupply -> UniqSM a -> a
899 newPolyBndrs :: Level -> LevelEnv -> [Var] -> [Id] -> UniqSM (LevelEnv, [Id])
900 newPolyBndrs dest_lvl env abs_vars bndrs = do
902 let new_bndrs = zipWith mk_poly_bndr bndrs uniqs
903 return (extendPolyLvlEnv dest_lvl env abs_vars (bndrs `zip` new_bndrs), new_bndrs)
905 mk_poly_bndr bndr uniq = transferPolyIdInfo bndr abs_vars $ -- Note [transferPolyIdInfo] in Id.lhs
906 mkSysLocal (mkFastString str) uniq poly_ty
908 str = "poly_" ++ occNameString (getOccName bndr)
909 poly_ty = mkPiTypes abs_vars (idType bndr)
911 newLvlVar :: [CoreBndr] -> Type -- Abstract wrt these bndrs
912 -> Maybe (Arity, StrictSig) -- Note [Bottoming floats]
914 newLvlVar vars body_ty mb_bot
915 = do { uniq <- getUniqueM
916 ; return (mkLocalIdWithInfo (mk_name uniq) (mkPiTypes vars body_ty) info) }
918 mk_name uniq = mkSystemVarName uniq (mkFastString "lvl")
919 arity = count isId vars
920 info = case mb_bot of
921 Nothing -> vanillaIdInfo
922 Just (bot_arity, sig) -> vanillaIdInfo
923 `setArityInfo` (arity + bot_arity)
924 `setStrictnessInfo` Just (increaseStrictSigArity arity sig)
926 -- The deeply tiresome thing is that we have to apply the substitution
927 -- to the rules inside each Id. Grr. But it matters.
929 cloneVar :: TopLevelFlag -> LevelEnv -> Id -> Level -> Level -> LvlM (LevelEnv, Id)
930 cloneVar TopLevel env v _ _
931 = return (extendInScopeEnv env v, v) -- Don't clone top level things
932 -- But do extend the in-scope env, to satisfy the in-scope invariant
934 cloneVar NotTopLevel env@(_,_,subst,_) v ctxt_lvl dest_lvl
935 = ASSERT( isId v ) do
936 us <- getUniqueSupplyM
938 (subst', v1) = cloneIdBndr subst us v
939 v2 = zap_demand ctxt_lvl dest_lvl v1
940 env' = extendCloneLvlEnv dest_lvl env subst' [(v,v2)]
943 cloneRecVars :: TopLevelFlag -> LevelEnv -> [Id] -> Level -> Level -> LvlM (LevelEnv, [Id])
944 cloneRecVars TopLevel env vs _ _
945 = return (extendInScopeEnvList env vs, vs) -- Don't clone top level things
946 cloneRecVars NotTopLevel env@(_,_,subst,_) vs ctxt_lvl dest_lvl
947 = ASSERT( all isId vs ) do
948 us <- getUniqueSupplyM
950 (subst', vs1) = cloneRecIdBndrs subst us vs
951 vs2 = map (zap_demand ctxt_lvl dest_lvl) vs1
952 env' = extendCloneLvlEnv dest_lvl env subst' (vs `zip` vs2)
955 -- VERY IMPORTANT: we must zap the demand info
956 -- if the thing is going to float out past a lambda,
957 -- or if it's going to top level (where things can't be strict)
958 zap_demand :: Level -> Level -> Id -> Id
959 zap_demand dest_lvl ctxt_lvl id
960 | ctxt_lvl == dest_lvl,
961 not (isTopLvl dest_lvl) = id -- Stays, and not going to top level
962 | otherwise = zapDemandIdInfo id -- Floats out