Multiple improvements to CPS algorithm.
authorMichael D. Adams <t-madams@microsoft.com>
Mon, 2 Jul 2007 08:42:21 +0000 (08:42 +0000)
committerMichael D. Adams <t-madams@microsoft.com>
Mon, 2 Jul 2007 08:42:21 +0000 (08:42 +0000)
These include:
 - Stack size detection now includes function arguments.
 - Stack size detection now avoids stack checks just because of
   the GC block.
 - A CmmCall followed by a CmmBranch will no longer generate an extra
   continuation consisting just of the brach.
 - Multiple CmmCall/CmmBranch pairs that all go to the same place
   will try to use the same continuation.  If they can't (because
   the return value signature is different), adaptor block are built.
 - Function entry statements are now in a separate block.
   (Fixed bug with branches to the entry block having unintended effects.)
 - Other changes that I can't recall right now.

compiler/cmm/Cmm.hs
compiler/cmm/CmmBrokenBlock.hs
compiler/cmm/CmmCPS.hs
compiler/cmm/CmmProcPoint.hs
compiler/codeGen/ClosureInfo.lhs

index 530fab5..1f7161b 100644 (file)
@@ -10,7 +10,7 @@ module Cmm (
        GenCmm(..), Cmm, RawCmm,
        GenCmmTop(..), CmmTop, RawCmmTop,
        CmmInfo(..), ClosureTypeInfo(..), ProfilingInfo(..), ClosureTypeTag,
-       GenBasicBlock(..), CmmBasicBlock, blockId, blockStmts,
+       GenBasicBlock(..), CmmBasicBlock, blockId, blockStmts, mapBlockStmts,
        CmmStmt(..), CmmActuals, CmmFormal, CmmFormals, CmmHintFormals,
         CmmSafety(..),
        CmmCallTarget(..),
@@ -104,6 +104,8 @@ blockId (BasicBlock blk_id _ ) = blk_id
 blockStmts :: GenBasicBlock i -> [i]
 blockStmts (BasicBlock _ stmts) = stmts
 
+mapBlockStmts f (BasicBlock id bs) = BasicBlock id (map f bs)
+
 -----------------------------------------------------------------------------
 --     Info Tables
 -----------------------------------------------------------------------------
index af47e8d..14259c6 100644 (file)
@@ -5,17 +5,24 @@ module CmmBrokenBlock (
   breakBlock,
   cmmBlockFromBrokenBlock,
   blocksToBlockEnv,
+  adaptBlockToFormat,
+  selectContinuations,
+  ContFormat,
+  makeContinuationEntries,
   ) where
 
 #include "HsVersions.h"
 
 import Cmm
 import CLabel
+import MachOp (MachHint(..))
 
 import ClosureInfo
 
 import Maybes
+import List
 import Panic
+import UniqSupply
 import Unique
 import UniqFM
 
@@ -59,6 +66,7 @@ data BlockEntryInfo
   | ContinuationEntry          -- ^ Return point of a function call
       CmmFormals                -- ^ return values (argument to continuation)
       C_SRT                     -- ^ SRT for the continuation's info table
+      Bool                      -- ^ True <=> GC block so ignore stack size
 
   | ControlEntry               -- ^ Any other kind of block.
                                 -- Only entered due to control flow.
@@ -67,6 +75,11 @@ data BlockEntryInfo
   -- no return values, but some live might end up as
   -- params or possibly in the frame
 
+data ContFormat = ContFormat
+      CmmHintFormals            -- ^ return values (argument to continuation)
+      C_SRT                     -- ^ SRT for the continuation's info table
+      Bool                      -- ^ True <=> GC block so ignore stack size
+  deriving (Eq)
 
 -- | Final statement in a 'BlokenBlock'.
 -- Constructors and arguments match those in 'Cmm',
@@ -90,6 +103,8 @@ data FinalStmt
       CmmHintFormals                -- ^ Results from call
                                 -- (redundant with ContinuationEntry)
       CmmActuals                -- ^ Arguments to call
+      C_SRT                     -- ^ SRT for the continuation's info table
+      Bool                      -- ^ True <=> GC block so ignore stack size
 
   | FinalSwitch                 -- ^ Same as a 'CmmSwitch'
       CmmExpr                   -- ^ Scrutinee (zero based)
@@ -99,17 +114,74 @@ data FinalStmt
 -- Operations for broken blocks
 -----------------------------------------------------------------------------
 
+-- Naively breaking at *every* CmmCall leads to sub-optimal code.
+-- In particular, a CmmCall followed by a CmmBranch would result
+-- in a continuation that has the single CmmBranch statement in it.
+-- It would be better have the CmmCall directly return to the block
+-- that the branch jumps to.
+--
+-- This requires the target of the branch to look like the parameter
+-- format that the CmmCall is expecting.  If other CmmCall/CmmBranch
+-- sequences go to the same place they might not be expecting the
+-- same format.  So this transformation uses the following solution.
+-- First the blocks are broken up but none of the blocks are marked
+-- as continuations yet.  This is the 'breakBlock' function.
+-- Second, the blocks "vote" on what other blocks need to be continuations
+-- and how they should be layed out.  Plurality wins, but other selection
+-- methods could be selected at a later time.
+-- This is the 'selectContinuations' function.
+-- Finally, the blocks are upgraded to 'ContEntry' continuations
+-- based on the results with the 'makeContinuationEntries' function,
+-- and the blocks that didn't get the format they wanted for their
+-- targets get a small adaptor block created for them by
+-- the 'adaptBlockToFormat' function.
+-- could be 
+
+breakProc ::
+    [BlockId]                   -- ^ Any GC blocks that should be special
+    -> [[Unique]]               -- ^ An infinite list of uniques
+                                -- to create names of the new blocks with
+    -> CmmInfo                  -- ^ Info table for the procedure
+    -> CLabel                   -- ^ Name of the procedure
+    -> CmmFormals               -- ^ Parameters of the procedure
+    -> [CmmBasicBlock]          -- ^ Blocks of the procecure
+                                -- (First block is the entry block)
+    -> [BrokenBlock]
+
+breakProc gc_block_idents uniques info ident params blocks =
+    let
+        (adaptor_uniques : block_uniques) = uniques
+
+        broken_blocks :: ([(BlockId, ContFormat)], [BrokenBlock])
+        broken_blocks =
+            let new_blocks =
+                    zipWith3 (breakBlock gc_block_idents)
+                             block_uniques
+                             blocks
+                             (FunctionEntry info ident params :
+                              repeat ControlEntry)
+            in (concatMap fst new_blocks, concatMap snd new_blocks)
+
+        selected = selectContinuations (fst broken_blocks)
+
+    in map (makeContinuationEntries selected) $
+       concat $
+       zipWith (adaptBlockToFormat selected)
+               adaptor_uniques
+               (snd broken_blocks)
+
 -----------------------------------------------------------------------------
 -- | Takes a 'CmmBasicBlock' and breaks it up into a list of 'BrokenBlock'
 -- by splitting on each 'CmmCall' in the 'CmmBasicBlock'.
 
 breakBlock ::
-    [Unique]                    -- ^ An infinite list of uniques
+    [BlockId]                   -- ^ Any GC blocks that should be special
+    -> [Unique]                 -- ^ An infinite list of uniques
                                 -- to create names of the new blocks with
     -> CmmBasicBlock            -- ^ Input block to break apart
     -> BlockEntryInfo           -- ^ Info for the first created 'BrokenBlock'
-    -> [BrokenBlock]
-breakBlock uniques (BasicBlock ident stmts) entry =
+    -> ([(BlockId, ContFormat)], [BrokenBlock])
+breakBlock gc_block_idents uniques (BasicBlock ident stmts) entry =
     breakBlock' uniques ident entry [] [] stmts
     where
       breakBlock' uniques current_id entry exits accum_stmts stmts =
@@ -118,21 +190,25 @@ breakBlock uniques (BasicBlock ident stmts) entry =
 
             -- Last statement.  Make the 'BrokenBlock'
             [CmmJump target arguments] ->
-                [BrokenBlock current_id entry accum_stmts
-                             exits
-                             (FinalJump target arguments)]
+                ([],
+                 [BrokenBlock current_id entry accum_stmts
+                              exits
+                              (FinalJump target arguments)])
             [CmmReturn arguments] ->
-                [BrokenBlock current_id entry accum_stmts
+                ([],
+                 [BrokenBlock current_id entry accum_stmts
                              exits
-                             (FinalReturn arguments)]
+                             (FinalReturn arguments)])
             [CmmBranch target] ->
-                [BrokenBlock current_id entry accum_stmts
+                ([],
+                 [BrokenBlock current_id entry accum_stmts
                              (target:exits)
-                             (FinalBranch target)]
+                             (FinalBranch target)])
             [CmmSwitch expr targets] ->
-                [BrokenBlock current_id entry accum_stmts
+                ([],
+                 [BrokenBlock current_id entry accum_stmts
                              (mapMaybe id targets ++ exits)
-                             (FinalSwitch expr targets)]
+                             (FinalSwitch expr targets)])
 
             -- These shouldn't happen in the middle of a block.
             -- They would cause dead code.
@@ -143,24 +219,28 @@ breakBlock uniques (BasicBlock ident stmts) entry =
 
             -- Detect this special case to remain an inverse of
             -- 'cmmBlockFromBrokenBlock'
-            {- TODO: Interferes with proc point detection
-            [CmmCall target results arguments,
-             CmmBranch next_id] -> [block]
-              where
-                block = do_call current_id entry accum_stmts exits next_id
-                                target results arguments
-             -}
+            [CmmCall target results arguments (CmmSafe srt),
+             CmmBranch next_id] ->
+                ([cont_info], [block])
+                where
+                  cont_info = (next_id,
+                               ContFormat results srt
+                                              (ident `elem` gc_block_idents))
+                  block = do_call current_id entry accum_stmts exits next_id
+                                target results arguments srt
 
             -- Break the block on safe calls (the main job of this function)
             (CmmCall target results arguments (CmmSafe srt):stmts) ->
-                block : rest
+                (cont_info : cont_infos, block : blocks)
                 where
                   next_id = BlockId $ head uniques
                   block = do_call current_id entry accum_stmts exits next_id
-                                  target results arguments
-                  rest = breakBlock' (tail uniques) next_id
-                                     (ContinuationEntry (map fst results) srt)
-                                     [] [] stmts
+                                  target results arguments srt
+                  cont_info = (next_id,
+                               ContFormat results srt
+                                              (ident `elem` gc_block_idents))
+                  (cont_infos, blocks) = breakBlock' (tail uniques) next_id
+                                         ControlEntry [] [] stmts
 
             -- Default case.  Just keep accumulating statements
             -- and branch targets.
@@ -171,14 +251,86 @@ breakBlock uniques (BasicBlock ident stmts) entry =
                             stmts
 
       do_call current_id entry accum_stmts exits next_id
-              target results arguments =
+              target results arguments srt =
           BrokenBlock current_id entry accum_stmts (next_id:exits)
-                      (FinalCall next_id target results arguments)
+                      (FinalCall next_id target results arguments srt
+                                     (current_id `elem` gc_block_idents))
 
       cond_branch_target (CmmCondBranch _ target) = [target]
       cond_branch_target _ = []
 
 -----------------------------------------------------------------------------
+
+selectContinuations :: [(BlockId, ContFormat)] -> [(BlockId, ContFormat)]
+selectContinuations needed_continuations = formats
+    where
+      formats = map select_format format_groups
+      format_groups = groupBy by_target needed_continuations
+      by_target x y = fst x == fst y
+
+      select_format formats = winner
+          where
+            winner = head $ head $ sortBy more_votes format_votes
+            format_votes = groupBy by_format formats
+            by_format x y = snd x == snd y
+            more_votes x y = compare (length y) (length x)
+              -- sort so the most votes goes *first*
+              -- (thus the order of x and y is reversed)
+
+makeContinuationEntries formats
+                        block@(BrokenBlock ident entry stmts targets exit) =
+    case lookup ident formats of
+      Nothing -> block
+      Just (ContFormat formals srt is_gc) ->
+          BrokenBlock ident (ContinuationEntry (map fst formals) srt is_gc)
+                      stmts targets exit
+
+adaptBlockToFormat :: [(BlockId, ContFormat)]
+                   -> Unique
+                   -> BrokenBlock
+                   -> [BrokenBlock]
+adaptBlockToFormat formats unique
+                   block@(BrokenBlock ident entry stmts targets
+                                      exit@(FinalCall next target formals
+                                                      actuals srt is_gc)) =
+    if format_formals == formals &&
+       format_srt == srt &&
+       format_is_gc == is_gc
+    then [block] -- Woohoo! This block got the continuation format it wanted
+    else [adaptor_block, revised_block]
+           -- This block didn't get the format it wanted for the
+           -- continuation, so we have to build an adaptor.
+    where
+      (ContFormat format_formals format_srt format_is_gc) =
+          maybe unknown_block id $ lookup next formats
+      unknown_block = panic "unknown block in adaptBlockToFormat"
+
+      revised_block = BrokenBlock ident entry stmts revised_targets revised_exit
+      revised_targets = adaptor_ident : delete next targets
+      revised_exit = FinalCall
+                       adaptor_ident -- ^ The only part that changed
+                       target formals actuals srt is_gc
+
+      adaptor_block = mk_adaptor_block adaptor_ident
+                  (ContinuationEntry (map fst formals) srt is_gc)
+                  next format_formals
+      adaptor_ident = BlockId unique
+
+      mk_adaptor_block :: BlockId -> BlockEntryInfo -> BlockId -> CmmHintFormals -> BrokenBlock
+      mk_adaptor_block ident entry next formals =
+          BrokenBlock ident entry [] [next] exit
+              where
+                exit = FinalJump
+                         (CmmLit (CmmLabel (mkReturnPtLabel (getUnique next))))
+                         (map formal_to_actual format_formals)
+
+                formal_to_actual (reg, hint) = ((CmmReg (CmmLocal reg)), hint)
+                -- TODO: Check if NoHint is right.  We're
+                -- jumping to a C-- function not a foreign one
+                -- so it might always be right.
+adaptBlockToFormat _ _ block = [block]
+
+-----------------------------------------------------------------------------
 -- | Convert from a BrokenBlock back to an equivalent CmmBasicBlock
 -- Needed by liveness analysis
 cmmBlockFromBrokenBlock :: BrokenBlock -> CmmBasicBlock
@@ -191,8 +343,8 @@ cmmBlockFromBrokenBlock (BrokenBlock ident _ stmts _ exit) =
             FinalReturn arguments -> [CmmReturn arguments]
             FinalJump target arguments -> [CmmJump target arguments]
             FinalSwitch expr targets -> [CmmSwitch expr targets]
-            FinalCall branch_target call_target results arguments ->
-                [CmmCall call_target results arguments (panic "needed SRT from cmmBlockFromBrokenBlock"),
+            FinalCall branch_target call_target results arguments srt _ ->
+                [CmmCall call_target results arguments (CmmSafe srt),
                  CmmBranch branch_target]
 
 -----------------------------------------------------------------------------
index afb55d5..e6d70d4 100644 (file)
@@ -79,23 +79,31 @@ make_gc_block block_id fun_label formals safety = BasicBlock block_id stmts
       actuals = map (\x -> (CmmReg (CmmLocal x), NoHint)) formals
       fun_expr = CmmLit (CmmLabel fun_label)
 
-force_gc_block old_info block_id fun_label formals =
+make_gc_check stack_use gc_block =
+    [CmmCondBranch
+     (CmmMachOp (MO_U_Lt $ cmmRegRep spReg)
+                    [CmmReg stack_use, CmmReg spLimReg])
+    gc_block]
+
+force_gc_block old_info stack_use block_id fun_label formals =
     case old_info of
-      CmmNonInfo (Just _) -> (old_info, [])
-      CmmInfo _ (Just _) _ _ -> (old_info, [])
+      CmmNonInfo (Just existing) -> (old_info, [], make_gc_check stack_use existing)
+      CmmInfo _ (Just existing) _ _ -> (old_info, [], make_gc_check stack_use existing)
       CmmNonInfo Nothing
           -> (CmmNonInfo (Just block_id),
-              [make_gc_block block_id fun_label formals (CmmSafe NoC_SRT)])
+              [make_gc_block block_id fun_label formals (CmmSafe NoC_SRT)],
+              make_gc_check stack_use block_id)
       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'    
+          -> (CmmInfo prof (Just block_id) type_tag type_info,
+              [make_gc_block block_id fun_label formals (CmmSafe srt)],
+              make_gc_check stack_use block_id)
+             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)
@@ -108,26 +116,46 @@ cpsProc :: UniqSupply
 cpsProc uniqSupply (CmmData sec dat) = [CmmData sec dat]
 cpsProc uniqSupply (CmmProc info ident params blocks) = info_procs
     where
