From 4e09e47e72d222e95b2c9cbc901a60f384e5aae7 Mon Sep 17 00:00:00 2001 From: alissa-tung 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 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 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