Simpify bitmap decoding

3 files changed, 77 insertions, 96 deletions

diff --git a/libraries/ghc-heap/GHC/Exts/Stack/Decode.hs b/libraries/ghc-heap/GHC/Exts/Stack/Decode.hs
index 5e3644a361..5a7a4e7cda 100644
--- a/libraries/ghc-heap/GHC/Exts/Stack/Decode.hs
+++ b/libraries/ghc-heap/GHC/Exts/Stack/Decode.hs
@@ -6,6 +6,7 @@
 {-# LANGUAGE GHCForeignImportPrim #-}
 {-# LANGUAGE MagicHash #-}
 {-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE RecordWildCards #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE TypeInType #-}
@@ -17,7 +18,6 @@ module GHC.Exts.Stack.Decode
   )
 where
 
-import Data.Array.Byte
 import Data.Bits
 import Data.Maybe
 import Foreign
@@ -32,6 +32,7 @@ import GHC.IO (IO (..))
 import GHC.Stack.CloneStack
 import GHC.Word
 import Prelude
+import Debug.Trace
 
 {- Note [Decoding the stack]
    ~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -145,7 +146,7 @@ getRetFunType stackSnapshot# index =
             (# s1, rft# #) -> (# s1, W# rft# #)
       )
 
-type LargeBitmapGetter = StackSnapshot# -> Word# -> State# RealWorld -> (# State# RealWorld, ByteArray#, Word# #)
+type LargeBitmapGetter = StackSnapshot# -> Word# -> State# RealWorld -> (# State# RealWorld, Addr#, Word# #)
 
 foreign import prim "getLargeBitmapzh" getLargeBitmap# :: LargeBitmapGetter
 
@@ -229,29 +230,60 @@ data StackFrameIter
   | -- | Represents a primitive word on the stack
     SfiPrimitive !StackSnapshot# !WordOffset
 
+data LargeBitmap = LargeBitmap
+    { largeBitmapSize :: Word
+    , largebitmapWords :: Ptr Word
+    }
+
+-- | Is a bitmap entry a closure pointer or a primitive non-pointer?
+data Pointerness = Pointer | NonPointer
+  deriving Show
+
 decodeLargeBitmap :: LargeBitmapGetter -> StackSnapshot# -> WordOffset -> WordOffset -> IO [Closure]
 decodeLargeBitmap getterFun# stackSnapshot# index relativePayloadOffset = do
-  (bitmapArray, size) <- IO $ \s ->
+  largeBitmap <- IO $ \s ->
     case getterFun# stackSnapshot# (wordOffsetToWord# index) s of
-      (# s1, ba#, s# #) -> (# s1, (ByteArray ba#, W# s#) #)
-  let bitmapWords :: [Word] = byteArrayToList bitmapArray
-  decodeBitmaps stackSnapshot# (index + relativePayloadOffset) bitmapWords size
+      (# s1, wordsAddr#, size# #) -> (# s1, LargeBitmap (W# size#) (Ptr wordsAddr#) #)
+  bitmapWords <-largeBitmapToList largeBitmap
+  decodeBitmaps stackSnapshot#
+    (index + relativePayloadOffset)
+    (bitmapWordsPointerness (largeBitmapSize largeBitmap) bitmapWords)
   where
-    byteArrayToList :: ByteArray -> [Word]
-    byteArrayToList (ByteArray bArray) = go 0
-      where
-        go i
-          | i < maxIndex = W# (indexWordArray# bArray (toInt# i)) : go (i + 1)
-          | otherwise = []
-        maxIndex = sizeofByteArray bArray `quot` sizeOf (undefined :: Word)
-
-    sizeofByteArray :: ByteArray# -> Int
-    sizeofByteArray arr# = I# (sizeofByteArray# arr#)
-
-decodeBitmaps :: StackSnapshot# -> WordOffset -> [Word] -> Word -> IO [Closure]
-decodeBitmaps stackSnapshot# index bitmapWords size =
-  let bes = wordsToBitmapEntries index bitmapWords size
-   in mapM toBitmapPayload bes
+    largeBitmapToList :: LargeBitmap -> IO [Word]
+    largeBitmapToList LargeBitmap {..} = cWordArrayToList largebitmapWords $
+      (usedBitmapWords.fromIntegral) largeBitmapSize
+
+    cWordArrayToList :: Ptr Word -> Int -> IO [Word]
+    cWordArrayToList ptr size = mapM (peekElemOff ptr) [0..(size-1)]
+
+    usedBitmapWords :: Int -> Int
+    usedBitmapWords 0 = error "Invalid large bitmap size 0."
+    usedBitmapWords size = (size `div` (fromIntegral wORD_SIZE_IN_BITS)) + 1
+
+    bitmapWordsPointerness :: Word -> [Word] -> [Pointerness]
+    bitmapWordsPointerness size _ | size <= 0 = []
+    bitmapWordsPointerness _ [] = []
+    bitmapWordsPointerness size (w:wds) =
+      bitmapWordPointerness (min size (fromIntegral wORD_SIZE_IN_BITS)) w ++
+        bitmapWordsPointerness (size - (fromIntegral wORD_SIZE_IN_BITS)) wds
+
+bitmapWordPointerness :: Word -> Word -> [Pointerness]
+bitmapWordPointerness 0 _ = []
+bitmapWordPointerness bSize bitmapWord =
+  ( if (bitmapWord .&. 1) /= 0
+      then NonPointer
+      else Pointer
+  )
+    : bitmapWordPointerness
+      (bSize - 1)
+      (bitmapWord `shiftR` 1)
+
+decodeBitmaps :: StackSnapshot# -> WordOffset -> [Pointerness] -> IO [Closure]
+decodeBitmaps stackSnapshot# index bitmapWords =
+  let bes = toEntries index bitmapWords
+   in do
+    traceM $ "decodeBitmaps - index: " ++ show index ++ " words: " ++ show bitmapWords
+    mapM toBitmapPayload bes
   where
     toBitmapPayload :: StackFrameIter -> IO Closure
     toBitmapPayload (SfiPrimitive stack# i) = do
@@ -259,42 +291,14 @@ decodeBitmaps stackSnapshot# index bitmapWords size =
       pure $ UnknownTypeWordSizedPrimitive w
     toBitmapPayload (SfiClosure stack# i) = getClosure stack# i
 
-    wordsToBitmapEntries :: WordOffset -> [Word] -> Word -> [StackFrameIter]
-    wordsToBitmapEntries _ [] 0 = []
-    wordsToBitmapEntries _ [] i = error $ "Invalid state: Empty list, size " ++ show i
-    wordsToBitmapEntries _ l 0 = error $ "Invalid state: Size 0, list " ++ show l
-    wordsToBitmapEntries index' (b : bs) bitmapSize =
-      let entries = toBitmapEntries index' b (min bitmapSize (fromIntegral wORD_SIZE_IN_BITS))
-          mbLastFrame = (listToMaybe . reverse) entries
-       in case mbLastFrame of
-            Just sfi' ->
-              entries
-                ++ wordsToBitmapEntries
-                  (getIndex sfi' + 1)
-                  bs
-                  subtractDecodedBitmapWord
-            _ -> error "This should never happen! Recursion ended not in base case."
-      where
-        subtractDecodedBitmapWord :: Word
-        subtractDecodedBitmapWord =
-          fromIntegral $
-            max 0 (fromIntegral bitmapSize - wORD_SIZE_IN_BITS)
-
-        toBitmapEntries :: WordOffset -> Word -> Word -> [StackFrameIter]
-        toBitmapEntries _ _ 0 = []
-        toBitmapEntries i bitmapWord bSize =
-          ( if (bitmapWord .&. 1) /= 0
-              then SfiPrimitive stackSnapshot# i
-              else SfiClosure stackSnapshot# i
-          )
-            : toBitmapEntries
-              (i + 1)
-              (bitmapWord `shiftR` 1)
-              (bSize - 1)
-
-        getIndex :: StackFrameIter -> WordOffset
-        getIndex (SfiClosure _ i) = i
-        getIndex (SfiPrimitive _ i) = i
+    toEntries :: WordOffset -> [Pointerness] -> [StackFrameIter]
+    toEntries _ [] = []
+    toEntries i (p:ps) =
+      let sn = case p of
+                NonPointer -> SfiPrimitive stackSnapshot# i
+                Pointer -> SfiClosure stackSnapshot# i
+      in
+        sn : toEntries (i + 1) ps
 
 decodeSmallBitmap :: SmallBitmapGetter -> StackSnapshot# -> WordOffset -> WordOffset -> IO [Closure]
 decodeSmallBitmap getterFun# stackSnapshot# index relativePayloadOffset =
@@ -302,8 +306,7 @@ decodeSmallBitmap getterFun# stackSnapshot# index relativePayloadOffset =
     (bitmap, size) <- IO $ \s ->
       case getterFun# stackSnapshot# (wordOffsetToWord# index) s of
         (# s1, b#, s# #) -> (# s1, (W# b#, W# s#) #)
-    let bitmapWords = [bitmap | size > 0]
-    decodeBitmaps stackSnapshot# (index + relativePayloadOffset) bitmapWords size
+    decodeBitmaps stackSnapshot# (index + relativePayloadOffset) (bitmapWordPointerness size bitmap)
 
 unpackStackFrame :: StackFrameLocation -> IO StackFrame
 unpackStackFrame (StackSnapshot stackSnapshot#, index) = do
diff --git a/libraries/ghc-heap/cbits/Stack.c b/libraries/ghc-heap/cbits/Stack.c
index 89f1e4e27a..a8a47605ff 100644
--- a/libraries/ghc-heap/cbits/Stack.c
+++ b/libraries/ghc-heap/cbits/Stack.c
@@ -110,52 +110,30 @@ StgWord getBCOLargeBitmapSize(StgClosure *c) {
   return BCO_BITMAP_SIZE(bco);
 }
 
-#define ROUNDUP_BITS_TO_WDS(n)                                                 \
-  (((n) + WORD_SIZE_IN_BITS - 1) / WORD_SIZE_IN_BITS)
-
-// Copied from Cmm.h
-#define SIZEOF_W SIZEOF_VOID_P
-#define WDS(n) ((n)*SIZEOF_W)
-
-static StgArrBytes *largeBitmapToStgArrBytes(Capability *cap,
-                                             StgLargeBitmap *bitmap) {
-  StgWord neededWords = ROUNDUP_BITS_TO_WDS(bitmap->size);
-  StgArrBytes *array =
-      (StgArrBytes *)allocate(cap, sizeofW(StgArrBytes) + neededWords);
-  SET_HDR(array, &stg_ARR_WORDS_info, CCS_SYSTEM);
-  array->bytes = WDS(ROUNDUP_BITS_TO_WDS(bitmap->size));
-
-  for (int i = 0; i < neededWords; i++) {
-    array->payload[i] = bitmap->bitmap[i];
-  }
-
-  return array;
-}
-
-StgArrBytes *getLargeBitmap(Capability *cap, StgClosure *c) {
+StgWord *getLargeBitmap(Capability *cap, StgClosure *c) {
   ASSERT(LOOKS_LIKE_CLOSURE_PTR(c));
   const StgInfoTable *info = get_itbl(c);
   StgLargeBitmap *bitmap = GET_LARGE_BITMAP(info);
 
-  return largeBitmapToStgArrBytes(cap, bitmap);
+  return bitmap->bitmap;
 }
 
-StgArrBytes *getRetFunLargeBitmap(Capability *cap, StgRetFun *ret_fun) {
+StgWord *getRetFunLargeBitmap(Capability *cap, StgRetFun *ret_fun) {
   ASSERT(LOOKS_LIKE_CLOSURE_PTR(ret_fun));
 
   const StgFunInfoTable *fun_info = get_fun_itbl(UNTAG_CLOSURE(ret_fun->fun));
   StgLargeBitmap *bitmap = GET_FUN_LARGE_BITMAP(fun_info);
 
-  return largeBitmapToStgArrBytes(cap, bitmap);
+  return bitmap->bitmap;
 }
 
-StgArrBytes *getBCOLargeBitmap(Capability *cap, StgClosure *c) {
+StgWord *getBCOLargeBitmap(Capability *cap, StgClosure *c) {
   ASSERT(LOOKS_LIKE_CLOSURE_PTR(c));
 
   StgBCO *bco = (StgBCO *)*c->payload;
   StgLargeBitmap *bitmap = BCO_BITMAP(bco);
 
-  return largeBitmapToStgArrBytes(cap, bitmap);
+  return bitmap->bitmap;
 }
 
 StgStack *getUnderflowFrameNextChunk(StgUnderflowFrame *frame) {
diff --git a/libraries/ghc-heap/cbits/Stack.cmm b/libraries/ghc-heap/cbits/Stack.cmm
index 88ae7d656e..16504c73b3 100644
--- a/libraries/ghc-heap/cbits/Stack.cmm
+++ b/libraries/ghc-heap/cbits/Stack.cmm
@@ -77,41 +77,41 @@ getRetFunSmallBitmapzh(P_ stack, W_ offsetWords) {
 
 // getLargeBitmapzh(StgStack* stack, StgWord offsetWords)
 getLargeBitmapzh(P_ stack, W_ offsetWords) {
-  P_ c, stgArrBytes;
+  P_ c, words;
   W_ size;
   c = StgStack_sp(stack) + WDS(offsetWords);
   ASSERT(LOOKS_LIKE_CLOSURE_PTR(c));
 
-  (stgArrBytes) = ccall getLargeBitmap(MyCapability(), c);
+  (words) = ccall getLargeBitmap(MyCapability(), c);
   (size) = ccall getLargeBitmapSize(c);
 
-  return (stgArrBytes, size);
+  return (words, size);
 }
 
 // getBCOLargeBitmapzh(StgStack* stack, StgWord offsetWords)
 getBCOLargeBitmapzh(P_ stack, W_ offsetWords) {
-  P_ c, stgArrBytes;
+  P_ c, words;
   W_ size;
   c = StgStack_sp(stack) + WDS(offsetWords);
   ASSERT(LOOKS_LIKE_CLOSURE_PTR(c));
 
-  (stgArrBytes) = ccall getBCOLargeBitmap(MyCapability(), c);
+  (words) = ccall getBCOLargeBitmap(MyCapability(), c);
   (size) = ccall getBCOLargeBitmapSize(c);
 
-  return (stgArrBytes, size);
+  return (words, size);
 }
 
 // getRetFunLargeBitmapzh(StgStack* stack, StgWord offsetWords)
 getRetFunLargeBitmapzh(P_ stack, W_ offsetWords) {
-  P_ c, stgArrBytes;
+  P_ c, words;
   W_ size;
   c = StgStack_sp(stack) + WDS(offsetWords);
   ASSERT(LOOKS_LIKE_CLOSURE_PTR(c));
 
-  (stgArrBytes) = ccall getRetFunLargeBitmap(MyCapability(), c);
+  (words) = ccall getRetFunLargeBitmap(MyCapability(), c);
   (size) = ccall getRetFunSize(c);
 
-  return (stgArrBytes, size);
+  return (words, size);
 }
 
 // getWordzh(StgStack* stack, StgWord offsetWords)