Super-monster patch implementing the new typechecker -- at last
[ghc-hetmet.git] / compiler / simplCore / FloatOut.lhs
index bbf6f0a..fba88e7 100644 (file)
@@ -6,30 +6,25 @@
 ``Long-distance'' floating of bindings towards the top level.
 
 \begin{code}
-{-# OPTIONS -w #-}
--- The above warning supression flag is a temporary kludge.
--- While working on this module you are encouraged to remove it and fix
--- any warnings in the module. See
---     http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#Warnings
--- for details
-
 module FloatOut ( floatOutwards ) where
 
-#include "HsVersions.h"
-
 import CoreSyn
 import CoreUtils
+import CoreArity       ( etaExpand )
+import CoreMonad       ( FloatOutSwitches(..) )
 
-import DynFlags        ( DynFlags, DynFlag(..), FloatOutSwitches(..) )
+import DynFlags                ( DynFlags, DynFlag(..) )
 import ErrUtils                ( dumpIfSet_dyn )
 import CostCentre      ( dupifyCC, CostCentre )
-import Id              ( Id, idType )
+import Id              ( Id, idType, idArity, isBottomingId )
 import Type            ( isUnLiftedType )
-import CoreLint                ( showPass, endPass )
 import SetLevels       ( Level(..), LevelledExpr, LevelledBind,
-                         setLevels, ltMajLvl, ltLvl, isTopLvl )
+                         setLevels, isTopLvl )
 import UniqSupply       ( UniqSupply )
-import List            ( partition )
+import Bag
+import Util
+import Maybes
+import UniqFM
 import Outputable
 import FastString
 \end{code}
@@ -106,10 +101,6 @@ vwhich might usefully be separated to
 @
 Well, maybe.  We don't do this at the moment.
 
-\begin{code}
-type FloatBind     = (Level, CoreBind) -- INVARIANT: a FloatBind is always lifted
-type FloatBinds    = [FloatBind]       
-\end{code}
 
 %************************************************************************
 %*                                                                     *
