From 40cc0f86803aea479f2e7ddc741aa7c3ad674a12 Mon Sep 17 00:00:00 2001 From: Oleg Grenrus Date: Fri, 20 Jan 2023 16:16:52 +0200 Subject: [PATCH] Add encodingToStrictByteString This runs Builder to produce Strict ByteString directly, by making a mutable buffer and growing it exponentially. This might be good or bad, better or worse than LBS.toStrict . encodingToLazyByteString. Latter allocates many small chunks, and copies once; encodingToStrictByteString makes a buffer exponentially, but copies data everytime. --- aeson.cabal | 1 + src/Data/Aeson/Encoding.hs | 1 + src/Data/Aeson/Encoding/Internal.hs | 14 +++++ src/Data/Aeson/Internal/StrictBuilder.hs | 77 ++++++++++++++++++++++++ tests/PropUtils.hs | 12 +++- 5 files changed, 104 insertions(+), 1 deletion(-) create mode 100644 src/Data/Aeson/Internal/StrictBuilder.hs diff --git a/aeson.cabal b/aeson.cabal index 87877928d..d535281af 100644 --- a/aeson.cabal +++ b/aeson.cabal @@ -81,6 +81,7 @@ library Data.Aeson.Internal.Functions Data.Aeson.Internal.Prelude Data.Aeson.Internal.Scientific + Data.Aeson.Internal.StrictBuilder Data.Aeson.Internal.Text Data.Aeson.Internal.TH Data.Aeson.Internal.Unescape diff --git a/src/Data/Aeson/Encoding.hs b/src/Data/Aeson/Encoding.hs index e1d17d892..5a0ce0433 100644 --- a/src/Data/Aeson/Encoding.hs +++ b/src/Data/Aeson/Encoding.hs @@ -11,6 +11,7 @@ module Data.Aeson.Encoding Encoding , Encoding' , encodingToLazyByteString + , encodingToStrictByteString , fromEncoding , unsafeToEncoding , Series diff --git a/src/Data/Aeson/Encoding/Internal.hs b/src/Data/Aeson/Encoding/Internal.hs index b8ce355a3..a1b914606 100644 --- a/src/Data/Aeson/Encoding/Internal.hs +++ b/src/Data/Aeson/Encoding/Internal.hs @@ -8,6 +8,7 @@ module Data.Aeson.Encoding.Internal Encoding' (..) , Encoding , encodingToLazyByteString + , encodingToStrictByteString , unsafeToEncoding , retagEncoding , Series (..) @@ -63,6 +64,7 @@ module Data.Aeson.Encoding.Internal import Data.Aeson.Internal.Prelude hiding (empty) import Data.Aeson.Types.Internal (Value, Key) +import Data.Aeson.Internal.StrictBuilder (toStrictByteString) import Data.ByteString.Builder (Builder, char7, toLazyByteString) import Data.ByteString.Short (ShortByteString) import qualified Data.Aeson.Key as Key @@ -70,6 +72,7 @@ import Data.Time (Day, LocalTime, TimeOfDay, ZonedTime) import Data.Time.Calendar.Month.Compat (Month) import Data.Time.Calendar.Quarter.Compat (Quarter) import qualified Data.Aeson.Encoding.Builder as EB +import qualified Data.ByteString as BS import qualified Data.ByteString.Builder as B import qualified Data.ByteString.Lazy as BSL import qualified Data.Text.Lazy as LT @@ -94,10 +97,21 @@ type Encoding = Encoding' Value unsafeToEncoding :: Builder -> Encoding' a unsafeToEncoding = Encoding +-- | Convert 'Encoding' to /lazy/ 'BSL.ByteString'. encodingToLazyByteString :: Encoding' a -> BSL.ByteString encodingToLazyByteString = toLazyByteString . fromEncoding {-# INLINE encodingToLazyByteString #-} +-- | Convert 'Encoding' to /strict/ 'BS.ByteString'. +-- +-- This might or might not be more efficient than @'BSL.toStrict' . 'encodingToLazyByteString'@ +-- +-- @since 2.1.2.0 +-- +encodingToStrictByteString :: Encoding' a -> BS.ByteString +encodingToStrictByteString = toStrictByteString . fromEncoding +{-# INLINE encodingToStrictByteString #-} + retagEncoding :: Encoding' a -> Encoding' b retagEncoding = Encoding . fromEncoding diff --git a/src/Data/Aeson/Internal/StrictBuilder.hs b/src/Data/Aeson/Internal/StrictBuilder.hs new file mode 100644 index 000000000..0935c4292 --- /dev/null +++ b/src/Data/Aeson/Internal/StrictBuilder.hs @@ -0,0 +1,77 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE UnboxedTuples #-} +module Data.Aeson.Internal.StrictBuilder ( + toStrictByteString, + toStrictByteStringIO, +) where + +import Data.ByteString.Builder.Internal (BufferRange (BufferRange), BuildStep, Builder, fillWithBuildStep, runBuilder) +import Data.ByteString.Internal (ByteString (..)) +import Data.Word (Word8) +import GHC.Exts (Addr#, Ptr (..), minusAddr#, plusAddr#) +import GHC.Exts (Int (I#), Int#, orI#, (+#)) +import GHC.Exts (MutableByteArray#, RealWorld, newPinnedByteArray#, resizeMutableByteArray#, shrinkMutableByteArray#) +import GHC.ForeignPtr (ForeignPtr (ForeignPtr), ForeignPtrContents (PlainPtr)) +import GHC.IO (IO (IO), unIO, unsafePerformIO) + +#if MIN_VERSION_base(4,16,0) +import GHC.Exts (mutableByteArrayContents#) +#else +import GHC.Exts (byteArrayContents#, unsafeCoerce#) + +mutableByteArrayContents# :: MutableByteArray# s -> Addr# +mutableByteArrayContents# mba = byteArrayContents# (unsafeCoerce# mba) +#endif + +toStrictByteString :: Builder -> ByteString +toStrictByteString b = unsafePerformIO (toStrictByteStringIO b) +{-# NOINLINE toStrictByteString #-} + +toStrictByteStringIO :: Builder -> IO ByteString +toStrictByteStringIO b = IO $ \s -> + case newPinnedByteArray# 4096# s of + (# s', mba #) -> case mutableByteArrayContents# mba of + start -> unIO (toStrictByteStringWorker mba 4096# start start (plusAddr# start 4096#) (runBuilder b)) s' + +-- Progressively double the buffer size if it's reported to be full. +-- (convertion to lazy bytestring allocates new buffer chunks). +toStrictByteStringWorker + :: MutableByteArray# RealWorld -- ^ the buffer bytearray + -> Int# -- ^ size of the bytearray + -> Addr# -- ^ beginning of the bytearray + -> Addr# -- ^ current write position + -> Addr# -- ^ end of the bytearray + -> BuildStep () + -> IO ByteString +toStrictByteStringWorker mba size start begin end !curr = + fillWithBuildStep curr kDone kFull kChunk (BufferRange (Ptr begin) (Ptr end)) + where + kDone :: Ptr Word8 -> () -> IO ByteString + kDone (Ptr pos) _ = IO $ \s1 -> + case minusAddr# pos start of { len -> + case shrinkMutableByteArray# mba len s1 of { s2 -> +#if MIN_VERSION_bytestring(0,11,0) + (# s2 , BS (ForeignPtr start (PlainPtr mba)) (I# len) #) +#else + (# s2 , PS (ForeignPtr start (PlainPtr mba)) 0 (I# len) #) +#endif + }} + + kFull :: Ptr Word8 -> Int -> BuildStep () -> IO ByteString + kFull (Ptr pos) (I# nsize) next = IO $ \s1 -> + -- orI# is an approximation of max + case size +# orI# size nsize of { size' -> + case resizeMutableByteArray# mba size' s1 of { (# s2, mba' #) -> + case mutableByteArrayContents# mba' of { start' -> + unIO (toStrictByteStringWorker mba' size' start' (plusAddr# start' (minusAddr# pos start)) (plusAddr# start' size') next) s2 + }}} + + kChunk :: Ptr Word8 -> ByteString -> BuildStep () -> IO ByteString +#if MIN_VERSION_bytestring(0,11,0) + kChunk (Ptr pos) (BS _ 0) next = toStrictByteStringWorker mba size start pos end next +#else + kChunk (Ptr pos) (PS _ _ 0) next = toStrictByteStringWorker mba size start pos end next +#endif + kChunk _ _ _ = fail "TODO: non-empty chunk" diff --git a/tests/PropUtils.hs b/tests/PropUtils.hs index 19b54fe7f..c3c1fffbb 100644 --- a/tests/PropUtils.hs +++ b/tests/PropUtils.hs @@ -30,7 +30,7 @@ module PropUtils ( import Prelude.Compat import Data.Aeson (eitherDecode, encode) -import Data.Aeson.Encoding (encodingToLazyByteString) +import Data.Aeson.Encoding (encodingToLazyByteString, encodingToStrictByteString) import Data.Aeson.Types import qualified Data.Aeson.Key as Key import qualified Data.Aeson.KeyMap as KM @@ -42,6 +42,7 @@ import Instances () import Test.QuickCheck (Arbitrary(..), Property, Testable, (===), (.&&.), counterexample, property) import Types import Text.Read (readMaybe) +import qualified Data.Attoparsec.ByteString as S import qualified Data.ByteString.Lazy.Char8 as L import qualified Data.Vector as V import qualified Data.Aeson.Decoding as Dec @@ -89,6 +90,14 @@ roundTripDecEnc eq i = Right v -> v `eq` i Left err -> failure "parse" err i +roundTripStrictEnc :: (FromJSON a, ToJSON a, Show a) => + (a -> a -> Property) -> a -> a -> Property +roundTripStrictEnc eq _ i = + case fmap ifromJSON . S.parseOnly value . encodingToStrictByteString . toEncoding $ i of + Right (ISuccess v) -> v `eq` i + Right (IError path err) -> failure "fromJSON" (formatError path err) i + Left err -> failure "parse" err i + roundTripNoEnc :: (FromJSON a, ToJSON a, Show a) => (a -> a -> Property) -> a -> Property roundTripNoEnc eq i = @@ -112,6 +121,7 @@ roundTripEq y = roundTripEnc (===) y .&&. roundTripNoEnc (===) y .&&. roundTripDecEnc (===) y .&&. + roundTripStrictEnc (===) x y .&&. roundTripOmit (===) y roundtripReadShow :: Value -> Property