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