@@ -125,8 +116,6 @@ floatOutwards :: FloatOutSwitches
 
 floatOutwards float_sws dflags us pgm
   = do {
-       showPass dflags float_msg ;
-
        let { annotated_w_levels = setLevels float_sws pgm us ;
              (fss, binds_s')    = unzip (map floatTopBind annotated_w_levels)
            } ;
@@ -137,24 +126,17 @@ floatOutwards float_sws dflags us pgm
        let { (tlets, ntlets, lams) = get_stats (sum_stats fss) };
 
        dumpIfSet_dyn dflags Opt_D_dump_simpl_stats "FloatOut stats:"
-               (hcat [ int tlets,  ptext SLIT(" Lets floated to top level; "),
-                       int ntlets, ptext SLIT(" Lets floated elsewhere; from "),
-                       int lams,   ptext SLIT(" Lambda groups")]);
+               (hcat [ int tlets,  ptext (sLit " Lets floated to top level; "),
+                       int ntlets, ptext (sLit " Lets floated elsewhere; from "),
+                       int lams,   ptext (sLit " Lambda groups")]);
 
-       endPass dflags float_msg  Opt_D_verbose_core2core (concat binds_s')
-                       {- no specific flag for dumping float-out -} 
+       return (concat binds_s')
     }
-  where
-    float_msg = showSDoc (text "Float out" <+> parens (sws float_sws))
-    sws (FloatOutSw lam const) = pp_not lam   <+> text "lambdas" <> comma <+>
-                                pp_not const <+> text "constants"
-    pp_not True  = empty
-    pp_not False = text "not"
 
+floatTopBind :: LevelledBind -> (FloatStats, [CoreBind])
 floatTopBind bind
   = case (floatBind bind) of { (fs, floats) ->
-    (fs, floatsToBinds floats)
-    }
+    (fs, bagToList (flattenFloats floats)) }
 \end{code}
 
 %************************************************************************
@@ -163,52 +145,56 @@ floatTopBind bind
 %*                                                                     *
 %************************************************************************
 
-
 \begin{code}
 floatBind :: LevelledBind -> (FloatStats, FloatBinds)
-
-floatBind (NonRec (TB name level) rhs)
+floatBind (NonRec (TB var level) rhs)
   = case (floatRhs level rhs) of { (fs, rhs_floats, rhs') ->
-    (fs, rhs_floats ++ [(level, NonRec name rhs')]) }
 
-floatBind bind@(Rec pairs)
-  = case (unzip3 (map do_pair pairs)) of { (fss, rhss_floats, new_pairs) ->
-    let rhs_floats = concat rhss_floats in
+       -- A tiresome hack: 
+       -- see Note [Bottoming floats: eta expansion] in SetLevels
+    let rhs'' | isBottomingId var = etaExpand (idArity var) rhs'
+             | otherwise         = rhs'
 
-    if not (isTopLvl bind_dest_lvl) then
+    in (fs, rhs_floats `plusFloats` unitFloat level (NonRec var rhs'')) }
+
+floatBind (Rec pairs)
+  = case floatList do_pair pairs of { (fs, rhs_floats, new_pairs) ->
+        -- NB: the rhs floats may contain references to the 
+       -- bound things.  For example
+       --      f = ...(let v = ...f... in b) ...
+    if not (isTopLvl dest_lvl) then
        -- Find which bindings float out at least one lambda beyond this one
        -- These ones can't mention the binders, because they couldn't 
        -- be escaping a major level if so.
        -- The ones that are not going further can join the letrec;
        -- they may not be mutually recursive but the occurrence analyser will
-       -- find that out.
-       case (partitionByMajorLevel bind_dest_lvl rhs_floats) of { (floats', heres) ->
-       (sum_stats fss, floats' ++ [(bind_dest_lvl, Rec (floatsToBindPairs heres ++ new_pairs))]) }
-    else
-       -- In a recursive binding, *destined for* the top level
-       -- (only), the rhs floats may contain references to the 
-       -- bound things.  For example
-       --      f = ...(let v = ...f... in b) ...
-       --  might get floated to
+       -- find that out. In our example we make a Rec thus:
        --      v = ...f...
        --      f = ... b ...
-       -- and hence we must (pessimistically) make all the floats recursive
-       -- with the top binding.  Later dependency analysis will unravel it.
-       --
-       -- This can only happen for bindings destined for the top level,
-       -- because only then will partitionByMajorLevel allow through a binding
-       -- that only differs in its minor level
-       (sum_stats fss, [(bind_dest_lvl, Rec (new_pairs ++ floatsToBindPairs rhs_floats))])
-    }
+       case (partitionByMajorLevel dest_lvl rhs_floats) of { (floats', heres) ->
+       (fs, floats' `plusFloats` unitFloat dest_lvl 
+                (Rec (floatsToBindPairs heres new_pairs))) }
+    else
+       -- For top level, no need to partition; just make them all recursive
+       -- (And the partition wouldn't work because they'd all end up in floats')
+       (fs, unitFloat dest_lvl
+                (Rec (floatsToBindPairs (flattenFloats rhs_floats) new_pairs)))  }
   where
-    bind_dest_lvl = getBindLevel bind
+    (((TB _ dest_lvl), _) : _) = pairs
 
     do_pair (TB name level, rhs)
       = case (floatRhs level rhs) of { (fs, rhs_floats, rhs') ->
-       (fs, rhs_floats, (name, rhs'))
-       }
+       (fs, rhs_floats, (name, rhs')) }
+
+---------------
+floatList :: (a -> (FloatStats, FloatBinds, b)) -> [a] -> (FloatStats, FloatBinds, [b])
+floatList _ [] = (zeroStats, emptyFloats, [])
+floatList f (a:as) = case f a           of { (fs_a,  binds_a,  b)  ->
+                    case floatList f as of { (fs_as, binds_as, bs) ->
+                    (fs_a `add_stats` fs_as, binds_a `plusFloats`  binds_as, b:bs) }}
 \end{code}
 
+
 %************************************************************************
 
 \subsection[FloatOut-Expr]{Floating in expressions}
@@ -229,53 +215,26 @@ floatCaseAlt lvl arg      -- Used rec rhss, and case-alternative rhss
        -- the rec or case alternative
     (fsa, floats', install heres arg') }}
 
+-----------------
 floatRhs lvl arg       -- Used for nested non-rec rhss, and fn args
                        -- See Note [Floating out of RHS]
-  = case (floatExpr lvl arg) of { (fsa, floats, arg') ->
-    if exprIsCheap arg' then   
-       (fsa, floats, arg')
-    else
-    case (partitionByMajorLevel lvl floats) of { (floats', heres) ->
-    (fsa, floats', install heres arg') }}
+  = floatExpr lvl arg
 
--- Note [Floating out of RHSs]
--- ~~~~~~~~~~~~~~~~~~~~~~~~~~~
--- Dump bindings that aren't going to escape from a lambda
--- This isn't a scoping issue (the binder isn't in scope in the RHS 
---     of a non-rec binding)
--- Rather, it is to avoid floating the x binding out of
---     f (let x = e in b)
--- unnecessarily.  But we first test for values or trival rhss,
--- because (in particular) we don't want to insert new bindings between
--- the "=" and the "\".  E.g.
---     f = \x -> let <bind> in <body>
--- We do not want
---     f = let <bind> in \x -> <body>
--- (a) The simplifier will immediately float it further out, so we may
---     as well do so right now; in general, keeping rhss as manifest 
---     values is good
--- (b) If a float-in pass follows immediately, it might add yet more
---     bindings just after the '='.  And some of them might (correctly)
---     be strict even though the 'let f' is lazy, because f, being a value,
---     gets its demand-info zapped by the simplifier.
---
--- We use exprIsCheap because that is also what's used by the simplifier
--- to decide whether to float a let out of a let
-
-floatExpr _ (Var v)   = (zeroStats, [], Var v)
-floatExpr _ (Type ty) = (zeroStats, [], Type ty)
-floatExpr _ (Lit lit) = (zeroStats, [], Lit lit)
+-----------------
+floatExpr _ (Var v)   = (zeroStats, emptyFloats, Var v)
+floatExpr _ (Type ty) = (zeroStats, emptyFloats, Type ty)
+floatExpr _ (Lit lit) = (zeroStats, emptyFloats, Lit lit)
          
 floatExpr lvl (App e a)
   = case (floatExpr      lvl e) of { (fse, floats_e, e') ->
     case (floatRhs lvl a)      of { (fsa, floats_a, a') ->
-    (fse `add_stats` fsa, floats_e ++ floats_a, App e' a') }}
+    (fse `add_stats` fsa, floats_e `plusFloats` floats_a, App e' a') }}
 
-floatExpr lvl lam@(Lam _ _)
+floatExpr _ lam@(Lam _ _)
   = let
        (bndrs_w_lvls, body) = collectBinders lam
        bndrs                = [b | TB b _ <- bndrs_w_lvls]
-       lvls                 = [l | TB b l <- bndrs_w_lvls]
+       lvls                 = [l | TB _ l <- bndrs_w_lvls]
 
        -- For the all-tyvar case we are prepared to pull 
        -- the lets out, to implement the float-out-of-big-lambda
@@ -283,7 +242,7 @@ floatExpr lvl lam@(Lam _ _)
        -- going to escape a value lambda.
        -- In particular, for one-shot lambdas we don't float things
        -- out; we get no saving by so doing.
-       partition_fn | all isTyVar bndrs = partitionByLevel
+       partition_fn | all isTyCoVar bndrs = partitionByLevel
                     | otherwise         = partitionByMajorLevel
     in
     case (floatExpr (last lvls) body) of { (fs, floats, body') ->
@@ -300,27 +259,9 @@ floatExpr lvl (Note note@(SCC cc) expr)
        -- Annotate bindings floated outwards past an scc expression
        -- with the cc.  We mark that cc as "duplicated", though.
 
-       annotated_defns = annotate (dupifyCC cc) floating_defns
+       annotated_defns = wrapCostCentre (dupifyCC cc) floating_defns
     in
     (fs, annotated_defns, Note note expr') }
-  where
-    annotate :: CostCentre -> FloatBinds -> FloatBinds
-
-    annotate dupd_cc defn_groups
-      = [ (level, ann_bind floater) | (level, floater) <- defn_groups ]
-      where
-       ann_bind (NonRec binder rhs)
-         = NonRec binder (mkSCC dupd_cc rhs)
-
-       ann_bind (Rec pairs)
-         = Rec [(binder, mkSCC dupd_cc rhs) | (binder, rhs) <- pairs]
-
-floatExpr lvl (Note InlineMe expr)     -- Other than SCCs
-  = (zeroStats, [], Note InlineMe (unTag expr))
-       -- Do no floating at all inside INLINE. [_$_]
-       -- The SetLevels pass did not clone the bindings, so it's
-       -- unsafe to do any floating, even if we dump the results
-       -- inside the Note (which is what we used to do).
 
 floatExpr lvl (Note note expr) -- Other than SCCs
   = case (floatExpr lvl expr)    of { (fs, floating_defns, expr') ->
@@ -331,23 +272,24 @@ floatExpr lvl (Cast expr co)
     (fs, floating_defns, Cast expr' co) }
 
 floatExpr lvl (Let (NonRec (TB bndr bndr_lvl) rhs) body)
-  | isUnLiftedType (idType bndr)       -- Treat unlifted lets just like a case
-                               -- I.e. floatExpr for rhs, floatCaseAlt for body
-  = case floatExpr lvl rhs         of { (fs, rhs_floats, rhs') ->
+  | isUnLiftedType (idType bndr)  -- Treat unlifted lets just like a case
+                                 -- I.e. floatExpr for rhs, floatCaseAlt for body
+  = case floatExpr lvl rhs         of { (_, rhs_floats, rhs') ->
     case floatCaseAlt bndr_lvl body of { (fs, body_floats, body') ->
-    (fs, rhs_floats ++ body_floats, Let (NonRec bndr rhs') body') }}
+    (fs, rhs_floats `plusFloats` body_floats, Let (NonRec bndr rhs') body') }}
 
 floatExpr lvl (Let bind body)
   = case (floatBind bind)     of { (fsb, bind_floats) ->
     case (floatExpr lvl body) of { (fse, body_floats, body') ->
-    (add_stats fsb fse,
-     bind_floats ++ body_floats,
-     body')  }}
+    case partitionByMajorLevel lvl (bind_floats `plusFloats` body_floats) 
+                              of { (floats, heres) ->
+       -- See Note [Avoiding unnecessary floating]
+    (add_stats fsb fse, floats, install heres body')  } } }
 
 floatExpr lvl (Case scrut (TB case_bndr case_lvl) ty alts)
   = case floatExpr lvl scrut   of { (fse, fde, scrut') ->
     case floatList float_alt alts      of { (fsa, fda, alts')  ->
-    (add_stats fse fsa, fda ++ fde, Case scrut' case_bndr ty alts')
+    (add_stats fse fsa, fda `plusFloats` fde, Case scrut' case_bndr ty alts')
     }}
   where
        -- Use floatCaseAlt for the alternatives, so that we
@@ -355,31 +297,43 @@ floatExpr lvl (Case scrut (TB case_bndr case_lvl) ty alts)
     float_alt (con, bs, rhs)
        = case (floatCaseAlt case_lvl rhs)      of { (fs, rhs_floats, rhs') ->
          (fs, rhs_floats, (con, [b | TB b _ <- bs], rhs')) }
-
-
-floatList :: (a -> (FloatStats, FloatBinds, b)) -> [a] -> (FloatStats, FloatBinds, [b])
-floatList f [] = (zeroStats, [], [])
-floatList f (a:as) = case f a           of { (fs_a,  binds_a,  b)  ->
-                    case floatList f as of { (fs_as, binds_as, bs) ->
-                    (fs_a `add_stats` fs_as, binds_a ++ binds_as, b:bs) }}
-
-unTagBndr :: TaggedBndr tag -> CoreBndr
-unTagBndr (TB b _) = b
-
-unTag :: TaggedExpr tag -> CoreExpr
-unTag (Var v)            = Var v
-unTag (Lit l)            = Lit l
-unTag (Type ty)   = Type ty
-unTag (Note n e)  = Note n (unTag e)
-unTag (App e1 e2) = App (unTag e1) (unTag e2)
-unTag (Lam b e)   = Lam (unTagBndr b) (unTag e)
-unTag (Cast e co) = Cast (unTag e) co
-unTag (Let (Rec prs) e)    = Let (Rec [(unTagBndr b,unTag r) | (b, r) <- prs]) (unTag e)
-unTag (Let (NonRec b r) e) = Let (NonRec (unTagBndr b) (unTag r)) (unTag e)
-unTag (Case e b ty alts)   = Case (unTag e) (unTagBndr b) ty
-                                 [(c, map unTagBndr bs, unTag r) | (c,bs,r) <- alts]
 \end{code}
 
+Note [Avoiding unnecessary floating]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+In general we want to avoid floating a let unnecessarily, because
+it might worsen strictness:
+    let 
+       x = ...(let y = e in y+y)....
+Here y is demanded.  If we float it outside the lazy 'x=..' then
+we'd have to zap its demand info, and it may never be restored.
+
+So at a 'let' we leave the binding right where the are unless
+the binding will escape a value lambda.  That's what the 
+partitionByMajorLevel does in the floatExpr (Let ...) case.
+
+Notice, though, that we must take care to drop any bindings
+from the body of the let that depend on the staying-put bindings.
+
+We used instead to do the partitionByMajorLevel on the RHS of an '=',
+in floatRhs.  But that was quite tiresome.  We needed to test for
+values or trival rhss, because (in particular) we don't want to insert
+new bindings between the "=" and the "\".  E.g.
+       f = \x -> let <bind> in <body>
+We do not want
+       f = let <bind> in \x -> <body>
+(a) The simplifier will immediately float it further out, so we may
+       as well do so right now; in general, keeping rhss as manifest 
+       values is good
+(b) If a float-in pass follows immediately, it might add yet more
+       bindings just after the '='.  And some of them might (correctly)
+       be strict even though the 'let f' is lazy, because f, being a value,
+       gets its demand-info zapped by the simplifier.
+And even all that turned out to be very fragile, and broke
+altogether when profiling got in the way.
+
+So now we do the partition right at the (Let..) itself.
+
 %************************************************************************
 %*                                                                     *
 \subsection{Utility bits for floating stats}
@@ -395,21 +349,22 @@ data FloatStats
        Int  -- Number of non-top-floats * lambda groups they've been past
        Int  -- Number of lambda (groups) seen
 
+get_stats :: FloatStats -> (Int, Int, Int)
 get_stats (FlS a b c) = (a, b, c)
 
+zeroStats :: FloatStats
 zeroStats = FlS 0 0 0
 
+sum_stats :: [FloatStats] -> FloatStats
 sum_stats xs = foldr add_stats zeroStats xs
 
+add_stats :: FloatStats -> FloatStats -> FloatStats
 add_stats (FlS a1 b1 c1) (FlS a2 b2 c2)
   = FlS (a1 + a2) (b1 + b2) (c1 + c2)
 
-add_to_stats (FlS a b c) floats
-  = FlS (a + length top_floats) (b + length other_floats) (c + 1)
-  where
-    (top_floats, other_floats) = partition to_very_top floats
-
-    to_very_top (my_lvl, _) = isTopLvl my_lvl
+add_to_stats :: FloatStats -> FloatBinds -> FloatStats
+add_to_stats (FlS a b c) (FB tops others)
+  = FlS (a + lengthBag tops) (b + lengthBag (flattenMajor others)) (c + 1)
 \end{code}
 
 
@@ -419,55 +374,114 @@ add_to_stats (FlS a b c) floats
 %*                                                                     *
 %************************************************************************
 
-\begin{code}
-getBindLevel (NonRec (TB _ lvl) _)      = lvl
-getBindLevel (Rec (((TB _ lvl), _) : _)) = lvl
-\end{code}
+Note [Representation of FloatBinds]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+The FloatBinds types is somewhat important.  We can get very large numbers
+of floating bindings, often all destined for the top level.  A typical example
+is     x = [4,2,5,2,5, .... ]
+Then we get lots of small expressions like (fromInteger 4), which all get
+lifted to top level.  
+
+The trouble is that  
+  (a) we partition these floating bindings *at every binding site* 
+  (b) SetLevels introduces a new bindings site for every float
+So we had better not look at each binding at each binding site!
+
+That is why MajorEnv is represented as a finite map.
+
+We keep the bindings destined for the *top* level separate, because
+we float them out even if they don't escape a *value* lambda; see
+partitionByMajorLevel.
+
 
 \begin{code}
-partitionByMajorLevel, partitionByLevel
-       :: Level                -- Partitioning level
+type FloatBind = CoreBind      -- INVARIANT: a FloatBind is always lifted
 
-       -> FloatBinds           -- Defns to be divided into 2 piles...
+data FloatBinds  = FB !(Bag FloatBind)         -- Destined for top level
+                     !MajorEnv                 -- Levels other than top
+     -- See Note [Representation of FloatBinds]
 
-       -> (FloatBinds, -- Defns  with level strictly < partition level,
-           FloatBinds) -- The rest
+type MajorEnv = UniqFM MinorEnv                        -- Keyed by major level
+type MinorEnv = UniqFM (Bag FloatBind)         -- Keyed by minor level
 
+flattenFloats :: FloatBinds -> Bag FloatBind
+flattenFloats (FB tops others) = tops `unionBags` flattenMajor others
 
-partitionByMajorLevel ctxt_lvl defns
-  = partition float_further defns
-  where
-       -- Float it if we escape a value lambda, or if we get to the top level
-    float_further (my_lvl, bind) = my_lvl `ltMajLvl` ctxt_lvl || isTopLvl my_lvl
-       -- The isTopLvl part says that if we can get to the top level, say "yes" anyway
-       -- This means that 
-       --      x = f e
-       -- transforms to 
-       --    lvl = e
-       --    x = f lvl
-       -- which is as it should be
-
-partitionByLevel ctxt_lvl defns
-  = partition float_further defns
-  where
-    float_further (my_lvl, _) = my_lvl `ltLvl` ctxt_lvl
-\end{code}
+flattenMajor :: MajorEnv -> Bag FloatBind
+flattenMajor = foldUFM (unionBags . flattenMinor) emptyBag
 
-\begin{code}
-floatsToBinds :: FloatBinds -> [CoreBind]
-floatsToBinds floats = map snd floats
+flattenMinor :: MinorEnv -> Bag FloatBind
+flattenMinor = foldUFM unionBags emptyBag
 
-floatsToBindPairs :: FloatBinds -> [(Id,CoreExpr)]
+emptyFloats :: FloatBinds
+emptyFloats = FB emptyBag emptyUFM
 
-floatsToBindPairs floats = concat (map mk_pairs floats)
-  where
-   mk_pairs (_, Rec pairs)         = pairs
-   mk_pairs (_, NonRec binder rhs) = [(binder,rhs)]
+unitFloat :: Level -> FloatBind -> FloatBinds
+unitFloat lvl@(Level major minor) b 
+  | isTopLvl lvl = FB (unitBag b) emptyUFM
+  | otherwise    = FB emptyBag (unitUFM major (unitUFM minor (unitBag b)))
+
+plusFloats :: FloatBinds -> FloatBinds -> FloatBinds
+plusFloats (FB t1 b1) (FB t2 b2) = FB (t1 `unionBags` t2) (b1 `plusMajor` b2)
 
-install :: FloatBinds -> CoreExpr -> CoreExpr
+plusMajor :: MajorEnv -> MajorEnv -> MajorEnv
+plusMajor = plusUFM_C plusMinor
 
+plusMinor :: MinorEnv -> MinorEnv -> MinorEnv
+plusMinor = plusUFM_C unionBags
+
+floatsToBindPairs :: Bag FloatBind -> [(Id,CoreExpr)] -> [(Id,CoreExpr)]
+floatsToBindPairs floats binds = foldrBag add binds floats
+  where
+   add (Rec pairs)         binds = pairs ++ binds
+   add (NonRec binder rhs) binds = (binder,rhs) : binds
+
+install :: Bag FloatBind -> CoreExpr -> CoreExpr
 install defn_groups expr
-  = foldr install_group expr defn_groups
+  = foldrBag install_group expr defn_groups
+  where
+    install_group defns body = Let defns body
+
+partitionByMajorLevel, partitionByLevel
+       :: Level                -- Partitioning level
+       -> FloatBinds           -- Defns to be divided into 2 piles...
+       -> (FloatBinds,         -- Defns  with level strictly < partition level,
+           Bag FloatBind)      -- The rest
+
+--      ---- partitionByMajorLevel ----
+-- Float it if we escape a value lambda, *or* if we get to the top level
+-- If we can get to the top level, say "yes" anyway. This means that 
+--     x = f e
+-- transforms to 
+--    lvl = e
+--    x = f lvl
+-- which is as it should be
+
+partitionByMajorLevel (Level major _) (FB tops defns)
+  = (FB tops outer, heres `unionBags` flattenMajor inner)
+  where
+    (outer, mb_heres, inner) = splitUFM defns major
+    heres = case mb_heres of 
+               Nothing -> emptyBag
+               Just h  -> flattenMinor h
+
+partitionByLevel (Level major minor) (FB tops defns)
+  = (FB tops (outer_maj `plusMajor` unitUFM major outer_min),
+     here_min `unionBags` flattenMinor inner_min 
+              `unionBags` flattenMajor inner_maj)
+
+  where
+    (outer_maj, mb_here_maj, inner_maj) = splitUFM defns major
+    (outer_min, mb_here_min, inner_min) = case mb_here_maj of
+                                            Nothing -> (emptyUFM, Nothing, emptyUFM)
+                                            Just min_defns -> splitUFM min_defns minor
+    here_min = mb_here_min `orElse` emptyBag
+
+wrapCostCentre :: CostCentre -> FloatBinds -> FloatBinds
+wrapCostCentre cc (FB tops defns)
+  = FB (wrap_defns tops) (mapUFM (mapUFM wrap_defns) defns)
   where
-    install_group (_, defns) body = Let defns body
+    wrap_defns = mapBag wrap_one 
+    wrap_one (NonRec binder rhs) = NonRec binder (mkSCC cc rhs)
+    wrap_one (Rec pairs)         = Rec (mapSnd (mkSCC cc) pairs)
 \end{code}