Better underflow frames

10 files changed, 94 insertions, 37 deletions

diff --git a/libraries/ghc-heap/GHC/Exts/DecodeHeap.hs b/libraries/ghc-heap/GHC/Exts/DecodeHeap.hs
index a0eca760ef..cce4c35885 100644
--- a/libraries/ghc-heap/GHC/Exts/DecodeHeap.hs
+++ b/libraries/ghc-heap/GHC/Exts/DecodeHeap.hs
@@ -234,6 +234,7 @@ getClosureDataFromHeapRepPrim getConDesc decodeCCS itbl heapRep pts = do
 #if __GLASGOW_HASKELL__ >= 811
                                 , stack_marking = FFIClosures.stack_marking fields
 #endif
+                                , stack = []
                                 })
             | otherwise
                 -> fail $ "Expected 0 ptr argument to STACK, found "
diff --git a/libraries/ghc-heap/GHC/Exts/DecodeStack.hs b/libraries/ghc-heap/GHC/Exts/DecodeStack.hs
index 0815db8764..8f714594e0 100644
--- a/libraries/ghc-heap/GHC/Exts/DecodeStack.hs
+++ b/libraries/ghc-heap/GHC/Exts/DecodeStack.hs
@@ -37,6 +37,7 @@ import GHC.Stack.CloneStack
 import Prelude
 import GHC.IO (IO (..))
 import Data.Array.Byte
