diff --git a/BENCHMARKS.md b/BENCHMARKS.md index 474bf7c..b2f08e1 100644 --- a/BENCHMARKS.md +++ b/BENCHMARKS.md @@ -31,18 +31,18 @@ This runs with 2 concurrent queries, 10 times over: This benchmark is unfair towards both hpgsql and postgresql-simple because the row decoder is Generically derived for them while it is hand-written for hasql. ```csv -postgresql-simple Record List (100000 rows),12.23,142.52M,91.3 -hasql Record List (100000 rows),6.119,142.48M,78.5 -hpgsql Record List (100000 rows),3.719,72.07M,120.2 +postgresql-simple Record List (100000 rows),11.69,143.04M,91.9 +hasql Record List (100000 rows),6.086,142.48M,80.2 +hpgsql Record List (100000 rows),3.548,72.07M,119.9 ``` ### Materializing 100_000 rows with 13 columns each into a List of Tuples This runs with 2 concurrent queries, 10 times over: ```csv -postgresql-simple Tuple List (100000 rows),14.37,142.54M,137.7 -hasql Tuple List (100000 rows),8.305,142.48M,201.3 -hpgsql Tuple List (100000 rows),4.541,72.07M,150.5 +postgresql-simple Tuple List (100000 rows),14.31,143.18M,144.9 +hasql Tuple List (100000 rows),8.769,142.48M,203.9 +hpgsql Tuple List (100000 rows),4.256,72.07M,150.4 ``` ### Streaming 100_000 rows with 13 columns as Records @@ -55,9 +55,9 @@ However, Hpgsql's implementation streams directly from the socket while the othe it might not be a fair comparison in terms of implementation (e.g. you can advance multiple cursors simultaneously, but not hpgsql's Streamed-from-socket streams). ```csv -streaming-postgresql-simple Record Stream (100000 rows),13.05,73.30M,0.0 -postgresql-simple Record fold (100000 rows),12.23,76.22M,0.0 -hpgsql Record Stream (100000 rows),1.117,72.07M,0.0 +streaming-postgresql-simple Record Stream (100000 rows),13.15,73.32M,0.0 +postgresql-simple Record fold (100000 rows),12.58,76.49M,0.0 +hpgsql Record Stream (100000 rows),1.125,72.07M,0.0 ``` ### Streaming 100_000 rows with 13 columns as Tuples @@ -67,9 +67,9 @@ Hpgsql's implementation streams directly from the socket while the others use cu it might not be a fair comparison in terms of implementation (e.g. you can advance multiple cursors simultaneously, but not hpgsql's Streamed-from-socket streams). ```csv -streaming-postgresql-simple Tuple Stream (100000 rows),14.04,73.37M,0.0 -postgresql-simple Tuple fold (100000 rows),13.53,82.45M,0.0 -hpgsql Tuple Stream (100000 rows),880.0,72.07M,0.0 +streaming-postgresql-simple Tuple Stream (100000 rows),13.75,73.23M,0.0 +postgresql-simple Tuple fold (100000 rows),12.60,79.92M,0.0 +hpgsql Tuple Stream (100000 rows),799.7,72.07M,0.0 ``` ### COPY FROM STDIN @@ -77,6 +77,6 @@ hpgsql Tuple Stream (100000 rows),880.0,72.07M,0.0 This compares hpgsql's binary copy to a `forM` loop writing text rows. ```csv -postgresql-simple text COPY (100000 rows),1.373,72.10M,3.8 -hpgsql copyFromS binary COPY (100000 rows),652.0,72.07M,10.8 +postgresql-simple text COPY (100000 rows),1.324,72.09M,3.8 +hpgsql copyFromS binary COPY (100000 rows),631.8,72.07M,10.8 ``` diff --git a/hpgsql-tests/RowDecoderGhcCore.hs b/hpgsql-tests/RowDecoderGhcCore.hs index e53d8a5..e9f272b 100644 --- a/hpgsql-tests/RowDecoderGhcCore.hs +++ b/hpgsql-tests/RowDecoderGhcCore.hs @@ -1,38 +1,60 @@ -{-# OPTIONS_GHC -ddump-simpl -ddump-to-file #-} +{-# OPTIONS_GHC -ddump-simpl -dno-typeable-binds -dsuppress-coercions -dsuppress-module-prefixes -dsuppress-type-applications -ddump-to-file #-} -- | -- This is not a real test module. It's just a type deriving `FromPgRow` --- so we can look at GHC Core output. +-- so we can look at GHC Core output. It's as small as we can make it to +-- facilitate reading GHC Core. module RowDecoderGhcCore where import Data.Int (Int64) import Data.Text (Text) import Data.Time (Day, UTCTime) import GHC.Generics (Generic) -import Hpgsql.Encoding (FromPgRow (..), fieldDecoder, genericFromPgRow, singleField) - -data BenchRow = BenchRow - { brId :: !Int, - brDate1 :: !Day, - brDate2 :: !Day, - brTimestamp1 :: !UTCTime, - brTimestamp2 :: !UTCTime, - brText1 :: !Text, - brText2 :: !Text, - brDouble1 :: !Double, - brDouble2 :: !Double, - brMaybeInt :: !(Maybe Int), - brMaybeText :: !(Maybe Text), - brMaybeDouble :: !(Maybe Double), - brMaybeDay :: !(Maybe Day) +import Hpgsql.Encoding (FromPgField (..), FromPgRow (..), genericFromPgRow, singleField) + +-- | BestCaseScenarioRecord's purpose is to have a very small row decoder in GHC Core +-- for my own understanding/comprehension of what a RowDecoder gets compiled to +-- in the best case scenario. Also, we expect one day to maybe reach a fully inlined +-- row decoder that only peeks at bytes and allocates 3 values per row (one +-- for each field), plus one `BestCaseScenarioRecord` per row. +-- In the GHC Core of this module (use `run ghc-core` to output it), it helps to: +-- - Look for the Record constructor and grep for it to find where the RowDecoder +-- invokes it, only to find where the RowDecoder is. +-- - Grep for numbers that exist in the decoders' implementation, such as 8#, 13#, 4#. +-- These are strong indicators that each decoder was inlined into the RowDecoder. +-- There still are unnecessary allocations/boxing even with full inlining, but maybe +-- one day we'll find a way to get rid of all of them. +data BestCaseScenarioRecord = BestCaseScenarioRecord + { bcsId :: !Int, + bcsDate :: !Day, + bcsText :: !(Maybe Int) } +instance FromPgRow BestCaseScenarioRecord where + rowDecoder = BestCaseScenarioRecord <$> singleField fieldDecoder <*> singleField fieldDecoder <*> singleField fieldDecoder + +-- data BenchRow = BenchRow +-- { brId :: !Int, +-- brDate1 :: !Day, +-- brDate2 :: !Day, +-- brTimestamp1 :: !UTCTime, +-- brTimestamp2 :: !UTCTime, +-- brText1 :: !Text, +-- brText2 :: !Text, +-- brDouble1 :: !Double, +-- brDouble2 :: !Double, +-- brMaybeInt :: !(Maybe Int), +-- brMaybeText :: !(Maybe Text), +-- brMaybeDouble :: !(Maybe Double), +-- brMaybeDay :: !(Maybe Day) +-- } + -- Generically deriving section. -deriving instance Generic BenchRow +-- deriving instance Generic BenchRow -instance FromPgRow BenchRow where - rowDecoder = genericFromPgRow +-- instance FromPgRow BenchRow where +-- rowDecoder = genericFromPgRow -- Hand-written applicative style deriving section. -- instance FromPgRow BenchRow where diff --git a/hpgsql/src/Hpgsql/Encoding.hs b/hpgsql/src/Hpgsql/Encoding.hs index 6a1e771..dd5378f 100644 --- a/hpgsql/src/Hpgsql/Encoding.hs +++ b/hpgsql/src/Hpgsql/Encoding.hs @@ -735,9 +735,9 @@ binaryIntDecoder typOid = \bs -> maxBoundPgType :: Integer intDecoder :: ByteString -> Either String a (maxBoundPgType, intDecoder) - | typOid == int8Oid = (fromIntegral $ maxBound @Int64, fmap fromIntegral . BinSer.decodeInt64BE) - | typOid == int4Oid = (fromIntegral $ maxBound @Int32, fmap fromIntegral . BinSer.decodeInt32BE) - | typOid == int2Oid = (fromIntegral $ maxBound @Int16, fmap fromIntegral . BinSer.decodeInt16BE) + | typOid == int8Oid = (fromIntegral $ maxBound @Int64, fmap fromIntegral . BinSer.decodeInt64BE 0) + | typOid == int4Oid = (fromIntegral $ maxBound @Int32, fmap fromIntegral . BinSer.decodeInt32BE 0) + | typOid == int2Oid = (fromIntegral $ maxBound @Int16, fmap fromIntegral . BinSer.decodeInt16BE 0) | otherwise = error "Bug in Hpgsql. Decoding binary integral type not an int2, int4 or int8" doesFit = maxBoundPgType <= fromIntegral (maxBound @a) @@ -991,7 +991,7 @@ instance FromPgField UTCTime where fieldDecoder = parsePgType [timestamptzOid] $ \case Just bs -> do -- See https://github.com/postgres/postgres/blob/50cb7505b3010736b9a7922e903931534785f3aa/src/backend/utils/adt/timestamp.c#L1909 - totalusecs <- BinSer.decodeInt64BE bs + totalusecs <- BinSer.decodeInt64BE 0 bs let (day, timeusecs) = totalusecs `divMod` 86_400_000_000 -- USECS per day parsedDate = addJulianDurationClip (CalendarDiffDays 0 (fromIntegral day)) $ fromJulian 1999 12 19 Right $ UTCTime parsedDate (picosecondsToDiffTime $ fromIntegral timeusecs * 1_000_000) @@ -1001,7 +1001,7 @@ instance FromPgField (Unbounded UTCTime) where fieldDecoder = parsePgType [timestamptzOid] $ \case Just bs -> do -- See https://github.com/postgres/postgres/blob/50cb7505b3010736b9a7922e903931534785f3aa/src/backend/utils/adt/timestamp.c#L1909 - totalusecs <- BinSer.decodeInt64BE bs + totalusecs <- BinSer.decodeInt64BE 0 bs Right $ if totalusecs == minBound then NegInfinity @@ -1018,7 +1018,7 @@ instance FromPgField ZonedTime where fieldDecoder = parsePgType [timestamptzOid] $ \case Just bs -> do -- See https://github.com/postgres/postgres/blob/50cb7505b3010736b9a7922e903931534785f3aa/src/backend/utils/adt/timestamp.c#L1909 - totalusecs <- BinSer.decodeInt64BE bs + totalusecs <- BinSer.decodeInt64BE 0 bs let (day, timeusecs) = totalusecs `divMod` 86_400_000_000 -- USECS per day parsedDate = addJulianDurationClip (CalendarDiffDays 0 (fromIntegral day)) $ fromJulian 1999 12 19 Right $ utcToZonedTime utc $ UTCTime parsedDate (picosecondsToDiffTime $ fromIntegral timeusecs * 1_000_000) @@ -1028,7 +1028,7 @@ instance FromPgField (Unbounded ZonedTime) where fieldDecoder = parsePgType [timestamptzOid] $ \case Just bs -> do -- See https://github.com/postgres/postgres/blob/50cb7505b3010736b9a7922e903931534785f3aa/src/backend/utils/adt/timestamp.c#L1909 - totalusecs <- BinSer.decodeInt64BE bs + totalusecs <- BinSer.decodeInt64BE 0 bs Right $ if totalusecs == minBound then NegInfinity @@ -1044,7 +1044,7 @@ instance FromPgField (Unbounded ZonedTime) where instance FromPgField LocalTime where fieldDecoder = parsePgType [timestampOid] $ \case Just bs -> do - totalusecs <- BinSer.decodeInt64BE bs + totalusecs <- BinSer.decodeInt64BE 0 bs let (day, timeusecs) = totalusecs `divMod` 86_400_000_000 -- USECS per day parsedDate = addJulianDurationClip (CalendarDiffDays 0 (fromIntegral day)) $ fromJulian 1999 12 19 Right $ LocalTime parsedDate (timeToTimeOfDay $ picosecondsToDiffTime $ fromIntegral timeusecs * 1_000_000) @@ -1053,7 +1053,7 @@ instance FromPgField LocalTime where instance FromPgField TimeOfDay where fieldDecoder = parsePgType [timeOid] $ \case Just bs -> do - usecs <- BinSer.decodeInt64BE bs + usecs <- BinSer.decodeInt64BE 0 bs Right $ timeToTimeOfDay $ picosecondsToDiffTime $ fromIntegral usecs * 1_000_000 Nothing -> Left "Cannot decode SQL null as the Haskell TimeOfDay type. Use a `Maybe TimeOfDay`" @@ -1063,7 +1063,7 @@ instance FromPgField Day where -- There is a very specific conversion function for these, which I poorly translated to Haskell -- https://github.com/postgres/postgres/blob/799959dc7cf0e2462601bea8d07b6edec3fa0c4f/src/backend/utils/adt/datetime.c#L321 -- But I found a simpler way to do this. Let's see if it works in our property based tests - jd <- BinSer.decodeInt32BE bs + jd <- BinSer.decodeInt32BE 0 bs Right $ addJulianDurationClip (CalendarDiffDays 0 (fromIntegral jd - 13)) $ fromJulian 2000 01 01 Nothing -> Left "Cannot decode SQL null as the Haskell Day type. Use a `Maybe Day`" @@ -1073,7 +1073,7 @@ instance FromPgField (Unbounded Day) where -- There is a very specific conversion function for these, which I poorly translated to Haskell -- https://github.com/postgres/postgres/blob/799959dc7cf0e2462601bea8d07b6edec3fa0c4f/src/backend/utils/adt/datetime.c#L321 -- But I found a simpler way to do this. Let's see if it works in our property based tests - jd <- BinSer.decodeInt32BE bs + jd <- BinSer.decodeInt32BE 0 bs Right $ if jd == minBound then NegInfinity @@ -1087,9 +1087,9 @@ instance FromPgField (Unbounded Day) where instance FromPgField CalendarDiffTime where fieldDecoder = parsePgType [intervalOid] $ \case Just bs -> do - nMicrosecs <- BinSer.decodeInt64BE bs - nDays <- BinSer.decodeInt32BE (BS.drop 8 bs) - nMonths <- BinSer.decodeInt32BE (BS.drop 12 bs) + nMicrosecs <- BinSer.decodeInt64BE 0 bs + nDays <- BinSer.decodeInt32BE 8 bs + nMonths <- BinSer.decodeInt32BE 12 bs Right $ CalendarDiffTime {ctMonths = fromIntegral nMonths, ctTime = secondsToNominalDiffTime (fromIntegral nDays * 86400) + realToFrac (picosecondsToDiffTime (fromIntegral nMicrosecs * 1_000_000))} Nothing -> Left "Cannot decode SQL null as the Haskell CalendarDiffTime type. Use a `Maybe CalendarDiffTime`" diff --git a/hpgsql/src/Hpgsql/Encoding/BinarySerializer.hs b/hpgsql/src/Hpgsql/Encoding/BinarySerializer.hs index a21dd27..cfaa2c8 100644 --- a/hpgsql/src/Hpgsql/Encoding/BinarySerializer.hs +++ b/hpgsql/src/Hpgsql/Encoding/BinarySerializer.hs @@ -9,7 +9,8 @@ -- The caveat is that this module makes unaligned memory access. For the target -- CPU architectures of this library, this should be fine. module Hpgsql.Encoding.BinarySerializer - ( decodeInt16BE, + ( ByteStringIdx (..), + decodeInt16BE, decodeInt32BE, decodeInt64BE, decodeWord32BE, @@ -31,13 +32,11 @@ import Prelude hiding (encodeFloat) #if WORDS_BIGENDIAN import Data.Word (Word16, Word32, Word64) #else -import Data.Word (Word16, Word32, Word64, byteSwap16, byteSwap32, byteSwap64) +import Data.Word (Word16, Word32, Word64, byteSwap16, byteSwap32, byteSwap64, Word8) #endif import Data.Bits (Bits (unsafeShiftR)) -import qualified Data.ByteString as BS import Data.Coerce (coerce) -import Data.Maybe (fromMaybe) -import Foreign (Storable (..), peek, (.&.)) +import Foreign (Storable (..), (.&.)) import Foreign.ForeignPtr (withForeignPtr) import GHC.Float (castDoubleToWord64, castFloatToWord32) import System.IO.Unsafe (unsafeDupablePerformIO) @@ -63,15 +62,20 @@ fromBigEndian16 = Prelude.id fromBigEndian16 = byteSwap16 #endif -{-# INLINE unsafeDecodeWord #-} -unsafeDecodeWord :: (Storable a) => ByteString -> Int -> (a -> a) -> Either String a -unsafeDecodeWord (InternalBS.BS bytesPtr len) minLen endianConvert = - if len >= minLen - then - -- A bang (strictness) in `decodedWord` makes our benchmarks allocate more memory and run slower! - let decodedWord = endianConvert $ unsafeDupablePerformIO $ withForeignPtr bytesPtr $ \ptr -> peek (coerce ptr) - in Right decodedWord - else Left "Less than enough bytes to decode" +data CoolWordDec a where + CWord8 :: CoolWordDec Word8 + CWord16 :: CoolWordDec Word16 + CWord32 :: CoolWordDec Word32 + CWord64 :: CoolWordDec Word64 + +{-# INLINE decodeWord #-} +decodeWord :: CoolWordDec a -> ByteStringIdx -> ByteString -> (a -> a) -> Either String a +decodeWord wdec idx (InternalBS.BS bytesPtr len) endianConvert = + case wdec of + CWord8 -> if len < 1 + idx.idx then Left "Less than enough bytes to decode" else Right $ endianConvert $ unsafeDupablePerformIO $ withForeignPtr bytesPtr $ \ptr -> peekByteOff (coerce ptr) idx.idx + CWord16 -> if len < 2 + idx.idx then Left "Less than enough bytes to decode" else Right $ endianConvert $ unsafeDupablePerformIO $ withForeignPtr bytesPtr $ \ptr -> peekByteOff (coerce ptr) idx.idx + CWord32 -> if len < 4 + idx.idx then Left "Less than enough bytes to decode" else Right $ endianConvert $ unsafeDupablePerformIO $ withForeignPtr bytesPtr $ \ptr -> peekByteOff (coerce ptr) idx.idx + CWord64 -> if len < 8 + idx.idx then Left "Less than enough bytes to decode" else Right $ endianConvert $ unsafeDupablePerformIO $ withForeignPtr bytesPtr $ \ptr -> peekByteOff (coerce ptr) idx.idx {-# INLINE unsafeEncodeWord #-} unsafeEncodeWord :: (Storable a) => a -> (a -> a) -> Int -> ByteString @@ -79,33 +83,40 @@ unsafeEncodeWord n endianConvert len = InternalBS.unsafeCreate len $ \bufferPtr -> poke (coerce bufferPtr) $ endianConvert n +newtype ByteStringIdx = ByteStringIdx {idx :: Int} + deriving newtype (Num) + {-# INLINE decodeInt16BE #-} -decodeInt16BE :: ByteString -> Either String Int16 -decodeInt16BE bs = fromIntegral <$> unsafeDecodeWord bs 2 fromBigEndian16 +decodeInt16BE :: ByteStringIdx -> ByteString -> Either String Int16 +decodeInt16BE idx bs = fromIntegral <$> decodeWord CWord16 idx bs fromBigEndian16 {-# INLINE encodeInt16BE #-} encodeInt16BE :: Int16 -> ByteString encodeInt16BE n = unsafeEncodeWord (fromIntegral n) fromBigEndian16 2 +{-# INLINE decodeWord8 #-} +decodeWord8 :: ByteStringIdx -> ByteString -> Either String Word8 +decodeWord8 idx bs = decodeWord CWord8 idx bs Prelude.id + {-# INLINE decodeWord32BE #-} decodeWord32BE :: ByteString -> Either String Word32 -decodeWord32BE bs = unsafeDecodeWord bs 4 fromBigEndian32 +decodeWord32BE bs = decodeWord CWord32 0 bs fromBigEndian32 {-# INLINE decodeWord64BE #-} decodeWord64BE :: ByteString -> Either String Word64 -decodeWord64BE bs = unsafeDecodeWord bs 8 fromBigEndian64 +decodeWord64BE bs = decodeWord CWord64 0 bs fromBigEndian64 {-# INLINE decodeInt32BE #-} -decodeInt32BE :: ByteString -> Either String Int32 -decodeInt32BE bs = fromIntegral <$> unsafeDecodeWord bs 4 fromBigEndian32 +decodeInt32BE :: ByteStringIdx -> ByteString -> Either String Int32 +decodeInt32BE idx bs = fromIntegral <$> decodeWord CWord32 idx bs fromBigEndian32 {-# INLINE encodeInt32BE #-} encodeInt32BE :: Int32 -> ByteString encodeInt32BE n = unsafeEncodeWord (fromIntegral n) fromBigEndian32 4 {-# INLINE decodeInt64BE #-} -decodeInt64BE :: ByteString -> Either String Int64 -decodeInt64BE bs = fromIntegral <$> unsafeDecodeWord bs 8 fromBigEndian64 +decodeInt64BE :: ByteStringIdx -> ByteString -> Either String Int64 +decodeInt64BE idx bs = fromIntegral <$> decodeWord CWord64 idx bs fromBigEndian64 {-# INLINE encodeInt64BE #-} encodeInt64BE :: Int64 -> ByteString @@ -127,17 +138,16 @@ encodePgBoolean v = if v then "\SOH" else "\NUL" -- | A super specialized decoder to decode a postgres DataRow message -- more quickly than a naive implementation. --- Returns first the parsed DataRow (only column sizes and values) and second --- the left-unparsed original bytestring. -decodeDataRow :: ByteString -> Either String (ByteString, ByteString) -decodeDataRow bs@(InternalBS.BS _bytesPtr len) = +-- Returns the index into the left-unparsed contents of the supplied bytestring. +decodeDataRow :: ByteStringIdx -> ByteString -> Either String ByteStringIdx +decodeDataRow idx bs@(InternalBS.BS _bytesPtr len) = -- We have a fast path when rows are at least 8 bytes long (should be the case -- for all but 0-column query results or bytestring chunks "cut in the middle of the message") -- by playing with bitwise operations. -- Whether this is worth keeping is sort of questionable. It's complex -- (even if I think it's safe and well tested) and reduces runtime of one of -- our benchmarks by 2% compared to not having it. - case unsafeDecodeWord bs 8 fromBigEndian64 of + case decodeWord CWord64 idx bs fromBigEndian64 of Right (w64 :: Word64) -> -- After fromBigEndian64, the Word64 has bytes in big-endian order: -- byte 0 (msg type) in MSB, bytes 1-4 (length) next, bytes 5-6 (col count), byte 7 in LSB. @@ -151,15 +161,15 @@ decodeDataRow bs@(InternalBS.BS _bytesPtr len) = Left _ -> -- It is possible the DataRow has length less than 8 bytes, so -- we still have to try to parse that. - if len >= 5 - then - let (InternalBS.w2c -> msgIdentChar, lenbs) = fromMaybe (error "impossible") $ BS.uncons bs - lenFullMsg = fromIntegral $ either error id (decodeInt32BE lenbs) - in if msgIdentChar == 'D' - then toResult lenFullMsg - else Left "Not a DataRow" + if len >= 5 + idx.idx + then do + msgIdentChar <- decodeWord8 idx bs + lenFullMsg <- decodeInt32BE (1 + idx) bs + if msgIdentChar == 68 -- Letter 'D' + then toResult (fromIntegral lenFullMsg) + else Left "Not a DataRow" else Left "Less than enough bytes to decode a DataRow" where toResult lenFullMsg - | len >= 1 + lenFullMsg = let (a, rest) = BS.splitAt (1 + lenFullMsg) bs in Right (BS.drop 7 a, rest) + | len >= 1 + lenFullMsg + idx.idx = Right $ ByteStringIdx $ 1 + lenFullMsg + idx.idx | otherwise = Left "Less than enough bytes to decode a full DataRow" diff --git a/hpgsql/src/Hpgsql/Internal.hs b/hpgsql/src/Hpgsql/Internal.hs index 71d6de0..f986b59 100644 --- a/hpgsql/src/Hpgsql/Internal.hs +++ b/hpgsql/src/Hpgsql/Internal.hs @@ -115,6 +115,7 @@ import qualified Control.Concurrent.STM as STM import Control.Exception.Safe (Exception (..), MonadThrow, SomeException, bracket, bracketOnError, finally, handleJust, mask, mask_, onException, throw, toException, tryJust) import Control.Monad (forM, forM_, join, unless, void, when) import Data.ByteString (ByteString) +import qualified Data.ByteString as BS import Data.ByteString.Internal (w2c) import qualified Data.ByteString.Lazy as LBS import Data.Data (Proxy (..)) @@ -506,7 +507,7 @@ receiveNextMsgWithMaskedContinuation conn parser f = Left (msgIdentChar, mPgError) -> throw IrrecoverableHpgsqlError {hpgsqlDetails = "Could not parse postgres message with ident char " <> Text.pack (show msgIdentChar) <> ". This is an internal error in Hpgsql. Please report it.", innerException = toException <$> mPgError, relatedStatement = Nothing} data ReceiveWhat a b where - ReceiveDataRows :: ReceiveWhat DataRow [DataRow] + ReceiveDataRows :: ReceiveWhat DataRow ByteString ReceiveArbitraryMsg :: PgMsgParser a -> (Either (Char, Maybe PostgresError) a -> STM b) -> ReceiveWhat a b -- | Masks asynchronous exceptions in between the moment the message is extracted from @@ -534,11 +535,11 @@ receiveNextMsgGeneric conn@HPgConnection {socket, recvBuffer} receiveWhat = do (initialBuf, initialBufLen) <- receiveUntilBufferHasAtLeast 5 let charAndLength = LBS.take 5 initialBuf let (w2c -> msgIdentChar, lenbs) = fromMaybe (error "impossible") $ LBS.uncons charAndLength - lenLeftToFetch :: Int64 = fromIntegral $ either error id (BinSer.decodeInt32BE $ LBS.toStrict lenbs) - 4 + lenLeftToFetch :: Int64 = fromIntegral $ either error id (BinSer.decodeInt32BE 0 $ LBS.toStrict lenbs) - 4 fullMessageLen = 5 + lenLeftToFetch (nowBuf, _nowBufLen) <- if initialBufLen >= fullMessageLen then pure (initialBuf, initialBufLen) else receiveUntilBufferHasAtLeast fullMessageLen - let restOfMsg = LBS.drop 5 $ LBS.take fullMessageLen nowBuf - receivedNoticeOrParameterSoTryAgain <- go msgIdentChar restOfMsg fullMessageLen nowBuf + let fullMsg = LBS.take fullMessageLen nowBuf + receivedNoticeOrParameterSoTryAgain <- go msgIdentChar fullMsg fullMessageLen nowBuf case receivedNoticeOrParameterSoTryAgain of Nothing -> receiveNextMsgGeneric conn receiveWhat Just res -> pure res @@ -551,12 +552,12 @@ receiveNextMsgGeneric conn@HPgConnection {socket, recvBuffer} receiveWhat = do -- the recvBuffer, then we _must_ remove that message from recvBuffer. -- Ideally we'd have non-retriable STM at the type-level here. Maybe later. -- Make sure to do very little work inside `go`! - go msgIdentChar restOfMsg fullMessageLen nowBuf = mask_ $ modifyIORefIO recvBuffer $ do + go msgIdentChar fullMsg fullMessageLen nowBuf = mask_ $ modifyIORefIO recvBuffer $ do let bufferWithoutMsg = LBS.drop fullMessageLen nowBuf handleUnexpectedMsg onNotAnyReasonableMsg = -- This could be a Notification, NOTICE or a ParameterStatus message, since these -- can be received _at any time_ according to the docs. - case parsePgMessage msgIdentChar restOfMsg (Left3 <$> msgParser @NotificationResponse <|> Middle3 <$> msgParser @NoticeResponse <|> Right3 <$> msgParser @ParameterStatus) of + case parsePgMessage msgIdentChar fullMsg (Left3 <$> msgParser @NotificationResponse <|> Middle3 <$> msgParser @NoticeResponse <|> Right3 <$> msgParser @ParameterStatus) of Just (Left3 notifResponse) -> do debugPrint "Received notification. Will add it to internal queue." STM.atomically $ do @@ -578,25 +579,25 @@ receiveNextMsgGeneric conn@HPgConnection {socket, recvBuffer} receiveWhat = do Nothing -> do -- Just in case this is a postgres error, it might include useful information, -- so we spit that out - let mPgError = mkPostgresError "" <$> parsePgMessage msgIdentChar restOfMsg (msgParser @ErrorResponse) + let mPgError = mkPostgresError "" <$> parsePgMessage msgIdentChar fullMsg (msgParser @ErrorResponse) fmap (nowBuf,) $ Just <$> STM.atomically (onNotAnyReasonableMsg (msgIdentChar, mPgError)) case receiveWhat of ReceiveDataRows -> -- Parse as many DataRows as we can to do as much work as we can per buffer "churn" - case Parser.parseOnly (Parser.matchLeftUnconsumed (Parser.parseMany customDataRowParser)) (LBS.toStrict nowBuf) of - Parser.ParseOk (unconsumedBuffer, msgs@(_ : _)) -> do - debugPrint $ "Received " ++ show msgs - pure (LBS.fromStrict unconsumedBuffer, Just msgs) - _ -> handleUnexpectedMsg $ const $ pure [] -- No error when we stop receiving DataRows, only emptiness + let fullBuf = LBS.toStrict nowBuf + in case Parser.parseOnly Parser.parseManyRows fullBuf of + Parser.ParseOk unconsumedBufferBegin | unconsumedBufferBegin.idx > 0 -> do + let (msgs, unconsumedBuffer) = BS.splitAt unconsumedBufferBegin.idx fullBuf + debugPrint $ "Received one or more messages with total length " ++ show (BS.length msgs) + pure (LBS.fromStrict unconsumedBuffer, Just msgs) + _ -> handleUnexpectedMsg $ const $ pure "" -- No error when we stop receiving DataRows, only emptiness ReceiveArbitraryMsg parser f -> - case parsePgMessage msgIdentChar restOfMsg parser of + case parsePgMessage msgIdentChar fullMsg parser of Just msg -> do debugPrint $ "Received " ++ show msg fmap (bufferWithoutMsg,) $ Just <$> STM.atomically (f (Right msg)) Nothing -> handleUnexpectedMsg (f . Left) - customDataRowParser = DataRow <$> Parser.takeDataRow - -- \| Appends into the internal buffer by reading from the socket -- until the buffer has at least N bytes. -- Returns the current buffer and its length. @@ -843,6 +844,8 @@ receiveOutstandingResponseMsgsAtomically thisThreadId conn qryId = do } pure (Just respMsg, newState) +newtype DataRows = DataRows ByteString + -- | After sending one or more queries to the backend, run this function for each query to fetch that query's results. -- You must call the returned IO function and consume the returned Stream completely until you get to the -- `Either ErrorResponse CommandComplete` object. @@ -855,7 +858,7 @@ receiveOutstandingResponseMsgsAtomically thisThreadId conn qryId = do consumeResults :: HPgConnection -> QueryId -> - IO (Maybe (Either3 NoData RowDescription CopyInResponse), Stream (Of DataRow) IO (Either ErrorResponse CommandComplete)) + IO (Maybe (Either3 NoData RowDescription CopyInResponse), Stream (Of DataRows) IO (Either ErrorResponse CommandComplete)) consumeResults conn qryId = do -- debugPrint "++++ Inside consumeResults" -- We assume it's possible to receive a DataRow here even in the first call because `consumeResults` @@ -886,29 +889,28 @@ consumeResults conn qryId = do pure (mERowDesc, pure $ Right cmd) (mERowDesc, Middle3 mDataRow) -> do let allOtherRows = - S.concat $ - S.unfold - ( \() -> do - mRow <- receiveNextMsgGeneric conn ReceiveDataRows - case mRow of - rows@(_ : _) -> pure $ Right (rows :> ()) - [] -> do - stateAfterNextMsg <- snd <$> receiveOutstandingResponseMsgsAtomically thisThreadId conn qryId - case stateAfterNextMsg of - ErrorResponseReceived _ err -> do - receiveReadyForQueryIfNecessary thisThreadId - pure $ Left $ Left err - CommandCompleteReceived _ cmd -> do - receiveReadyForQueryIfNecessary thisThreadId - pure $ Left $ Right cmd - ReadyForQueryReceived errOrCmd _ -> pure $ Left errOrCmd - st -> throwIrrecoverableError $ "Internal error in Hpgsql. After the last DataRow we should get either an ErrorResponse or a CommandComplete message. State: " <> Text.pack (show st) - ) - () + S.unfold + ( \() -> do + mRow <- receiveNextMsgGeneric conn ReceiveDataRows + case mRow of + rows | not (BS.null rows) -> pure $ Right (DataRows rows :> ()) + _ -> do + stateAfterNextMsg <- snd <$> receiveOutstandingResponseMsgsAtomically thisThreadId conn qryId + case stateAfterNextMsg of + ErrorResponseReceived _ err -> do + receiveReadyForQueryIfNecessary thisThreadId + pure $ Left $ Left err + CommandCompleteReceived _ cmd -> do + receiveReadyForQueryIfNecessary thisThreadId + pure $ Left $ Right cmd + ReadyForQueryReceived errOrCmd _ -> pure $ Left errOrCmd + st -> throwIrrecoverableError $ "Internal error in Hpgsql. After the last DataRow we should get either an ErrorResponse or a CommandComplete message. State: " <> Text.pack (show st) + ) + () finalStream = case mDataRow of Nothing -> allOtherRows Just dr -> - dr `S.cons` allOtherRows + DataRows dr.fullDataRow `S.cons` allOtherRows pure (mERowDesc, finalStream) where receiveReadyForQueryIfNecessary :: WeakThreadId -> IO () @@ -1409,17 +1411,18 @@ consumeStreamingResults rp conn qryId = S.effect $ do let typecheckedColInfos = rtypecheck colInfos unless (numResultColumns == expectedNumCols) $ throwIrrecoverableErrorWithStatement qText $ "Query result contains " <> Text.pack (show numResultColumns) <> " columns but row parser expected " <> Text.pack (show expectedNumCols) unless (all snd typecheckedColInfos) $ throwIrrecoverableErrorWithStatement qText "Query result column types do not match expected column types" - pure $ rparser colInfos <* Parser.endOfInput - MonadicRowDecoder (RowDecoderMonadic rparser) -> pure $ fmap fst $ rparser ConversionState {colsLeftToParse = colInfos} <* Parser.endOfInput + pure $ Parser.skip 7 *> rparser colInfos -- Skip msg ident., length, number of columns, then parse fields + MonadicRowDecoder (RowDecoderMonadic rparser) -> pure $ Parser.skip 7 *> fmap fst (rparser ConversionState {colsLeftToParse = colInfos}) pure $ do errOrCmdComplete <- - S.mapM - ( \(DataRow rowColumnData) -> - case Parser.parseOnly rowparser rowColumnData of - Parser.ParseOk row -> pure row - Parser.ParseFail err -> throwIrrecoverableErrorWithStatement qText $ "Failed parsing a row: " <> Text.pack (show err) - ) - rowsStream + S.concat $ + S.mapM + ( \(DataRows rowColumnData) -> + case Parser.parseOnly (Parser.parseMany rowparser <* Parser.endOfInput) rowColumnData of + Parser.ParseOk rows -> pure rows + Parser.ParseFail err -> throwIrrecoverableErrorWithStatement qText $ "Failed parsing a row: " <> Text.pack (show err) + ) + rowsStream S.effect $ case errOrCmdComplete of Left err -> throwPostgresError qText err Right _cmdComplete -> pure mempty diff --git a/hpgsql/src/Hpgsql/InternalTypes.hs b/hpgsql/src/Hpgsql/InternalTypes.hs index 7b9534e..af81141 100644 --- a/hpgsql/src/Hpgsql/InternalTypes.hs +++ b/hpgsql/src/Hpgsql/InternalTypes.hs @@ -368,7 +368,9 @@ newtype ErrorResponse = ErrorResponse (Map ErrorDetail LBS.ByteString) newtype CommandComplete = CommandComplete {numRows :: Int64} deriving stock (Show) -newtype DataRow = DataRow {rowColumnData :: ByteString} +-- | A DataRow with its leading identifying character ('D'), the 32bits self-length, +-- the 2 bytes for the number of fields and the fields' lengths and values themselves. +newtype DataRow = DataRow {fullDataRow :: ByteString} instance Show DataRow where show _ = "DataRow" diff --git a/hpgsql/src/Hpgsql/Msgs.hs b/hpgsql/src/Hpgsql/Msgs.hs index 9e1f6cf..3446e15 100644 --- a/hpgsql/src/Hpgsql/Msgs.hs +++ b/hpgsql/src/Hpgsql/Msgs.hs @@ -34,7 +34,7 @@ class ToPgMessage a where newtype PgMsgParser a = PgMsgParser ( Char -> - -- \| Message contents after the Int32 length attribute + -- \| Full PG message contents, including the message identifier byte, the rest-of-message length and the message contents. LBS.ByteString -> Maybe a ) @@ -54,7 +54,7 @@ colParser = do colName <- nulTerminatedCStringParser -- Column name as C string void $ Parsec.take (4 + 2) -- TODO: OIDs are unsigned integers! Try `select (-1)::oid` to see. Change to UInt32 somehow - typOid <- either fail pure . BinSer.decodeInt32BE =<< Parsec.take 4 + typOid <- either fail pure . BinSer.decodeInt32BE 0 =<< Parsec.take 4 void $ Parsec.take (2 + 4 + 2) pure (colName, Oid (fromIntegral typOid)) @@ -137,8 +137,8 @@ data Terminate = Terminate deriving stock (Show) instance FromPgMessage AuthenticationResponse where - msgParser = PgMsgParser $ \c restOfMsg -> case c of - 'R' -> case first (BinSer.decodeInt32BE . LBS.toStrict) $ LBS.splitAt 4 restOfMsg of + msgParser = PgMsgParser $ \c (LBS.drop 5 -> restOfMsg) -> case c of + 'R' -> case first (BinSer.decodeInt32BE 0 . LBS.toStrict) $ LBS.splitAt 4 restOfMsg of (Right 0, _) -> Just $ AuthenticationResponse AuthOk (Right 2, _) -> Just $ AuthenticationResponse AuthKerberosV5 (Right 3, _) -> Just $ AuthenticationResponse AuthCleartextPassword @@ -154,19 +154,19 @@ instance FromPgMessage AuthenticationResponse where _ -> Nothing instance FromPgMessage BackendKeyData where - msgParser = PgMsgParser $ \c (LBS.splitAt 4 -> (pidBS, backendSecretKey)) -> case c of - 'K' -> case BinSer.decodeInt32BE $ LBS.toStrict pidBS of + msgParser = PgMsgParser $ \c (LBS.splitAt 4 . LBS.drop 5 -> (pidBS, backendSecretKey)) -> case c of + 'K' -> case BinSer.decodeInt32BE 0 $ LBS.toStrict pidBS of Right pid -> Just $ BackendKeyData {backendPid = pid, backendSecretKey = LBS.toStrict backendSecretKey} Left _ -> Nothing _ -> Nothing instance FromPgMessage BindComplete where - msgParser = PgMsgParser $ \c _restOfMsg -> case c of + msgParser = PgMsgParser $ \c _ -> case c of '2' -> Just BindComplete _ -> Nothing instance FromPgMessage CommandComplete where - msgParser = PgMsgParser $ \c restOfMsg -> case c of + msgParser = PgMsgParser $ \c (LBS.drop 5 -> restOfMsg) -> case c of 'C' -> let astext = decodeASCII $ LBS.toStrict $ LBS.dropEnd 1 restOfMsg in case TextParsec.parseOnly ((ins <|> del <|> upd <|> merge <|> sel <|> move <|> fetch <|> copy) <* TextParsec.endOfInput) astext of @@ -224,8 +224,8 @@ instance FromPgMessage CopyInResponse where _ -> Nothing instance FromPgMessage DataRow where - msgParser = PgMsgParser $ \c !restOfMsg -> case c of - 'D' -> Just $ DataRow {rowColumnData = LBS.toStrict $ LBS.drop 2 restOfMsg} + msgParser = PgMsgParser $ \c !fullDataRow -> case c of + 'D' -> Just $ DataRow {fullDataRow = LBS.toStrict fullDataRow} _ -> Nothing instance FromPgMessage NoData where @@ -234,7 +234,7 @@ instance FromPgMessage NoData where _ -> Nothing instance FromPgMessage ParameterStatus where - msgParser = PgMsgParser $ \c !restOfMsg -> case c of + msgParser = PgMsgParser $ \c !(LBS.drop 5 -> restOfMsg) -> case c of 'S' -> case LazyParsec.parseOnly (((,) <$> nulTerminatedCStringParser <*> nulTerminatedCStringParser) <* Parsec.endOfInput) restOfMsg of Left _ -> error "Failed parsing ParameterStatus" Right (parameterName, parameterValue) -> Just $ ParameterStatus {..} @@ -348,18 +348,18 @@ instance ToPgMessage StartupMessage where Builder.int32BE (4 + contentsLen) <> contents instance FromPgMessage ReadyForQuery where - msgParser = PgMsgParser $ \c restOfMsg -> case (c, restOfMsg) of + msgParser = PgMsgParser $ \c (LBS.drop 5 -> restOfMsg) -> case (c, restOfMsg) of ('Z', "I") -> Just $ ReadyForQuery TransIdle ('Z', "T") -> Just $ ReadyForQuery TransInTrans ('Z', "E") -> Just $ ReadyForQuery TransInError _ -> Nothing instance FromPgMessage RowDescription where - msgParser = PgMsgParser $ \c restOfMsg -> + msgParser = PgMsgParser $ \c (LBS.drop 5 -> restOfMsg) -> if c == 'T' then let (numColsBS, colContents) = LBS.splitAt 2 restOfMsg - numCols = either error id $ BinSer.decodeInt16BE $ LBS.toStrict numColsBS + numCols = either error id $ BinSer.decodeInt16BE 0 $ LBS.toStrict numColsBS allColOidsParser :: Parsec.Parser [(Text, Oid)] allColOidsParser = replicateM (fromIntegral numCols) colParser in case LazyParsec.parseOnly (allColOidsParser <* Parsec.endOfInput) colContents of @@ -368,7 +368,7 @@ instance FromPgMessage RowDescription where else Nothing instance FromPgMessage ErrorResponse where - msgParser = PgMsgParser $ \c restOfMsg -> + msgParser = PgMsgParser $ \c (LBS.drop 5 -> restOfMsg) -> if c /= 'E' then Nothing else @@ -384,7 +384,7 @@ instance FromPgMessage ErrorResponse where in Just $ ErrorResponse $ Map.fromList $ mapMaybe parseSingleErrorField errorFields instance FromPgMessage NoticeResponse where - msgParser = PgMsgParser $ \c restOfMsg -> + msgParser = PgMsgParser $ \c (LBS.drop 5 -> restOfMsg) -> if c /= 'N' then Nothing else @@ -401,12 +401,12 @@ instance FromPgMessage NoticeResponse where in Just $ NoticeResponse $ Map.fromList $ mapMaybe parseSingleErrorField errorFields instance FromPgMessage NotificationResponse where - msgParser = PgMsgParser $ \c restOfMsg -> + msgParser = PgMsgParser $ \c (LBS.drop 5 -> restOfMsg) -> if c /= 'A' then Nothing else let (notifierPidBs, channelNameAndPayload) = LBS.splitAt 4 restOfMsg - notifierPid = either error id $ BinSer.decodeInt32BE $ LBS.toStrict notifierPidBs + notifierPid = either error id $ BinSer.decodeInt32BE 0 $ LBS.toStrict notifierPidBs in case LazyParsec.parseOnly ((NotificationResponse notifierPid <$> nulTerminatedCStringParser <*> nulTerminatedCStringParser) <* Parsec.endOfInput) channelNameAndPayload of diff --git a/hpgsql/src/Hpgsql/SimpleParser.hs b/hpgsql/src/Hpgsql/SimpleParser.hs index 64a13bd..74940f0 100644 --- a/hpgsql/src/Hpgsql/SimpleParser.hs +++ b/hpgsql/src/Hpgsql/SimpleParser.hs @@ -22,12 +22,15 @@ module Hpgsql.SimpleParser takeInt32BE, takeInt64BE, takeDataRow, + parseManyRows, + skip, ) where import Data.ByteString (ByteString) import qualified Data.ByteString as BS import Data.Int (Int16, Int32, Int64) +import Hpgsql.Encoding.BinarySerializer (ByteStringIdx (..)) import qualified Hpgsql.Encoding.BinarySerializer as BinSer import Prelude hiding (take) @@ -40,128 +43,152 @@ data ParseResult a newtype Parser a = Parser { unParser :: forall r. + ByteStringIdx -> ByteString -> (String -> r) -> -- \^ failure continuation - (a -> ByteString -> r) -> - -- \^ success continuation, taking left-unparsed ByteString and parsed value + (a -> ByteStringIdx -> ByteString -> r) -> + -- \^ success continuation, taking original or new ByteString, the index into the original/new bytestring of the first yet-unparsed byte, and parsed value r } instance Functor Parser where - fmap f (Parser p) = Parser $ \bs kf ks -> - p bs kf (\a bs' -> ks (f a) bs') + fmap f (Parser p) = Parser $ \idx bs kf ks -> + p idx bs kf (\a bs' -> ks (f a) bs') {-# INLINE fmap #-} instance Applicative Parser where - pure a = Parser $ \bs _ ks -> ks a bs + pure a = Parser $ \idx bs _ ks -> ks a idx bs {-# INLINE pure #-} - Parser pf <*> Parser pa = Parser $ \bs kf ks -> - pf bs kf (\f bs' -> pa bs' kf (\a bs'' -> ks (f a) bs'')) + Parser pf <*> Parser pa = Parser $ \idx bs kf ks -> + pf idx bs kf (\f bs' idx' -> pa bs' idx' kf (\a bs'' idx'' -> ks (f a) bs'' idx'')) {-# INLINE (<*>) #-} instance Monad Parser where return = pure {-# INLINE return #-} - Parser p >>= k = Parser $ \bs kf ks -> - p bs kf (\a bs' -> unParser (k a) bs' kf ks) + Parser p >>= k = Parser $ \idx bs kf ks -> + p idx bs kf (\a bs' idx' -> unParser (k a) bs' idx' kf ks) {-# INLINE (>>=) #-} instance MonadFail Parser where - fail msg = Parser $ \_ kf _ -> kf msg + fail msg = Parser $ \_ _ kf _ -> kf msg {-# INLINE fail #-} -- | Run a parser and return either an error message or the parsed value, -- using the strict 'ParseResult' type. Any unconsumed trailing input is -- discarded. parseOnly :: Parser a -> ByteString -> ParseResult a -parseOnly (Parser p) bs = p bs ParseFail (\a _ -> ParseOk a) +parseOnly p = parseOnlyOffset p 0 {-# INLINE parseOnly #-} +-- | Run a parser and return either an error message or the parsed value, +-- using the strict 'ParseResult' type. Any unconsumed trailing input is +-- discarded. +parseOnlyOffset :: Parser a -> ByteStringIdx -> ByteString -> ParseResult a +parseOnlyOffset (Parser p) idx bs = p idx bs ParseFail (\a _ _ -> ParseOk a) +{-# INLINE parseOnlyOffset #-} + -- | Consume exactly @n@ bytes of input, failing if fewer than @n@ bytes -- remain. take :: Int -> Parser ByteString -take n = Parser $ \bs kf ks -> - -- Special-casing n>0 helps reduce memory usage - -- by ~1.5% in our benchmarks without a measurable - -- difference in run time - if n > 0 - then - if BS.length bs >= n - then case BS.splitAt n bs of - (!h, !t) -> ks h t - else kf ("take: wanted " <> show n <> " bytes but only " <> show (BS.length bs) <> " remain") - else - ks mempty bs +take n = Parser $ \idx bs kf ks -> + let skip' = n + idx.idx + in if BS.length bs >= skip' + then case BS.take n $ BS.drop idx.idx bs of + -- Strict on the bytestring because we're pretty sure + -- the field decoder will need to evaluate this anyway, + -- so no need for an extra thunk + !h -> ks h (ByteStringIdx skip') bs + else kf "take: insufficient bytes" {-# INLINE take #-} +-- | Consume exactly @n@ bytes of input, failing if fewer than @n@ bytes +-- remain. +skip :: Int -> Parser () +skip n = Parser $ \idx bs _ ks -> + ks () (ByteStringIdx $ idx.idx + n) bs +{-# INLINE skip #-} + {-# INLINE takeInt16BE #-} takeInt16BE :: Parser Int16 -takeInt16BE = Parser $ \bs kf ks -> - case BinSer.decodeInt16BE bs of +takeInt16BE = Parser $ \idx bs kf ks -> + case BinSer.decodeInt16BE idx bs of Left err -> kf err - Right v -> ks v (BS.drop 2 bs) + Right v -> ks v (idx + 2) bs {-# INLINE takeInt32BE #-} takeInt32BE :: Parser Int32 -takeInt32BE = Parser $ \bs kf ks -> - case BinSer.decodeInt32BE bs of +takeInt32BE = Parser $ \idx bs kf ks -> + case BinSer.decodeInt32BE idx bs of Left err -> kf err - Right v -> ks v (BS.drop 4 bs) + Right v -> ks v (idx + 4) bs {-# INLINE takeInt64BE #-} takeInt64BE :: Parser Int64 -takeInt64BE = Parser $ \bs kf ks -> - case BinSer.decodeInt64BE bs of +takeInt64BE = Parser $ \idx bs kf ks -> + case BinSer.decodeInt64BE idx bs of Left err -> kf err - Right v -> ks v (BS.drop 8 bs) + Right v -> ks v (idx + 8) bs {-# INLINE takeDataRow #-} --- | A specialized parser to parse a postgres DataRow. -takeDataRow :: Parser ByteString -takeDataRow = Parser $ \bs kf ks -> - case BinSer.decodeDataRow bs of +-- | A specialized parser to parse a postgres DataRow, +-- returning the index of the byte after this DataRow's last. +takeDataRow :: Parser ByteStringIdx +takeDataRow = Parser $ \idx bs kf ks -> + case BinSer.decodeDataRow idx bs of Left err -> kf err - Right (thisDataRow, rest) -> ks thisDataRow rest + Right idxRest -> ks idxRest idxRest bs parseMany :: Parser a -> Parser [a] -parseMany p = Parser $ \bs' _kf ks -> let (vs, rest) = go bs' in ks vs rest +parseMany p = Parser $ \idx' bs' _kf ks -> let (vs, restIdx) = go idx' bs' in ks vs restIdx bs' where - go bs = case parseOnly (matchLeftUnconsumed p) bs of - ParseOk (unconsumed, v) -> let (vs, rest) = go unconsumed in (v : vs, rest) - ParseFail _ -> ([], bs) + go idx bs = case parseOnlyOffset (matchLeftUnconsumed p) idx bs of + ParseOk (unconsumedIdx, v) -> let (vs, rest) = go unconsumedIdx bs in (v : vs, rest) + ParseFail _ -> ([], idx) {-# INLINE parseMany #-} +parseManyRows :: Parser ByteStringIdx +parseManyRows = Parser $ \idx' bs' _kf ks -> let restIdx = go idx' bs' in ks restIdx restIdx bs' + where + go idx bs = case parseOnlyOffset takeDataRow idx bs of + ParseOk unconsumedIdx -> go unconsumedIdx bs + ParseFail _ -> idx +{-# INLINE parseManyRows #-} + -- | Succeeds only when the input has been fully consumed. endOfInput :: Parser () -endOfInput = Parser $ \bs kf ks -> - if BS.null bs then ks () bs else kf "endOfInput: input remaining" +endOfInput = Parser $ \idx bs kf ks -> + if BS.length bs <= idx.idx then ks () idx bs else kf "endOfInput: input remaining" {-# INLINE endOfInput #-} -- | Run a parser and additionally return the slice of input it consumed. -- Because the input is a strict 'ByteString', the returned slice is a view -- over the original buffer and allocates no extra memory. match :: Parser a -> Parser (ByteString, a) -match (Parser p) = Parser $ \bs kf ks -> +match (Parser p) = Parser $ \idx bs kf ks -> p + idx bs kf - ( \a bs' -> - let !consumed = BS.take (BS.length bs - BS.length bs') bs - in ks (consumed, a) bs' + ( \a idx' bs' -> + let !consumed = BS.take (idx'.idx - idx.idx) $ BS.drop idx.idx bs + in ks (consumed, a) idx' bs' ) {-# INLINE match #-} --- | Run a parser and additionally return the unconsumed/unparsed ByteString. -matchLeftUnconsumed :: Parser a -> Parser (ByteString, a) -matchLeftUnconsumed (Parser p) = Parser $ \bs kf ks -> +-- | Run a parser and additionally return the index to the first unconsumed/unparsed byte +-- in the supplied ByteString. +matchLeftUnconsumed :: Parser a -> Parser (ByteStringIdx, a) +matchLeftUnconsumed (Parser p) = Parser $ \idx bs kf ks -> p + idx bs kf - ( \a bs' -> - ks (bs', a) bs' + ( \a idx' bs' -> + ks (idx', a) idx' bs' ) {-# INLINE matchLeftUnconsumed #-}