+      (uniqSupply1, uniqSupply2) = splitUniqSupply uniqSupply
       uniques :: [[Unique]]
-      uniques = map uniqsFromSupply $ listSplitUniqSupply uniqSupply
-      (gc_unique:info_uniques):block_uniques = uniques
+      uniques = map uniqsFromSupply $ listSplitUniqSupply uniqSupply1
+      (gc_unique:stack_use_unique:info_uniques):adaptor_uniques:block_uniques = uniques
+      proc_uniques = map uniqsFromSupply $ listSplitUniqSupply uniqSupply2
+
+      stack_use = CmmLocal (LocalReg stack_use_unique (cmmRegRep spReg) KindPtr)
 
-      -- Ensure that 
-      forced_gc :: (CmmInfo, [CmmBasicBlock])
-      forced_gc = force_gc_block info (BlockId gc_unique) ident params
+      -- TODO: doc
+      forced_gc :: (CmmInfo, [CmmBasicBlock], [CmmStmt])
+      forced_gc = force_gc_block info stack_use (BlockId gc_unique) ident params
+      (forced_info, gc_blocks, check_stmts) = forced_gc
+
+      forced_blocks =
+          case blocks of
+            (BasicBlock id stmts) : bs ->
+                (BasicBlock id (check_stmts ++ stmts)) : (bs ++ gc_blocks)
+            [] -> [] -- If there is no code then we don't need a stack check
 
