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