Implemented and fixed bugs in CmmInfo handling
[ghc-hetmet.git] / compiler / cmm / CmmCPS.hs
index 88d2585..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,8 +10,15 @@ import Cmm
 import CmmLint
 import PprCmm
 
-import Dataflow -- (fixedpoint, cmmLivenessComment, cmmLiveness, CmmLive)
+import CmmLive
+import CmmBrokenBlock
+import CmmProcPoint
+import CmmCallConv
+import CmmInfo
+import CmmUtils
 
+import Bitmap
+import ClosureInfo
 import MachOp
 import ForeignCall
 import CLabel
@@ -25,36 +36,152 @@ import Unique
 
 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"
 
---------------------------------------------------------------------------------
--- Monad for the CPSer
--- Contains:
---  * State for the uniqSupply
-
-data CPSState = CPSState { cps_uniqs :: UniqSupply }
-
-data CPS a = CPS { runCPS :: CPSState -> (CPSState, a) }
+  -- TODO: more lint checking
+  --        check for use of branches to non-existant blocks
+  --        check for use of Sp, SpLim, R1, R2, etc.
 
-instance Monad CPS where
-  return a = CPS $ \s -> (s, a)
-  (CPS m) >>= f = CPS $ \s ->
-    let (s', m') = m s
-    in runCPS (f m') s'
+  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
 
---------------------------------------------------------------------------------
--- Utility functions
+  dumpIfSet_dyn dflags Opt_D_dump_cps_cmm "CPS Cmm" (pprCmms continuationC)
 
-getState = CPS $ \s -> (s, s)
-putState s = CPS $ \_ -> (s, ())
+  -- TODO: add option to dump Cmm to file
 
-newLabelCPS = do
-  state <- getState
-  let (us1, us2) = splitUniqSupply (cps_uniqs state)
-  putState $ state { cps_uniqs = us1 }
-  return $ BlockId (uniqFromSupply us2)
+  return continuationC
 
-mapMCmmTop :: (Monad m) => (CmmTop -> m [CmmTop]) -> Cmm -> m Cmm
-mapMCmmTop f (Cmm xs) = liftM Cmm $ liftM concat $ mapM f xs
+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)
 
 --------------------------------------------------------------------------------
 
@@ -73,417 +200,344 @@ mapMCmmTop f (Cmm xs) = liftM Cmm $ liftM concat $ mapM f xs
 -- and heap memory (not sure if that's usefull at all though, but it may
 -- be worth exploring the design space).
 
-data BrokenBlock
-  = BrokenBlock                
-       BlockId                 -- Like a CmmBasicBlock
-       BlockEntryInfo          -- How this block can be entered
-       [CmmStmt]               -- Like a CmmBasicBlock (but without
-                               --      the last statement)
-       BlockExitInfo           -- How the block can be left
-
-data BlockEntryInfo
-  = FunctionEntry              -- Beginning of function
-
-  | ContinuationEntry          -- Return point of a call
-       CmmFormals {- return values -}
-  -- TODO | ProcPointEntry {- no return values, but some live might end up as params -}
-
-  | ControlEntry               -- A label in the input
-
-data BlockExitInfo
-  = ControlExit [BlockId] -- blocks branched to conditionally 
-    BlockId -- next block (must be a ControlEntry)
-
-  | ReturnExit [BlockId] -- blocks branched to conditionally 
-    CmmActuals -- return values
-
-  | TailCallExit [BlockId] -- blocks branched to conditionally 
-    CmmExpr -- the function to call
-    CmmActuals -- arguments to call
-
-  | CallExit [BlockId] -- blocks branched to conditionally 
-    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)
-  -- TODO: | ProcPointExit (needed?)
+continuationLabel (Continuation _ l _ _) = l
+data Continuation info =
+  Continuation
+     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
+                       -- the entry point.  The order is otherwise initially 
+                       -- unimportant, but at some point the code gen will
+                       -- fix the order.
+
+                      -- the BlockId of the first block does not give rise
+                      -- to a label.  To jump to the first block in a Proc,
+                      -- use the appropriate CLabel.
 
