Fewer ByteString allocations in CPS parser with ByteStringIdx

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,16 +67,16 @@ 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
 
 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
 ```

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

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`"
 

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,49 +62,61 @@ 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
 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"