-      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 :: ([(BlockId, ContFormat)], [BrokenBlock])
       broken_blocks =
-          concat $ zipWith3 breakBlock block_uniques forced_blocks
-                     (FunctionEntry forced_info ident params:repeat ControlEntry)
+          (\x -> (concatMap fst x, concatMap snd x)) $
+          zipWith3 (breakBlock [forced_gc_id])
+                     block_uniques
+                     forced_blocks
+                     (FunctionEntry forced_info ident params :
+                      repeat ControlEntry)
+
+      f' = selectContinuations (fst broken_blocks)
+      broken_blocks' = map (makeContinuationEntries f') $
+                       concat $
+                       zipWith (adaptBlockToFormat f')
+                               adaptor_uniques
+                               (snd broken_blocks)
 
       -- Calculate live variables for each broken block.
       --
@@ -135,22 +163,22 @@ cpsProc uniqSupply (CmmProc info ident params blocks) = info_procs
       -- 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
+      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
+      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
+      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)
+                        (gatherBlocksIntoContinuation live proc_points block_env)
                         (uniqSetToList proc_points)
-                        (Just forced_gc_id : repeat Nothing)
+                        (Just forced_gc_id : repeat Nothing) {-dead-}
 
       -- Select the stack format on entry to each continuation.
       -- Return the max stack offset and an association list
