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