Make CPS account for on-stack arguments when doing the stack check

[ghc-hetmet.git] / compiler / cmm / CmmCPS.hs

diff --git a/compiler/cmm/CmmCPS.hs b/compiler/cmm/CmmCPS.hs
index b6c57ee..e4a17a9 100644 (file)
--- a/compiler/cmm/CmmCPS.hs
+++ b/compiler/cmm/CmmCPS.hs
@@ -157,17 +157,15 @@ 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, (Maybe CLabel, [Maybe LocalReg]))]
+      formats :: [(CLabel,              -- key
+                   (Maybe CLabel,       -- label in top slot
+                    [Maybe LocalReg]))] -- slots
       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.
+      formats' = processFormats formats continuations
 
       -- Update the info table data on the continuations with
       -- the selected stack formats.
@@ -203,7 +201,7 @@ cpsProc uniqSupply (CmmProc info ident params blocks) = info_procs
 continuationLabel (Continuation _ l _ _) = l
 data Continuation info =
   Continuation
-     info --(Either C_SRT CmmInfo)   -- Left <=> Continuation created by the CPS
+     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)
@@ -308,10 +306,12 @@ selectStackFormat live continuations =
       unknown_block = panic "unknown BlockId in selectStackFormat"
 
 processFormats :: [(CLabel, (Maybe CLabel, [Maybe LocalReg]))]
+               -> [Continuation (Either C_SRT CmmInfo)]
                -> (WordOff, [(CLabel, StackFormat)])
-processFormats formats = (max_size, formats')
+processFormats formats continuations = (max_size, formats')
     where
-      max_size = foldl max 0 (map (stack_frame_size . snd) formats')
+      max_size = maximum $
+                 0 : map (continuationMaxStack formats') continuations
       formats' = map make_format formats
       make_format (label, format) =
           (label,
@@ -333,6 +333,44 @@ processFormats formats = (max_size, formats')
             width = machRepByteWidth (localRegRep reg) `quot` wORD_SIZE
             -- TODO: it would be better if we had a machRepWordWidth
 
+continuationMaxStack :: [(CLabel, StackFormat)]
+                     -> Continuation a
+                     -> WordOff
+continuationMaxStack formats (Continuation _ label _ blocks) =
+    max_arg_size + stack_frame_size stack_format
+    where
+      stack_format = maybe unknown_format id $ lookup label formats
+      unknown_format = panic "Unknown format in continuationMaxStack"
+
+      max_arg_size = maximum $ 0 : map block_max_arg_size blocks
+
+      block_max_arg_size block =
+          maximum (final_arg_size (brokenBlockExit block) :
+                   map stmt_arg_size (brokenBlockStmts block))
+
+      final_arg_size (FinalReturn args) =
+          argumentsSize (cmmExprRep . fst) args
+      final_arg_size (FinalJump _ args) =
+          argumentsSize (cmmExprRep . fst) args
+      final_arg_size (FinalCall next _ _ args) =
+          -- 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
+          where 
+            next_format = maybe unknown_format id $ lookup next' formats
+            next' = mkReturnPtLabel $ getUnique next
+
+      final_arg_size _ = 0
+
+      stmt_arg_size (CmmJump _ args) =
+          argumentsSize (cmmExprRep . fst) args
+      stmt_arg_size (CmmCall _ _ _ (CmmSafe _)) =
+          panic "Safe call in processFormats"
+      stmt_arg_size (CmmReturn _) =
+          panic "CmmReturn in processFormats"
+      stmt_arg_size _ = 0
+
 -----------------------------------------------------------------------------
 applyStackFormat :: [(CLabel, StackFormat)]
                  -> Continuation (Either C_SRT CmmInfo)
@@ -361,7 +399,7 @@ applyStackFormat formats (Continuation (Left srt) label formals blocks) =
       -- 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
+      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"