From 4e09e47e72d222e95b2c9cbc901a60f384e5aae7 Mon Sep 17 00:00:00 2001
From: alissa-tung <alissa-tung@outlook.com>
Date: Tue, 15 Mar 2022 11:43:46 +0800
Subject: [PATCH 1/3] Add strict wrapper for `toLazyByteString`
---
src/Proto3/Wire/Encode.hs | 7 +++++--
src/Proto3/Wire/Reverse.hs | 4 ++++
src/Proto3/Wire/Reverse/Internal.hs | 4 ++++
3 files changed, 13 insertions(+), 2 deletions(-)
diff --git a/src/Proto3/Wire/Encode.hs b/src/Proto3/Wire/Encode.hs
index 7104fc3..335b4c0 100644
--- a/src/Proto3/Wire/Encode.hs
+++ b/src/Proto3/Wire/Encode.hs
@@ -38,10 +38,8 @@
-- > 1 `strings` Just "some string" <>
-- > 2 `strings` [ "foo", "bar", "baz" ]
-{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
-{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
@@ -49,6 +47,7 @@
{-# 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
--
diff --git a/src/Proto3/Wire/Reverse.hs b/src/Proto3/Wire/Reverse.hs
index dcf87cf..ccd4c29 100644
--- a/src/Proto3/Wire/Reverse.hs
+++ b/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
diff --git a/src/Proto3/Wire/Reverse/Internal.hs b/src/Proto3/Wire/Reverse/Internal.hs
index 0600d7b..8dac918 100644
--- a/src/Proto3/Wire/Reverse/Internal.hs
+++ b/src/Proto3/Wire/Reverse/Internal.hs
@@ -30,6 +30,7 @@ module Proto3.Wire.Reverse.Internal
, fromBuildR
, etaBuildR
, runBuildR
+ , runBuildRStrict
, withUnused
, withTotal
, withLengthOf
@@ -481,6 +482,9 @@ runBuildR f = unsafePerformIO $ do
SealedState { sealedSB = bytes, totalSB = total } <- sealBuffer v1 u1
pure (total, bytes)
+runBuildRStrict :: BuildR -> (Int, B.ByteString)
+runBuildRStrict f = (\(x, y) -> (x, BL.toStrict y)) (runBuildR f)
+
-- | 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
From 373e2d6d3707bbe2ca4c5f0424591ad7e07bf122 Mon Sep 17 00:00:00 2001
From: alissa-tung <alissa-tung@outlook.com>
Date: Wed, 16 Mar 2022 14:10:28 +0800
Subject: [PATCH 2/3] conv strict
---
src/Proto3/Wire/Reverse/Internal.hs | 53 +++++++++++++++++++++++++++--
1 file changed, 50 insertions(+), 3 deletions(-)
diff --git a/src/Proto3/Wire/Reverse/Internal.hs b/src/Proto3/Wire/Reverse/Internal.hs
index 8dac918..70ab5a4 100644
--- a/src/Proto3/Wire/Reverse/Internal.hs
+++ b/src/Proto3/Wire/Reverse/Internal.hs
@@ -222,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
@@ -282,6 +287,9 @@ readState 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
@@ -351,6 +359,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# ->
@@ -375,6 +391,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
@@ -483,7 +523,14 @@ runBuildR f = unsafePerformIO $ do
pure (total, bytes)
runBuildRStrict :: BuildR -> (Int, B.ByteString)
-runBuildRStrict f = (\(x, y) -> (x, BL.toStrict y)) (runBuildR f)
+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
+ SealedState { sealedSB = bytes, totalSB = total } <- sealBuffer v1 u1
+ pure (total, BL.toStrict bytes)
-- | First reads the number of unused bytes in the current buffer.
withUnused :: (Int -> BuildR) -> BuildR
@@ -491,7 +538,7 @@ 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
@@ -685,7 +732,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 ->
From 19caaafb9b6310137224a1078fea799bd2f1ddd6 Mon Sep 17 00:00:00 2001
From: alissa-tung <alissa-tung@outlook.com>
Date: Wed, 16 Mar 2022 15:51:34 +0800
Subject: [PATCH 3/3] Remove type cast
---
src/Proto3/Wire/Reverse/Internal.hs | 69 ++++++++++++++++++++++++++++-
1 file changed, 67 insertions(+), 2 deletions(-)
diff --git a/src/Proto3/Wire/Reverse/Internal.hs b/src/Proto3/Wire/Reverse/Internal.hs
index 70ab5a4..ecded5f 100644
--- a/src/Proto3/Wire/Reverse/Internal.hs
+++ b/src/Proto3/Wire/Reverse/Internal.hs
@@ -284,6 +284,9 @@ 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
@@ -436,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.
--
@@ -462,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# ->
@@ -509,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.
@@ -529,8 +594,8 @@ runBuildRStrict f = unsafePerformIO $ do
let u0 = smallChunkSize
v0 <- newBufferStrict B.empty 0 stateVar statePtr u0
(v1, u1) <- fromBuildR f v0 u0
- SealedState { sealedSB = bytes, totalSB = total } <- sealBuffer v1 u1
- pure (total, BL.toStrict bytes)
+ 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