@@ -158,23 +186,24 @@ cpsProc uniqSupply (CmmProc info ident params blocks) = info_procs
       -- This is an association list instead of a UniqFM because
       -- CLabel's don't have a 'Uniqueable' instance.
       formats :: [(CLabel,              -- key
-                   (Maybe CLabel,       -- label in top slot
+                   (CmmFormals,         -- arguments
+                    Maybe CLabel,       -- label in top slot
                     [Maybe LocalReg]))] -- slots
-      formats = selectStackFormat live continuations
+      formats = selectContinuationFormat 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' :: (WordOff, [(CLabel, ContinuationFormat)])
       formats' = processFormats formats continuations
 
       -- Update the info table data on the continuations with
       -- the selected stack formats.
       continuations' :: [Continuation CmmInfo]
-      continuations' = map (applyStackFormat (snd formats')) continuations
+      continuations' = map (applyContinuationFormat (snd formats')) continuations
 
       -- Do the actual CPS transform.
       cps_procs :: [CmmTop]
-      cps_procs = map (continuationToProc formats') continuations'
+      cps_procs = zipWith (continuationToProc formats' stack_use) proc_uniques continuations'
 
       -- Convert the info tables from CmmInfo to [CmmStatic]
       -- We might want to put this in another pass eventually
@@ -198,13 +227,14 @@ cpsProc uniqSupply (CmmProc info ident params blocks) = info_procs
 -- and heap memory (not sure if that's usefull at all though, but it may
 -- be worth exploring the design space).
 
-continuationLabel (Continuation _ l _ _) = l
+continuationLabel (Continuation _ l _ _ _) = l
 data Continuation info =
   Continuation
      info              -- 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)
+     Bool              -- ^ True <=> GC block so ignore stack size
      [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
@@ -214,11 +244,12 @@ data Continuation info =
                       -- to a label.  To jump to the first block in a Proc,
                       -- use the appropriate CLabel.
 
-data StackFormat
-    = 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
+data ContinuationFormat
+    = ContinuationFormat {
+        continuation_formals :: CmmFormals,
+        continuation_label :: Maybe CLabel,    -- The label occupying the top slot
+        continuation_frame_size :: WordOff,    -- Total frame size in words (not including arguments)
+        continuation_stack :: [Maybe LocalReg] -- local reg offsets from stack top
       }
 
 -- A block can be a continuation of a call
@@ -254,74 +285,90 @@ collectNonProcPointTargets proc_points blocks current_targets new_blocks =
 --  * Proc points might get some live variables passed as arguments
 
 gatherBlocksIntoContinuation ::
-    UniqSet BlockId -> BlockEnv BrokenBlock
+    BlockEntryLiveness -> UniqSet BlockId -> BlockEnv BrokenBlock
     -> BlockId -> Maybe BlockId -> Continuation (Either C_SRT CmmInfo)
-gatherBlocksIntoContinuation proc_points blocks start gc =
-  Continuation info_table clabel params body
+gatherBlocksIntoContinuation live proc_points blocks start gc =
+  Continuation info_table clabel params is_gc_cont 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_and_gc = [start] -- : maybeToList gc
+      --children = (collectNonProcPointTargets proc_points blocks (mkUniqSet start_and_gc) start_and_gc) `minusUniqSet` (mkUniqSet start_and_gc)
+      children = (collectNonProcPointTargets proc_points blocks (unitUniqSet start) [start]) `minusUniqSet` (unitUniqSet start)
       start_block = lookupWithDefaultUFM blocks (panic "TODO") start
-      gc_block = map (lookupWithDefaultUFM blocks (panic "TODO)")) (maybeToList gc)
+      unknown_block = panic "unknown block in gatherBlocksIntoContinuation"
+      --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
+      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
+                     ContinuationEntry _ srt -> Left srt
                      ControlEntry -> Right (CmmNonInfo Nothing)
 
+      is_gc_cont = case start_block_entry of
+                     FunctionEntry _ _ _ -> False
+                     ContinuationEntry _ _ gc_cont -> gc_cont
+                     ControlEntry -> False
+
       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
+                 ContinuationEntry args _ _ -> args
+                 ControlEntry ->
+                     uniqSetToList $
+                     lookupWithDefaultUFM live unknown_block start
+                     -- it's a proc-point, 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
+selectContinuationFormat :: BlockEnv CmmLive
                   -> [Continuation (Either C_SRT CmmInfo)]
-                  -> [(CLabel, (Maybe CLabel, [Maybe LocalReg]))]
-selectStackFormat live continuations =
-    map (\c -> (continuationLabel c, selectStackFormat' c)) continuations
+                  -> [(CLabel, (CmmFormals, Maybe CLabel, [Maybe LocalReg]))]
+selectContinuationFormat live continuations =
+    map (\c -> (continuationLabel c, selectContinuationFormat' c)) continuations
     where
-      selectStackFormat' (Continuation
+      selectContinuationFormat' (Continuation
                           (Right (CmmInfo _ _ _ (ContInfo format srt)))
-                          label _ _) = (Just label, format)
-      selectStackFormat' (Continuation (Right _) _ _ _) = (Nothing, [])
-      selectStackFormat' (Continuation (Left srt) label _ blocks) =
+                          label formals _ _) =
+          (formals, Just label, format)
+      selectContinuationFormat' (Continuation (Right _) _ formals _ _) =
+          (formals, Nothing, [])
+      -- CPS generated continuations
+      selectContinuationFormat' (Continuation (Left srt) label formals _ 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,
+          in (formals,
+              Just label,
               map Just $ uniqSetToList $
               lookupWithDefaultUFM live unknown_block ident)
 
-      unknown_block = panic "unknown BlockId in selectStackFormat"
+      unknown_block = panic "unknown BlockId in selectContinuationFormat"
 
-processFormats :: [(CLabel, (Maybe CLabel, [Maybe LocalReg]))]
+processFormats :: [(CLabel, (CmmFormals, Maybe CLabel, [Maybe LocalReg]))]
                -> [Continuation (Either C_SRT CmmInfo)]
-               -> (WordOff, [(CLabel, StackFormat)])
+               -> (WordOff, [(CLabel, ContinuationFormat)])
 processFormats formats continuations = (max_size, formats')
     where
       max_size = maximum $
                  0 : map (continuationMaxStack formats') continuations
       formats' = map make_format formats
-      make_format (label, format) =
+      make_format (label, (formals, top, stack)) =
           (label,
-           StackFormat {
-             stack_label = fst format,
-             stack_frame_size = stack_size (snd format) +
-                                if isJust (fst format)
+           ContinuationFormat {
+             continuation_formals = formals,
+             continuation_label = top,
+             continuation_frame_size = stack_size stack +
+                                if isJust top
                                 then label_size
                                 else 0,
-             stack_live = snd format })
+             continuation_stack = stack })
 
       -- TODO: get rid of "+ 1" etc.
       label_size = 1 :: WordOff
@@ -333,11 +380,12 @@ processFormats formats continuations = (max_size, formats')
             width = machRepByteWidth (localRegRep reg) `quot` wORD_SIZE
             -- TODO: it would be better if we had a machRepWordWidth
 
-continuationMaxStack :: [(CLabel, StackFormat)]
+continuationMaxStack :: [(CLabel, ContinuationFormat)]
                      -> Continuation a
                      -> WordOff
-continuationMaxStack formats (Continuation _ label _ blocks) =
-    max_arg_size + stack_frame_size stack_format
+continuationMaxStack _ (Continuation _ _ _ True _) = 0
+continuationMaxStack formats (Continuation _ label _ False blocks) =
+    max_arg_size + continuation_frame_size stack_format
     where
       stack_format = maybe unknown_format id $ lookup label formats
       unknown_format = panic "Unknown format in continuationMaxStack"
@@ -352,11 +400,12 @@ continuationMaxStack formats (Continuation _ label _ blocks) =
           argumentsSize (cmmExprRep . fst) args
       final_arg_size (FinalJump _ args) =
           argumentsSize (cmmExprRep . fst) args
-      final_arg_size (FinalCall next _ _ args) =
+      final_arg_size (FinalCall next _ _ args _ True) = 0
+      final_arg_size (FinalCall next _ _ args _ False) =
           -- We have to account for the stack used when we build a frame
           -- for the *next* continuation from *this* continuation
           argumentsSize (cmmExprRep . fst) args +
-          stack_frame_size next_format
+          continuation_frame_size next_format
           where 
             next_format = maybe unknown_format id $ lookup next' formats
             next' = mkReturnPtLabel $ getUnique next
@@ -372,28 +421,30 @@ continuationMaxStack formats (Continuation _ label _ blocks) =
       stmt_arg_size _ = 0
 
 -----------------------------------------------------------------------------
-applyStackFormat :: [(CLabel, StackFormat)]
+applyContinuationFormat :: [(CLabel, ContinuationFormat)]
                  -> Continuation (Either C_SRT CmmInfo)
                  -> Continuation CmmInfo
 
 -- User written continuations
-applyStackFormat formats (Continuation
+applyContinuationFormat formats (Continuation
                           (Right (CmmInfo prof gc tag (ContInfo _ srt)))
-                          label formals blocks) =
+                          label formals is_gc blocks) =
     Continuation (CmmInfo prof gc tag (ContInfo format srt))
-                 label formals blocks
+                 label formals is_gc blocks
     where
-      format = stack_live $ maybe unknown_block id $ lookup label formats
-      unknown_block = panic "unknown BlockId in applyStackFormat"
+      format = continuation_stack $ maybe unknown_block id $ lookup label formats
+      unknown_block = panic "unknown BlockId in applyContinuationFormat"
 
 -- User written non-continuation code
-applyStackFormat formats (Continuation (Right info) label formals blocks) =
-    Continuation info label formals blocks
+applyContinuationFormat formats (Continuation
+                          (Right info) label formals is_gc blocks) =
+    Continuation info label formals is_gc 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
+applyContinuationFormat formats (Continuation
+                          (Left srt) label formals is_gc blocks) =
+    Continuation (CmmInfo prof gc tag (ContInfo (continuation_stack $ format) srt))
+                 label formals is_gc blocks
     where
       gc = Nothing -- Generated continuations never need a stack check
       -- TODO prof: this is the same as the current implementation
@@ -401,40 +452,69 @@ applyStackFormat formats (Continuation (Left srt) label formals blocks) =
       prof = ProfilingInfo zeroCLit zeroCLit
       tag = rET_SMALL -- cmmToRawCmm may convert it to rET_BIG
       format = maybe unknown_block id $ lookup label formats
-      unknown_block = panic "unknown BlockId in applyStackFormat"
+      unknown_block = panic "unknown BlockId in applyContinuationFormat"
 
 -----------------------------------------------------------------------------
-continuationToProc :: (WordOff, [(CLabel, StackFormat)])
+continuationToProc :: (WordOff, [(CLabel, ContinuationFormat)])
+                   -> CmmReg
+                   -> [Unique]
                    -> Continuation CmmInfo
                    -> CmmTop
-continuationToProc (max_stack, formats)
-                   (Continuation info label formals blocks) =
-    CmmProc info label formals (map continuationToProc' blocks)
+continuationToProc (max_stack, formats) stack_use uniques
+                   (Continuation info label formals blocks) =
+    CmmProc info label formals (concat $ zipWith3 continuationToProc' uniques blocks (True : repeat False))
     where
       curr_format = maybe unknown_block id $ lookup label formats
       unknown_block = panic "unknown BlockId in continuationToProc"
-      curr_stack = stack_frame_size curr_format
-
-      continuationToProc' :: BrokenBlock -> CmmBasicBlock
-      continuationToProc' (BrokenBlock ident entry stmts _ exit) =
-          BasicBlock ident (prefix++stmts++postfix)
+      curr_stack = continuation_frame_size curr_format
+      arg_stack = argumentsSize localRegRep formals
+
+      param_stmts :: [CmmStmt]
+      param_stmts = function_entry curr_format
+
+      gc_stmts :: [CmmStmt]
+      gc_stmts =
+          case info of
+            CmmInfo _ (Just gc_block) _ _ ->
+                gc_stack_check' stack_use arg_stack (max_stack - curr_stack)
+            CmmInfo _ Nothing _ _ ->
+                panic "continuationToProc: missing GC block"
+            CmmNonInfo (Just gc_block) ->
+                gc_stack_check' stack_use arg_stack (max_stack - curr_stack)
+            CmmNonInfo Nothing ->
+                panic "continuationToProc: missing non-info GC block"
+
+      continuationToProc' :: Unique -> BrokenBlock -> Bool -> [CmmBasicBlock]
+      continuationToProc' unique (BrokenBlock ident entry stmts _ exit) is_entry =
+          case gc_prefix ++ param_prefix of
+            [] -> [main_block]
+            stmts -> [BasicBlock prefix_id (gc_prefix ++ param_prefix ++ [CmmBranch ident]),
+                      main_block]
           where
-            prefix = case entry of
+            main_block = BasicBlock ident (stmts ++ postfix)
+            prefix_id = BlockId unique
+            gc_prefix = case entry of
+                       FunctionEntry _ _ _ -> gc_stmts
                        ControlEntry -> []
-                       FunctionEntry (CmmInfo _ (Just gc_block) _ _) _ formals ->
-                           gc_stack_check gc_block (max_stack - curr_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 - curr_stack) ++
-                           function_entry formals curr_format
-                       FunctionEntry (CmmNonInfo Nothing) _ formals ->
-                           panic "continuationToProc: missing non-info GC block"
-                       ContinuationEntry formals _ ->
-                           function_entry formals curr_format
+                       ContinuationEntry _ _ _ -> []
+            param_prefix = if is_entry
+                           then param_stmts
+                           else []
             postfix = case exit of
-                        FinalBranch next -> [CmmBranch next]
+                        FinalBranch next ->
+                            if (mkReturnPtLabel $ getUnique next) == label
+                            then [CmmBranch next]
+                            else case lookup (mkReturnPtLabel $ getUnique next) formats of
+                              Nothing -> [CmmBranch next]
+                              Just cont_format ->
+                                pack_continuation False curr_format cont_format ++
+                                tail_call (curr_stack - cont_stack)
+                                          (CmmLit $ CmmLabel $ mkReturnPtLabel $ getUnique next)
+                                          arguments
+                                where
+                                  cont_stack = continuation_frame_size cont_format
+                                  arguments = map formal_to_actual (continuation_formals cont_format)
+                                  formal_to_actual reg = (CmmReg (CmmLocal reg), NoHint)
                         FinalSwitch expr targets -> [CmmSwitch expr targets]
                         FinalReturn arguments ->
                             tail_call curr_stack
@@ -443,15 +523,15 @@ continuationToProc (max_stack, formats)
                         FinalJump target arguments ->
                             tail_call curr_stack target arguments
                         FinalCall next (CmmForeignCall target CmmCallConv)
-                            results arguments ->
-                                pack_continuation curr_format cont_format ++
+                            results arguments _ _ ->
+                                pack_continuation True curr_format cont_format ++
                                 tail_call (curr_stack - cont_stack)
                                               target arguments
                             where
                               cont_format = maybe unknown_block id $
                                             lookup (mkReturnPtLabel $ getUnique next) formats
-                              cont_stack = stack_frame_size cont_format
-                        FinalCall next _ results arguments -> panic "unimplemented CmmCall"
+                              cont_stack = continuation_frame_size cont_format
+                        FinalCall next _ results arguments _ _ -> panic "unimplemented CmmCall"
 
 -----------------------------------------------------------------------------
 -- Functions that generate CmmStmt sequences
@@ -474,6 +554,12 @@ tail_call spRel target arguments
 
     argument_formats = assignArguments (cmmExprRep . fst) arguments
 
+gc_stack_check' stack_use arg_stack max_frame_size =
+    if max_frame_size > arg_stack
+    then [CmmAssign stack_use (CmmRegOff spReg (-max_frame_size*wORD_SIZE))]
+    else [CmmAssign stack_use (CmmReg spLimReg)]
+         -- Trick the optimizer into eliminating the branch for us
+  
 gc_stack_check :: BlockId -> WordOff -> [CmmStmt]
 gc_stack_check gc_block max_frame_size
   = check_stack_limit where
@@ -484,11 +570,13 @@ gc_stack_check gc_block max_frame_size
                      CmmReg spLimReg])
      gc_block]
 
+
 -- 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 _)
-                       (StackFormat cont_id cont_frame_size live_regs)
+pack_continuation :: Bool -> ContinuationFormat -> ContinuationFormat -> [CmmStmt]
+pack_continuation allow_header_set
+                      (ContinuationFormat _ curr_id curr_frame_size _)
+                      (ContinuationFormat _ cont_id cont_frame_size live_regs)
   = store_live_values ++ set_stack_header where
     -- TODO: only save variables when actually needed
     -- (may be handled by latter pass)
@@ -496,7 +584,7 @@ pack_continuation (StackFormat curr_id curr_frame_size _)
         [stack_put spRel (CmmReg (CmmLocal reg)) offset
          | (reg, offset) <- cont_offsets]
     set_stack_header =
-        if needs_header_set
+        if needs_header_set && allow_header_set
         then [stack_put spRel continuation_function 0]
         else []
 
@@ -523,8 +611,8 @@ pack_continuation (StackFormat curr_id curr_frame_size _)
 -- that could branch to eachother (i.e. control blocks)
 -- have the same stack format (this causes a problem
 -- only for proc-point).
-function_entry :: CmmFormals -> StackFormat -> [CmmStmt]
-function_entry formals (StackFormat _ _ live_regs)
+function_entry :: ContinuationFormat -> [CmmStmt]
+function_entry (ContinuationFormat formals _ _ live_regs)
   = load_live_values ++ load_args where
     -- TODO: only save variables when actually needed
     -- (may be handled by latter pass)
index 5a159a6..36c02ff 100644 (file)
@@ -46,7 +46,7 @@ calculateProcPoints blocks =
       always_proc_point BrokenBlock {
                               brokenBlockEntry = FunctionEntry _ _ _ } = True
       always_proc_point BrokenBlock {
-                              brokenBlockEntry = ContinuationEntry _ _ } = True
+                              brokenBlockEntry = ContinuationEntry _ _ } = True
       always_proc_point _ = False
 
 calculateProcPoints' :: UniqSet BlockId -> [BrokenBlock] -> UniqSet BlockId
index db46368..d0d2ed9 100644 (file)
@@ -124,6 +124,7 @@ data ClosureInfo
 
 data C_SRT = NoC_SRT
           | C_SRT !CLabel !WordOff !StgHalfWord {-bitmap or escape-}
+           deriving (Eq)
 
 needsSRT :: C_SRT -> Bool
 needsSRT NoC_SRT       = False