-data CPSBlockInfo
-  = ControlBlock -- Consider whether a proc-point might want arguments on stack
-  | ContinuationBlock [(CmmReg,MachHint)] {- params -}
-  | EntryBlock
-
---type StackFormat = [Maybe LocalReg] -- TODO: consider params as part of format
 data StackFormat
-    = StackFormat
-        BlockId {- block that is the start of the continuation. may or may not be the current block -}
-        WordOff {- total frame size -}
-        [(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
 
+-----------------------------------------------------------------------------
+
+collectNonProcPointTargets ::
+    UniqSet BlockId -> BlockEnv BrokenBlock
+    -> 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
+                (map (:[]) targets)
+    where
+      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 $ (unionManyUniqSets $ map (mkUniqSet . brokenBlockTargets) blocks')
+                          `minusUniqSet` proc_points
+        -- TODO: remove redundant uniqSetToList
+      new_targets = current_targets `unionUniqSets` (mkUniqSet targets)
+
+-- 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 -> Maybe BlockId -> Continuation (Either C_SRT CmmInfo)
+gatherBlocksIntoContinuation proc_points blocks start gc =
+  Continuation info_table clabel params body
+    where
+      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 : 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
+      clabel = case start_block_entry of
+                 FunctionEntry _ label _ -> label
+                 _ -> mkReturnPtLabel $ getUnique start
+      params = case start_block_entry of
+                 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
--- TODO: select a format that keeps blocks that can jump to each other the same
--- Assumed that jumps, calls 
-selectStackFormat :: UniqFM {-BlockId-} CmmFormals -> UniqFM {-BlockId-} CmmLive -> UniqFM {-BlockId-} [(CPSBlockInfo, CmmBasicBlock)] -> UniqFM {-BlockId-} StackFormat
-selectStackFormat = undefined
-{-
-selectStackFormat param live blocks = fixedpoint 
-listToUFM $ map live_to_format $ ufmToList live
+
+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
+                          (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 (Just label,
+              map Just $ uniqSetToList $
+              lookupWithDefaultUFM live unknown_block ident)
+
+      unknown_block = panic "unknown BlockId in selectStackFormat"
+
+processFormats :: [(CLabel, (Maybe CLabel, [Maybe LocalReg]))]
+               -> (WordOff, [(CLabel, StackFormat)])
+processFormats formats = (max_size, formats')
     where
-      live_to_format (unique, live) = (unique, format) where
-          format = foldl extend_format
-                    (StackFormat (BlockId unique) 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 :: UniqFM {-BlockId-} CmmLive -> [BrokenBlock] -> UniqFM {-BlockId-} StackFormat
-selectStackFormat2 live blocks = fixedpoint dependants update (map brokenBlockId blocks) emptyUFM where
-  blocks_ufm = listToUFM $ map (\b -> (brokenBlockId b, b)) blocks
-  dependants ident =
-      case lookupWithDefaultUFM blocks_ufm (panic "TODO") ident of
-        (BrokenBlock _ _ _ (ControlExit exits next)) -> next:exits
-        (BrokenBlock _ _ _ (ReturnExit exits _)) -> exits
-        (BrokenBlock _ _ _ (TailCallExit exits _ _)) -> exits
-        (BrokenBlock _ _ _ (CallExit exits _ _ _ _ _)) -> exits
-  update ident cause formats =
-    let BrokenBlock _ entry _ _ = lookupWithDefaultUFM blocks_ufm (panic "unknown BlockId in selectStackFormat:live") ident in
-    case cause of
-      -- Propagate only to blocks entered by branches (not function entry blocks or continuation entry blocks)
-      Just cause_name ->
-          let cause_format = lookupWithDefaultUFM formats (panic "update signaled for block not in format") cause_name
-          in case entry of
-            ControlEntry -> Just $ addToUFM formats ident cause_format
-            FunctionEntry -> Nothing
-            ContinuationEntry _ -> Nothing
-      -- Do initial calculates for function blocks
-      Nothing ->
-          case entry of
-            ControlEntry -> Nothing
-            FunctionEntry -> Just $ addToUFM formats ident $ StackFormat ident 0 []
-            ContinuationEntry _ -> Just $ addToUFM formats ident $ live_to_format ident $ lookupWithDefaultUFM live (panic "TODO") ident
-  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)
-
-slot_size reg = ((machRepByteWidth (localRegRep reg) - 1) `div` wORD_SIZE) + 1
-
-transformReturn :: UniqFM {-BlockId-} CPSBlockInfo -> UniqFM {-BlockId-} StackFormat -> CmmBasicBlock -> CmmBasicBlock
-transformReturn block_infos formats (BasicBlock ident stmts) =
-  -- NOTE: assumes that return/jump can *only* appear at end of block
-  case last stmts of
-    CmmReturn arguments ->
-        BasicBlock ident $
-                  (init stmts) ++
-                  exit_function curr_format (CmmLoad (CmmReg spReg) wordRep) arguments
-    CmmJump target arguments ->
-        BasicBlock ident $
-                  (init stmts) ++
-                  exit_function curr_format target arguments
-    _ -> BasicBlock ident stmts
-  where
-  curr_format = lookupWithDefaultUFM formats (panic $ "format: unknown block " ++ (showSDoc $ ppr $ getUnique ident)) ident
-
-destructContinuation :: UniqFM {-BlockId-} CPSBlockInfo -> UniqFM {-BlockId-} StackFormat -> CmmBasicBlock -> CmmBasicBlock
-destructContinuation block_infos formats (BasicBlock ident stmts) =
-  case info of
-    ControlBlock -> BasicBlock ident stmts
-    ContinuationBlock _ -> BasicBlock ident (unpack_continuation curr_format ++ stmts)
-  where
-  info = lookupWithDefaultUFM block_infos (panic $ "info: unknown block " ++ (showSDoc $ ppr $ getUnique ident)) ident
-  curr_format = lookupWithDefaultUFM formats (panic $ "format: unknown block " ++ (showSDoc $ ppr $ getUnique ident)) ident
-
-constructContinuation2 :: UniqFM {-BlockId-} StackFormat -> BrokenBlock -> CmmBasicBlock
-constructContinuation2 formats (BrokenBlock ident entry stmts exit) =
-    BasicBlock ident (prefix++stmts++postfix)
+      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
+
+-----------------------------------------------------------------------------
+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
-      curr_format = lookupWithDefaultUFM formats (panic $ "format: unknown block " ++ (showSDoc $ ppr $ getUnique ident)) ident
-      prefix = case entry of
-                 ControlEntry -> []
-                 FunctionEntry -> []
-                 ContinuationEntry formals -> unpack_continuation curr_format
-      postfix = case exit of
-                  ControlExit _ next -> [CmmBranch next]
-                  ReturnExit _ arguments -> exit_function curr_format (CmmLoad (CmmReg spReg) wordRep) arguments
-                  TailCallExit _ target arguments -> exit_function curr_format target arguments
-                  -- TODO: do something about global saves
-                  CallExit _ next (CmmForeignCall target CmmCallConv) results arguments saves ->
-                      let cont_format = lookupWithDefaultUFM formats (panic $ "format: unknown block " ++ (showSDoc $ ppr $ getUnique next)) next
-                      in pack_continuation curr_format cont_format ++
-                             [CmmJump target arguments]
-                  CallExit _ next _ results arguments saves -> panic "unimplemented CmmCall"
-
-constructContinuation :: UniqFM {-BlockId-} CPSBlockInfo -> UniqFM {-BlockId-} StackFormat -> CmmBasicBlock -> CmmBasicBlock
-constructContinuation block_infos formats (BasicBlock ident stmts) =
-  case last $ init stmts of
-    -- TODO: global_saves
-    --CmmCall (CmmForeignCall target CmmCallConv) results arguments (Just []) -> --TODO: handle globals
-    CmmCall (CmmForeignCall target CmmCallConv) results arguments _ ->
-        BasicBlock ident $
-                   init (init stmts) ++
-                   pack_continuation curr_format cont_format ++
-                   [CmmJump target arguments]
-    CmmCall target results arguments _ -> panic "unimplemented CmmCall"
-    -- TODO: branches for proc-points
-    -- _ -> BasicBlock ident $ (init stmts) ++ build_block_branch
-    _ -> BasicBlock ident stmts
-  where
-  info = lookupWithDefaultUFM block_infos (panic $ "info: unknown block " ++ (showSDoc $ ppr $ getUnique next_block)) next_block
-  cont_format = lookupWithDefaultUFM formats (panic $ "format: unknown block " ++ (showSDoc $ ppr $ getUnique next_block)) next_block
-  curr_format = lookupWithDefaultUFM formats (panic $ "format: unknown block " ++ (showSDoc $ ppr $ getUnique next_block)) ident
-  next_block = case last stmts of
-    CmmBranch next -> next
-    -- TODO: blocks with jump at end
-    -- TODO: blocks with return at end
-    _ -> panic $ "basic block without a branch at the end (unimplemented) " ++ (showSDoc $ ppr $ stmts)
-  next_block_as_proc_expr = CmmLit $ CmmLabel $ mkReturnPtLabel $ getUnique next_block
-  block_needs_call = True -- TODO: use a table (i.e. proc-point)
-  build_block_branch =
-    if block_needs_call
-       then [CmmJump next_block_as_proc_expr [] {- TODO: pass live -}] {- NOTE: a block can never be both a continuation and a controll block -}
-       else [CmmBranch next_block]
+      format = stack_live $ maybe unknown_block id $ lookup label formats
+      unknown_block = panic "unknown BlockId in applyStackFormat"
 
---------------------------------------------------------------------------------
+-- User written non-continuation code
+applyStackFormat formats (Continuation (Right info) label formals blocks) =
+    Continuation info label formals 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"
+
+-----------------------------------------------------------------------------
+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 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 ++
+                                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 -> 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
-    adjust_spReg = [
-     CmmAssign spReg
-     (CmmRegOff spReg (curr_frame_size*wORD_SIZE))]
+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 spRel == 0
+        then []
+        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]
-    set_stack_header = -- TODO: only set when needed
-        [CmmStore (CmmRegOff spReg (wORD_SIZE*(curr_frame_size - cont_frame_size))) continuation_function]
-    continuation_function = CmmLit $ CmmLabel $ mkReturnPtLabel $ getUnique 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))]
+    set_stack_header =
+        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
+    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]
+
+    argument_formats = assignArguments (localRegRep) formals
+
+    -- TODO: eliminate copy/paste with pack_continuation
+    curr_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
+
+-----------------------------------------------------------------------------
+-- 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
+
+--------------------------------
+-- |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))
 
