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 ()