0b81f29a7de8519b91369fc7314827b88f76413b
[ghc.git] / compiler / simplCore / SetLevels.hs
1 {-
2 (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
3
4 \section{SetLevels}
5
6 ***************************
7 Overview
8 ***************************
9
10 1. We attach binding levels to Core bindings, in preparation for floating
11 outwards (@FloatOut@).
12
13 2. We also let-ify many expressions (notably case scrutinees), so they
14 will have a fighting chance of being floated sensible.
15
16 3. Note [Need for cloning during float-out]
17 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
18 We clone the binders of any floatable let-binding, so that when it is
19 floated out it will be unique. Example
20 (let x=2 in x) + (let x=3 in x)
21 we must clone before floating so we get
22 let x1=2 in
23 let x2=3 in
24 x1+x2
25
26 NOTE: this can't be done using the uniqAway idea, because the variable
27 must be unique in the whole program, not just its current scope,
28 because two variables in different scopes may float out to the
29 same top level place
30
31 NOTE: Very tiresomely, we must apply this substitution to
32 the rules stored inside a variable too.
33
34 We do *not* clone top-level bindings, because some of them must not change,
35 but we *do* clone bindings that are heading for the top level
36
37 4. Note [Binder-swap during float-out]
38 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
39 In the expression
40 case x of wild { p -> ...wild... }
41 we substitute x for wild in the RHS of the case alternatives:
42 case x of wild { p -> ...x... }
43 This means that a sub-expression involving x is not "trapped" inside the RHS.
44 And it's not inconvenient because we already have a substitution.
45
46 Note that this is EXACTLY BACKWARDS from the what the simplifier does.
47 The simplifier tries to get rid of occurrences of x, in favour of wild,
48 in the hope that there will only be one remaining occurrence of x, namely
49 the scrutinee of the case, and we can inline it.
50 -}
51
52 {-# LANGUAGE CPP #-}
53 module SetLevels (
54 setLevels,
55
56 Level(..), tOP_LEVEL,
57 LevelledBind, LevelledExpr, LevelledBndr,
58 FloatSpec(..), floatSpecLevel,
59
60 incMinorLvl, ltMajLvl, ltLvl, isTopLvl
61 ) where
62
63 #include "HsVersions.h"
64
65 import CoreSyn
66 import CoreMonad ( FloatOutSwitches(..) )
67 import CoreUtils ( exprType
68 , isExprLevPoly
69 , exprOkForSpeculation
70 , collectMakeStaticArgs
71 )
72 import CoreArity ( exprBotStrictness_maybe )
73 import CoreFVs -- all of it
74 import CoreSubst
75 import MkCore ( sortQuantVars )
76 import Id
77 import IdInfo
78 import Var
79 import VarSet
80 import VarEnv
81 import Literal ( litIsTrivial )
82 import Demand ( StrictSig, increaseStrictSigArity )
83 import Name ( getOccName, mkSystemVarName )
84 import OccName ( occNameString )
85 import Type ( isUnliftedType, Type, mkLamTypes, splitTyConApp_maybe )
86 import BasicTypes ( Arity, RecFlag(..) )
87 import DataCon ( dataConOrigResTy )
88 import TysWiredIn
89 import UniqSupply
90 import Util
91 import Outputable
92 import FastString
93 import UniqDFM
94 import FV
95 import Data.Maybe
96
97 {-
98 ************************************************************************
99 * *
100 \subsection{Level numbers}
101 * *
102 ************************************************************************
103 -}
104
105 type LevelledExpr = TaggedExpr FloatSpec
106 type LevelledBind = TaggedBind FloatSpec
107 type LevelledBndr = TaggedBndr FloatSpec
108
109 data Level = Level Int -- Major level: number of enclosing value lambdas
110 Int -- Minor level: number of big-lambda and/or case
111 -- expressions between here and the nearest
112 -- enclosing value lambda
113
114 data FloatSpec
115 = FloatMe Level -- Float to just inside the binding
116 -- tagged with this level
117 | StayPut Level -- Stay where it is; binding is
118 -- tagged with tihs level
119
120 floatSpecLevel :: FloatSpec -> Level
121 floatSpecLevel (FloatMe l) = l
122 floatSpecLevel (StayPut l) = l
123
124 {-
125 The {\em level number} on a (type-)lambda-bound variable is the
126 nesting depth of the (type-)lambda which binds it. The outermost lambda
127 has level 1, so (Level 0 0) means that the variable is bound outside any lambda.
128
129 On an expression, it's the maximum level number of its free
130 (type-)variables. On a let(rec)-bound variable, it's the level of its
131 RHS. On a case-bound variable, it's the number of enclosing lambdas.
132
133 Top-level variables: level~0. Those bound on the RHS of a top-level
134 definition but ``before'' a lambda; e.g., the \tr{x} in (levels shown
135 as ``subscripts'')...
136 \begin{verbatim}
137 a_0 = let b_? = ... in
138 x_1 = ... b ... in ...
139 \end{verbatim}
140
141 The main function @lvlExpr@ carries a ``context level'' (@ctxt_lvl@).
142 That's meant to be the level number of the enclosing binder in the
143 final (floated) program. If the level number of a sub-expression is
144 less than that of the context, then it might be worth let-binding the
145 sub-expression so that it will indeed float.
146
147 If you can float to level @Level 0 0@ worth doing so because then your
148 allocation becomes static instead of dynamic. We always start with
149 context @Level 0 0@.
150
151
152 Note [FloatOut inside INLINE]
153 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
154 @InlineCtxt@ very similar to @Level 0 0@, but is used for one purpose:
155 to say "don't float anything out of here". That's exactly what we
156 want for the body of an INLINE, where we don't want to float anything
157 out at all. See notes with lvlMFE below.
158
159 But, check this out:
160
161 -- At one time I tried the effect of not float anything out of an InlineMe,
162 -- but it sometimes works badly. For example, consider PrelArr.done. It
163 -- has the form __inline (\d. e)
164 -- where e doesn't mention d. If we float this to
165 -- __inline (let x = e in \d. x)
166 -- things are bad. The inliner doesn't even inline it because it doesn't look
167 -- like a head-normal form. So it seems a lesser evil to let things float.
168 -- In SetLevels we do set the context to (Level 0 0) when we get to an InlineMe
169 -- which discourages floating out.
170
171 So the conclusion is: don't do any floating at all inside an InlineMe.
172 (In the above example, don't float the {x=e} out of the \d.)
173
174 One particular case is that of workers: we don't want to float the
175 call to the worker outside the wrapper, otherwise the worker might get
176 inlined into the floated expression, and an importing module won't see
177 the worker at all.
178 -}
179
180 instance Outputable FloatSpec where
181 ppr (FloatMe l) = char 'F' <> ppr l
182 ppr (StayPut l) = ppr l
183
184 tOP_LEVEL :: Level
185 tOP_LEVEL = Level 0 0
186
187 incMajorLvl :: Level -> Level
188 incMajorLvl (Level major _) = Level (major + 1) 0
189
190 incMinorLvl :: Level -> Level
191 incMinorLvl (Level major minor) = Level major (minor+1)
192
193 maxLvl :: Level -> Level -> Level
194 maxLvl l1@(Level maj1 min1) l2@(Level maj2 min2)
195 | (maj1 > maj2) || (maj1 == maj2 && min1 > min2) = l1
196 | otherwise = l2
197
198 ltLvl :: Level -> Level -> Bool
199 ltLvl (Level maj1 min1) (Level maj2 min2)
200 = (maj1 < maj2) || (maj1 == maj2 && min1 < min2)
201
202 ltMajLvl :: Level -> Level -> Bool
203 -- Tells if one level belongs to a difft *lambda* level to another
204 ltMajLvl (Level maj1 _) (Level maj2 _) = maj1 < maj2
205
206 isTopLvl :: Level -> Bool
207 isTopLvl (Level 0 0) = True
208 isTopLvl _ = False
209
210 instance Outputable Level where
211 ppr (Level maj min) = hcat [ char '<', int maj, char ',', int min, char '>' ]
212
213 instance Eq Level where
214 (Level maj1 min1) == (Level maj2 min2) = maj1 == maj2 && min1 == min2
215
216 {-
217 ************************************************************************
218 * *
219 \subsection{Main level-setting code}
220 * *
221 ************************************************************************
222 -}
223
224 setLevels :: FloatOutSwitches
225 -> CoreProgram
226 -> UniqSupply
227 -> [LevelledBind]
228
229 setLevels float_lams binds us
230 = initLvl us (do_them init_env binds)
231 where
232 init_env = initialEnv float_lams
233
234 do_them :: LevelEnv -> [CoreBind] -> LvlM [LevelledBind]
235 do_them _ [] = return []
236 do_them env (b:bs)
237 = do { (lvld_bind, env') <- lvlTopBind env b
238 ; lvld_binds <- do_them env' bs
239 ; return (lvld_bind : lvld_binds) }
240
241 lvlTopBind :: LevelEnv -> Bind Id -> LvlM (LevelledBind, LevelEnv)
242 lvlTopBind env (NonRec bndr rhs)
243 = do { rhs' <- lvlExpr env (freeVars rhs)
244 ; let (env', [bndr']) = substAndLvlBndrs NonRecursive env tOP_LEVEL [bndr]
245 ; return (NonRec bndr' rhs', env') }
246
247 lvlTopBind env (Rec pairs)
248 = do let (bndrs,rhss) = unzip pairs
249 (env', bndrs') = substAndLvlBndrs Recursive env tOP_LEVEL bndrs
250 rhss' <- mapM (lvlExpr env' . freeVars) rhss
251 return (Rec (bndrs' `zip` rhss'), env')
252
253 {-
254 ************************************************************************
255 * *
256 \subsection{Setting expression levels}
257 * *
258 ************************************************************************
259
260 Note [Floating over-saturated applications]
261 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
262 If we see (f x y), and (f x) is a redex (ie f's arity is 1),
263 we call (f x) an "over-saturated application"
264
265 Should we float out an over-sat app, if can escape a value lambda?
266 It is sometimes very beneficial (-7% runtime -4% alloc over nofib -O2).
267 But we don't want to do it for class selectors, because the work saved
268 is minimal, and the extra local thunks allocated cost money.
269
270 Arguably we could float even class-op applications if they were going to
271 top level -- but then they must be applied to a constant dictionary and
272 will almost certainly be optimised away anyway.
273 -}
274
275 lvlExpr :: LevelEnv -- Context
276 -> CoreExprWithFVs -- Input expression
277 -> LvlM LevelledExpr -- Result expression
278
279 {-
280 The @ctxt_lvl@ is, roughly, the level of the innermost enclosing
281 binder. Here's an example
282
283 v = \x -> ...\y -> let r = case (..x..) of
284 ..x..
285 in ..
286
287 When looking at the rhs of @r@, @ctxt_lvl@ will be 1 because that's
288 the level of @r@, even though it's inside a level-2 @\y@. It's
289 important that @ctxt_lvl@ is 1 and not 2 in @r@'s rhs, because we
290 don't want @lvlExpr@ to turn the scrutinee of the @case@ into an MFE
291 --- because it isn't a *maximal* free expression.
292
293 If there were another lambda in @r@'s rhs, it would get level-2 as well.
294 -}
295
296 lvlExpr env (_, AnnType ty) = return (Type (CoreSubst.substTy (le_subst env) ty))
297 lvlExpr env (_, AnnCoercion co) = return (Coercion (substCo (le_subst env) co))
298 lvlExpr env (_, AnnVar v) = return (lookupVar env v)
299 lvlExpr _ (_, AnnLit lit) = return (Lit lit)
300
301 lvlExpr env (_, AnnCast expr (_, co)) = do
302 expr' <- lvlExpr env expr
303 return (Cast expr' (substCo (le_subst env) co))
304
305 lvlExpr env (_, AnnTick tickish expr) = do
306 expr' <- lvlExpr env expr
307 return (Tick tickish expr')
308
309 lvlExpr env expr@(_, AnnApp _ _) = do
310 let
311 (fun, args) = collectAnnArgs expr
312 --
313 case fun of
314 (_, AnnVar f) | floatOverSat env -- See Note [Floating over-saturated applications]
315 , arity > 0
316 , arity < n_val_args
317 , Nothing <- isClassOpId_maybe f ->
318 do
319 let (lapp, rargs) = left (n_val_args - arity) expr []
320 rargs' <- mapM (lvlMFE False env) rargs
321 lapp' <- lvlMFE False env lapp
322 return (foldl App lapp' rargs')
323 where
324 n_val_args = count (isValArg . deAnnotate) args
325 arity = idArity f
326
327 -- separate out the PAP that we are floating from the extra
328 -- arguments, by traversing the spine until we have collected
329 -- (n_val_args - arity) value arguments.
330 left 0 e rargs = (e, rargs)
331 left n (_, AnnApp f a) rargs
332 | isValArg (deAnnotate a) = left (n-1) f (a:rargs)
333 | otherwise = left n f (a:rargs)
334 left _ _ _ = panic "SetLevels.lvlExpr.left"
335
336 -- No PAPs that we can float: just carry on with the
337 -- arguments and the function.
338 _otherwise -> do
339 args' <- mapM (lvlMFE False env) args
340 fun' <- lvlExpr env fun
341 return (foldl App fun' args')
342
343 -- We don't split adjacent lambdas. That is, given
344 -- \x y -> (x+1,y)
345 -- we don't float to give
346 -- \x -> let v = x+1 in \y -> (v,y)
347 -- Why not? Because partial applications are fairly rare, and splitting
348 -- lambdas makes them more expensive.
349
350 lvlExpr env expr@(_, AnnLam {})
351 = do { new_body <- lvlMFE True new_env body
352 ; return (mkLams new_bndrs new_body) }
353 where
354 (bndrs, body) = collectAnnBndrs expr
355 (env1, bndrs1) = substBndrsSL NonRecursive env bndrs
356 (new_env, new_bndrs) = lvlLamBndrs env1 (le_ctxt_lvl env) bndrs1
357 -- At one time we called a special verion of collectBinders,
358 -- which ignored coercions, because we don't want to split
359 -- a lambda like this (\x -> coerce t (\s -> ...))
360 -- This used to happen quite a bit in state-transformer programs,
361 -- but not nearly so much now non-recursive newtypes are transparent.
362 -- [See SetLevels rev 1.50 for a version with this approach.]
363
364 lvlExpr env (_, AnnLet bind body)
365 = do { (bind', new_env) <- lvlBind env bind
366 ; body' <- lvlExpr new_env body
367 -- No point in going via lvlMFE here. If the binding is alive
368 -- (mentioned in body), and the whole let-expression doesn't
369 -- float, then neither will the body
370 ; return (Let bind' body') }
371
372 lvlExpr env (_, AnnCase scrut case_bndr ty alts)
373 = do { scrut' <- lvlMFE True env scrut
374 ; lvlCase env (freeVarsOf scrut) scrut' case_bndr ty alts }
375
376 -------------------------------------------
377 lvlCase :: LevelEnv -- Level of in-scope names/tyvars
378 -> DVarSet -- Free vars of input scrutinee
379 -> LevelledExpr -- Processed scrutinee
380 -> Id -> Type -- Case binder and result type
381 -> [CoreAltWithFVs] -- Input alternatives
382 -> LvlM LevelledExpr -- Result expression
383 lvlCase env scrut_fvs scrut' case_bndr ty alts
384 | [(con@(DataAlt {}), bs, body)] <- alts
385 , exprOkForSpeculation scrut' -- See Note [Check the output scrutinee for okForSpec]
386 , not (isTopLvl dest_lvl) -- Can't have top-level cases
387 , not (floatTopLvlOnly env) -- Can float anywhere
388 = -- See Note [Floating cases]
389 -- Always float the case if possible
390 -- Unlike lets we don't insist that it escapes a value lambda
391 do { (env1, (case_bndr' : bs')) <- cloneCaseBndrs env dest_lvl (case_bndr : bs)
392 ; let rhs_env = extendCaseBndrEnv env1 case_bndr scrut'
393 ; body' <- lvlMFE True rhs_env body
394 ; let alt' = (con, [TB b (StayPut dest_lvl) | b <- bs'], body')
395 ; return (Case scrut' (TB case_bndr' (FloatMe dest_lvl)) ty [alt']) }
396
397 | otherwise -- Stays put
398 = do { let (alts_env1, [case_bndr']) = substAndLvlBndrs NonRecursive env incd_lvl [case_bndr]
399 alts_env = extendCaseBndrEnv alts_env1 case_bndr scrut'
400 ; alts' <- mapM (lvl_alt alts_env) alts
401 ; return (Case scrut' case_bndr' ty alts') }
402 where
403 incd_lvl = incMinorLvl (le_ctxt_lvl env)
404 dest_lvl = maxFvLevel (const True) env scrut_fvs
405 -- Don't abstact over type variables, hence const True
406
407 lvl_alt alts_env (con, bs, rhs)
408 = do { rhs' <- lvlMFE True new_env rhs
409 ; return (con, bs', rhs') }
410 where
411 (new_env, bs') = substAndLvlBndrs NonRecursive alts_env incd_lvl bs
412
413 {-
414 Note [Floating cases]
415 ~~~~~~~~~~~~~~~~~~~~~
416 Consider this:
417 data T a = MkT !a
418 f :: T Int -> blah
419 f x vs = case x of { MkT y ->
420 let f vs = ...(case y of I# w -> e)...f..
421 in f vs
422 Here we can float the (case y ...) out , because y is sure
423 to be evaluated, to give
424 f x vs = case x of { MkT y ->
425 caes y of I# w ->
426 let f vs = ...(e)...f..
427 in f vs
428
429 That saves unboxing it every time round the loop. It's important in
430 some DPH stuff where we really want to avoid that repeated unboxing in
431 the inner loop.
432
433 Things to note
434 * We can't float a case to top level
435 * It's worth doing this float even if we don't float
436 the case outside a value lambda. Example
437 case x of {
438 MkT y -> (case y of I# w2 -> ..., case y of I# w2 -> ...)
439 If we floated the cases out we could eliminate one of them.
440 * We only do this with a single-alternative case
441
442 Note [Check the output scrutinee for okForSpec]
443 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
444 Consider this:
445 case x of y {
446 A -> ....(case y of alts)....
447 }
448 Because of the binder-swap, the inner case will get substituted to
449 (case x of ..). So when testing whether the scrutinee is
450 okForSpecuation we must be careful to test the *result* scrutinee ('x'
451 in this case), not the *input* one 'y'. The latter *is* ok for
452 speculation here, but the former is not -- and indeed we can't float
453 the inner case out, at least not unless x is also evaluated at its
454 binding site.
455
456 That's why we apply exprOkForSpeculation to scrut' and not to scrut.
457 -}
458
459 lvlMFE :: Bool -- True <=> strict context [body of case or let]
460 -> LevelEnv -- Level of in-scope names/tyvars
461 -> CoreExprWithFVs -- input expression
462 -> LvlM LevelledExpr -- Result expression
463 -- lvlMFE is just like lvlExpr, except that it might let-bind
464 -- the expression, so that it can itself be floated.
465
466 lvlMFE _ env (_, AnnType ty)
467 = return (Type (CoreSubst.substTy (le_subst env) ty))
468
469 -- No point in floating out an expression wrapped in a coercion or note
470 -- If we do we'll transform lvl = e |> co
471 -- to lvl' = e; lvl = lvl' |> co
472 -- and then inline lvl. Better just to float out the payload.
473 lvlMFE strict_ctxt env (_, AnnTick t e)
474 = do { e' <- lvlMFE strict_ctxt env e
475 ; return (Tick t e') }
476
477 lvlMFE strict_ctxt env (_, AnnCast e (_, co))
478 = do { e' <- lvlMFE strict_ctxt env e
479 ; return (Cast e' (substCo (le_subst env) co)) }
480
481 -- Note [Case MFEs]
482 lvlMFE True env e@(_, AnnCase {})
483 = lvlExpr env e -- Don't share cases
484
485 lvlMFE strict_ctxt env ann_expr
486 | floatTopLvlOnly env && not (isTopLvl dest_lvl)
487 -- Only floating to the top level is allowed.
488 || isExprLevPoly expr
489 -- We can't let-bind levity polymorphic expressions
490 -- See Note [Levity polymorphism invariants] in CoreSyn
491 || notWorthFloating expr abs_vars
492 || not float_me
493 = -- Don't float it out
494 lvlExpr env ann_expr
495
496 | Just (wrap_float, wrap_use)
497 <- canFloat_maybe rhs_env strict_ctxt float_is_lam expr_ty
498 = do { expr1 <- lvlExpr rhs_env ann_expr
499 ; let abs_expr = mkLams abs_vars_w_lvls (wrap_float expr1)
500 ; var <- newLvlVar abs_expr
501 ; let var2 = annotateBotStr var float_n_lams mb_bot_str
502 ; return (Let (NonRec (TB var2 (FloatMe dest_lvl)) abs_expr)
503 (wrap_use (mkVarApps (Var var2) abs_vars))) }
504
505 | otherwise
506 = lvlExpr env ann_expr
507
508 where
509 expr = deAnnotate ann_expr
510 expr_ty = exprType expr
511 fvs = freeVarsOf ann_expr
512 is_bot = isJust mb_bot_str
513 mb_bot_str = exprBotStrictness_maybe expr
514 -- See Note [Bottoming floats]
515 -- esp Bottoming floats (2)
516 dest_lvl = destLevel env fvs (isFunction ann_expr) is_bot
517 abs_vars = abstractVars dest_lvl env fvs
518 float_is_lam = float_n_lams > 0 -- The floated thing will be a value lambda
519 float_n_lams = count isId abs_vars -- so nothing is shared; the only benefit
520 -- is getting it to the top level
521 (rhs_env, abs_vars_w_lvls) = lvlLamBndrs env dest_lvl abs_vars
522
523 -- A decision to float entails let-binding this thing, and we only do
524 -- that if we'll escape a value lambda, or will go to the top level.
525 float_me = (dest_lvl `ltMajLvl` (le_ctxt_lvl env) -- Escapes a value lambda
526 && not float_is_lam) -- See Note [Escaping a value lambda]
527
528 || (isTopLvl dest_lvl -- Only float if we are going to the top level
529 && floatConsts env -- and the floatConsts flag is on
530 && not strict_ctxt) -- Don't float from a strict context
531 -- We are keen to float something to the top level, even if it does not
532 -- escape a lambda, because then it needs no allocation. But it's controlled
533 -- by a flag, because doing this too early loses opportunities for RULES
534 -- which (needless to say) are important in some nofib programs
535 -- (gcd is an example).
536 --
537 -- Beware:
538 -- concat = /\ a -> foldr ..a.. (++) []
539 -- was getting turned into
540 -- lvl = /\ a -> foldr ..a.. (++) []
541 -- concat = /\ a -> lvl a
542 -- which is pretty stupid. Hence the strict_ctxt test
543
544 canFloat_maybe :: LevelEnv
545 -> Bool -- Strict context
546 -> Bool -- The float has a value lambda
547 -> Type
548 -> Maybe ( LevelledExpr -> LevelledExpr -- Wrep the flaot
549 , LevelledExpr -> LevelledExpr) -- Wrap the use
550 -- See Note [Floating MFEs of unlifted type]
551 canFloat_maybe env strict_ctxt float_is_lam expr_ty
552 | float_is_lam || not (isUnliftedType expr_ty)
553 = Just (id, id) -- No wrapping needed if the type is lifted, or
554 -- if we are wrapping it in one or more value lambdas
555
556 -- OK, so the float has an unlifted type and no value lambdas
557 | strict_ctxt
558 , Just (tc, _) <- splitTyConApp_maybe expr_ty
559 , Just dc <- boxingDataCon_maybe tc
560 , let dc_res_ty = dataConOrigResTy dc -- No free type variables
561 [bx_bndr, ubx_bndr] = mkTemplateLocals [dc_res_ty, expr_ty]
562 l1 = incMinorLvl (le_ctxt_lvl env)
563 l2 = incMinorLvl l1
564 = Just ( \e -> Case e (TB ubx_bndr (StayPut l1)) dc_res_ty
565 [(DEFAULT, [], mkConApp dc [Var ubx_bndr])]
566 , \e -> Case e (TB bx_bndr (StayPut l1)) expr_ty
567 [(DataAlt dc, [TB ubx_bndr (StayPut l2)], Var ubx_bndr)] )
568
569 | otherwise -- e.g. do not float unboxed tuples
570 = Nothing
571
572 {- Note [Floating MFEs of unlifted type]
573 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
574 Suppose we have
575 case f x of (r::Int#) -> blah
576 we'd like to float (f x). But it's not trivial because it has type
577 Int#, and we don't want to evaluate it to early. But we can instead
578 float a boxed version
579 y = case f x of r -> I# r
580 and replace the original (f x) with
581 case (case y of I# r -> r) of r -> blah
582
583 Being able to float unboxed expressions is sometimes important; see
584 Trac #12603. I'm not sure how /often/ it is important, but it's
585 not hard to achieve.
586
587 We only do it for a fixed collection of types for which we have a
588 convenient boxing constructor (see boxingDataCon_maybe). In
589 particular we /don't/ do it for unboxed tuples; it's better to float
590 the components of the tuple individually.
591
592 The work is done by canFloat_maybe, which constructs both the code
593 that wraps the floating binding, and the code to appear at the
594 original use site.
595
596 I did experiment with a form of boxing that works for any type, namely
597 wrapping in a function. In our example
598
599 let y = case f x of r -> \v. f x
600 in case y void of r -> blah
601
602 It works fine, but it's 50% slower (based on some crude benchmarking).
603 I suppose we could do it for types not covered by boxingDataCon_maybe,
604 but it's more code and I'll wait to see if anyone wants it.
605
606 Note [Bottoming floats]
607 ~~~~~~~~~~~~~~~~~~~~~~~
608 If we see
609 f = \x. g (error "urk")
610 we'd like to float the call to error, to get
611 lvl = error "urk"
612 f = \x. g lvl
613
614 * Bottoming floats (1): Furthermore, we want to float a bottoming
615 expression even if it has free variables:
616 f = \x. g (let v = h x in error ("urk" ++ v))
617 Then we'd like to abstact over 'x' can float the whole arg of g:
618 lvl = \x. let v = h x in error ("urk" ++ v)
619 f = \x. g (lvl x)
620 To achieve this we pass is_bot to destLevel
621
622 * Bottoming floats (2): And we'd like to do this even if it's a
623 function that guarantees to return bottom:
624 f = \x. ....(\y z. if x then error y else error z)....
625 ===>
626 lvl = \x y z. if b then error y else error z
627 f = \x. ...(lvl x)...
628 To achieve this we use exprBotStrictness_maybe, which spots
629 an expression that diverges after applying some arguments
630
631 See Maessen's paper 1999 "Bottom extraction: factoring error handling out
632 of functional programs" (unpublished I think).
633
634 When we do this, we set the strictness and arity of the new bottoming
635 Id, *immediately*, for three reasons:
636
637 * To prevent the abstracted thing being immediately inlined back in again
638 via preInlineUnconditionally. The latter has a test for bottoming Ids
639 to stop inlining them, so we'd better make sure it *is* a bottoming Id!
640
641 * So that it's properly exposed as such in the interface file, even if
642 this is all happening after strictness analysis.
643
644 * In case we do CSE with the same expression that *is* marked bottom
645 lvl = error "urk"
646 x{str=bot) = error "urk"
647 Here we don't want to replace 'x' with 'lvl', else we may get Lint
648 errors, e.g. via a case with empty alternatives: (case x of {})
649 Lint complains unless the scrutinee of such a case is clearly bottom.
650
651 This was reported in Trac #11290. But since the whole bottoming-float
652 thing is based on the cheap-and-cheerful exprIsBottom, I'm not sure
653 that it'll nail all such cases.
654
655 Note [Bottoming floats: eta expansion] c.f Note [Bottoming floats]
656 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
657 Tiresomely, though, the simplifier has an invariant that the manifest
658 arity of the RHS should be the same as the arity; but we can't call
659 etaExpand during SetLevels because it works over a decorated form of
660 CoreExpr. So we do the eta expansion later, in FloatOut.
661
662 Note [Case MFEs]
663 ~~~~~~~~~~~~~~~~
664 We don't float a case expression as an MFE from a strict context. Why not?
665 Because in doing so we share a tiny bit of computation (the switch) but
666 in exchange we build a thunk, which is bad. This case reduces allocation
667 by 7% in spectral/puzzle (a rather strange benchmark) and 1.2% in real/fem.
668 Doesn't change any other allocation at all.
669 -}
670
671 annotateBotStr :: Id -> Arity -> Maybe (Arity, StrictSig) -> Id
672 -- See Note [Bottoming floats] for why we want to add
673 -- bottoming information right now
674 --
675 -- n_extra are the number of extra value arguments added during floating
676 annotateBotStr id n_extra mb_str
677 = case mb_str of
678 Nothing -> id
679 Just (arity, sig) -> id `setIdArity` (arity + n_extra)
680 `setIdStrictness` (increaseStrictSigArity n_extra sig)
681
682 notWorthFloating :: CoreExpr -> [Var] -> Bool
683 -- Returns True if the expression would be replaced by
684 -- something bigger than it is now. For example:
685 -- abs_vars = tvars only: return True if e is trivial,
686 -- but False for anything bigger
687 -- abs_vars = [x] (an Id): return True for trivial, or an application (f x)
688 -- but False for (f x x)
689 --
690 -- One big goal is that floating should be idempotent. Eg if
691 -- we replace e with (lvl79 x y) and then run FloatOut again, don't want
692 -- to replace (lvl79 x y) with (lvl83 x y)!
693
694 notWorthFloating e abs_vars
695 = go e (count isId abs_vars)
696 where
697 go (Var {}) n = n >= 0
698 go (Lit lit) n = ASSERT( n==0 )
699 litIsTrivial lit -- Note [Floating literals]
700 go (Tick t e) n = not (tickishIsCode t) && go e n
701 go (Cast e _) n = go e n
702 go (App e arg) n
703 | (Type {}) <- arg = go e n
704 | (Coercion {}) <- arg = go e n
705 | n==0 = False
706 | is_triv arg = go e (n-1)
707 | otherwise = False
708 go _ _ = False
709
710 is_triv (Lit {}) = True -- Treat all literals as trivial
711 is_triv (Var {}) = True -- (ie not worth floating)
712 is_triv (Cast e _) = is_triv e
713 is_triv (App e (Type {})) = is_triv e
714 is_triv (App e (Coercion {})) = is_triv e
715 is_triv (Tick t e) = not (tickishIsCode t) && is_triv e
716 is_triv _ = False
717
718 {-
719 Note [Floating literals]
720 ~~~~~~~~~~~~~~~~~~~~~~~~
721 It's important to float Integer literals, so that they get shared,
722 rather than being allocated every time round the loop.
723 Hence the litIsTrivial.
724
725 We'd *like* to share MachStr literal strings too, mainly so we could
726 CSE them, but alas can't do so directly because they are unlifted.
727
728
729 Note [Escaping a value lambda]
730 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
731 We want to float even cheap expressions out of value lambdas,
732 because that saves allocation. Consider
733 f = \x. .. (\y.e) ...
734 Then we'd like to avoid allocating the (\y.e) every time we call f,
735 (assuming e does not mention x). An example where this really makes a
736 difference is simplrun009.
737
738 Another reason it's good is because it makes SpecContr fire on functions.
739 Consider
740 f = \x. ....(f (\y.e))....
741 After floating we get
742 lvl = \y.e
743 f = \x. ....(f lvl)...
744 and that is much easier for SpecConstr to generate a robust
745 specialisation for.
746
747 However, if we are wrapping the thing in extra value lambdas (in
748 abs_vars), then nothing is saved. E.g.
749 f = \xyz. ...(e1[y],e2)....
750 If we float
751 lvl = \y. (e1[y],e2)
752 f = \xyz. ...(lvl y)...
753 we have saved nothing: one pair will still be allocated for each
754 call of 'f'. Hence the (not float_is_lam) in float_me.
755
756
757 ************************************************************************
758 * *
759 \subsection{Bindings}
760 * *
761 ************************************************************************
762
763 The binding stuff works for top level too.
764 -}
765
766 lvlBind :: LevelEnv
767 -> CoreBindWithFVs
768 -> LvlM (LevelledBind, LevelEnv)
769
770 lvlBind env (AnnNonRec bndr rhs)
771 | isTyVar bndr -- Don't do anything for TyVar binders
772 -- (simplifier gets rid of them pronto)
773 || isCoVar bndr -- Difficult to fix up CoVar occurrences (see extendPolyLvlEnv)
774 -- so we will ignore this case for now
775 || not (profitableFloat env dest_lvl)
776 || (isTopLvl dest_lvl && isUnliftedType (idType bndr))
777 -- We can't float an unlifted binding to top level, so we don't
778 -- float it at all. It's a bit brutal, but unlifted bindings
779 -- aren't expensive either
780 = -- No float
781 do { rhs' <- lvlExpr env rhs
782 ; let bind_lvl = incMinorLvl (le_ctxt_lvl env)
783 (env', [bndr']) = substAndLvlBndrs NonRecursive env bind_lvl [bndr]
784 ; return (NonRec bndr' rhs', env') }
785
786 -- Otherwise we are going to float
787 | null abs_vars
788 = do { -- No type abstraction; clone existing binder
789 rhs' <- lvlExpr (setCtxtLvl env dest_lvl) rhs
790 ; (env', [bndr']) <- cloneLetVars NonRecursive env dest_lvl [bndr]
791 ; let bndr2 = annotateBotStr bndr' 0 mb_bot_str
792 ; return (NonRec (TB bndr2 (FloatMe dest_lvl)) rhs', env') }
793
794 | otherwise
795 = do { -- Yes, type abstraction; create a new binder, extend substitution, etc
796 rhs' <- lvlFloatRhs abs_vars dest_lvl env rhs
797 ; (env', [bndr']) <- newPolyBndrs dest_lvl env abs_vars [bndr]
798 ; let bndr2 = annotateBotStr bndr' n_extra mb_bot_str
799 ; return (NonRec (TB bndr2 (FloatMe dest_lvl)) rhs', env') }
800
801 where
802 rhs_fvs = freeVarsOf rhs
803 bind_fvs = rhs_fvs `unionDVarSet` dIdFreeVars bndr
804 abs_vars = abstractVars dest_lvl env bind_fvs
805 dest_lvl = destLevel env bind_fvs (isFunction rhs) is_bot
806 mb_bot_str = exprBotStrictness_maybe (deAnnotate rhs)
807 -- See Note [Bottoming floats]
808 -- esp Bottoming floats (2)
809 is_bot = isJust mb_bot_str
810 n_extra = count isId abs_vars
811
812 lvlBind env (AnnRec pairs)
813 | floatTopLvlOnly env && not (isTopLvl dest_lvl)
814 -- Only floating to the top level is allowed.
815 || not (profitableFloat env dest_lvl)
816 = do { let bind_lvl = incMinorLvl (le_ctxt_lvl env)
817 (env', bndrs') = substAndLvlBndrs Recursive env bind_lvl bndrs
818 ; rhss' <- mapM (lvlExpr env') rhss
819 ; return (Rec (bndrs' `zip` rhss'), env') }
820
821 | null abs_vars
822 = do { (new_env, new_bndrs) <- cloneLetVars Recursive env dest_lvl bndrs
823 ; new_rhss <- mapM (lvlExpr (setCtxtLvl new_env dest_lvl)) rhss
824 ; return ( Rec ([TB b (FloatMe dest_lvl) | b <- new_bndrs] `zip` new_rhss)
825 , new_env) }
826
827 -- ToDo: when enabling the floatLambda stuff,
828 -- I think we want to stop doing this
829 | [(bndr,rhs)] <- pairs
830 , count isId abs_vars > 1
831 = do -- Special case for self recursion where there are
832 -- several variables carried around: build a local loop:
833 -- poly_f = \abs_vars. \lam_vars . letrec f = \lam_vars. rhs in f lam_vars
834 -- This just makes the closures a bit smaller. If we don't do
835 -- this, allocation rises significantly on some programs
836 --
837 -- We could elaborate it for the case where there are several
838 -- mutually functions, but it's quite a bit more complicated
839 --
840 -- This all seems a bit ad hoc -- sigh
841 let (rhs_env, abs_vars_w_lvls) = lvlLamBndrs env dest_lvl abs_vars
842 rhs_lvl = le_ctxt_lvl rhs_env
843
844 (rhs_env', [new_bndr]) <- cloneLetVars Recursive rhs_env rhs_lvl [bndr]
845 let
846 (lam_bndrs, rhs_body) = collectAnnBndrs rhs
847 (body_env1, lam_bndrs1) = substBndrsSL NonRecursive rhs_env' lam_bndrs
848 (body_env2, lam_bndrs2) = lvlLamBndrs body_env1 rhs_lvl lam_bndrs1
849 new_rhs_body <- lvlExpr body_env2 rhs_body
850 (poly_env, [poly_bndr]) <- newPolyBndrs dest_lvl env abs_vars [bndr]
851 return (Rec [(TB poly_bndr (FloatMe dest_lvl)
852 , mkLams abs_vars_w_lvls $
853 mkLams lam_bndrs2 $
854 Let (Rec [( TB new_bndr (StayPut rhs_lvl)
855 , mkLams lam_bndrs2 new_rhs_body)])
856 (mkVarApps (Var new_bndr) lam_bndrs1))]
857 , poly_env)
858
859 | otherwise -- Non-null abs_vars
860 = do { (new_env, new_bndrs) <- newPolyBndrs dest_lvl env abs_vars bndrs
861 ; new_rhss <- mapM (lvlFloatRhs abs_vars dest_lvl new_env) rhss
862 ; return ( Rec ([TB b (FloatMe dest_lvl) | b <- new_bndrs] `zip` new_rhss)
863 , new_env) }
864
865 where
866 (bndrs,rhss) = unzip pairs
867
868 -- Finding the free vars of the binding group is annoying
869 bind_fvs = ((unionDVarSets [ freeVarsOf rhs | (_, rhs) <- pairs])
870 `unionDVarSet`
871 (fvDVarSet $ unionsFV [ idFVs bndr
872 | (bndr, (_,_)) <- pairs]))
873 `delDVarSetList`
874 bndrs
875
876 dest_lvl = destLevel env bind_fvs (all isFunction rhss) False
877 abs_vars = abstractVars dest_lvl env bind_fvs
878
879 profitableFloat :: LevelEnv -> Level -> Bool
880 profitableFloat env dest_lvl
881 = (dest_lvl `ltMajLvl` le_ctxt_lvl env) -- Escapes a value lambda
882 || isTopLvl dest_lvl -- Going all the way to top level
883
884 ----------------------------------------------------
885 -- Three help functions for the type-abstraction case
886
887 lvlFloatRhs :: [OutVar] -> Level -> LevelEnv -> CoreExprWithFVs
888 -> UniqSM (Expr LevelledBndr)
889 lvlFloatRhs abs_vars dest_lvl env rhs
890 = do { body' <- lvlExpr rhs_env body
891 ; return (mkLams all_bndrs_w_lvls body') }
892 where
893 (bndrs, body) = collectAnnBndrs rhs
894 (env1, bndrs1) = substBndrsSL NonRecursive env bndrs
895 all_bndrs = abs_vars ++ bndrs1
896 (rhs_env, all_bndrs_w_lvls) = lvlLamBndrs env1 dest_lvl all_bndrs
897 -- The important thing here is that we call lvlLamBndrs on
898 -- all these binders at once (abs_vars and bndrs), so they
899 -- all get the same major level. Otherwise we create stupid
900 -- let-bindings inside, joyfully thinking they can float; but
901 -- in the end they don't because we never float bindings in
902 -- between lambdas
903
904 {-
905 ************************************************************************
906 * *
907 \subsection{Deciding floatability}
908 * *
909 ************************************************************************
910 -}
911
912 substAndLvlBndrs :: RecFlag -> LevelEnv -> Level -> [InVar] -> (LevelEnv, [LevelledBndr])
913 substAndLvlBndrs is_rec env lvl bndrs
914 = lvlBndrs subst_env lvl subst_bndrs
915 where
916 (subst_env, subst_bndrs) = substBndrsSL is_rec env bndrs
917
918 substBndrsSL :: RecFlag -> LevelEnv -> [InVar] -> (LevelEnv, [OutVar])
919 -- So named only to avoid the name clash with CoreSubst.substBndrs
920 substBndrsSL is_rec env@(LE { le_subst = subst, le_env = id_env }) bndrs
921 = ( env { le_subst = subst'
922 , le_env = foldl add_id id_env (bndrs `zip` bndrs') }
923 , bndrs')
924 where
925 (subst', bndrs') = case is_rec of
926 NonRecursive -> substBndrs subst bndrs
927 Recursive -> substRecBndrs subst bndrs
928
929 lvlLamBndrs :: LevelEnv -> Level -> [OutVar] -> (LevelEnv, [LevelledBndr])
930 -- Compute the levels for the binders of a lambda group
931 lvlLamBndrs env lvl bndrs
932 = lvlBndrs env new_lvl bndrs
933 where
934 new_lvl | any is_major bndrs = incMajorLvl lvl
935 | otherwise = incMinorLvl lvl
936
937 is_major bndr = isId bndr && not (isProbablyOneShotLambda bndr)
938 -- The "probably" part says "don't float things out of a
939 -- probable one-shot lambda"
940 -- See Note [Computing one-shot info] in Demand.hs
941
942
943 lvlBndrs :: LevelEnv -> Level -> [CoreBndr] -> (LevelEnv, [LevelledBndr])
944 -- The binders returned are exactly the same as the ones passed,
945 -- apart from applying the substitution, but they are now paired
946 -- with a (StayPut level)
947 --
948 -- The returned envt has ctxt_lvl updated to the new_lvl
949 --
950 -- All the new binders get the same level, because
951 -- any floating binding is either going to float past
952 -- all or none. We never separate binders.
953 lvlBndrs env@(LE { le_lvl_env = lvl_env }) new_lvl bndrs
954 = ( env { le_ctxt_lvl = new_lvl
955 , le_lvl_env = addLvls new_lvl lvl_env bndrs }
956 , lvld_bndrs)
957 where
958 lvld_bndrs = [TB bndr (StayPut new_lvl) | bndr <- bndrs]
959
960 -- Destination level is the max Id level of the expression
961 -- (We'll abstract the type variables, if any.)
962 destLevel :: LevelEnv -> DVarSet
963 -> Bool -- True <=> is function
964 -> Bool -- True <=> is bottom
965 -> Level
966 destLevel env fvs is_function is_bot
967 | is_bot = tOP_LEVEL -- Send bottoming bindings to the top
968 -- regardless; see Note [Bottoming floats]
969 -- Esp Bottoming floats (1)
970 | Just n_args <- floatLams env
971 , n_args > 0 -- n=0 case handled uniformly by the 'otherwise' case
972 , is_function
973 , countFreeIds fvs <= n_args
974 = tOP_LEVEL -- Send functions to top level; see
975 -- the comments with isFunction
976
977 | otherwise = maxFvLevel isId env fvs -- Max over Ids only; the tyvars
978 -- will be abstracted
979
980 isFunction :: CoreExprWithFVs -> Bool
981 -- The idea here is that we want to float *functions* to
982 -- the top level. This saves no work, but
983 -- (a) it can make the host function body a lot smaller,
984 -- and hence inlinable.
985 -- (b) it can also save allocation when the function is recursive:
986 -- h = \x -> letrec f = \y -> ...f...y...x...
987 -- in f x
988 -- becomes
989 -- f = \x y -> ...(f x)...y...x...
990 -- h = \x -> f x x
991 -- No allocation for f now.
992 -- We may only want to do this if there are sufficiently few free
993 -- variables. We certainly only want to do it for values, and not for
994 -- constructors. So the simple thing is just to look for lambdas
995 isFunction (_, AnnLam b e) | isId b = True
996 | otherwise = isFunction e
997 -- isFunction (_, AnnTick _ e) = isFunction e -- dubious
998 isFunction _ = False
999
1000 countFreeIds :: DVarSet -> Int
1001 countFreeIds = nonDetFoldUDFM add 0
1002 -- It's OK to use nonDetFoldUDFM here because we're just counting things.
1003 where
1004 add :: Var -> Int -> Int
1005 add v n | isId v = n+1
1006 | otherwise = n
1007
1008 {-
1009 ************************************************************************
1010 * *
1011 \subsection{Free-To-Level Monad}
1012 * *
1013 ************************************************************************
1014 -}
1015
1016 data LevelEnv
1017 = LE { le_switches :: FloatOutSwitches
1018 , le_ctxt_lvl :: Level -- The current level
1019 , le_lvl_env :: VarEnv Level -- Domain is *post-cloned* TyVars and Ids
1020 , le_subst :: Subst -- Domain is pre-cloned TyVars and Ids
1021 -- The Id -> CoreExpr in the Subst is ignored
1022 -- (since we want to substitute a LevelledExpr for
1023 -- an Id via le_env) but we do use the Co/TyVar substs
1024 , le_env :: IdEnv ([OutVar], LevelledExpr) -- Domain is pre-cloned Ids
1025 }
1026 -- We clone let- and case-bound variables so that they are still
1027 -- distinct when floated out; hence the le_subst/le_env.
1028 -- (see point 3 of the module overview comment).
1029 -- We also use these envs when making a variable polymorphic
1030 -- because we want to float it out past a big lambda.
1031 --
1032 -- The le_subst and le_env always implement the same mapping, but the
1033 -- le_subst maps to CoreExpr and the le_env to LevelledExpr
1034 -- Since the range is always a variable or type application,
1035 -- there is never any difference between the two, but sadly
1036 -- the types differ. The le_subst is used when substituting in
1037 -- a variable's IdInfo; the le_env when we find a Var.
1038 --
1039 -- In addition the le_env records a list of tyvars free in the
1040 -- type application, just so we don't have to call freeVars on
1041 -- the type application repeatedly.
1042 --
1043 -- The domain of the both envs is *pre-cloned* Ids, though
1044 --
1045 -- The domain of the le_lvl_env is the *post-cloned* Ids
1046
1047 initialEnv :: FloatOutSwitches -> LevelEnv
1048 initialEnv float_lams
1049 = LE { le_switches = float_lams
1050 , le_ctxt_lvl = tOP_LEVEL
1051 , le_lvl_env = emptyVarEnv
1052 , le_subst = emptySubst
1053 , le_env = emptyVarEnv }
1054
1055 addLvl :: Level -> VarEnv Level -> OutVar -> VarEnv Level
1056 addLvl dest_lvl env v' = extendVarEnv env v' dest_lvl
1057
1058 addLvls :: Level -> VarEnv Level -> [OutVar] -> VarEnv Level
1059 addLvls dest_lvl env vs = foldl (addLvl dest_lvl) env vs
1060
1061 floatLams :: LevelEnv -> Maybe Int
1062 floatLams le = floatOutLambdas (le_switches le)
1063
1064 floatConsts :: LevelEnv -> Bool
1065 floatConsts le = floatOutConstants (le_switches le)
1066
1067 floatOverSat :: LevelEnv -> Bool
1068 floatOverSat le = floatOutOverSatApps (le_switches le)
1069
1070 floatTopLvlOnly :: LevelEnv -> Bool
1071 floatTopLvlOnly le = floatToTopLevelOnly (le_switches le)
1072
1073 setCtxtLvl :: LevelEnv -> Level -> LevelEnv
1074 setCtxtLvl env lvl = env { le_ctxt_lvl = lvl }
1075
1076 -- extendCaseBndrEnv adds the mapping case-bndr->scrut-var if it can
1077 -- See Note [Binder-swap during float-out]
1078 extendCaseBndrEnv :: LevelEnv
1079 -> Id -- Pre-cloned case binder
1080 -> Expr LevelledBndr -- Post-cloned scrutinee
1081 -> LevelEnv
1082 extendCaseBndrEnv le@(LE { le_subst = subst, le_env = id_env })
1083 case_bndr (Var scrut_var)
1084 = le { le_subst = extendSubstWithVar subst case_bndr scrut_var
1085 , le_env = add_id id_env (case_bndr, scrut_var) }
1086 extendCaseBndrEnv env _ _ = env
1087
1088 maxFvLevel :: (Var -> Bool) -> LevelEnv -> DVarSet -> Level
1089 maxFvLevel max_me (LE { le_lvl_env = lvl_env, le_env = id_env }) var_set
1090 = foldDVarSet max_in tOP_LEVEL var_set
1091 where
1092 max_in in_var lvl
1093 = foldr max_out lvl (case lookupVarEnv id_env in_var of
1094 Just (abs_vars, _) -> abs_vars
1095 Nothing -> [in_var])
1096
1097 max_out out_var lvl
1098 | max_me out_var = case lookupVarEnv lvl_env out_var of
1099 Just lvl' -> maxLvl lvl' lvl
1100 Nothing -> lvl
1101 | otherwise = lvl -- Ignore some vars depending on max_me
1102
1103 lookupVar :: LevelEnv -> Id -> LevelledExpr
1104 lookupVar le v = case lookupVarEnv (le_env le) v of
1105 Just (_, expr) -> expr
1106 _ -> Var v
1107
1108 abstractVars :: Level -> LevelEnv -> DVarSet -> [OutVar]
1109 -- Find the variables in fvs, free vars of the target expression,
1110 -- whose level is greater than the destination level
1111 -- These are the ones we are going to abstract out
1112 --
1113 -- Note that to get reproducible builds, the variables need to be
1114 -- abstracted in deterministic order, not dependent on the values of
1115 -- Uniques. This is achieved by using DVarSets, deterministic free
1116 -- variable computation and deterministic sort.
1117 -- See Note [Unique Determinism] in Unique for explanation of why
1118 -- Uniques are not deterministic.
1119 abstractVars dest_lvl (LE { le_subst = subst, le_lvl_env = lvl_env }) in_fvs
1120 = -- NB: sortQuantVars might not put duplicates next to each other
1121 map zap $ sortQuantVars $ uniq
1122 [out_var | out_fv <- dVarSetElems (substDVarSet subst in_fvs)
1123 , out_var <- dVarSetElems (close out_fv)
1124 , abstract_me out_var ]
1125 -- NB: it's important to call abstract_me only on the OutIds the
1126 -- come from substDVarSet (not on fv, which is an InId)
1127 where
1128 uniq :: [Var] -> [Var]
1129 -- Remove duplicates, preserving order
1130 uniq = dVarSetElems . mkDVarSet
1131
1132 abstract_me v = case lookupVarEnv lvl_env v of
1133 Just lvl -> dest_lvl `ltLvl` lvl
1134 Nothing -> False
1135
1136 -- We are going to lambda-abstract, so nuke any IdInfo,
1137 -- and add the tyvars of the Id (if necessary)
1138 zap v | isId v = WARN( isStableUnfolding (idUnfolding v) ||
1139 not (isEmptyRuleInfo (idSpecialisation v)),
1140 text "absVarsOf: discarding info on" <+> ppr v )
1141 setIdInfo v vanillaIdInfo
1142 | otherwise = v
1143
1144 close :: Var -> DVarSet -- Close over variables free in the type
1145 -- Result includes the input variable itself
1146 close v = foldDVarSet (unionDVarSet . close)
1147 (unitDVarSet v)
1148 (fvDVarSet $ varTypeTyCoFVs v)
1149
1150 type LvlM result = UniqSM result
1151
1152 initLvl :: UniqSupply -> UniqSM a -> a
1153 initLvl = initUs_
1154
1155 newPolyBndrs :: Level -> LevelEnv -> [OutVar] -> [InId] -> UniqSM (LevelEnv, [OutId])
1156 -- The envt is extended to bind the new bndrs to dest_lvl, but
1157 -- the ctxt_lvl is unaffected
1158 newPolyBndrs dest_lvl
1159 env@(LE { le_lvl_env = lvl_env, le_subst = subst, le_env = id_env })
1160 abs_vars bndrs
1161 = ASSERT( all (not . isCoVar) bndrs ) -- What would we add to the CoSubst in this case. No easy answer.
1162 do { uniqs <- getUniquesM
1163 ; let new_bndrs = zipWith mk_poly_bndr bndrs uniqs
1164 bndr_prs = bndrs `zip` new_bndrs
1165 env' = env { le_lvl_env = addLvls dest_lvl lvl_env new_bndrs
1166 , le_subst = foldl add_subst subst bndr_prs
1167 , le_env = foldl add_id id_env bndr_prs }
1168 ; return (env', new_bndrs) }
1169 where
1170 add_subst env (v, v') = extendIdSubst env v (mkVarApps (Var v') abs_vars)
1171 add_id env (v, v') = extendVarEnv env v ((v':abs_vars), mkVarApps (Var v') abs_vars)
1172
1173 mk_poly_bndr bndr uniq = transferPolyIdInfo bndr abs_vars $ -- Note [transferPolyIdInfo] in Id.hs
1174 mkSysLocalOrCoVar (mkFastString str) uniq poly_ty
1175 where
1176 str = "poly_" ++ occNameString (getOccName bndr)
1177 poly_ty = mkLamTypes abs_vars (CoreSubst.substTy subst (idType bndr))
1178
1179 newLvlVar :: LevelledExpr -- The RHS of the new binding
1180 -> LvlM Id
1181 newLvlVar lvld_rhs
1182 = do { uniq <- getUniqueM
1183 ; return (mk_id uniq rhs_ty) }
1184 where
1185 de_tagged_rhs = deTagExpr lvld_rhs
1186 rhs_ty = exprType de_tagged_rhs
1187
1188 mk_id uniq rhs_ty
1189 -- See Note [Grand plan for static forms] in StaticPtrTable.
1190 | isJust $ collectMakeStaticArgs $ snd $
1191 collectTyBinders de_tagged_rhs
1192 = mkExportedVanillaId (mkSystemVarName uniq (mkFastString "static_ptr"))
1193 rhs_ty
1194 | otherwise
1195 = mkLocalIdOrCoVar (mkSystemVarName uniq (mkFastString "lvl")) rhs_ty
1196
1197 cloneCaseBndrs :: LevelEnv -> Level -> [Var] -> LvlM (LevelEnv, [Var])
1198 cloneCaseBndrs env@(LE { le_subst = subst, le_lvl_env = lvl_env, le_env = id_env })
1199 new_lvl vs
1200 = do { us <- getUniqueSupplyM
1201 ; let (subst', vs') = cloneBndrs subst us vs
1202 env' = env { le_ctxt_lvl = new_lvl
1203 , le_lvl_env = addLvls new_lvl lvl_env vs'
1204 , le_subst = subst'
1205 , le_env = foldl add_id id_env (vs `zip` vs') }
1206
1207 ; return (env', vs') }
1208
1209 cloneLetVars :: RecFlag -> LevelEnv -> Level -> [Var] -> LvlM (LevelEnv, [Var])
1210 -- See Note [Need for cloning during float-out]
1211 -- Works for Ids bound by let(rec)
1212 -- The dest_lvl is attributed to the binders in the new env,
1213 -- but cloneVars doesn't affect the ctxt_lvl of the incoming env
1214 cloneLetVars is_rec
1215 env@(LE { le_subst = subst, le_lvl_env = lvl_env, le_env = id_env })
1216 dest_lvl vs
1217 = do { us <- getUniqueSupplyM
1218 ; let (subst', vs1) = case is_rec of
1219 NonRecursive -> cloneBndrs subst us vs
1220 Recursive -> cloneRecIdBndrs subst us vs
1221 vs2 = map zap_demand_info vs1 -- See Note [Zapping the demand info]
1222 prs = vs `zip` vs2
1223 env' = env { le_lvl_env = addLvls dest_lvl lvl_env vs2
1224 , le_subst = subst'
1225 , le_env = foldl add_id id_env prs }
1226
1227 ; return (env', vs2) }
1228
1229 add_id :: IdEnv ([Var], LevelledExpr) -> (Var, Var) -> IdEnv ([Var], LevelledExpr)
1230 add_id id_env (v, v1)
1231 | isTyVar v = delVarEnv id_env v
1232 | otherwise = extendVarEnv id_env v ([v1], ASSERT(not (isCoVar v1)) Var v1)
1233
1234 zap_demand_info :: Var -> Var
1235 zap_demand_info v
1236 | isId v = zapIdDemandInfo v
1237 | otherwise = v
1238
1239 {-
1240 Note [Zapping the demand info]
1241 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
1242 VERY IMPORTANT: we must zap the demand info if the thing is going to
1243 float out, because it may be less demanded than at its original
1244 binding site. Eg
1245 f :: Int -> Int
1246 f x = let v = 3*4 in v+x
1247 Here v is strict; but if we float v to top level, it isn't any more.
1248 -}