From b0823d03c3ceb6f32edfc705298a9b42e7da0450 Mon Sep 17 00:00:00 2001
From: Andrew Martin <andrew.thaddeus@gmail.com>
Date: Fri, 2 Aug 2019 16:08:13 -0400
Subject: [PATCH] make pasteGrowST accept an offset. correct the implementation
 of bytes

---
 src/Data/ByteArray/Builder/Small.hs        | 14 +++++----
 src/Data/ByteArray/Builder/Small/Unsafe.hs | 33 ++++++++++++++++++++++
 2 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/src/Data/ByteArray/Builder/Small.hs b/src/Data/ByteArray/Builder/Small.hs
index ffc367c..8762826 100644
--- a/src/Data/ByteArray/Builder/Small.hs
+++ b/src/Data/ByteArray/Builder/Small.hs
@@ -55,6 +55,8 @@ import qualified Data.ByteArray.Builder.Small.Unsafe as Unsafe
 -- | An unmaterialized sequence of bytes that may be pasted
 -- into a mutable byte array.
 newtype Builder = Builder
+  -- This functions takes an offset and a number of remaining bytes
+  -- and returns the new offset.
   (forall s. MutableByteArray# s -> Int# -> Int# -> State# s -> (# State# s, Int# #))
 
 instance Semigroup Builder where
@@ -115,15 +117,15 @@ pasteArrayST (MutableBytes arr off0 len0) f !xs0 = do
 pasteGrowST ::
      Int -- ^ How many bytes to grow by at a time
   -> Builder
-  -> MutableByteArray s
+  -> MutableByteArrayOffset s
      -- ^ Initial buffer, used linearly. Do not reuse this argument.
   -> ST s (MutableByteArrayOffset s)
      -- ^ Final buffer that accomodated the builder.
-pasteGrowST !n b !arr0 = do
-  let go !arr !sz = pasteST b (MutableBytes arr 0 sz) >>= \case
+pasteGrowST !n b !(MutableByteArrayOffset arr0 off0) = do
+  let go !arr !sz = pasteST b (MutableBytes arr off0 (sz - off0)) >>= \case
         Nothing -> do
           let szNext = sz + n
-          arrNext <- PM.newByteArray szNext
+          arrNext <- PM.resizeMutableByteArray arr szNext
           go arrNext szNext
         Just ix -> pure (MutableByteArrayOffset{array=arr,offset=ix})
   go arr0 =<< PM.getSizeofMutableByteArray arr0
@@ -132,7 +134,7 @@ pasteGrowST !n b !arr0 = do
 pasteGrowIO ::
      Int -- ^ How many bytes to grow by at a time
   -> Builder
-  -> MutableByteArray RealWorld
+  -> MutableByteArrayOffset RealWorld
      -- ^ Initial buffer, used linearly. Do not reuse this argument.
   -> IO (MutableByteArrayOffset RealWorld)
      -- ^ Final buffer that accomodated the builder.
@@ -183,7 +185,7 @@ bytes :: Bytes -> Builder
 bytes (Bytes src soff slen) = construct $ \(MutableBytes arr off len) -> if len >= slen
   then do
     copyByteArray arr off src soff slen
-    pure (Just (len - slen))
+    pure (Just (off + slen))
   else pure Nothing
 
 -- | Encodes an unsigned 64-bit integer as decimal.
diff --git a/src/Data/ByteArray/Builder/Small/Unsafe.hs b/src/Data/ByteArray/Builder/Small/Unsafe.hs
index a1a5f0c..f4feabc 100644
--- a/src/Data/ByteArray/Builder/Small/Unsafe.hs
+++ b/src/Data/ByteArray/Builder/Small/Unsafe.hs
@@ -19,6 +19,7 @@ module Data.ByteArray.Builder.Small.Unsafe
     -- * Execute
   , run
   , pasteST
+  , pasteGrowST
   , pasteIO
     -- * Combine
   , append
@@ -34,6 +35,7 @@ module Data.ByteArray.Builder.Small.Unsafe
   , word64BE
   , word32BE
   , word16BE
+  , word8
   ) where
 
 import Control.Monad.Primitive
@@ -47,6 +49,9 @@ import GHC.Word
 import GHC.Int
 import Data.Kind
 import GHC.TypeLits (KnownNat,Nat,type (+),natVal')
+import Data.Primitive.ByteArray.Offset (MutableByteArrayOffset(..))
+
+import qualified Data.Primitive as PM
 
 -- | A builder parameterized by the maximum number of bytes it uses
 -- when executed.
@@ -74,6 +79,29 @@ pasteST (Builder f) (MutableByteArray arr) (I# off) =
   ST $ \s0 -> case f arr off s0 of
     (# s1, r #) -> (# s1, (I# r) #)
 
+-- | Paste the builder into the byte array starting at offset zero.
+-- This reallocates the byte array if it cannot accomodate the builder,
+-- growing it by the minimum amount necessary.
+pasteGrowST :: forall n s. KnownNat n
+  => Builder n
+  -> MutableByteArrayOffset s
+     -- ^ Initial buffer, used linearly. Do not reuse this argument.
+  -> ST s (MutableByteArrayOffset s)
+     -- ^ Final buffer that accomodated the builder.
+{-# inline pasteGrowST #-}
+pasteGrowST b !(MutableByteArrayOffset{array=arr0,offset=off0}) = do
+  sz0 <- PM.getSizeofMutableByteArray arr0
+  let req = fromIntegral (natVal' (proxy# :: Proxy# n))
+  let sz1 = off0 + req
+  if sz1 <= sz0
+    then do
+      off1 <- pasteST b arr0 off0
+      pure (MutableByteArrayOffset arr0 off1)
+    else do
+      arr1 <- PM.resizeMutableByteArray arr0 sz1
+      off1 <- pasteST b arr1 off0
+      pure (MutableByteArrayOffset arr1 off1)
+
 -- | This function does not enforce the known upper bound on the
 -- size. It is up to the user to do this.
 pasteIO :: Builder n -> MutableByteArray RealWorld -> Int -> IO Int
@@ -283,6 +311,11 @@ word16BE w = construct $ \arr off -> do
   writeByteArray arr (off + 1) (fromIntegral @Word16 @Word8 w)
   pure (off + 2)
 
+word8 :: Word8 -> Builder 1
+word8 w = construct $ \arr off -> do
+  writeByteArray arr off w
+  pure (off + 1)
+
 -- Reverse the bytes in the designated slice. This takes
 -- an inclusive start offset and an inclusive end offset.
 reverseBytes :: MutableByteArray s -> Int -> Int -> ST s ()