add -fsimpleopt-before-flatten
[ghc-hetmet.git] / compiler / cmm / CmmStackLayout.hs
index 3518df8..01543c4 100644 (file)
@@ -1,3 +1,16 @@
+{-# OPTIONS_GHC -XGADTs -XNoMonoLocalBinds #-}
+-- Norman likes local bindings
+-- If this module lives on I'd like to get rid of this flag in due course
+
+-- Todo: remove
+{-# OPTIONS_GHC -fno-warn-warnings-deprecations #-}
+
+{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-}
+#if __GLASGOW_HASKELL__ >= 701
+-- GHC 7.0.1 improved incomplete pattern warnings with GADTs
+{-# OPTIONS_GHC -fwarn-incomplete-patterns #-}
+#endif
+
 module CmmStackLayout
     ( SlotEnv, liveSlotAnal, liveSlotTransfers, removeLiveSlotDefs
     , layout, manifestSP, igraph, areaBuilder
@@ -5,24 +18,24 @@ module CmmStackLayout
 where
 
 import Constants
-import qualified Prelude as P
-import Prelude hiding (zip, unzip, last)
+import Prelude hiding (succ, zip, unzip, last)
 
 import BlockId
+import Cmm
 import CmmExpr
-import CmmProcPointZ
-import CmmTx
-import DFMonad
-import FiniteMap
+import CmmProcPoint
 import Maybes
-import MkZipCfg
-import MkZipCfgCmm hiding (CmmBlock, CmmGraph)
-import Monad
+import MkGraph (stackStubExpr)
+import Control.Monad
+import OptimizationFuel
 import Outputable
-import Panic
-import ZipCfg
-import ZipCfgCmmRep
-import ZipDataflow
+import SMRep (ByteOff)
+
+import Compiler.Hoopl
+
+import Data.Map (Map)
+import qualified Data.Map as Map
+import qualified FiniteMap as Map
 
 ------------------------------------------------------------------------
 --                    Stack Layout                                    --
@@ -57,22 +70,23 @@ import ZipDataflow
 -- a single slot, on insertion.
 
 slotLattice :: DataflowLattice SubAreaSet
-slotLattice = DataflowLattice "live slots" emptyFM add False
-  where add new old = case foldFM addArea (False, old) new of
-                        (True,  x) -> aTx  x
-                        (False, x) -> noTx x
+slotLattice = DataflowLattice "live slots" Map.empty add
+  where add _ (OldFact old) (NewFact new) = case Map.foldRightWithKey addArea (False, old) new of
+                                              (change, x) -> (changeIf change, x)
         addArea a newSlots z = foldr (addSlot a) z newSlots
         addSlot a slot (changed, map) =
-          let (c, live) = liveGen slot $ lookupWithDefaultFM map [] a
-          in (c || changed, addToFM map a live)
+          let (c, live) = liveGen slot $ Map.findWithDefault [] a map
+          in (c || changed, Map.insert a live map)
+
+slotLatticeJoin :: [SubAreaSet] -> SubAreaSet
+slotLatticeJoin facts = foldr extend (fact_bot slotLattice) facts
+  where extend fact res = snd $ fact_join slotLattice undefined (OldFact fact) (NewFact res)
 
 type SlotEnv   = BlockEnv SubAreaSet
-type SlotFix a = FuelMonad (BackwardFixedPoint Middle Last SubAreaSet a)
+  -- The sub-areas live on entry to the block
 
-liveSlotAnal :: LGraph Middle Last -> FuelMonad SlotEnv
-liveSlotAnal g = liftM zdfFpFacts (res :: SlotFix ())
-  where res = zdfSolveFromL emptyBlockEnv "live slot analysis" slotLattice
-                            liveSlotTransfers (fact_bot slotLattice) g
+liveSlotAnal :: CmmGraph -> FuelUniqSM SlotEnv
+liveSlotAnal g = liftM snd $ dataflowPassBwd g [] $ analBwd slotLattice liveSlotTransfers
 
 -- Add the subarea s to the subareas in the list-set (possibly coalescing it with
 -- adjacent subareas), and also return whether s was a new addition.
@@ -111,29 +125,41 @@ liveKill (a, hi, w) set = -- pprTrace "killing slots in area" (ppr a) $
 -- considered live in to the block -- we treat the first node as a definition site.
 -- BEWARE?: Am I being a little careless here in failing to check for the
 -- entry Id (which would use the CallArea Old).
-liveSlotTransfers :: BackwardTransfers Middle Last SubAreaSet
-liveSlotTransfers =
-  BackwardTransfers first liveInSlots liveLastIn
-    where first live id = delFromFM live (CallArea (Young id))
+liveSlotTransfers :: BwdTransfer CmmNode SubAreaSet
+liveSlotTransfers = mkBTransfer3 frt mid lst
+  where frt :: CmmNode C O -> SubAreaSet -> SubAreaSet
+        frt (CmmEntry l) f = Map.delete (CallArea (Young l)) f
+
+        mid :: CmmNode O O -> SubAreaSet -> SubAreaSet
+        mid n f = foldSlotsUsed addSlot (removeLiveSlotDefs f n) n
+        lst :: CmmNode O C -> FactBase SubAreaSet -> SubAreaSet
+        lst n f = liveInSlots n $ case n of
+          CmmCall {cml_cont=Nothing, cml_args=args} -> add_area (CallArea Old) args out
+          CmmCall {cml_cont=Just k, cml_args=args}  -> add_area (CallArea Old) args (add_area (CallArea (Young k)) args out)
+          CmmForeignCall {succ=k, updfr=oldend}     -> add_area (CallArea Old) oldend (add_area (CallArea (Young k)) wORD_SIZE out)
+          _                                         -> out
+         where out = joinOutFacts slotLattice n f
+               add_area _ n live | n == 0 = live
+               add_area a n live = Map.insert a (snd $ liveGen (a, n, n) $ Map.findWithDefault [] a live) live
 
 -- Slot sets: adding slots, removing slots, and checking for membership.
 liftToArea :: Area -> ([SubArea] -> [SubArea]) -> SubAreaSet -> SubAreaSet 
 addSlot, removeSlot :: SubAreaSet -> SubArea -> SubAreaSet
 elemSlot            :: SubAreaSet -> SubArea -> Bool
-liftToArea a f map = addToFM map a $ f (lookupWithDefaultFM map [] a)
+liftToArea a f map = Map.insert a (f (Map.findWithDefault [] a map)) map
 addSlot    live (a, i, w) = liftToArea a (snd . liveGen  (a, i, w)) live
 removeSlot live (a, i, w) = liftToArea a       (liveKill (a, i, w)) live
 elemSlot   live (a, i, w) =
-  not $ fst $ liveGen  (a, i, w) (lookupWithDefaultFM live [] a)
+  not $ fst $ liveGen  (a, i, w) (Map.findWithDefault [] a live)
 
 removeLiveSlotDefs :: (DefinerOfSlots s, UserOfSlots s) => SubAreaSet -> s -> SubAreaSet
 removeLiveSlotDefs = foldSlotsDefd removeSlot
 
-liveInSlots :: (DefinerOfSlots s, UserOfSlots s) => SubAreaSet -> s -> SubAreaSet
-liveInSlots live x = foldSlotsUsed addSlot (removeLiveSlotDefs live x) x
+liveInSlots :: (DefinerOfSlots s, UserOfSlots s) => s -> SubAreaSet -> SubAreaSet
+liveInSlots x live = foldSlotsUsed addSlot (removeLiveSlotDefs live x) x
 
-liveLastIn :: (BlockId -> SubAreaSet) -> Last -> SubAreaSet
-liveLastIn env l = liveInSlots (liveLastOut env l) l
+liveLastIn :: CmmNode O C -> (BlockId -> SubAreaSet) -> SubAreaSet
+liveLastIn l env = liveInSlots l (liveLastOut env l)
 
 -- Don't forget to keep the outgoing parameters in the CallArea live,
 -- as well as the update frame.
@@ -142,20 +168,20 @@ liveLastIn env l = liveInSlots (liveLastOut env l) l
 -- be a return to keep the update frame live. We'd still better keep the
 -- info pointer in the update frame live at any call site;
 -- otherwise we could screw up the garbage collector.
-liveLastOut :: (BlockId -> SubAreaSet) -> Last -> SubAreaSet
+liveLastOut :: (BlockId -> SubAreaSet) -> CmmNode O C -> SubAreaSet
 liveLastOut env l =
   case l of
-    LastCall _ Nothing n _ -> 
+    CmmCall _ Nothing n _ _ -> 
       add_area (CallArea Old) n out -- add outgoing args (includes upd frame)
-    LastCall _ (Just k) n (Just upd_n) ->
+    CmmCall _ (Just k) n _ _ ->
       add_area (CallArea Old) n (add_area (CallArea (Young k)) n out)
-    LastCall _ (Just k) n Nothing ->
-      add_area (CallArea (Young k)) n out
+    CmmForeignCall { succ = k, updfr = oldend } ->
+      add_area (CallArea Old) oldend (add_area (CallArea (Young k)) wORD_SIZE out)
     _ -> out
-  where out = joinOuts slotLattice env l
+  where out = slotLatticeJoin $ map env $ successors l
         add_area _ n live | n == 0 = live
         add_area a n live =
-          addToFM live a $ snd $ liveGen (a, n, n) $ lookupWithDefaultFM live [] a
+          Map.insert a (snd $ liveGen (a, n, n) $ Map.findWithDefault [] a live) live
 
 -- The liveness analysis must be precise: otherwise, we won't know if a definition
 -- should really kill a live-out stack slot.
@@ -166,7 +192,7 @@ liveLastOut env l =
 -- every time, I provide a function to fold over the nodes, which should be a
 -- reasonably efficient approach for the implementations we envision.
 -- Of course, it will probably be much easier to program if we just return a list...
-type Set x = FiniteMap x ()
+type Set x = Map x ()
 data IGraphBuilder n =
   Builder { foldNodes     :: forall z. SubArea -> (n -> z -> z) -> z -> z
           , _wordsOccupied :: AreaMap -> AreaMap -> n -> [Int]
@@ -176,9 +202,9 @@ areaBuilder :: IGraphBuilder Area
 areaBuilder = Builder fold words
   where fold (a, _, _) f z = f a z
         words areaSize areaMap a =
-          case lookupFM areaMap a of
-            Just addr -> [addr .. addr + (lookupFM areaSize a `orElse`
-                                          pprPanic "wordsOccupied: unknown area" (ppr a))]
+          case Map.lookup a areaMap of
+            Just addr -> [addr .. addr + (Map.lookup a areaSize `orElse`
+                                          pprPanic "wordsOccupied: unknown area" (ppr areaSize <+> ppr a))]
             Nothing   -> []
 
 --slotBuilder :: IGraphBuilder (Area, Int)
@@ -187,69 +213,77 @@ areaBuilder = Builder fold words
 -- Now, we can build the interference graph.
 -- The usual story: a definition interferes with all live outs and all other
 -- definitions.
-type IGraph x = FiniteMap x (Set x)
+type IGraph x = Map x (Set x)
 type IGPair x = (IGraph x, IGraphBuilder x)
-igraph :: (Ord x) => IGraphBuilder x -> SlotEnv -> LGraph Middle Last -> IGraph x
-igraph builder env g = foldr interfere emptyFM (postorder_dfs g)
+igraph :: (Ord x) => IGraphBuilder x -> SlotEnv -> CmmGraph -> IGraph x
+igraph builder env g = foldr interfere Map.empty (postorderDfs g)
   where foldN = foldNodes builder
-        interfere block igraph =
-          let (h, l) = goto_end (unzip block)
-              --heads :: ZHead Middle -> (IGraph x, SubAreaSet) -> IGraph x
-              heads (ZFirst _ _) (igraph, _)       = igraph
-              heads (ZHead h m)    (igraph, liveOut) =
-                heads h (addEdges igraph m liveOut, liveInSlots liveOut m)
-              -- add edges between a def and the other defs and liveouts
-              addEdges igraph i out = fst $ foldSlotsDefd addDef (igraph, out) i
-              addDef (igraph, out) def@(a, _, _) =
-                (foldN def (addDefN out) igraph,
-                 addToFM out a (snd $ liveGen def (lookupWithDefaultFM out [] a)))
-              addDefN out n igraph =
-                let addEdgeNO o igraph = foldN o addEdgeNN igraph
-                    addEdgeNN n' igraph = addEdgeNN' n n' $ addEdgeNN' n' n igraph
-                    addEdgeNN' n n' igraph = addToFM igraph n (addToFM set n' ())
-                      where set = lookupWithDefaultFM igraph emptyFM n
-                in foldFM (\ _ os igraph -> foldr addEdgeNO igraph os) igraph out
-              env' bid = lookupBlockEnv env bid `orElse` panic "unknown blockId in igraph"
-          in heads h $ case l of LastExit    -> (igraph, emptyFM)
-                                 LastOther l -> (addEdges igraph l $ liveLastOut env' l,
-                                                 liveLastIn env' l)
+        interfere block igraph = foldBlockNodesB3 (first, middle, last) block igraph
+          where first _ (igraph, _) = igraph
+                middle node (igraph, liveOut) =
+                  (addEdges igraph node liveOut, liveInSlots node liveOut)
+                last node igraph =
+                  (addEdges igraph node $ liveLastOut env' node, liveLastIn node env')
+
+                -- add edges between a def and the other defs and liveouts
+                addEdges igraph i out = fst $ foldSlotsDefd addDef (igraph, out) i
+                addDef (igraph, out) def@(a, _, _) =
+                  (foldN def (addDefN out) igraph,
+                   Map.insert a (snd $ liveGen def (Map.findWithDefault [] a out)) out)
+                addDefN out n igraph =
+                  let addEdgeNO o igraph = foldN o addEdgeNN igraph
+                      addEdgeNN n' igraph = addEdgeNN' n n' $ addEdgeNN' n' n igraph
+                      addEdgeNN' n n' igraph = Map.insert n (Map.insert n' () set) igraph
+                        where set = Map.findWithDefault Map.empty n igraph
+                  in Map.foldRightWithKey (\ _ os igraph -> foldr addEdgeNO igraph os) igraph out
+                env' bid = mapLookup bid env `orElse` panic "unknown blockId in igraph"
 
 -- Before allocating stack slots, we need to collect one more piece of information:
 -- what's the highest offset (in bytes) used in each Area?
 -- We'll need to allocate that much space for each Area.
-getAreaSize :: LGraph Middle Last -> AreaMap
-getAreaSize g@(LGraph _ off _) =
-  fold_blocks (fold_fwd_block first add_regslots last)
-              (unitFM (CallArea Old) off) g
-  where first id (StackInfo {argBytes = Just off}) z = add z (CallArea (Young id)) off
-        first _  _          z = z
+
+-- JD: WHY CAN'T THIS COME FROM THE slot-liveness info?
+getAreaSize :: ByteOff -> CmmGraph -> AreaMap
+  -- The domain of the returned mapping consists only of Areas
+  -- used for (a) variable spill slots, and (b) parameter passing ares for calls
+getAreaSize entry_off g =
+  foldGraphBlocks (foldBlockNodesF3 (first, add_regslots, last))
+              (Map.singleton (CallArea Old) entry_off) g
+  where first _  z = z
+        last :: CmmNode O C -> Map Area Int -> Map Area Int
+        last l@(CmmCall _ Nothing args res _) z  =  add_regslots l (add (add z area args) area res)
+          where area = CallArea Old
+        last l@(CmmCall _ (Just k) args res _) z =  add_regslots l (add (add z area args) area res)
+          where area = CallArea (Young k)
+        last l@(CmmForeignCall {succ = k}) z     =  add_regslots l (add z area wORD_SIZE)
+          where area = CallArea (Young k)
+        last l z                                 =  add_regslots l z
         add_regslots i z = foldSlotsUsed addSlot (foldSlotsDefd addSlot z i) i
-        last l@(LastOther (LastCall _ Nothing off _)) z =
-          add_regslots l (add z (CallArea Old) off)
-        last l@(LastOther (LastCall _ (Just k) off _)) z =
-          add_regslots l (add z (CallArea (Young k)) off)
-        last l z = add_regslots l z
-        addSlot z (a@(RegSlot _), off, _) = add z a off
+        addSlot z (a@(RegSlot (LocalReg _ ty)), _, _) =
+          add z a $ widthInBytes $ typeWidth ty
         addSlot z _ = z
-        add z a off = addToFM z a (max off (lookupWithDefaultFM z 0 a))
+        add z a off = Map.insert a (max off (Map.findWithDefault 0 a z)) z
+       -- The 'max' is important.  Two calls, to f and g, might share a common
+       -- continuation (and hence a common CallArea), but their number of overflow
+       -- parameters might differ.
 
 
 -- Find the Stack slots occupied by the subarea's conflicts
 conflictSlots :: Ord x => IGPair x -> AreaMap -> AreaMap -> SubArea -> Set Int
 conflictSlots (ig, Builder foldNodes wordsOccupied) areaSize areaMap subarea =
-  foldNodes subarea foldNode emptyFM
-  where foldNode n set = foldFM conflict set $ lookupWithDefaultFM ig emptyFM n
+  foldNodes subarea foldNode Map.empty
+  where foldNode n set = Map.foldRightWithKey conflict set $ Map.findWithDefault Map.empty n ig
         conflict n' () set = liveInSlots areaMap n' set
         -- Add stack slots occupied by igraph node n
         liveInSlots areaMap n set = foldr setAdd set (wordsOccupied areaSize areaMap n)
-        setAdd w s = addToFM s w ()
+        setAdd w s = Map.insert w () s
 
 -- Find any open space on the stack, starting from the offset.
 -- If the area is a CallArea or a spill slot for a pointer, then it must
 -- be word-aligned.
 freeSlotFrom :: Ord x => IGPair x -> AreaMap -> Int -> AreaMap -> Area -> Int
 freeSlotFrom ig areaSize offset areaMap area =
-  let size = lookupFM areaSize area `orElse` 0
+  let size = Map.lookup area areaSize `orElse` 0
       conflicts = conflictSlots ig areaSize areaMap (area, size, size)
       -- CallAreas and Ptrs need to be word-aligned (round up!)
       align = case area of CallArea _                                -> align'
@@ -259,7 +293,7 @@ freeSlotFrom ig areaSize offset areaMap area =
       -- Find a space big enough to hold the area
       findSpace curr 0 = curr
       findSpace curr cnt = -- part of target slot, # of bytes left to check
-        if elemFM curr conflicts then
+        if Map.member curr conflicts then
           findSpace (align (curr + size)) size -- try the next (possibly) open space
         else findSpace (curr - 1) (cnt - 1)
   in findSpace (align (offset + size)) size
@@ -267,69 +301,86 @@ freeSlotFrom ig areaSize offset areaMap area =
 -- Find an open space on the stack, and assign it to the area.
 allocSlotFrom :: Ord x => IGPair x -> AreaMap -> Int -> AreaMap -> Area -> AreaMap
 allocSlotFrom ig areaSize from areaMap area =
-  if elemFM area areaMap then areaMap
-  else addToFM areaMap area $ freeSlotFrom ig areaSize from areaMap area
+  if Map.member area areaMap then areaMap
+  else Map.insert area (freeSlotFrom ig areaSize from areaMap area) areaMap
 
 -- | Greedy stack layout.
 -- Compute liveness, build the interference graph, and allocate slots for the areas.
 -- We visit each basic block in a (generally) forward order.
+
 -- At each instruction that names a register subarea r, we immediately allocate
 -- any available slot on the stack by the following procedure:
---  1. Find the nodes N' that conflict with r
---  2. Find the stack slots used for N'
---  3. Choose a contiguous stack space s not in N' (s must be large enough to hold r)
+--  1. Find the sub-areas S that conflict with r
+--  2. Find the stack slots used for S
+--  3. Choose a contiguous stack space s not in S (s must be large enough to hold r)
+
 -- For a CallArea, we allocate the stack space only when we reach a function
 -- call that returns to the CallArea's blockId.
--- We use a similar procedure, with one exception: the stack space
--- must be allocated below the youngest stack slot that is live out.
+-- Then, we allocate the Area subject to the following constraints:
+--   a) It must be younger than all the sub-areas that are live on entry to the block
+--         This constraint is only necessary for the successor of a call
+--   b) It must not overlap with any already-allocated Area with which it conflicts
+--        (ie at some point, not necessarily now, is live at the same time)
+--   Part (b) is just the 1,2,3 part above
 
 -- Note: The stack pointer only has to be younger than the youngest live stack slot
 -- at proc points. Otherwise, the stack pointer can point anywhere.
-layout :: ProcPointSet -> SlotEnv -> LGraph Middle Last -> AreaMap
-layout procPoints env g@(LGraph _ entrySp _) =
-  let builder = areaBuilder
-      ig = (igraph builder env g, builder)
-      env' bid = lookupBlockEnv env bid `orElse` panic "unknown blockId in igraph"
-      areaSize = getAreaSize g
-      -- Find the slots that are live-in to the block
-      live_in (ZTail m l) = liveInSlots (live_in l) m
-      live_in (ZLast (LastOther l)) = liveLastIn env' l
-      live_in (ZLast LastExit) = emptyFM 
-      -- Find the youngest live stack slot
+
+layout :: ProcPointSet -> SlotEnv -> ByteOff -> CmmGraph -> AreaMap
+-- The domain of the returned map includes an Area for EVERY block
+-- including each block that is not the successor of a call (ie is not a proc-point)
+-- That's how we return the info of what the SP should be at the entry of every block
+
+layout procPoints env entry_off g =
+  let ig = (igraph areaBuilder env g, areaBuilder)
+      env' bid = mapLookup bid env `orElse` panic "unknown blockId in igraph"
+      areaSize = getAreaSize entry_off g
+
+      -- Find the youngest live stack slot that has already been allocated
+      youngest_live :: AreaMap            -- Already allocated
+                    -> SubAreaSet  -- Sub-areas live here
+                   -> ByteOff     -- Offset of the youngest byte of any 
+                                  --    already-allocated, live sub-area
       youngest_live areaMap live = fold_subareas young_slot live 0
-        where young_slot (a, o, _) z = case lookupFM areaMap a of
+        where young_slot (a, o, _) z = case Map.lookup a areaMap of
                                          Just top -> max z $ top + o
                                          Nothing  -> z
-      fold_subareas :: (SubArea -> z -> z) -> SubAreaSet -> z -> z
-      fold_subareas f m z = foldFM (\_ s z -> foldr f z s) z m
+              fold_subareas f m z = Map.foldRightWithKey (\_ s z -> foldr f z s) z m
+
       -- Allocate space for spill slots and call areas
       allocVarSlot = allocSlotFrom ig areaSize 0
-      allocCallSlot areaMap (Block id stackInfo t)
-        | elemBlockSet id procPoints =
-        let young  = youngest_live areaMap $ live_in t
-            start = case returnOff stackInfo of Just b  -> max b young
-                                                Nothing -> young
-            z = allocSlotFrom ig areaSize start areaMap (CallArea (Young id))
-        in -- pprTrace "allocCallSlot for" (ppr id <+> ppr young <+> ppr (live_in t) <+> ppr z) 
-           z
-      allocCallSlot areaMap _ = areaMap
-      -- mid foreign calls need to have info tables placed on the stack
-      allocMidCall m@(MidForeignCall (Safe bid _) _ _ _) t areaMap =
-        let young     = youngest_live areaMap $ removeLiveSlotDefs (live_in t) m
-            area      = CallArea (Young bid)
-            areaSize' = addToFM areaSize area (widthInBytes (typeWidth gcWord))
-        in  allocSlotFrom ig areaSize' young areaMap area
-      allocMidCall _ _ areaMap = areaMap
-      alloc m t areaMap =
-          foldSlotsDefd alloc' (foldSlotsUsed alloc' (allocMidCall m t areaMap) m) m
-        where alloc' areaMap (a@(RegSlot _), _, _) = allocVarSlot areaMap a
-              alloc' areaMap _ = areaMap
-      layoutAreas areaMap b@(Block _ _ t) = layout areaMap t
-        where layout areaMap (ZTail m t) = layout (alloc m t areaMap) t
-              layout areaMap (ZLast _)   = allocCallSlot areaMap b
-      areaMap = foldl layoutAreas (addToFM emptyFM (CallArea Old) 0) (postorder_dfs g)
+
+      -- Update the successor's incoming SP.
+      setSuccSPs inSp bid areaMap =
+        case (Map.lookup area areaMap , mapLookup bid (toBlockMap g)) of
+          (Just _, _) -> areaMap -- succ already knows incoming SP
+          (Nothing, Just _) ->
+            if setMember bid procPoints then
+              let young = youngest_live areaMap $ env' bid
+                  -- start = case returnOff stackInfo of Just b  -> max b young
+                  --                                     Nothing -> young
+                  start = young -- maybe wrong, but I don't understand
+                                -- why the preceding is necessary...
+              in  allocSlotFrom ig areaSize start areaMap area
+            else Map.insert area inSp areaMap
+          (_, Nothing) -> panic "Block not found in cfg"
+        where area = CallArea (Young bid)
+
+      layoutAreas areaMap block = foldBlockNodesF3 (flip const, allocMid, allocLast (entryLabel block)) block areaMap
+      allocMid m areaMap = foldSlotsDefd alloc' (foldSlotsUsed alloc' areaMap m) m
+      allocLast bid l areaMap =
+        foldr (setSuccSPs inSp) areaMap' (successors l)
+        where inSp = expectJust "sp in" $ Map.lookup (CallArea (Young bid)) areaMap
+              areaMap' = foldSlotsDefd alloc' (foldSlotsUsed alloc' areaMap l) l
+      alloc' areaMap (a@(RegSlot _), _, _) = allocVarSlot areaMap a
+      alloc' areaMap _ = areaMap
+
+      initMap = Map.insert (CallArea (Young (g_entry g))) 0 $
+                  Map.insert (CallArea Old) 0 Map.empty
+                        
+      areaMap = foldl layoutAreas initMap (postorderDfs g)
   in -- pprTrace "ProcPoints" (ppr procPoints) $
-       -- pprTrace "Area SizeMap" (ppr areaSize) $
+        -- pprTrace "Area SizeMap" (ppr areaSize) $
          -- pprTrace "Entry SP" (ppr entrySp) $
            -- pprTrace "Area Map" (ppr areaMap) $
      areaMap
@@ -343,83 +394,76 @@ layout procPoints env g@(LGraph _ entrySp _) =
 --    stack pointer to be younger than the live values on the stack at proc points.
 -- 3. Compute the maximum stack offset used in the procedure and replace
 --    the stack high-water mark with that offset.
-manifestSP :: ProcPointSet -> BlockEnv Status -> AreaMap ->
-                LGraph Middle Last -> FuelMonad (LGraph Middle Last)
-manifestSP procPoints procMap areaMap g@(LGraph entry args blocks) =
-  liftM (LGraph entry args) blocks'
-  where blocks' = foldl replB (return emptyBlockEnv) (postorder_dfs g)
-        slot a = -- pprTrace "slot" (ppr a) $
-                   lookupFM areaMap a `orElse` panic "unallocated Area"
+manifestSP :: AreaMap -> ByteOff -> CmmGraph -> FuelUniqSM CmmGraph
+manifestSP areaMap entry_off g@(CmmGraph {g_entry=entry}) =
+  ofBlockMap entry `liftM` foldl replB (return mapEmpty) (postorderDfs g)
+  where slot a = -- pprTrace "slot" (ppr a) $
+                   Map.lookup a areaMap `orElse` panic "unallocated Area"
         slot' (Just id) = slot $ CallArea (Young id)
         slot' Nothing   = slot $ CallArea Old
         sp_high = maxSlot slot g
-        proc_entry_sp = slot (CallArea Old) + args
+        proc_entry_sp = slot (CallArea Old) + entry_off
+
+        add_sp_off :: CmmBlock -> BlockEnv Int -> BlockEnv Int
+        add_sp_off b env =
+          case lastNode b of
+            CmmCall {cml_cont=Just succ, cml_ret_args=off} -> mapInsert succ off env
+            CmmForeignCall {succ=succ}                     -> mapInsert succ wORD_SIZE env
+            _                                              -> env
+        spEntryMap = foldGraphBlocks add_sp_off (mapInsert entry entry_off emptyBlockMap) g
+        spOffset id = mapLookup id spEntryMap `orElse` 0
+
         sp_on_entry id | id == entry = proc_entry_sp
-        sp_on_entry id =
-          case lookupBlockEnv blocks id of
-            Just (Block _ (StackInfo {argBytes = Just o}) _) -> slot' (Just id) + o
-            _ -> 
-             case expectJust "sp_on_entry" (lookupBlockEnv procMap id) of
-               ReachedBy pp ->
-                 case blockSetToList pp of
-                   [id] -> sp_on_entry id
-                   _    -> panic "block not reached by one proc point"
-               ProcPoint -> pprPanic "procpoint doesn't take any arguments?"
-                               (ppr id <+> ppr g <+> ppr procPoints <+> ppr procMap)
+        sp_on_entry id = slot' (Just id) + spOffset id
 
         -- On entry to procpoints, the stack pointer is conventional;
         -- otherwise, we check the SP set by predecessors.
-        replB :: FuelMonad (BlockEnv CmmBlock) -> CmmBlock -> FuelMonad (BlockEnv CmmBlock)
-        replB blocks (Block id o t) =
-          do bs <- replTail (Block id o) spIn t
-             -- pprTrace "spIn" (ppr id <+> ppr spIn) $ do
-             liftM (flip (foldr insertBlock) bs) blocks
-          where spIn = sp_on_entry id
-        replTail :: (ZTail Middle Last -> CmmBlock) -> Int -> (ZTail Middle Last) -> 
-                    FuelMonad ([CmmBlock])
-        replTail h spOff (ZTail m@(MidForeignCall (Safe bid _) _ _ _) t) =
-          replTail (\t' -> h (setSp spOff spOff' (ZTail (middle spOff m) t'))) spOff' t
-            where spOff' = slot' (Just bid) + widthInBytes (typeWidth gcWord)
-        replTail h spOff (ZTail m t) = replTail (h . ZTail (middle spOff m)) spOff t
-        replTail h spOff (ZLast (LastOther l)) = fixSp h spOff l
-        replTail h _   l@(ZLast LastExit) = return [h l]
-        middle spOff m = mapExpDeepMiddle (replSlot spOff) m
-        last   spOff l = mapExpDeepLast   (replSlot spOff) l
-        replSlot spOff (CmmStackSlot a i) = CmmRegOff (CmmGlobal Sp) (spOff - (slot a + i))
-        replSlot spOff (CmmLit CmmHighStackMark) = -- replacing the high water mark
-          CmmLit (CmmInt (toInteger (max 0 (sp_high - proc_entry_sp))) (typeWidth bWord))
-        replSlot _ e = e
-        -- The block must establish the SP expected at each successsor.
-        fixSp :: (ZTail Middle Last -> CmmBlock) -> Int -> Last -> FuelMonad ([CmmBlock])
-        fixSp h spOff l@(LastCall _ k n _) = updSp h spOff (slot' k + n) l
-        fixSp h spOff l@(LastBranch k) =
-          let succSp = sp_on_entry k in
-          if succSp /= spOff then
-               -- pprTrace "updSp" (ppr k <> ppr spOff <> ppr (sp_on_entry k)) $
-               updSp h spOff succSp l
-          else return $ [h (ZLast (LastOther (last spOff l)))]
-        fixSp h spOff l = liftM (uncurry (:)) $ fold_succs succ l $ return (b, [])
-          where b = h (ZLast (LastOther (last spOff l)))
-                succ succId z =
-                  let succSp = sp_on_entry succId in
-                  if succSp /= spOff then
-                    do (b,  bs)  <- z
-                       (b', bs') <- insertBetween b [setSpMid spOff succSp] succId
-                       return (b', bs ++ bs')
-                  else z
-        updSp h old new l = return [h $ setSp old new $ ZLast $ LastOther (last new l)]
-        setSpMid sp sp' = MidAssign (CmmGlobal Sp) e
-          where e = CmmMachOp (MO_Add wordWidth) [CmmReg (CmmGlobal Sp), off]
-                off = CmmLit $ CmmInt (toInteger $ sp - sp') wordWidth
-        setSp sp sp' t = if sp == sp' then t else ZTail (setSpMid sp sp') t
+        replB :: FuelUniqSM (BlockEnv CmmBlock) -> CmmBlock -> FuelUniqSM (BlockEnv CmmBlock)
+        replB blocks block =
+          do let (head, middles, JustC tail :: MaybeC C (CmmNode O C)) = blockToNodeList block
+                 middles' = map (middle spIn) middles
+             bs <- replLast head middles' tail
+             flip (foldr insertBlock) bs `liftM` blocks
+          where spIn = sp_on_entry (entryLabel block)
+
+                middle spOff m = mapExpDeep (replSlot spOff) m
+                last   spOff l = mapExpDeep (replSlot spOff) l
+                replSlot spOff (CmmStackSlot a i) = CmmRegOff (CmmGlobal Sp) (spOff - (slot a + i))
+                replSlot _ (CmmLit CmmHighStackMark) = -- replacing the high water mark
+                  CmmLit (CmmInt (toInteger (max 0 (sp_high - proc_entry_sp))) (typeWidth bWord))
+                replSlot _ e = e
+
+                replLast :: MaybeC C (CmmNode C O) -> [CmmNode O O] -> CmmNode O C -> FuelUniqSM [CmmBlock]
+                replLast h m l@(CmmCall _ k n _ _)       = updSp (slot' k + n) h m l
+                -- JD: LastForeignCall probably ought to have an outgoing
+                --     arg size, just like LastCall
+                replLast h m l@(CmmForeignCall {succ=k}) = updSp (slot' (Just k) + wORD_SIZE) h m l
+                replLast h m l@(CmmBranch k)             = updSp (sp_on_entry k) h m l
+                replLast h m l                           = uncurry (:) `liftM` foldr succ (return (b, [])) (successors l)
+                  where b :: CmmBlock
+                        b = updSp' spIn h m l
+                        succ succId z =
+                          let succSp = sp_on_entry succId in
+                          if succSp /= spIn then
+                            do (b,  bs)  <- z
+                               (b', bs') <- insertBetween b (adjustSp succSp) succId
+                               return (b', bs' ++ bs)
+                          else z
+
+                updSp sp h m l = return [updSp' sp h m l]
+                updSp' sp h m l | sp == spIn = blockOfNodeList (h, m, JustC $ last sp l)
+                                | otherwise  = blockOfNodeList (h, m ++ adjustSp sp, JustC $ last sp l)
+                adjustSp sp = [CmmAssign (CmmGlobal Sp) e]
+                  where e = CmmMachOp (MO_Add wordWidth) [CmmReg (CmmGlobal Sp), off]
+                        off = CmmLit $ CmmInt (toInteger $ spIn - sp) wordWidth
 
 
 -- To compute the stack high-water mark, we fold over the graph and
 -- compute the highest slot offset.
 maxSlot :: (Area -> Int) -> CmmGraph -> Int
-maxSlot slotOff g = fold_blocks (fold_fwd_block (\ _ _ x -> x) highSlot highSlot) 0 g
+maxSlot slotOff g = foldGraphBlocks (foldBlockNodesF3 (flip const, highSlot, highSlot)) 0 g
   where highSlot i z = foldSlotsUsed add (foldSlotsDefd add z i) i
-        add z (a, i, w) = max z (slotOff a + i)
+        add z (a, i, _) = max z (slotOff a + i)
 
 -----------------------------------------------------------------------------
 -- | Sanity check: stub pointers immediately after they die
@@ -427,19 +471,17 @@ maxSlot slotOff g = fold_blocks (fold_fwd_block (\ _ _ x -> x) highSlot highSlot
 -- This will miss stack slots that are last used in a Last node,
 -- but it should do pretty well...
 
-type StubPtrFix = FuelMonad (BackwardFixedPoint Middle Last SubAreaSet CmmGraph)
-
-stubSlotsOnDeath :: (LGraph Middle Last) -> FuelMonad (LGraph Middle Last)
-stubSlotsOnDeath g = liftM zdfFpContents $ (res :: StubPtrFix)
-    where res = zdfBRewriteFromL RewriteShallow emptyBlockEnv "stub ptrs" slotLattice
-                                 liveSlotTransfers rewrites (fact_bot slotLattice) g
-          rewrites = BackwardRewrites first middle last Nothing
-          first _ _ = Nothing
-          last  _ _ = Nothing
-          middle liveSlots m = foldSlotsUsed (stub liveSlots m) Nothing m
+stubSlotsOnDeath :: CmmGraph -> FuelUniqSM CmmGraph
+stubSlotsOnDeath g = liftM fst $ dataflowPassBwd g [] $ analRewBwd slotLattice
+                                                                   liveSlotTransfers
+                                                                   rewrites
+    where rewrites = mkBRewrite3 frt mid lst
+          frt _ _ = return Nothing
+          mid m liveSlots = return $ foldSlotsUsed (stub liveSlots m) Nothing m
+          lst _ _ = return Nothing
           stub liveSlots m rst subarea@(a, off, w) =
             if elemSlot liveSlots subarea then rst
-            else let store = mkStore (CmmStackSlot a off)
-                                     (stackStubExpr (widthFromBytes w))
+            else let store = mkMiddle $ CmmStore (CmmStackSlot a off)
+                                                 (stackStubExpr (widthFromBytes w))
                  in case rst of Nothing -> Just (mkMiddle m <*> store)
                                 Just g  -> Just (g <*> store)