Introduce toByteString API

src/Proto3/Wire/Encode.hs

@@ -38,17 +38,16 @@
 -- > 1 `strings` Just "some string" <>
 -- > 2 `strings` [ "foo", "bar", "baz" ]
 
-{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE FlexibleInstances #-}
-{-# LANGUAGE FunctionalDependencies #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE PolyKinds #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeFamilies #-}
 {-# LANGUAGE TypeOperators #-}
 {-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
 
 module Proto3.Wire.Encode
     ( -- * `MessageBuilder` type
@@ -57,6 +56,7 @@ module Proto3.Wire.Encode
     , vectorMessageBuilder
     , messageLength
     , toLazyByteString
+    , toByteString
     , unsafeFromLazyByteString
 
       -- * Standard Integers
@@ -172,6 +172,9 @@ messageLength = fromIntegral . fst . RB.runBuildR . unMessageBuilder
 toLazyByteString :: MessageBuilder -> BL.ByteString
 toLazyByteString = RB.toLazyByteString . unMessageBuilder
 
+toByteString :: MessageBuilder -> B.ByteString
+toByteString = RB.toByteString . unMessageBuilder
+
 -- | This lets you cast an arbitrary `ByteString` to a `MessageBuilder`, whether
 -- or not the `ByteString` corresponds to a valid serialized protobuf message
 --

src/Proto3/Wire/Reverse.hs

@@ -74,6 +74,7 @@ module Proto3.Wire.Reverse
     -- * Consume `BuildR`s
     , runBuildR
     , toLazyByteString
+    , toByteString
 
     -- * Helpful combinators
     , foldlRVector
@@ -117,6 +118,9 @@ import qualified Proto3.Wire.Reverse.Prim      as Prim
 toLazyByteString :: BuildR -> BL.ByteString
 toLazyByteString = snd . runBuildR
 
+toByteString :: BuildR -> B.ByteString
+toByteString = snd . runBuildRStrict
+
 -- | Convert a strict `B.ByteString` to a `BuildR`
 --
 -- > byteString (x <> y) = byteString x <> byteString y

src/Proto3/Wire/Reverse/Internal.hs

@@ -30,6 +30,7 @@ module Proto3.Wire.Reverse.Internal
     , fromBuildR
     , etaBuildR
     , runBuildR
+    , runBuildRStrict
     , withUnused
     , withTotal
     , withLengthOf
@@ -221,6 +222,11 @@ data BuildRState = BuildRState
       -- But to avoid redundant evaluation we do not mark it strict.
   }
 
+data BuildRStateStrict = BuildRStateStrict {
+  currentBuffer' :: {-# UNPACK #-}!(P.MutableByteArray RealWorld),
+  sealedBuffers' :: B.ByteString
+}
+
 -- | Allocates fields backward from offset 0 relative to some hypothetical
 -- address, yielding the total size and alignment requirements, respectively,
 -- along with the monadic return value.  The total size includes any padding
@@ -278,9 +284,15 @@ metaPtr v = plusPtr v . negate
 readState :: Ptr MetaData -> IO (StablePtr (IORef BuildRState))
 readState m = castPtrToStablePtr <$> peekByteOff m stateOffset
 
+readStateStrict :: Ptr MetaData -> IO (StablePtr (IORef BuildRStateStrict))
+readStateStrict m = castPtrToStablePtr <$> peekByteOff m stateOffset
+
 writeState :: Ptr MetaData -> StablePtr (IORef BuildRState) -> IO ()
 writeState m = pokeByteOff m stateOffset . castStablePtrToPtr
 
+writeState' :: Ptr MetaData -> StablePtr (IORef BuildRStateStrict) -> IO ()
+writeState' m = pokeByteOff m stateOffset . castStablePtrToPtr
+
 readSpace :: Ptr MetaData -> IO Int
 readSpace m = peekByteOff m spaceOffset
 
@@ -350,6 +362,14 @@ newBuffer sealed (I# total) (IORef (STRef stateVar)) (StablePtr stateSP)
     case newBuffer# sealed total stateVar stateSP unused s0 of
       (# s1, addr #) -> (# s1, Ptr addr #)
 
+newBufferStrict :: B.ByteString -> Int -> IORef BuildRStateStrict
+                -> StablePtr (IORef BuildRStateStrict) -> Int
+                -> IO (Ptr Word8)
+newBufferStrict sealed (I# total) (IORef (STRef stateVar)) (StablePtr stateSP) (I# unused) = do
+  IO $ \s0 ->
+    case newBufferStrict# sealed total stateVar stateSP unused s0 of
+      (# s1, addr #) -> (# s1, Ptr addr #)
+
 newBuffer# ::
   BL.ByteString ->
   Int# ->
@@ -374,6 +394,30 @@ newBuffer# sealed total stateVar stateSP unused s0 =
       writeIORef (IORef (STRef stateVar)) nextState
       pure v
 
+newBufferStrict# ::
+  B.ByteString ->
+  Int# ->
+  MutVar# RealWorld BuildRStateStrict ->
+  StablePtr# (IORef BuildRStateStrict) ->
+  Int# ->
+  State# RealWorld ->
+  (# State# RealWorld, Addr# #)
+newBufferStrict# sealed total stateVar stateSP unused s0 =
+    case go s0 of
+      (# s1, Ptr addr #) -> (# s1, addr #)
+  where
+    IO go = do
+      let allocation = metaDataSize + I# unused
+      buf <- P.newAlignedPinnedByteArray allocation metaDataAlign
+      let !(PTR base) = P.mutableByteArrayContents buf
+          !v = plusPtr (Ptr base) (metaDataSize + I# unused)
+          !m = plusPtr (Ptr base) metaDataSize
+      writeState' m (StablePtr stateSP)
+      writeSpace m (I# unused + I# total)
+      let !nextState = BuildRStateStrict{currentBuffer' = buf, sealedBuffers' = sealed}
+      writeIORef (IORef (STRef stateVar)) nextState
+      pure v
+
 -- | The result of a call to 'sealBuffer'.
 data SealedState = SealedState
   { sealedSB :: BL.ByteString
@@ -395,6 +439,14 @@ data SealedState = SealedState
       -- but please be sure to update the "space" metadatum.
   }
 
+data SealedStateStrict = SealedStateStrict
+  { sealedSB'   :: B.ByteString
+  , totalSB'    :: {-# UNPACK #-}!Int
+  , stateVarSB' :: {-# UNPACK #-}!(IORef BuildRStateStrict)
+  , statePtrSB' :: {-# UNPACK #-}!(StablePtr (IORef BuildRStateStrict))
+  , recycledSB' :: Maybe (P.MutableByteArray RealWorld)
+  }
+
 -- | Takes ownership of the current buffer,
 -- but sometimes hands it back for reuse.
 --
@@ -421,6 +473,20 @@ sealBuffer (Ptr addr) (I# u) = IO $ \s0 ->
            }
        #)
 
+sealBufferStrict :: Ptr Word8 -> Int -> IO SealedStateStrict
+sealBufferStrict (Ptr addr) (I# u) = IO $ \s0 ->
+  case sealBufferStrict# addr u s0 of
+    (# s1, sealed, total, stateVar, statePtr, recycled #) ->
+      (# s1
+       , SealedStateStrict
+           { sealedSB'   = sealed
+           , totalSB'    = I# total
+           , stateVarSB' = IORef (STRef stateVar)
+           , statePtrSB' = StablePtr statePtr
+           , recycledSB' = recycled
+           }
+       #)
+
 sealBuffer# ::
   Addr# ->
   Int# ->
@@ -468,6 +534,46 @@ sealBuffer# addr unused s0 =
             then finish untrimmed Nothing
             else finish (B.copy untrimmed) (Just buffer)
 
+sealBufferStrict# ::
+  Addr# ->
+  Int# ->
+  State# RealWorld ->
+  (# State# RealWorld
+   , B.ByteString
+   , Int#
+   , MutVar# RealWorld BuildRStateStrict
+   , StablePtr# (IORef BuildRStateStrict)
+   , Maybe (P.MutableByteArray RealWorld)
+   #)
+sealBufferStrict# addr unused s0 =
+    case go s0 of
+      (# s1, (sealed, I# total, IORef (STRef sv), StablePtr sp, re) #) ->
+        (# s1, sealed, total, sv, sp, re #)
+  where
+    IO go = do
+      let v = Ptr addr
+      statePtr <- readStateStrict (metaPtr v (I# unused))
+      stateVar <- deRefStablePtr statePtr
+      BuildRStateStrict { currentBuffer' = buffer, sealedBuffers' = oldSealed } <-
+       readIORef stateVar
+      total <- readTotal v (I# unused)
+      let allocation = P.sizeofMutableByteArray buffer - metaDataSize
+      if allocation <= I# unused
+        then
+          pure (oldSealed, total, stateVar, statePtr, Just buffer)
+        else do
+          let !(PTR base) = P.mutableByteArrayContents buffer
+              !(P.MutableByteArray mba) = buffer
+              fp = ForeignPtr base (PlainPtr mba)
+              offset = metaDataSize + I# unused
+              finish trimmed recycled = do
+                let !newSealed = B.append trimmed oldSealed
+                pure (newSealed, total, stateVar, statePtr, recycled)
+              untrimmed = BI.fromForeignPtr fp offset (allocation - I# unused)
+          if offset <= B.length untrimmed
+            then finish untrimmed Nothing
+            else finish (B.copy untrimmed) (Just buffer)
+
 -- | Like `Proto3.Wire.Reverse.toLazyByteString` but also
 -- returns the total length of the lazy 'BL.ByteString',
 -- which is computed as a side effect of encoding.
@@ -481,13 +587,23 @@ runBuildR f = unsafePerformIO $ do
     SealedState { sealedSB = bytes, totalSB = total } <- sealBuffer v1 u1
     pure (total, bytes)
 
+runBuildRStrict :: BuildR -> (Int, B.ByteString)
+runBuildRStrict f = unsafePerformIO $ do
+  stateVar <- newIORef undefined
+  bracket (newStablePtr stateVar) freeStablePtr $ \statePtr -> do
+    let u0 = smallChunkSize
+    v0 <- newBufferStrict B.empty 0 stateVar statePtr u0
+    (v1, u1) <- fromBuildR f v0 u0
+    SealedStateStrict { sealedSB' = bytes, totalSB' = total } <- sealBufferStrict v1 u1
+    pure (total, bytes)
+
 -- | First reads the number of unused bytes in the current buffer.
 withUnused :: (Int -> BuildR) -> BuildR
 withUnused f = toBuildR $ \v u -> fromBuildR (f u) v u
 
 -- | First reads the number of bytes previously written.
 withTotal :: (Int -> BuildR) -> BuildR
-withTotal f = withTotal# (\total -> f (I# total))
+withTotal f = withTotal# $ \total -> f (I# total)
 
 -- | First reads the number of bytes previously written.
 withTotal# :: (Int# -> BuildR) -> BuildR
@@ -681,7 +797,7 @@ prependReverseChunks# v0 u0 s0 ad ct off len cs0 = go v0 u0 s0
 -- | Ensures that the current buffer has at least the given
 -- number of bytes before executing the given builder.
 ensure :: Int -> BuildR -> BuildR
-ensure (I# required) f = ensure# required f
+ensure (I# required) = ensure# required
 
 ensure# :: Int# -> BuildR -> BuildR
 ensure# required (BuildR f) = BuildR $ \v u s ->