--- TODO: TBD when to adjust the stack
-
-cpsProc :: CmmTop -> CPS [CmmTop]
-cpsProc x@(CmmData _ _) = return [x]
-cpsProc x@(CmmProc info_table ident params blocks) = do
-
-  broken_blocks <- liftM concat $ mapM breakBlock blocks
-  broken_blocks2 <- liftM concat (zipWithM breakBlock2 blocks (FunctionEntry:repeat ControlEntry))
-       -- broken_blocks :: [BrokenBlock]
-
-   let live = cmmLiveness (map snd broken_blocks)
-  let live2 :: BlockEntryLiveness
-      live2 = cmmLiveness2 broken_blocks2
-
-  let blocks_with_live = map (cmmLivenessComment live . snd) broken_blocks
-
-  let formats = selectStackFormat (panic "params to selectStackFormat" {-TODO-}) live (undefined)
-  let formats2 :: BlockEnv StackFormat -- Stack format on entry
-      formats2 = selectStackFormat2 live2 broken_blocks2
-
-  let block_infos = listToUFM $ map (\(info, block) -> (blockId block, info)) broken_blocks
-  --let blocks_with_live' = map (constructContinuation block_infos formats) blocks_with_live
-  --let blocks_with_live'' = map (destructContinuation block_infos formats) blocks_with_live'
-  --let blocks_with_live''' = map (transformReturn block_infos formats) blocks_with_live''
-
-  return $ [CmmProc info_table ident params $ map (constructContinuation2 formats2) broken_blocks2]
-{-  
-  return $ [CmmProc info_table ident params $
-            map (constructContinuation block_infos formats .
-                 destructContinuation block_infos formats .
-                 transformReturn block_infos formats)
-            blocks_with_live]
--}
-
---------------------------------------------------------------------------------
--- 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.
-
-cmmBlockifyCalls :: [CmmBasicBlock] -> CPS [(CPSBlockInfo, CmmBasicBlock)]
-cmmBlockifyCalls blocks = liftM concat $ mapM breakBlock blocks
-
--- [(CmmReg,MachHint)] is the results from the previous block that are expected as parameters
---breakBlock :: CmmBasicBlock -> CPS [(Maybe BlockId, CmmBasicBlock)]
-breakBlock :: CmmBasicBlock -> CPS [(CPSBlockInfo, CmmBasicBlock)]
-breakBlock (BasicBlock ident stmts) = breakBlock' ident ControlBlock [] stmts
-
-breakBlock' current_id block_info accum_stmts [] =
-  return [(block_info, BasicBlock current_id accum_stmts)]
--- TODO: notice a call just before a branch, jump, call, etc.
-breakBlock' current_id block_info accum_stmts (stmt@(CmmCall _ results _ _):stmts) = do
-  new_id <- newLabelCPS
-  let new_block = (block_info, BasicBlock current_id (accum_stmts ++ [stmt, CmmBranch new_id]))
-  rest <- breakBlock' new_id (ContinuationBlock results) [] stmts
-  return $ (new_block:rest)
-breakBlock' current_id arguments accum_stmts (stmt:stmts) =
-  breakBlock' current_id arguments (accum_stmts ++ [stmt]) stmts
-
-breakBlock2 (BasicBlock ident stmts) entry = breakBlock2' ident entry [] [] stmts
-
-breakBlock2' current_id block_info exits accum_stmts [] =
-    panic "block doesn't end in jump, goto or return"
-breakBlock2' current_id entry exits accum_stmts [CmmJump target arguments] =
-    return [BrokenBlock current_id entry accum_stmts (TailCallExit exits target arguments)]
-breakBlock2' current_id entry exits accum_stmts [CmmReturn arguments] =
-    return [BrokenBlock current_id entry accum_stmts (ReturnExit exits arguments)]
-breakBlock2' current_id entry exits accum_stmts [CmmBranch target] =
-    return [BrokenBlock current_id entry accum_stmts (ControlExit exits target)]
-breakBlock2' _ _ _ _ (CmmJump _ _:_) = panic "jump in middle of block"
-breakBlock2' _ _ _ _ (CmmReturn _:_) = panic "return in middle of block"
-breakBlock2' _ _ _ _ (CmmBranch _:_) = panic "branch in middle of block"
-breakBlock2' _ _ _ _ (CmmSwitch _ _:_) = panic "switch in block not implemented"
-breakBlock2' current_id entry exits accum_stmts (CmmCall target results arguments saves:stmts) = do
-  new_id <- newLabelCPS
-  rest <- breakBlock2' new_id (ContinuationEntry results) [] [] stmts
-  return $ BrokenBlock current_id entry accum_stmts (CallExit exits new_id target results arguments saves) : rest
-breakBlock2' current_id entry exits accum_stmts (s@(CmmCondBranch test target):stmts) =
-    breakBlock2' current_id entry (target:exits) (accum_stmts++[s]) stmts
-breakBlock2' current_id entry exits accum_stmts (s:stmts) =
-    breakBlock2' current_id entry exits (accum_stmts++[s]) stmts
-
-brokenBlockTargets (BrokenBlock _ _ _ (TailCallExit exits _ _)) = exits
-brokenBlockTargets (BrokenBlock _ _ _ (ReturnExit exits _)) = exits
-brokenBlockTargets (BrokenBlock _ _ _ (ControlExit exits target)) = target:exits
-brokenBlockTargets (BrokenBlock _ _ _ (CallExit exits next _ _ _ _)) = next:exits
-
-brokenBlockId (BrokenBlock ident _ _ _) = ident
-
-cmmBrokenBlockSources ::
-    [BrokenBlock] -> UniqFM {-BlockId-} (UniqSet BlockId)
-cmmBrokenBlockSources blocks = foldr aux emptyUFM blocks where
-    aux block sourcesUFM  =
-        foldr add_source_edges sourcesUFM targets where
-            add_source_edges t ufm =
-                addToUFM_Acc (flip addOneToUniqSet) unitUniqSet ufm t ident
-            targets = brokenBlockTargets block
-            ident = brokenBlockId block
-
-cmmBrokenBlockNames :: [BrokenBlock] -> UniqFM {-BlockId-} BrokenBlock
-cmmBrokenBlockNames blocks = listToUFM $ map block_name blocks where
-    block_name b = (brokenBlockId b, b)
-
-cmmBrokenBlockDependants :: UniqFM {-BlockId-} (UniqSet BlockId) -> BlockId -> [BlockId]
-cmmBrokenBlockDependants sources ident =
-    uniqSetToList $ lookupWithDefaultUFM sources emptyUniqSet ident
-
-cmmBrokenBlockLive :: UniqFM {-BlockId-} CmmLive -> BrokenBlock -> CmmLive
-cmmBrokenBlockLive other_live (BrokenBlock _ _ stmts exit) =
-    foldr ((.) . (cmmStmtLive other_live)) id stmts live_at_end
-    where
-      live_at_end =
-          case exit of
-            ControlExit _ _ -> emptyUniqSet
-            ReturnExit _ actuals -> foldr ((.) . cmmExprLive) id (map fst actuals) emptyUniqSet
-            TailCallExit _ target actuals -> 
-                cmmExprLive target $ foldr ((.) . cmmExprLive) id (map fst actuals) $ emptyUniqSet
-            CallExit _ _ target _ actuals live ->
-                target_liveness $
-                foldr ((.) . cmmExprLive) id (map fst actuals) $
-                emptyUniqSet
-                where
-                  only_local_regs [] = []
-                  only_local_regs ((CmmGlobal _,_):args) = only_local_regs args
-                  only_local_regs ((CmmLocal r,_):args) = r:only_local_regs args
-                  target_liveness =
-                    case target of
-                      (CmmForeignCall target _) -> cmmExprLive target
-                      (CmmPrim _) -> id
-
-
-cmmBrokenBlockUpdate ::
-    UniqFM {-BlockId-} BrokenBlock
-    -> BlockId
-    -> Maybe BlockId
-    -> UniqFM {-BlockId-} CmmLive
-    -> Maybe (UniqFM {-BlockId-} CmmLive)
-cmmBrokenBlockUpdate blocks node _ state =
-    let old_live = lookupWithDefaultUFM state (panic "unknown block id during liveness analysis") node
-        block = lookupWithDefaultUFM blocks (panic "unknown block id during liveness analysis") node
-        new_live = cmmBrokenBlockLive state block
-    in if (sizeUniqSet old_live) == (sizeUniqSet new_live)
-       then Nothing
-       else Just $ addToUFM state node new_live
-
-
-cmmLiveness2 :: [BrokenBlock] -> UniqFM {-BlockId-} CmmLive
-cmmLiveness2 blocks =
-    fixedpoint (cmmBrokenBlockDependants sources) (cmmBrokenBlockUpdate blocks')
-               (map brokenBlockId blocks) (listToUFM [(brokenBlockId b, emptyUniqSet) | b <- blocks]) where
-                   sources = cmmBrokenBlockSources blocks
-                   blocks' = cmmBrokenBlockNames blocks
-
---------------------------------------------------------------------------------
-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.
-  -- continuationC <- return abstractC
-  -- TODO: find out if it is valid to create a new unique source like this
-  uniqSupply <- mkSplitUniqSupply 'p'
-  let (_, continuationC) = runCPS (mapM (mapMCmmTop cpsProc) abstractC) (CPSState uniqSupply)
-
-  dumpIfSet_dyn dflags Opt_D_dump_cps_cmm "CPS Cmm" (pprCmms continuationC)
-  -- TODO: add option to dump Cmm to file
-  return continuationC