+import GHC.Word
 
 {- Note [Decoding the stack]
    ~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -156,27 +157,34 @@ foreign import prim "advanceStackFrameIterzh" advanceStackFrameIter# :: StackSna
 
 foreign import prim "getInfoTableAddrzh" getInfoTableAddr# :: StackSnapshot# -> Word# -> Addr#
 
+foreign import prim "getStackInfoTableAddrzh" getStackInfoTableAddr# :: StackSnapshot# -> Addr#
+
 getInfoTable :: StackFrameIter -> IO StgInfoTable
-getInfoTable StackFrameIter {..} =
+getInfoTable StackFrameIter {..} | sfiKind == SfiClosure =
   let infoTablePtr = Ptr (getInfoTableAddr# stackSnapshot# (wordOffsetToWord# index))
    in peekItbl infoTablePtr
+getInfoTable StackFrameIter {..} | sfiKind == SfiStack = peekItbl $ Ptr (getStackInfoTableAddr# stackSnapshot#)
+getInfoTable StackFrameIter {..} | sfiKind == SfiPrimitive = error "Primitives have no info table!"
 
 foreign import prim "getBoxedClosurezh" getBoxedClosure# :: StackSnapshot# -> Word# -> State# RealWorld -> (# State# RealWorld, Any #)
 
--- -- TODO: Remove this instance (debug only)
--- instance Show StackFrameIter where
---   show (StackFrameIter {..}) = "StackFrameIter " ++ "(StackSnapshot _" ++ " " ++ show index
+foreign import prim "getStackFieldszh" getStackFields# :: StackSnapshot# -> State# RealWorld -> (# State# RealWorld, Word32#, Word8#, Word8# #)
+
+getStackFields :: StackFrameIter -> IO (Word32, Word8, Word8)
+getStackFields StackFrameIter {..} = IO $ \s ->
+  case getStackFields# stackSnapshot# s of (# s1, sSize#, sDirty#, sMarking# #)
+                                             -> (# s1, (W32# sSize#, W8# sDirty#, W8# sMarking#) #)
 
 -- | Get an interator starting with the top-most stack frame
 stackHead :: StackSnapshot -> StackFrameIter
-stackHead (StackSnapshot s) = StackFrameIter s 0 False -- GHC stacks are never empty
+stackHead (StackSnapshot s) = StackFrameIter s 0 SfiClosure -- GHC stacks are never empty
 
 -- | Advance iterator to the next stack frame (if any)
 advanceStackFrameIter :: StackFrameIter -> Maybe StackFrameIter
 advanceStackFrameIter (StackFrameIter {..}) =
   let !(# s', i', hasNext #) = advanceStackFrameIter# stackSnapshot# (wordOffsetToWord# index)
    in if (I# hasNext) > 0
-        then Just $ StackFrameIter s' (primWordToWordOffset i') False
+        then Just $ StackFrameIter s' (primWordToWordOffset i') SfiClosure
         else Nothing
 
 primWordToWordOffset :: Word# -> WordOffset
@@ -191,7 +199,7 @@ wordsToBitmapEntries sfi (b : bs) bitmapSize =
       mbLastFrame = (listToMaybe . reverse) entries
    in case mbLastFrame of
         Just (StackFrameIter {..}) ->
-          entries ++ wordsToBitmapEntries (StackFrameIter stackSnapshot# (index + 1) False) bs (subtractDecodedBitmapWord bitmapSize)
+          entries ++ wordsToBitmapEntries (StackFrameIter stackSnapshot# (index + 1) SfiClosure) bs (subtractDecodedBitmapWord bitmapSize)
         Nothing -> error "This should never happen! Recursion ended not in base case."
   where
     subtractDecodedBitmapWord :: Word -> Word
@@ -202,12 +210,12 @@ toBitmapEntries _ _ 0 = []
 toBitmapEntries sfi@(StackFrameIter {..}) bitmapWord bSize =
   -- TODO: overriding isPrimitive field is a bit weird. Could be calculated before
     sfi {
-        isPrimitive = (bitmapWord .&. 1) /= 0
+        sfiKind = if (bitmapWord .&. 1) /= 0 then SfiPrimitive else SfiClosure
         }
-    : toBitmapEntries (StackFrameIter stackSnapshot# (index + 1) False) (bitmapWord `shiftR` 1) (bSize - 1)
+    : toBitmapEntries (StackFrameIter stackSnapshot# (index + 1) SfiClosure) (bitmapWord `shiftR` 1) (bSize - 1)
 
 toBitmapPayload :: StackFrameIter -> IO Box
-toBitmapPayload sfi | isPrimitive sfi = pure (StackFrameBox sfi)
+toBitmapPayload sfi | sfiKind sfi == SfiPrimitive = pure (StackFrameBox sfi)
 toBitmapPayload sfi = getClosure sfi 0
 
 getClosure :: StackFrameIter -> WordOffset -> IO Box
@@ -226,7 +234,7 @@ decodeLargeBitmap getterFun# sfi@(StackFrameIter {..}) relativePayloadOffset = d
 
 decodeBitmaps :: StackFrameIter -> WordOffset -> [Word] -> Word -> IO [Box]
 decodeBitmaps (StackFrameIter {..}) relativePayloadOffset bitmapWords size =
-  let bes = wordsToBitmapEntries (StackFrameIter stackSnapshot# (index + relativePayloadOffset) False) bitmapWords size
+  let bes = wordsToBitmapEntries (StackFrameIter stackSnapshot# (index + relativePayloadOffset) SfiClosure) bitmapWords size
    in mapM toBitmapPayload bes
 
 decodeSmallBitmap :: (StackSnapshot# -> Word# -> State# RealWorld -> (# State# RealWorld, Word#, Word# #)) -> StackFrameIter -> WordOffset -> IO [Box]
@@ -249,7 +257,21 @@ wordOffsetToWord# :: WordOffset -> Word#
 wordOffsetToWord# wo = intToWord# (fromIntegral wo)
 
 unpackStackFrameIter :: StackFrameIter -> IO Closure
-unpackStackFrameIter sfi | isPrimitive sfi = UnknownTypeWordSizedPrimitive <$> (getWord sfi 0)
+unpackStackFrameIter sfi | sfiKind sfi == SfiPrimitive = UnknownTypeWordSizedPrimitive <$> (getWord sfi 0)
+unpackStackFrameIter sfi | sfiKind sfi == SfiStack = do
+  info <- getInfoTable sfi
+  (stack_size', stack_dirty', stack_marking') <- getStackFields sfi
+  case tipe info of
+    STACK -> do
+      let stack' = decodeStack' (StackSnapshot (stackSnapshot# sfi))
+      pure $ StackClosure {
+                            info = info,
+                            stack_size = stack_size',
+                            stack_dirty = stack_dirty',
+                            stack_marking = stack_marking',
+                            stack = stack'
+                          }
+    _ -> error $ "Expected STACK closure, got " ++ show info
 unpackStackFrameIter sfi = do
   traceM $ "unpackStackFrameIter - sfi " ++ show sfi
   info <- getInfoTable sfi
@@ -316,10 +338,14 @@ unpackStackFrameIter sfi = do
               handler = handler'
             }
         UNDERFLOW_FRAME -> do
-          nextChunk' <- getUnderflowFrameNextChunk sfi
+          (StackSnapshot nextChunk') <- getUnderflowFrameNextChunk sfi
           pure $ UnderflowFrame
             { info = info,
-              nextChunk = nextChunk'
+              nextChunk = StackFrameBox $ StackFrameIter {
+                                          stackSnapshot# = nextChunk',
+                                          index = 0,
+                                          sfiKind = SfiStack
+                                         }
             }
         STOP_FRAME -> pure $ StopFrame {info = info}
         ATOMICALLY_FRAME -> do
@@ -363,9 +389,12 @@ toInt# (I# i) = i
 intToWord# :: Int -> Word#
 intToWord# i = int2Word# (toInt# i)
 
-decodeStack :: StackSnapshot -> Closure
-decodeStack = SimpleStack . decodeStack'
-
+decodeStack :: StackSnapshot -> IO Closure
+decodeStack (StackSnapshot stack#) = unpackStackFrameIter $ StackFrameIter {
+  stackSnapshot# = stack#,
+  index = 0,
+  sfiKind = SfiStack
+                                   }
 decodeStack' :: StackSnapshot -> [Box]
 decodeStack' s = StackFrameBox (stackHead s) : go (advanceStackFrameIter (stackHead s))
   where
diff --git a/libraries/ghc-heap/GHC/Exts/Heap.hs b/libraries/ghc-heap/GHC/Exts/Heap.hs
index 23c3747869..0b451cc6fa 100644
--- a/libraries/ghc-heap/GHC/Exts/Heap.hs
+++ b/libraries/ghc-heap/GHC/Exts/Heap.hs
@@ -143,7 +143,7 @@ instance Double# ~ a => HasHeapRep (a :: TYPE 'DoubleRep) where
 
 #if MIN_TOOL_VERSION_ghc(9,5,0)
 instance {-# OVERLAPPING #-} HasHeapRep StackSnapshot# where
-    getClosureData s# = pure $ decodeStack (StackSnapshot s#)
+    getClosureData s# = decodeStack (StackSnapshot s#)
 #endif
 
 -- | Get the heap representation of a closure _at this moment_, even if it is
diff --git a/libraries/ghc-heap/GHC/Exts/Heap/Closures.hs b/libraries/ghc-heap/GHC/Exts/Heap/Closures.hs
index dec560c1c5..d7032f50d5 100644
--- a/libraries/ghc-heap/GHC/Exts/Heap/Closures.hs
+++ b/libraries/ghc-heap/GHC/Exts/Heap/Closures.hs
@@ -25,6 +25,7 @@ module GHC.Exts.Heap.Closures (
     , areBoxesEqual
     , asBox
 #if MIN_VERSION_base(4,17,0)
+    , SfiKind(..)
     , StackFrameIter(..)
 #endif
     ) where
@@ -78,10 +79,13 @@ foreign import prim "eqStackSnapshotszh" eqStackSnapshots# :: StackSnapshot# ->
 -- required, e.g. in 'getBoxedClosureData', the function knows how far it has
 -- to evaluate the argument.
 #if MIN_VERSION_base(4,17,0)
+data SfiKind = SfiClosure | SfiPrimitive | SfiStack
+  deriving (Eq, Show)
+
 data StackFrameIter = StackFrameIter
   { stackSnapshot# :: !StackSnapshot#,
     index :: !WordOffset,
-    isPrimitive :: !Bool
+    sfiKind :: !SfiKind
   }
 
 instance Show StackFrameIter where
@@ -360,14 +364,12 @@ data GenClosure b
 #if __GLASGOW_HASKELL__ >= 811
       , stack_marking   :: !Word8
 #endif
+      -- | The frames of the stack. Only available if a cloned stack was
+      -- decoded, otherwise empty.
+      , stack           :: ![b]
       }
 
 #if MIN_TOOL_VERSION_ghc(9,5,0)
-    -- TODO: I could model stack chunks here (much better). However, I need the
-    -- code to typecheck, now.
-  | SimpleStack {
-      stackClosures :: ![b]
-                }
   | UpdateFrame
       { info            :: !StgInfoTable
       , knownUpdateFrameType :: !UpdateFrameType
@@ -402,7 +404,7 @@ data GenClosure b
     -- TODO: nextChunk could be a CL.Closure, too! (StackClosure)
   | UnderflowFrame
       { info            :: !StgInfoTable
-      , nextChunk:: !StackSnapshot
+      , nextChunk:: !b
       }
 
   | StopFrame
@@ -621,7 +623,7 @@ allClosures (BlockingQueueClosure {..}) = [link, blackHole, owner, queue]
 allClosures (WeakClosure {..}) = [cfinalizers, key, value, finalizer] ++ Data.Foldable.toList weakLink
 allClosures (OtherClosure {..}) = hvalues
 #if MIN_TOOL_VERSION_ghc(9,5,0)
-allClosures (SimpleStack {..}) = stackClosures
+allClosures (StackClosure {..}) = stack
 allClosures (UpdateFrame {..}) = [updatee]
 allClosures (CatchFrame {..}) = [handler]
 allClosures (CatchStmFrame {..}) = [catchFrameCode, handler]
diff --git a/libraries/ghc-heap/cbits/Stack.cmm b/libraries/ghc-heap/cbits/Stack.cmm
index 19a22fa84f..680bbcc251 100644
--- a/libraries/ghc-heap/cbits/Stack.cmm
+++ b/libraries/ghc-heap/cbits/Stack.cmm
@@ -3,6 +3,7 @@
 
 #include "Cmm.h"
 
+#if defined(StgStack_marking)
 advanceStackFrameIterzh (P_ stack, W_ offsetWords) {
   W_ frameSize;
   (frameSize) = ccall stackFrameSize(stack, offsetWords);
@@ -175,6 +176,12 @@ getInfoTableAddrzh(P_ stack, W_ offsetWords){
   return (info);
 }
 
+getStackInfoTableAddrzh(P_ stack){
+  P_ info;
+  info = %GET_STD_INFO(UNTAG(stack));
+  return (info);
+}
+
 // Just a cast
 stackSnapshotToWordzh(P_ stack) {
   return (stack);
@@ -199,5 +206,18 @@ getBoxedClosurezh(P_ stack, W_ offsetWords){
   return (box);
 }
 
+// TODO: Unused?
 INFO_TABLE_CONSTR(box, 1, 0, 0, CONSTR_1_0, "BOX", "BOX")
 { foreign "C" barf("BOX object (%p) entered!", R1) never returns; }
+
+getStackFieldszh(P_ stack){
+  bits32 size;
+  bits8 dirty, marking;
+
+  size = StgStack_stack_size(stack);
+  dirty = StgStack_dirty(stack);
+  marking = StgStack_marking(stack);
+
+  return (size, dirty, marking);
+}
+#endif
diff --git a/libraries/ghc-heap/tests/TestUtils.hs b/libraries/ghc-heap/tests/TestUtils.hs
index 8b5bfcf8d7..53c40fc041 100644
--- a/libraries/ghc-heap/tests/TestUtils.hs
+++ b/libraries/ghc-heap/tests/TestUtils.hs
@@ -30,8 +30,8 @@ import Unsafe.Coerce (unsafeCoerce)
 getDecodedStack :: IO (StackSnapshot, [Closure])
 getDecodedStack = do
   s@(StackSnapshot s#) <- cloneMyStack
-  (SimpleStack cs) <- getClosureData s#
-  unboxedCs <- mapM getBoxedClosureData cs
+  stackClosure <- getClosureData s#
+  unboxedCs <- mapM getBoxedClosureData (stack stackClosure)
   pure (s, unboxedCs)
 
 assertEqual :: (HasCallStack, Monad m, Show a, Eq a) => a -> a -> m ()
diff --git a/libraries/ghc-heap/tests/stack_big_ret.hs b/libraries/ghc-heap/tests/stack_big_ret.hs
index 392e04f52e..c7b23d8a0a 100644
--- a/libraries/ghc-heap/tests/stack_big_ret.hs
+++ b/libraries/ghc-heap/tests/stack_big_ret.hs
@@ -37,8 +37,8 @@ main = do
 
   mbStackSnapshot <- readIORef stackRef
   let stackSnapshot@(StackSnapshot s#) = fromJust mbStackSnapshot
-  (SimpleStack boxedFrames) <- getClosureData s#
-  stackFrames <- mapM getBoxedClosureData boxedFrames
+  stackClosure <- getClosureData s#
+  stackFrames <- mapM getBoxedClosureData (stack stackClosure)
 
   assertStackInvariants stackSnapshot stackFrames
   assertThat
diff --git a/libraries/ghc-heap/tests/stack_misc_closures.hs b/libraries/ghc-heap/tests/stack_misc_closures.hs
index 822bfa63ef..1f4abda458 100644
--- a/libraries/ghc-heap/tests/stack_misc_closures.hs
+++ b/libraries/ghc-heap/tests/stack_misc_closures.hs
@@ -326,10 +326,11 @@ test setup assertion = do
   -- Better fail early, here.
   performGC
   traceM $ "test - sn' " ++ show sn
-  ss@(SimpleStack boxedFrames) <- getClosureData sn#
-  traceM $ "test - ss" ++ show ss
+  stackClosure <- getClosureData sn#
+  traceM $ "test - ss" ++ show stackClosure
   performGC
   traceM $ "call getBoxedClosureData"
+  let boxedFrames = stack stackClosure
   stack <- mapM getBoxedClosureData boxedFrames
   performGC
   assert sn stack
@@ -338,8 +339,8 @@ test setup assertion = do
   let (StackSnapshot sn#) = sn
   stack' <- getClosureData sn#
   case stack' of
-    SimpleStack {..} -> do
-      !cs <- mapM getBoxedClosureData stackClosures
+    StackClosure {..} -> do
+      !cs <- mapM getBoxedClosureData stack
       assert sn cs
     _ -> error $ "Unexpected closure type : " ++ show stack'
   where
@@ -364,8 +365,8 @@ entertainGC x = show x ++ entertainGC (x -1)
 testSize :: HasCallStack => SetupFunction -> Int -> IO ()
 testSize setup expectedSize = do
   (StackSnapshot sn#) <- getStackSnapshot setup
-  (SimpleStack boxedFrames) <- getClosureData sn#
-  assertEqual expectedSize =<< closureSize (head boxedFrames)
+  stackClosure <- getClosureData sn#
+  assertEqual expectedSize =<< (closureSize . head . stack) stackClosure
 
 -- | Get a `StackSnapshot` from test setup
 --
diff --git a/libraries/ghc-heap/tests/stack_underflow.hs b/libraries/ghc-heap/tests/stack_underflow.hs
index 13dd6a9cec..74f7c9d637 100644
--- a/libraries/ghc-heap/tests/stack_underflow.hs
+++ b/libraries/ghc-heap/tests/stack_underflow.hs
@@ -5,6 +5,7 @@ module Main where
 
 import Data.Bool (Bool (True))
 import GHC.Exts.DecodeStack
+import GHC.Exts.Heap
 import GHC.Exts.Heap.ClosureTypes
 import GHC.Exts.Heap.Closures
 import GHC.Exts.Heap.InfoTable.Types
@@ -37,7 +38,9 @@ isUnderflowFrame _ = False
 assertStackChunksAreDecodable :: HasCallStack => [Closure] -> IO ()
 assertStackChunksAreDecodable s = do
   let underflowFrames = filter isUnderflowFrame s
-  let framesOfChunks = map (stackClosures . decodeStack . nextChunk) underflowFrames
+  stackClosures <- mapM (getBoxedClosureData . nextChunk) underflowFrames
+  let stackBoxes = map stack stackClosures
+  framesOfChunks <- sequence (map (mapM getBoxedClosureData) stackBoxes)
   assertThat
     "No empty stack chunks"
     (== True)
diff --git a/utils/deriveConstants/Main.hs b/utils/deriveConstants/Main.hs
index a157202a5a..639ce62155 100644
--- a/utils/deriveConstants/Main.hs
+++ b/utils/deriveConstants/Main.hs
@@ -476,6 +476,7 @@ wanteds os = concat
           ,closureFieldOffset Both "StgStack" "stack"
           ,closureField       C    "StgStack" "stack_size"
           ,closureField       C    "StgStack" "dirty"
+          ,closureField       C    "StgStack" "marking"
 
           ,structSize C "StgTSOProfInfo"