Implemented and fixed bugs in CmmInfo handling
[ghc-hetmet.git] / compiler / cmm / CmmCPS.hs
index 4c1d025..b6c57ee 100644 (file)
@@ -1,4 +1,8 @@
-module CmmCPS (cmmCPS) where
+module CmmCPS (
+  -- | Converts C-- with full proceedures and parameters
+  -- to a CPS transformed C-- with the stack made manifest.
+  cmmCPS
+) where
 
 #include "HsVersions.h"
 
@@ -6,9 +10,15 @@ import Cmm
 import CmmLint
 import PprCmm
 
-import Dataflow (fixedpoint)
 import CmmLive
-
+import CmmBrokenBlock
+import CmmProcPoint
+import CmmCallConv
+import CmmInfo
+import CmmUtils
+
+import Bitmap
+import ClosureInfo
 import MachOp
 import ForeignCall
 import CLabel
@@ -28,6 +38,151 @@ import Monad
 import IO
 import Data.List
 
+-----------------------------------------------------------------------------
+-- |Top level driver for the CPS pass
+-----------------------------------------------------------------------------
+cmmCPS :: DynFlags -- ^ Dynamic flags: -dcmm-lint -ddump-cps-cmm
+       -> [GenCmm CmmStatic CmmInfo CmmStmt]    -- ^ Input C-- with Proceedures
+       -> IO [GenCmm CmmStatic [CmmStatic] CmmStmt] -- ^ Output CPS transformed C--
+cmmCPS dflags abstractC = do
+  when (dopt Opt_DoCmmLinting dflags) $
+       do showPass dflags "CmmLint"
+         case firstJust $ map cmmLint abstractC of
+           Just err -> do printDump err
+                          ghcExit dflags 1
+           Nothing  -> return ()
+  showPass dflags "CPS"
+
+  -- TODO: more lint checking
+  --        check for use of branches to non-existant blocks
+  --        check for use of Sp, SpLim, R1, R2, etc.
+
+  uniqSupply <- mkSplitUniqSupply 'p'
+  let supplies = listSplitUniqSupply uniqSupply
+  let doCpsProc s (Cmm c) =
+          Cmm $ concat $ zipWith cpsProc (listSplitUniqSupply s) c
+  let continuationC = zipWith doCpsProc supplies abstractC
+
+  dumpIfSet_dyn dflags Opt_D_dump_cps_cmm "CPS Cmm" (pprCmms continuationC)
+
+  -- TODO: add option to dump Cmm to file
+
+  return continuationC
+
+stg_gc_gen = mkRtsApFastLabel SLIT("gen_cg_TODO") --panic "Need the label for gc"
+make_gc_block block_id fun_label formals safety = BasicBlock block_id stmts
+    where
+      stmts = [CmmCall stg_gc_gen_target [] [] safety,
+               CmmJump fun_expr actuals]
+      stg_gc_gen_target =
+          CmmForeignCall (CmmLit (CmmLabel stg_gc_gen)) CmmCallConv
+      actuals = map (\x -> (CmmReg (CmmLocal x), NoHint)) formals
+      fun_expr = CmmLit (CmmLabel fun_label)
+
+force_gc_block old_info block_id fun_label formals blocks =
+    case old_info of
+      CmmNonInfo (Just _) -> (old_info, [])
+      CmmInfo _ (Just _) _ _ -> (old_info, [])
+      CmmNonInfo Nothing
+          -> (CmmNonInfo (Just block_id),
+              [make_gc_block block_id fun_label formals (CmmSafe NoC_SRT)])
+      CmmInfo prof Nothing type_tag type_info
+        -> (CmmInfo prof (Just block_id) type_tag type_info,
+            [make_gc_block block_id fun_label formals (CmmSafe srt)])
+           where
+             srt = case type_info of
+                     ConstrInfo _ _ _ -> NoC_SRT
+                     FunInfo _ srt' _ _ _ _ -> srt'
+                     ThunkInfo _ srt' -> srt'
+                     ThunkSelectorInfo _ srt' -> srt'
+                     ContInfo _ srt' -> srt'    
+
+-----------------------------------------------------------------------------
+-- |CPS a single CmmTop (proceedure)
+-- Only 'CmmProc' are transformed 'CmmData' will be left alone.
+-----------------------------------------------------------------------------
+
+cpsProc :: UniqSupply 
+        -> GenCmmTop CmmStatic CmmInfo CmmStmt     -- ^Input proceedure
+        -> [GenCmmTop CmmStatic [CmmStatic] CmmStmt]   -- ^Output proceedure and continuations
+cpsProc uniqSupply (CmmData sec dat) = [CmmData sec dat]
+cpsProc uniqSupply (CmmProc info ident params blocks) = info_procs
+    where
+      uniques :: [[Unique]]
+      uniques = map uniqsFromSupply $ listSplitUniqSupply uniqSupply
+      (gc_unique:info_uniques):block_uniques = uniques
+
+      -- Ensure that 
+      forced_gc :: (CmmInfo, [CmmBasicBlock])
+      forced_gc = force_gc_block info (BlockId gc_unique) ident params blocks
+
+      forced_info = fst forced_gc
+      forced_blocks = blocks ++ snd forced_gc
+      forced_gc_id = case forced_info of
+                       CmmNonInfo (Just x) -> x
+                       CmmInfo _ (Just x) _ _ -> x
+
+      -- Break the block at each function call.
+      -- The part after the function call will have to become a continuation.
+      broken_blocks :: [BrokenBlock]
+      broken_blocks =
+          concat $ zipWith3 breakBlock block_uniques forced_blocks
+                     (FunctionEntry forced_info ident params:repeat ControlEntry)
+
+      -- Calculate live variables for each broken block.
+      --
+      -- Nothing can be live on entry to the first block
+      -- so we could take the tail, but for now we wont
+      -- to help future proof the code.
+      live :: BlockEntryLiveness
+      live = cmmLiveness $ map cmmBlockFromBrokenBlock broken_blocks
+
+      -- Calculate which blocks must be made into full fledged procedures.
+      proc_points :: UniqSet BlockId
+      proc_points = calculateProcPoints broken_blocks
+
+      -- Construct a map so we can lookup a broken block by its 'BlockId'.
+      block_env :: BlockEnv BrokenBlock
+      block_env = blocksToBlockEnv broken_blocks
+
+      -- Group the blocks into continuations based on the set of proc-points.
+      continuations :: [Continuation (Either C_SRT CmmInfo)]
+      continuations = zipWith
+                        (gatherBlocksIntoContinuation proc_points block_env)
+                        (uniqSetToList proc_points)
+                        (Just forced_gc_id : repeat Nothing)
+
+      -- Select the stack format on entry to each continuation.
+      -- Return the max stack offset and an association list
+      --
+      -- This is an association list instead of a UniqFM because
+      -- CLabel's don't have a 'Uniqueable' instance.
+      formats :: [(CLabel, (Maybe CLabel, [Maybe LocalReg]))]
+      formats = selectStackFormat live continuations
+
+      -- Do a little meta-processing on the stack formats such as
+      -- getting the individual frame sizes and the maximum frame size
+      formats' :: (WordOff, [(CLabel, StackFormat)])
+      formats' = processFormats formats
+
+      -- TODO FIXME NOW: calculate a real max stack (including function call args)
+      -- TODO: from the maximum frame size get the maximum stack size.
+      -- The difference is due to the size taken by function calls.
+
+      -- Update the info table data on the continuations with
+      -- the selected stack formats.
+      continuations' :: [Continuation CmmInfo]
+      continuations' = map (applyStackFormat (snd formats')) continuations
+
+      -- Do the actual CPS transform.
+      cps_procs :: [CmmTop]
+      cps_procs = map (continuationToProc formats') continuations'
+
+      -- Convert the info tables from CmmInfo to [CmmStatic]
+      -- We might want to put this in another pass eventually
+      info_procs :: [RawCmmTop]
+      info_procs = concat (zipWith mkInfoTable info_uniques cps_procs)
+
 --------------------------------------------------------------------------------
 
 -- The format for the call to a continuation
@@ -45,33 +200,14 @@ import Data.List
 -- and heap memory (not sure if that's usefull at all though, but it may
 -- be worth exploring the design space).
 
-data BrokenBlock
-  = BrokenBlock {
-      brokenBlockId :: BlockId, -- Like a CmmBasicBlock
-      brokenBlockEntry :: BlockEntryInfo,
-                                -- How this block can be entered
-
-      brokenBlockStmts :: [CmmStmt],
-                                -- Like a CmmBasicBlock
-                                -- (but without the last statement)
-
-      brokenBlockTargets :: [BlockId],
-                                -- Blocks that this block could
-                                -- branch to one either by conditional
-                                -- branches or via the last statement
-
-      brokenBlockExit :: FinalStmt
-                                -- How the block can be left
-    }
-
-continuationLabel (Continuation _ _ l _ _) = l
-data Continuation =
+continuationLabel (Continuation _ l _ _) = l
+data Continuation info =
   Continuation
-     Bool              -- True => Function entry, False => Continuation/return point
-     [CmmStatic]       -- Info table, may be empty
+     info --(Either C_SRT CmmInfo)   -- Left <=> Continuation created by the CPS
+                       -- Right <=> Function or Proc point
      CLabel            -- Used to generate both info & entry labels
      CmmFormals        -- Argument locals live on entry (C-- procedure params)
-     [BrokenBlock]   -- Code, may be empty.  The first block is
+     [BrokenBlock]     -- Code, may be empty.  The first block is
                        -- the entry point.  The order is otherwise initially 
                        -- unimportant, but at some point the code gen will
                        -- fix the order.
@@ -80,491 +216,328 @@ data Continuation =
                       -- to a label.  To jump to the first block in a Proc,
                       -- use the appropriate CLabel.
 
-data BlockEntryInfo
-  = FunctionEntry              -- Beginning of a function
-      CLabel                    -- The function name
-      CmmFormals                -- Aguments to function
-
-  | ContinuationEntry          -- Return point of a call
-      CmmFormals                -- return values (argument to continuation)
-  -- TODO:
-  -- | ProcPointEntry -- no return values, but some live might end up as params or possibly in the frame
-
-  | ControlEntry               -- A label in the input
-
--- Final statement in a BlokenBlock
--- Constructors and arguments match those in Cmm,
--- but are restricted to branches, returns, jumps, calls and switches
-data FinalStmt
-  = FinalBranch
-      BlockId -- next block (must be a ControlEntry)
-
-  | FinalReturn
-      CmmActuals -- return values
-
-  | FinalJump
-      CmmExpr -- the function to call
-      CmmActuals -- arguments to call
-
-  | FinalCall
-      BlockId -- next block after call (must be a ContinuationEntry)
-      CmmCallTarget -- the function to call
-      CmmFormals -- results from call (redundant with ContinuationEntry)
-      CmmActuals -- arguments to call
-      (Maybe [GlobalReg]) -- registers that must be saved (TODO)
-
-  | FinalSwitch
-      CmmExpr [Maybe BlockId]   -- Table branch
-
-  -- TODO: | ProcPointExit (needed?)
-
--- Describes the layout of a stack frame for a continuation
 data StackFormat
-    = StackFormat
-         (Maybe CLabel)                -- The label occupying the top slot
-         WordOff               -- Total frame size in words
-         [(CmmReg, WordOff)]   -- local reg offsets from stack top
+    = StackFormat {
+         stack_label :: Maybe CLabel,  -- The label occupying the top slot
+         stack_frame_size :: WordOff,  -- Total frame size in words (not including arguments)
+         stack_live :: [Maybe LocalReg]        -- local reg offsets from stack top
+      }
 
 -- A block can be a continuation of a call
 -- A block can be a continuation of another block (w/ or w/o joins)
 -- A block can be an entry to a function
 
-blocksToBlockEnv :: [BrokenBlock] -> BlockEnv BrokenBlock
-blocksToBlockEnv blocks = listToUFM $ map (\b -> (brokenBlockId b, b)) blocks
-
 -----------------------------------------------------------------------------
-calculateOwnership :: UniqSet BlockId -> [BrokenBlock] -> BlockEnv (UniqSet BlockId)
-calculateOwnership proc_points blocks =
-    fixedpoint dependants update (map brokenBlockId blocks) emptyUFM
-    where
-      blocks_ufm :: BlockEnv BrokenBlock
-      blocks_ufm = blocksToBlockEnv blocks
-
-      dependants :: BlockId -> [BlockId]
-      dependants ident =
-          brokenBlockTargets $ lookupWithDefaultUFM
-                                 blocks_ufm unknown_block ident
-
-      update :: BlockId -> Maybe BlockId
-             -> BlockEnv (UniqSet BlockId) -> Maybe (BlockEnv (UniqSet BlockId))
-      update ident cause owners =
-          case (cause, ident `elementOfUniqSet` proc_points) of
-            (Nothing, True) -> Just $ addToUFM owners ident (unitUniqSet ident)
-            (Nothing, False) -> Nothing
-            (Just cause', True) -> Nothing
-            (Just cause', False) ->
-                if (sizeUniqSet old) == (sizeUniqSet new)
-                   then Nothing
-                   else Just $ addToUFM owners ident new
-                where
-                  old = lookupWithDefaultUFM owners emptyUniqSet ident
-                  new = old `unionUniqSets` lookupWithDefaultUFM owners emptyUniqSet cause'
-
-      unknown_block = panic "unknown BlockId in selectStackFormat"
-
-calculateProcPoints :: [BrokenBlock] -> UniqSet BlockId
-calculateProcPoints blocks = calculateProcPoints' init_proc_points blocks
-    where
-      init_proc_points = mkUniqSet $
-                         map brokenBlockId $
-                         filter always_proc_point blocks
-      always_proc_point BrokenBlock {
-                              brokenBlockEntry = FunctionEntry _ _ } = True
-      always_proc_point BrokenBlock {
-                              brokenBlockEntry = ContinuationEntry _ } = True
-      always_proc_point _ = False
-
-calculateProcPoints' :: UniqSet BlockId -> [BrokenBlock] -> UniqSet BlockId
-calculateProcPoints' old_proc_points blocks =
-    if sizeUniqSet old_proc_points == sizeUniqSet new_proc_points
-      then old_proc_points
-      else calculateProcPoints' new_proc_points blocks
-    where
-      owners = calculateOwnership old_proc_points blocks
-      new_proc_points = unionManyUniqSets (old_proc_points:(map (calculateProcPoints'' owners) blocks))
-
-calculateProcPoints'' :: BlockEnv (UniqSet BlockId) -> BrokenBlock -> UniqSet BlockId
-calculateProcPoints''  owners block =
-    unionManyUniqSets (map (f parent_id) child_ids)
-    where
-      parent_id = brokenBlockId block
-      child_ids = brokenBlockTargets block
-      -- TODO: name for f
-      f parent_id child_id = 
-          if needs_proc_point
-            then unitUniqSet child_id
-            else emptyUniqSet
-          where
-            parent_owners = lookupWithDefaultUFM owners emptyUniqSet parent_id
-            child_owners = lookupWithDefaultUFM owners emptyUniqSet child_id
-            needs_proc_point = not $ isEmptyUniqSet $ child_owners `minusUniqSet` parent_owners
 
 collectNonProcPointTargets ::
     UniqSet BlockId -> BlockEnv BrokenBlock
-    -> UniqSet BlockId -> BlockId -> UniqSet BlockId
-collectNonProcPointTargets proc_points blocks current_targets block =
+    -> UniqSet BlockId -> [BlockId] -> UniqSet BlockId
+collectNonProcPointTargets proc_points blocks current_targets new_blocks =
     if sizeUniqSet current_targets == sizeUniqSet new_targets
        then current_targets
-       else foldl (collectNonProcPointTargets proc_points blocks) new_targets targets
+       else foldl
+                (collectNonProcPointTargets proc_points blocks)
+                new_targets
+                (map (:[]) targets)
     where
-      block' = lookupWithDefaultUFM blocks (panic "TODO") block
+      blocks' = map (lookupWithDefaultUFM blocks (panic "TODO")) new_blocks
       targets =
         -- Note the subtlety that since the extra branch after a call
         -- will always be to a block that is a proc-point,
         -- this subtraction will always remove that case
-        uniqSetToList $ (mkUniqSet $ brokenBlockTargets block') `minusUniqSet` proc_points
+        uniqSetToList $ (unionManyUniqSets $ map (mkUniqSet . brokenBlockTargets) blocks')
+                          `minusUniqSet` proc_points
         -- TODO: remove redundant uniqSetToList
       new_targets = current_targets `unionUniqSets` (mkUniqSet targets)
 
-buildContinuation ::
+-- TODO: insert proc point code here
+--  * Branches and switches to proc points may cause new blocks to be created
+--    (or proc points could leave behind phantom blocks that just jump to them)
+--  * Proc points might get some live variables passed as arguments
+
+gatherBlocksIntoContinuation ::
     UniqSet BlockId -> BlockEnv BrokenBlock
-    -> BlockId -> Continuation
-buildContinuation proc_points blocks start =
-  Continuation is_entry info_table clabel params body
+    -> BlockId -> Maybe BlockId -> Continuation (Either C_SRT CmmInfo)
+gatherBlocksIntoContinuation proc_points blocks start gc =
+  Continuation info_table clabel params body
     where
-      children = (collectNonProcPointTargets proc_points blocks (unitUniqSet start) start) `delOneFromUniqSet` start
+      start_and_gc = start : maybeToList gc
+      children = (collectNonProcPointTargets proc_points blocks (mkUniqSet start_and_gc) start_and_gc) `minusUniqSet` (mkUniqSet start_and_gc)
       start_block = lookupWithDefaultUFM blocks (panic "TODO") start
+      gc_block = map (lookupWithDefaultUFM blocks (panic "TODO)")) (maybeToList gc)
       children_blocks = map (lookupWithDefaultUFM blocks (panic "TODO")) (uniqSetToList children)
-      body = start_block : children_blocks
-      info_table = [] -- TODO
+      body = start_block : gc_block ++ children_blocks
+
+      -- We can't properly annotate the continuation's stack parameters
+      -- at this point because this is before stack selection
+      -- but we want to keep the C_SRT around so we use 'Either'.
+      info_table = case start_block_entry of
+                     FunctionEntry info _ _ -> Right info
+                     ContinuationEntry _ srt -> Left srt
+                     ControlEntry -> Right (CmmNonInfo Nothing)
+
       start_block_entry = brokenBlockEntry start_block
-      is_entry = case start_block_entry of
-                   FunctionEntry _ _ -> True
-                   _ -> False
       clabel = case start_block_entry of
-                 FunctionEntry label _ -> label
+                 FunctionEntry _ label _ -> label
                  _ -> mkReturnPtLabel $ getUnique start
       params = case start_block_entry of
-                 FunctionEntry _ args -> args
-                 ContinuationEntry args -> args
+                 FunctionEntry _ _ args -> args
+                 ContinuationEntry args _ -> args
                  ControlEntry -> [] -- TODO: it's a proc-point, we could pass lives in parameter registers
 
 --------------------------------------------------------------------------------
 -- For now just select the continuation orders in the order they are in the set with no gaps
 
-selectStackFormat :: BlockEnv CmmLive -> [BrokenBlock] -> BlockEnv StackFormat
-selectStackFormat live blocks =
-    fixedpoint dependants update (map brokenBlockId blocks) emptyUFM
-    where
-      blocks_ufm :: BlockEnv BrokenBlock
-      blocks_ufm = listToUFM $ map (\b -> (brokenBlockId b, b)) blocks
-
-      dependants :: BlockId -> [BlockId]
-      dependants ident =
-          brokenBlockTargets $ lookupWithDefaultUFM
-                                 blocks_ufm unknown_block ident
-
-      update :: BlockId -> Maybe BlockId
-             -> BlockEnv StackFormat -> Maybe (BlockEnv StackFormat)
-      update ident cause formats =
-          if ident `elemUFM` formats
-             then Nothing -- Blocks only need to be updated once
-             else case (cause,
-                        brokenBlockEntry $ lookupWithDefaultUFM blocks_ufm
-                                             unknown_block ident) of
-                    -- Propagate only to blocks entered by branches
-                    -- (not function entry blocks or continuation entry blocks)
-                    (Just cause_name, ControlEntry) ->
-                        Just $ addToUFM formats ident cause_format
-                            where cause_format = lookupWithDefaultUFM
-                                                   formats unknown_block
-                                                   cause_name
-                    -- Do initial calculates for function blocks
-                    (Nothing, FunctionEntry _ _) ->
-                        Just $
-                             addToUFM formats ident $
-                             StackFormat ident 0 []
-                    -- Do initial calculates for continuation blocks
-                    (Nothing, ContinuationEntry _) ->
-                        Just $
-                             addToUFM formats ident $
-                             live_to_format ident $
-                             lookupWithDefaultUFM live unknown_block ident
-                    _ -> Nothing
-
-      unknown_block = panic "unknown BlockId in selectStackFormat"
-
-      live_to_format :: BlockId -> CmmLive -> StackFormat
-      live_to_format label live =
-          foldl extend_format
-                    (StackFormat label retAddrSizeW [])
-                    (uniqSetToList live)
-
-      extend_format :: StackFormat -> LocalReg -> StackFormat
-      extend_format (StackFormat block size offsets) reg =
-          StackFormat block (slot_size reg + size) ((CmmLocal reg, size) : offsets)
-
-selectStackFormat2 :: BlockEnv CmmLive -> [Continuation] -> [(CLabel, StackFormat)]
-selectStackFormat2 live continuations =
+selectStackFormat :: BlockEnv CmmLive
+                  -> [Continuation (Either C_SRT CmmInfo)]
+                  -> [(CLabel, (Maybe CLabel, [Maybe LocalReg]))]
+selectStackFormat live continuations =
     map (\c -> (continuationLabel c, selectStackFormat' c)) continuations
     where
-      selectStackFormat' (Continuation True info_table label formals blocks) =
-          --let ident = brokenBlockId $ head blocks -- TODO: CLabel isn't a uniquable, but we need a better way than this
-          --in
-          StackFormat (Just label) 0 []
-      selectStackFormat' (Continuation False info_table label formals blocks) =
+      selectStackFormat' (Continuation
+                          (Right (CmmInfo _ _ _ (ContInfo format srt)))
+                          label _ _) = (Just label, format)
+      selectStackFormat' (Continuation (Right _) _ _ _) = (Nothing, [])
+      selectStackFormat' (Continuation (Left srt) label _ blocks) =
           -- TODO: assumes the first block is the entry block
           let ident = brokenBlockId $ head blocks -- TODO: CLabel isn't a uniquable, but we need a better way than this
-          in live_to_format label formals $ lookupWithDefaultUFM live unknown_block ident
+          in (Just label,
+              map Just $ uniqSetToList $
+              lookupWithDefaultUFM live unknown_block ident)
 
-      live_to_format :: CLabel -> CmmFormals -> CmmLive -> StackFormat
-      live_to_format label formals live =
-          foldl extend_format
-                    (StackFormat (Just label) retAddrSizeW [])
-                    (uniqSetToList (live `minusUniqSet` mkUniqSet (cmmFormalsToLiveLocals formals)))
+      unknown_block = panic "unknown BlockId in selectStackFormat"
 
-      extend_format :: StackFormat -> LocalReg -> StackFormat
-      extend_format (StackFormat label size offsets) reg =
-          StackFormat label (slot_size reg + size) ((CmmLocal reg, size) : offsets)
+processFormats :: [(CLabel, (Maybe CLabel, [Maybe LocalReg]))]
+               -> (WordOff, [(CLabel, StackFormat)])
+processFormats formats = (max_size, formats')
+    where
+      max_size = foldl max 0 (map (stack_frame_size . snd) formats')
+      formats' = map make_format formats
+      make_format (label, format) =
+          (label,
+           StackFormat {
+             stack_label = fst format,
+             stack_frame_size = stack_size (snd format) +
+                                if isJust (fst format)
+                                then label_size
+                                else 0,
+             stack_live = snd format })
+
+      -- TODO: get rid of "+ 1" etc.
+      label_size = 1 :: WordOff
+
+      stack_size [] = 0
+      stack_size (Nothing:formats) = 1 + stack_size formats -- one dead word
+      stack_size (Just reg:formats) = width + stack_size formats
+          where
+            width = machRepByteWidth (localRegRep reg) `quot` wORD_SIZE
+            -- TODO: it would be better if we had a machRepWordWidth
 
-      unknown_block = panic "unknown BlockId in selectStackFormat"
+-----------------------------------------------------------------------------
+applyStackFormat :: [(CLabel, StackFormat)]
+                 -> Continuation (Either C_SRT CmmInfo)
+                 -> Continuation CmmInfo
+
+-- User written continuations
+applyStackFormat formats (Continuation
+                          (Right (CmmInfo prof gc tag (ContInfo _ srt)))
+                          label formals blocks) =
+    Continuation (CmmInfo prof gc tag (ContInfo format srt))
+                 label formals blocks
+    where
+      format = stack_live $ maybe unknown_block id $ lookup label formats
+      unknown_block = panic "unknown BlockId in applyStackFormat"
 
-slot_size reg = ((machRepByteWidth (localRegRep reg) - 1) `div` wORD_SIZE) + 1
+-- User written non-continuation code
+applyStackFormat formats (Continuation (Right info) label formals blocks) =
+    Continuation info label formals blocks
 
-constructContinuation :: [(CLabel, StackFormat)] -> Continuation -> CmmTop
-constructContinuation formats (Continuation is_entry info label formals blocks) =
-    CmmProc info label formals (map (constructContinuation2' label formats) blocks)
+-- CPS generated continuations
+applyStackFormat formats (Continuation (Left srt) label formals blocks) =
+    Continuation (CmmInfo prof gc tag (ContInfo (stack_live $ format) srt))
+                 label formals blocks
+    where
+      gc = Nothing -- Generated continuations never need a stack check
+      -- TODO prof: this is the same as the current implementation
+      -- but I think it could be improved
+      prof = ProfilingInfo zeroCLit zeroCLit
+      tag = rET_SMALL -- cmmToRawCmm will convert this to rET_BIG if needed
+      format = maybe unknown_block id $ lookup label formats
+      unknown_block = panic "unknown BlockId in applyStackFormat"
 
-constructContinuation2' :: CLabel -> [(CLabel, StackFormat)] -> BrokenBlock
-                       -> CmmBasicBlock
-constructContinuation2' curr_ident formats (BrokenBlock ident entry stmts _ exit) =
-    BasicBlock ident (prefix++stmts++postfix)
+-----------------------------------------------------------------------------
+continuationToProc :: (WordOff, [(CLabel, StackFormat)])
+                   -> Continuation CmmInfo
+                   -> CmmTop
+continuationToProc (max_stack, formats)
+                   (Continuation info label formals blocks) =
+    CmmProc info label formals (map continuationToProc' blocks)
     where
-      curr_format = maybe unknown_block id $ lookup curr_ident formats
-      unknown_block = panic "unknown BlockId in constructContinuation"
-      prefix = case entry of
-                 ControlEntry -> []
-                 FunctionEntry _ _ -> []
-                 ContinuationEntry formals ->
-                     unpack_continuation curr_format
-      postfix = case exit of
-                  FinalBranch next -> [CmmBranch next]
-                  FinalSwitch expr targets -> [CmmSwitch expr targets]
-                  FinalReturn arguments ->
-                      exit_function curr_format
-                                    (CmmLoad (CmmReg spReg) wordRep)
-                                    arguments
-                  FinalJump target arguments ->
-                      exit_function curr_format target arguments
-                  -- TODO: do something about global saves
-                  FinalCall next (CmmForeignCall target CmmCallConv)
-                            results arguments saves ->
+      curr_format = maybe unknown_block id $ lookup label formats
+      unknown_block = panic "unknown BlockId in continuationToProc"
+
+      continuationToProc' :: BrokenBlock -> CmmBasicBlock
+      continuationToProc' (BrokenBlock ident entry stmts _ exit) =
+          BasicBlock ident (prefix++stmts++postfix)
+          where
+            prefix = case entry of
+                       ControlEntry -> []
+                       FunctionEntry (CmmInfo _ (Just gc_block) _ _) _ formals ->
+                           gc_stack_check gc_block max_stack ++
+                           function_entry formals curr_format
+                       FunctionEntry (CmmInfo _ Nothing _ _) _ formals ->
+                           panic "continuationToProc: missing GC block"
+                       FunctionEntry (CmmNonInfo (Just gc_block)) _ formals ->
+                           gc_stack_check gc_block max_stack ++
+                           function_entry formals curr_format
+                       FunctionEntry (CmmNonInfo Nothing) _ formals ->
+                           panic "continuationToProc: missing non-info GC block"
+                       ContinuationEntry formals _ ->
+                           function_entry formals curr_format
+            postfix = case exit of
+                        FinalBranch next -> [CmmBranch next]
+                        FinalSwitch expr targets -> [CmmSwitch expr targets]
+                        FinalReturn arguments ->
+                            tail_call (stack_frame_size curr_format)
+                                (CmmLoad (CmmReg spReg) wordRep)
+                                arguments
+                        FinalJump target arguments ->
+                            tail_call (stack_frame_size curr_format) target arguments
+                        FinalCall next (CmmForeignCall target CmmCallConv)
+                            results arguments ->
                                 pack_continuation curr_format cont_format ++
-                                [CmmJump target arguments]
+                                tail_call (stack_frame_size curr_format - stack_frame_size cont_format)
+                                              target arguments
                             where
                               cont_format = maybe unknown_block id $
                                             lookup (mkReturnPtLabel $ getUnique next) formats
-                  FinalCall next _ results arguments saves -> panic "unimplemented CmmCall"
+                        FinalCall next _ results arguments -> panic "unimplemented CmmCall"
 
---------------------------------------------------------------------------------
+-----------------------------------------------------------------------------
 -- Functions that generate CmmStmt sequences
 -- for packing/unpacking continuations
 -- and entering/exiting functions
 
-exit_function :: StackFormat -> CmmExpr -> CmmActuals -> [CmmStmt]
-exit_function (StackFormat curr_id curr_frame_size curr_offsets) target arguments
-  = adjust_spReg ++ jump where
+tail_call :: WordOff -> CmmExpr -> CmmActuals -> [CmmStmt]
+tail_call spRel target arguments
+  = store_arguments ++ adjust_spReg ++ jump where
+    store_arguments =
+        [stack_put spRel expr offset
+         | ((expr, _), StackParam offset) <- argument_formats] ++
+        [global_put expr global
+         | ((expr, _), RegisterParam global) <- argument_formats]
     adjust_spReg =
-        if curr_frame_size == 0
+        if spRel == 0
         then []
-        else [CmmAssign spReg
-                 (CmmRegOff spReg (curr_frame_size*wORD_SIZE))]
+        else [CmmAssign spReg (CmmRegOff spReg (spRel*wORD_SIZE))]
     jump = [CmmJump target arguments]
 
-enter_function :: WordOff -> [CmmStmt]
-enter_function max_frame_size
+    argument_formats = assignArguments (cmmExprRep . fst) arguments
+
+gc_stack_check :: BlockId -> WordOff -> [CmmStmt]
+gc_stack_check gc_block max_frame_size
   = check_stack_limit where
     check_stack_limit = [
      CmmCondBranch
      (CmmMachOp (MO_U_Lt $ cmmRegRep spReg)
-                    [CmmRegOff spReg max_frame_size, CmmReg spLimReg])
+                    [CmmRegOff spReg (-max_frame_size*wORD_SIZE),
+                     CmmReg spLimReg])
      gc_block]
-    gc_block = undefined -- TODO: get stack and heap checks to go to same
 
--- TODO: fix branches to proc point (we have to insert a new block to marshel the continuation)
+-- TODO: fix branches to proc point
+-- (we have to insert a new block to marshel the continuation)
 pack_continuation :: StackFormat -> StackFormat -> [CmmStmt]
-pack_continuation (StackFormat curr_id curr_frame_size curr_offsets)
-                       (StackFormat cont_id cont_frame_size cont_offsets)
-  = save_live_values ++ set_stack_header ++ adjust_spReg where
+pack_continuation (StackFormat curr_id curr_frame_size _)
+                       (StackFormat cont_id cont_frame_size live_regs)
+  = store_live_values ++ set_stack_header where
     -- TODO: only save variables when actually needed
-    save_live_values =
-        [CmmStore
-         (CmmRegOff
-          spReg (wORD_SIZE*(curr_frame_size - cont_frame_size + offset)))
-         (CmmReg reg)
+    -- (may be handled by latter pass)
+    store_live_values =
+        [stack_put spRel (CmmReg (CmmLocal reg)) offset
          | (reg, offset) <- cont_offsets]
-    needs_header =
-      case (curr_id, cont_id) of
-        (Just x, Just y) -> x /= y
-        _ -> isJust cont_id
     set_stack_header =
-      if not needs_header
-         then []
-         else [CmmStore (CmmRegOff spReg (wORD_SIZE*(curr_frame_size - cont_frame_size))) continuation_function]
+        if needs_header_set
+        then [stack_put spRel continuation_function 0]
+        else []
+
+    -- TODO: factor with function_entry and CmmInfo.hs(?)
+    cont_offsets = mkOffsets label_size live_regs
+
+    label_size = 1 :: WordOff
+
+    mkOffsets size [] = []
+    mkOffsets size (Nothing:regs) = mkOffsets (size+1) regs
+    mkOffsets size (Just reg:regs) = (reg, size):mkOffsets (size + width) regs
+        where
+          width = machRepByteWidth (localRegRep reg) `quot` wORD_SIZE
+          -- TODO: it would be better if we had a machRepWordWidth
+
+    spRel = curr_frame_size - cont_frame_size
     continuation_function = CmmLit $ CmmLabel $ fromJust cont_id
-    adjust_spReg =
-        if curr_frame_size == cont_frame_size
-        then []
-        else [CmmAssign spReg (CmmRegOff spReg ((curr_frame_size - cont_frame_size)*wORD_SIZE))]
+    needs_header_set =
+        case (curr_id, cont_id) of
+          (Just x, Just y) -> x /= y
+          _ -> isJust cont_id
 
 -- Lazy adjustment of stack headers assumes all blocks
 -- that could branch to eachother (i.e. control blocks)
 -- have the same stack format (this causes a problem
 -- only for proc-point).
-unpack_continuation :: StackFormat -> [CmmStmt]
-unpack_continuation (StackFormat curr_id curr_frame_size curr_offsets)
-  = load_live_values where
+function_entry :: CmmFormals -> StackFormat -> [CmmStmt]
+function_entry formals (StackFormat _ _ live_regs)
+  = load_live_values ++ load_args where
     -- TODO: only save variables when actually needed
+    -- (may be handled by latter pass)
     load_live_values =
-        [CmmAssign
-         reg
-         (CmmLoad (CmmRegOff spReg (wORD_SIZE*offset)) (cmmRegRep reg))
+        [stack_get 0 reg offset
          | (reg, offset) <- curr_offsets]
+    load_args =
+        [stack_get 0 reg offset
+         | (reg, StackParam offset) <- argument_formats] ++
+        [global_get reg global
+         | (reg, RegisterParam global) <- argument_formats]
 
------------------------------------------------------------------------------
--- Breaking basic blocks on function calls
------------------------------------------------------------------------------
+    argument_formats = assignArguments (localRegRep) formals
 
------------------------------------------------------------------------------
--- Takes a basic block and breaks it up into a list of broken blocks
---
--- Takes a basic block and returns a list of basic blocks that
--- each have at most 1 CmmCall in them which must occur at the end.
--- Also returns with each basic block, the variables that will
--- be arguments to the continuation of the block once the call (if any)
--- returns.
-
-breakBlock :: [Unique] -> CmmBasicBlock -> BlockEntryInfo -> [BrokenBlock]
-breakBlock uniques (BasicBlock ident stmts) entry =
-    breakBlock' uniques ident entry [] [] stmts where
-        breakBlock' uniques current_id entry exits accum_stmts stmts =
-            case stmts of
-              [] -> panic "block doesn't end in jump, goto or return"
-              [CmmJump target arguments] ->
-                  [BrokenBlock current_id entry accum_stmts
-                               exits
-                               (FinalJump target arguments)]
-              [CmmReturn arguments] ->
-                  [BrokenBlock current_id entry accum_stmts
-                               exits
-                               (FinalReturn arguments)]
-              [CmmBranch target] ->
-                  [BrokenBlock current_id entry accum_stmts
-                               (target:exits)
-                               (FinalBranch target)]
-              [CmmSwitch expr targets] ->
-                  [BrokenBlock current_id entry accum_stmts
-                               (mapMaybe id targets ++ exits)
-                               (FinalSwitch expr targets)]
-              (CmmJump _ _:_) ->
-                  panic "jump in middle of block"
-              (CmmReturn _:_) ->
-                  panic "return in middle of block"
-              (CmmBranch _:_) ->
-                  panic "branch in middle of block"
-              (CmmSwitch _ _:_) ->
-                  panic ("switch in middle of block" ++ (showSDoc $ ppr stmts))
-              (CmmCall target results arguments saves:stmts) -> block : rest
-                  where
-                    new_id = BlockId $ head uniques
-                    block = BrokenBlock current_id entry accum_stmts
-                            (new_id:exits)
-                            (FinalCall new_id target results arguments saves)
-                    rest = breakBlock' (tail uniques) new_id
-                           (ContinuationEntry results) [] [] stmts
-              (s@(CmmCondBranch test target):stmts) ->
-                  breakBlock' uniques current_id entry
-                              (target:exits) (accum_stmts++[s]) stmts
-              (s:stmts) ->
-                  breakBlock' uniques current_id entry
-                              exits (accum_stmts++[s]) stmts
+    -- TODO: eliminate copy/paste with pack_continuation
+    curr_offsets = mkOffsets label_size live_regs
 
---------------------------------
--- Convert from a BrokenBlock
--- to a CmmBasicBlock so the
--- liveness analysis can run
--- on it.
---------------------------------
-cmmBlockFromBrokenBlock :: BrokenBlock -> CmmBasicBlock
-cmmBlockFromBrokenBlock (BrokenBlock ident _ stmts _ exit) =
-    BasicBlock ident (stmts++exit_stmt)
-    where
-      exit_stmt =
-          case exit of
-            FinalBranch target -> [CmmBranch target]
-            FinalReturn arguments -> [CmmReturn arguments]
-            FinalJump target arguments -> [CmmJump target arguments]
-            FinalSwitch expr targets -> [CmmSwitch expr targets]
-            FinalCall branch_target call_target results arguments saves ->
-                [CmmCall call_target results arguments saves,
-                 CmmBranch branch_target]
+    label_size = 1 :: WordOff
+
+    mkOffsets size [] = []
+    mkOffsets size (Nothing:regs) = mkOffsets (size+1) regs
+    mkOffsets size (Just reg:regs) = (reg, size):mkOffsets (size + width) regs
+        where
+          width = machRepByteWidth (localRegRep reg) `quot` wORD_SIZE
+          -- TODO: it would be better if we had a machRepWordWidth
 
 -----------------------------------------------------------------------------
--- CPS a single CmmTop (proceedure)
+-- Section: Stack and argument register puts and gets
 -----------------------------------------------------------------------------
+-- TODO: document
+
+-- |Construct a 'CmmStmt' that will save a value on the stack
+stack_put :: WordOff            -- ^ Offset from the real 'Sp' that 'offset'
+                                -- is relative to (added to offset)
+          -> CmmExpr            -- ^ What to store onto the stack
+          -> WordOff            -- ^ Where on the stack to store it
+                                -- (positive <=> higher addresses)
+          -> CmmStmt
+stack_put spRel expr offset =
+    CmmStore (CmmRegOff spReg (wORD_SIZE*(spRel + offset))) expr
 
-cpsProc :: UniqSupply -> CmmTop -> [CmmTop]
-cpsProc uniqSupply x@(CmmData _ _) = [x]
-cpsProc uniqSupply x@(CmmProc info_table ident params blocks) =
-    --[CmmProc info_table ident params cps_blocks]
-    cps_continuations
-    where
-      uniqes :: [[Unique]]
-      uniqes = map uniqsFromSupply $ listSplitUniqSupply uniqSupply
-
-      -- Break the block at each function call
-      broken_blocks :: [BrokenBlock]
-      broken_blocks = concat $ zipWith3 breakBlock uniqes blocks
-                                        (FunctionEntry ident params:repeat ControlEntry)
-
-      -- Calculate live variables for each broken block
-      live :: BlockEntryLiveness
-      live = cmmLiveness $ map cmmBlockFromBrokenBlock broken_blocks
-             -- nothing can be live on entry to the first block so we could take the tail
-
-      proc_points :: UniqSet BlockId
-      proc_points = calculateProcPoints broken_blocks
-
-      continuations :: [Continuation]
-      continuations = map (buildContinuation proc_points (blocksToBlockEnv broken_blocks)) (uniqSetToList proc_points)
-
-      -- TODO: insert proc point code here
-      --  * Branches and switches to proc points may cause new blocks to be created
-      --    (or proc points could leave behind phantom blocks that just jump to them)
-      --  * Proc points might get some live variables passed as arguments
-
-      -- TODO: let blocks_with_live = map (cmmLivenessComment live . snd) broken_blocks
-
-      --procs = groupBlocksIntoContinuations live broken_blocks
-
-      -- Select the stack format on entry to each block
-      formats2 :: [(CLabel, StackFormat)]
-      formats2 = selectStackFormat2 live continuations
-
-      -- Do the actual CPS transform
-      cps_continuations :: [CmmTop]
-      cps_continuations = map (constructContinuation formats2) continuations
-
---------------------------------------------------------------------------------
-cmmCPS :: DynFlags
-       -> [Cmm]                 -- C-- with Proceedures
-       -> IO [Cmm]             -- Output: CPS transformed C--
-
-cmmCPS dflags abstractC = do
-  when (dopt Opt_DoCmmLinting dflags) $
-       do showPass dflags "CmmLint"
-         case firstJust $ map cmmLint abstractC of
-           Just err -> do printDump err
-                          ghcExit dflags 1
-           Nothing  -> return ()
-  showPass dflags "CPS"
-  -- TODO: check for use of branches to non-existant blocks
-  -- TODO: check for use of Sp, SpLim, R1, R2, etc.
-  -- TODO: find out if it is valid to create a new unique source like this
-  uniqSupply <- mkSplitUniqSupply 'p'
-  let supplies = listSplitUniqSupply uniqSupply
-  let continuationC = zipWith (\s (Cmm c) -> Cmm $ concat $ zipWith (cpsProc) (listSplitUniqSupply s) c) supplies abstractC
+--------------------------------
+-- |Construct a 
+stack_get :: WordOff
+          -> LocalReg
+          -> WordOff
+          -> CmmStmt
+stack_get spRel reg offset =
+    CmmAssign (CmmLocal reg)
+              (CmmLoad (CmmRegOff spReg (wORD_SIZE*(spRel + offset)))
+                       (localRegRep reg))
+global_put :: CmmExpr -> GlobalReg -> CmmStmt
+global_put expr global = CmmAssign (CmmGlobal global) expr
+global_get :: LocalReg -> GlobalReg -> CmmStmt
+global_get reg global = CmmAssign (CmmLocal reg) (CmmReg (CmmGlobal global))
 
-  dumpIfSet_dyn dflags Opt_D_dump_cps_cmm "CPS Cmm" (pprCmms continuationC)
-  -- TODO: add option to dump Cmm to file
-  return continuationC