closureSize

3 files changed, 78 insertions, 20 deletions

diff --git a/libraries/ghc-heap/GHC/Exts/Heap/Closures.hs b/libraries/ghc-heap/GHC/Exts/Heap/Closures.hs
index a085f73013..03057f5f11 100644
--- a/libraries/ghc-heap/GHC/Exts/Heap/Closures.hs
+++ b/libraries/ghc-heap/GHC/Exts/Heap/Closures.hs
@@ -53,6 +53,7 @@ import Numeric
 
 #if MIN_VERSION_base(4,17,0)
 import GHC.Stack.CloneStack (StackSnapshot(..))
+import GHC.Exts.StackConstants
 import Unsafe.Coerce (unsafeCoerce)
 #endif
 
@@ -603,6 +604,20 @@ allClosures _ = []
 -- Includes header and payload. Does not follow pointers.
 --
 -- @since 8.10.1
--- TODO: Handle PrimitiveWordHolder
 closureSize :: Box -> Int
 closureSize (Box x) = I# (closureSize# x)
+#if MIN_VERSION_base(4,17,0)
+closureSize (DecodedBox c) = case c of
+  UpdateFrame {} -> sizeStgUpdateFrame
+  CatchFrame {} -> sizeStgCatchFrame
+  CatchStmFrame {} -> sizeStgCatchSTMFrame
+  CatchRetryFrame {} -> sizeStgCatchRetryFrame
+  AtomicallyFrame {} -> sizeStgAtomicallyFrame
+  RetSmall {..} -> sizeStgClosure + length payload
+  RetBig {..} -> sizeStgClosure + length payload
+  RetFun {..} -> sizeStgRetFunFrame + length retFunPayload
+  -- The one additional word is a pointer to the StgBCO in the closure's payload
+  RetBCO {..} -> sizeStgClosure + 1 + length bcoArgs
+  -- TODO: What to do about other closure types?
+  _ -> error "Unexpected closure type"
+#endif
diff --git a/libraries/ghc-heap/GHC/Exts/StackConstants.hsc b/libraries/ghc-heap/GHC/Exts/StackConstants.hsc
index cace9fa300..7a68f58935 100644
--- a/libraries/ghc-heap/GHC/Exts/StackConstants.hsc
+++ b/libraries/ghc-heap/GHC/Exts/StackConstants.hsc
@@ -26,21 +26,33 @@ offsetStgCatchFrameHandler = byteOffsetToWordOffset $ (#const OFFSET_StgCatchFra
 offsetStgCatchFrameExceptionsBlocked :: WordOffset
 offsetStgCatchFrameExceptionsBlocked = byteOffsetToWordOffset $ (#const OFFSET_StgCatchFrame_exceptions_blocked) + (#size StgHeader)
 
+sizeStgCatchFrame :: Int
+sizeStgCatchFrame = bytesToWords $ (#const SIZEOF_StgCatchFrame_NoHdr) + (#size StgHeader)
+
 offsetStgCatchSTMFrameCode :: WordOffset
 offsetStgCatchSTMFrameCode = byteOffsetToWordOffset $ (#const OFFSET_StgCatchSTMFrame_code) + (#size StgHeader)
 
 offsetStgCatchSTMFrameHandler :: WordOffset
 offsetStgCatchSTMFrameHandler = byteOffsetToWordOffset $ (#const OFFSET_StgCatchSTMFrame_handler) + (#size StgHeader)
 
+sizeStgCatchSTMFrame :: Int
+sizeStgCatchSTMFrame = bytesToWords $ (#const SIZEOF_StgCatchSTMFrame_NoHdr) + (#size StgHeader)
+
 offsetStgUpdateFrameUpdatee :: WordOffset
 offsetStgUpdateFrameUpdatee = byteOffsetToWordOffset $ (#const OFFSET_StgUpdateFrame_updatee) + (#size StgHeader)
 
+sizeStgUpdateFrame :: Int
+sizeStgUpdateFrame = bytesToWords $ (#const SIZEOF_StgUpdateFrame_NoHdr) + (#size StgHeader)
+
 offsetStgAtomicallyFrameCode :: WordOffset
 offsetStgAtomicallyFrameCode = byteOffsetToWordOffset $ (#const OFFSET_StgAtomicallyFrame_code) + (#size StgHeader)
 
 offsetStgAtomicallyFrameResult :: WordOffset
 offsetStgAtomicallyFrameResult = byteOffsetToWordOffset $ (#const OFFSET_StgAtomicallyFrame_result) + (#size StgHeader)
 
+sizeStgAtomicallyFrame :: Int
+sizeStgAtomicallyFrame = bytesToWords $ (#const SIZEOF_StgAtomicallyFrame_NoHdr) + (#size StgHeader)
+
 offsetStgCatchRetryFrameRunningAltCode :: WordOffset
 offsetStgCatchRetryFrameRunningAltCode = byteOffsetToWordOffset $ (#const OFFSET_StgCatchRetryFrame_running_alt_code) + (#size StgHeader)
 
@@ -50,6 +62,9 @@ offsetStgCatchRetryFrameRunningFirstCode = byteOffsetToWordOffset $ (#const OFFS
 offsetStgCatchRetryFrameAltCode :: WordOffset
 offsetStgCatchRetryFrameAltCode = byteOffsetToWordOffset $ (#const OFFSET_StgCatchRetryFrame_alt_code) + (#size StgHeader)
 
+sizeStgCatchRetryFrame :: Int
+sizeStgCatchRetryFrame = bytesToWords $ (#const SIZEOF_StgCatchRetryFrame_NoHdr) + (#size StgHeader)
+
 offsetStgRetFunFrameSize :: WordOffset
 -- StgRetFun has no header, but only a pointer to the info table at the beginning.
 offsetStgRetFunFrameSize = byteOffsetToWordOffset $ (#const OFFSET_StgRetFun_size)
@@ -60,6 +75,9 @@ offsetStgRetFunFrameFun = byteOffsetToWordOffset $ (#const OFFSET_StgRetFun_fun)
 offsetStgRetFunFramePayload :: WordOffset
 offsetStgRetFunFramePayload = byteOffsetToWordOffset $ (#const OFFSET_StgRetFun_payload)
 
+sizeStgRetFunFrame :: Int
+sizeStgRetFunFrame = bytesToWords (#const SIZEOF_StgRetFun)
+
 offsetStgBCOFrameInstrs :: ByteOffset
 offsetStgBCOFrameInstrs = (#const OFFSET_StgBCO_instrs) + (#size StgHeader)
 
@@ -78,12 +96,20 @@ offsetStgBCOFrameSize = (#const OFFSET_StgBCO_size) + (#size StgHeader)
 offsetStgClosurePayload :: WordOffset
 offsetStgClosurePayload = byteOffsetToWordOffset $ (#const OFFSET_StgClosure_payload) + (#size StgHeader)
 
+sizeStgClosure :: Int
+sizeStgClosure = bytesToWords (#size StgHeader)
+
 byteOffsetToWordOffset :: ByteOffset -> WordOffset
-byteOffsetToWordOffset bo = if bo `mod` bytesInWord == 0 then
-                              fromIntegral $ bo `div` bytesInWord
-                            else
-                              error "Unexpected struct alignment!"
-  where
-        bytesInWord = (#const SIZEOF_VOID_P)
+byteOffsetToWordOffset = WordOffset . bytesToWords . fromInteger . toInteger
+
+bytesToWords :: Int -> Int
+bytesToWords b =
+  if b `mod` bytesInWord == 0 then
+      fromIntegral $ b `div` bytesInWord
+    else
+      error "Unexpected struct alignment!"
+
+bytesInWord :: Int
+bytesInWord = (#const SIZEOF_VOID_P)
 
 #endif
diff --git a/libraries/ghc-heap/tests/stack_misc_closures.hs b/libraries/ghc-heap/tests/stack_misc_closures.hs
index 5113a5e242..18057c5ed3 100644
--- a/libraries/ghc-heap/tests/stack_misc_closures.hs
+++ b/libraries/ghc-heap/tests/stack_misc_closures.hs
@@ -62,8 +62,8 @@ foreign import ccall "maxSmallBitmapBits" maxSmallBitmapBits_c :: Word
 
 foreign import ccall "belchStack" belchStack# :: StackSnapshot# -> IO ()
 
-{-
-__Test stategy:__
+{- Test stategy
+   ~~~~~~~~~~~~
 
 - Create @StgStack@s in C that contain two closures (as they are on stack they
 may also be called "frames"). A stop frame and the frame which's decoding should
@@ -98,6 +98,7 @@ main = do
         assertEqual knownUpdateFrameType NormalUpdateFrame
         assertEqual 1 =<< (getWordFromBlackhole =<< getBoxedClosureData updatee)
       e -> error $ "Wrong closure type: " ++ show e
+  testSize any_update_frame# 2
   test any_catch_frame# $
     \case
       CatchFrame {..} -> do
@@ -105,6 +106,7 @@ main = do
         assertEqual exceptions_blocked 1
         assertConstrClosure 1 =<< getBoxedClosureData handler
       e -> error $ "Wrong closure type: " ++ show e
+  testSize any_catch_frame# 3
   test any_catch_stm_frame# $
     \case
       CatchStmFrame {..} -> do
@@ -112,6 +114,7 @@ main = do
         assertConstrClosure 1 =<< getBoxedClosureData catchFrameCode
         assertConstrClosure 2 =<< getBoxedClosureData handler
       e -> error $ "Wrong closure type: " ++ show e
+  testSize any_catch_stm_frame# 3
   test any_catch_retry_frame# $
     \case
       CatchRetryFrame {..} -> do
@@ -120,6 +123,7 @@ main = do
         assertConstrClosure 1 =<< getBoxedClosureData first_code
         assertConstrClosure 2 =<< getBoxedClosureData alt_code
       e -> error $ "Wrong closure type: " ++ show e
+  testSize any_catch_retry_frame# 4
   test any_atomically_frame# $
     \case
       AtomicallyFrame {..} -> do
@@ -127,6 +131,7 @@ main = do
         assertConstrClosure 1 =<< getBoxedClosureData atomicallyFrameCode
         assertConstrClosure 2 =<< getBoxedClosureData result
       e -> error $ "Wrong closure type: " ++ show e
+  testSize any_atomically_frame# 3
   -- TODO: Test for UnderflowFrame once it points to a Box payload
   test any_ret_small_prim_frame# $
     \case
@@ -137,6 +142,7 @@ main = do
         assertEqual (length pCs) 1
         assertUnknownTypeWordSizedPrimitive 1 (head pCs)
       e -> error $ "Wrong closure type: " ++ show e
+  testSize any_ret_small_prim_frame# 2
   test any_ret_small_closure_frame# $
     \case
       RetSmall {..} -> do
@@ -146,6 +152,7 @@ main = do
         assertEqual (length pCs) 1
         assertConstrClosure 1 (head pCs)
       e -> error $ "Wrong closure type: " ++ show e
+  testSize any_ret_small_closure_frame# 2
   test any_ret_small_closures_frame# $
     \case
       RetSmall {..} -> do
@@ -156,6 +163,7 @@ main = do
         let wds = map getWordFromConstr01 pCs
         assertEqual wds [1 .. 58]
       e -> error $ "Wrong closure type: " ++ show e
+  testSize any_ret_small_closures_frame# (1 + fromIntegral maxSmallBitmapBits_c)
   test any_ret_small_prims_frame# $
     \case
       RetSmall {..} -> do
@@ -166,6 +174,7 @@ main = do
         let wds = map getWordFromUnknownTypeWordSizedPrimitive pCs
         assertEqual wds [1 .. 58]
       e -> error $ "Wrong closure type: " ++ show e
+  testSize any_ret_small_prims_frame# (1 + fromIntegral maxSmallBitmapBits_c)
   test any_ret_big_prims_min_frame# $
     \case
       RetBig {..} -> do
@@ -175,15 +184,7 @@ main = do
         let wds = map getWordFromUnknownTypeWordSizedPrimitive pCs
         assertEqual wds [1 .. 59]
       e -> error $ "Wrong closure type: " ++ show e
-  test any_ret_big_prims_min_frame# $
-    \case
-      RetBig {..} -> do
-        assertEqual (tipe info) RET_BIG
-        pCs <- mapM getBoxedClosureData payload
-        assertEqual (length pCs) 59
-        let wds = map getWordFromUnknownTypeWordSizedPrimitive pCs
-        assertEqual wds [1 .. 59]
-      e -> error $ "Wrong closure type: " ++ show e
+  testSize any_ret_big_prims_min_frame# (minBigBitmapBits + 1)
   test any_ret_big_closures_min_frame# $
     \case
       RetBig {..} -> do
@@ -193,15 +194,18 @@ main = do
         let wds = map getWordFromConstr01 pCs
         assertEqual wds [1 .. 59]
       e -> error $ "Wrong closure type: " ++ show e
+  testSize any_ret_big_closures_min_frame# (minBigBitmapBits + 1)
   test any_ret_big_closures_two_words_frame# $
     \case
       RetBig {..} -> do
         assertEqual (tipe info) RET_BIG
         pCs <- mapM getBoxedClosureData payload
-        assertEqual (length pCs) 65
+        let closureCount = 64 + 1
+        assertEqual (length pCs) closureCount
         let wds = map getWordFromConstr01 pCs
-        assertEqual wds [1 .. 65]
+        assertEqual wds [1 .. (fromIntegral closureCount)]
       e -> error $ "Wrong closure type: " ++ show e
+  testSize any_ret_big_closures_two_words_frame# (64 + 1 + 1)
   test any_ret_fun_arg_n_prim_framezh# $
     \case
       RetFun {..} -> do
@@ -232,6 +236,7 @@ main = do
         let wds = map getWordFromConstr01 pCs
         assertEqual wds [1 .. 9]
       e -> error $ "Wrong closure type: " ++ show e
+  testSize any_ret_fun_arg_gen_framezh# (3 + 9)
   test any_ret_fun_arg_gen_big_framezh# $
     \case
       RetFun {..} -> do
@@ -249,6 +254,7 @@ main = do
         assertEqual (length pCs) 59
         let wds = map getWordFromConstr01 pCs
         assertEqual wds [1 .. 59]
+  testSize any_ret_fun_arg_gen_big_framezh# (3 + 59)
   test any_bco_frame# $
     \case
       RetBCO {..} -> do
@@ -271,6 +277,7 @@ main = do
               ] bitmap
           e -> error $ "Wrong closure type: " ++ show e
       e -> error $ "Wrong closure type: " ++ show e
+  testSize any_bco_frame# 3
 
 type SetupFunction = State# RealWorld -> (# State# RealWorld, StackSnapshot# #)
 
@@ -300,6 +307,7 @@ test setup assertion = do
     assert sn stack = do
       assertStackInvariants sn stack
       assertEqual (length stack) 2
+      -- TODO: Isn't this also a stack invariant? (assertStackInvariants)
       assertThat
         "Last frame is stop frame"
         ( \case
@@ -309,6 +317,12 @@ test setup assertion = do
         (last stack)
       assertion $ head stack
 
+testSize :: HasCallStack => SetupFunction -> Int -> IO ()
+testSize setup expectedSize = do
+  sn <- getStackSnapshot setup
+  (SimpleStack boxedFrames) <- decodeStack sn
+  assertEqual expectedSize (closureSize (head boxedFrames))
+
 -- | Get a `StackSnapshot` from test setup
 --
 -- This function mostly resembles `cloneStack`. Though, it doesn't clone, but
@@ -375,6 +389,9 @@ assertUnknownTypeWordSizedPrimitive w c = case c of
 unboxSingletonTuple :: (# StackSnapshot# #) -> StackSnapshot#
 unboxSingletonTuple (# s# #) = s#
 
+minBigBitmapBits :: Num a => a
+minBigBitmapBits = 1 + fromIntegral maxSmallBitmapBits_c
+
 -- | A function with 59 arguments
 --
 -- A small bitmap has @64 - 6 = 58@ entries on 64bit machines. On 32bit machines