Skip to content

Commit a0d1290

Browse files
VyacheslavHashovVyacheslavHashov
committed
Improved Encode type
1 parent 856d839 commit a0d1290

File tree

2 files changed

+75
-53
lines changed

2 files changed

+75
-53
lines changed

‎src/Database/PostgreSQL/Protocol/Encoders.hs

Lines changed: 23 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ module Database.PostgreSQL.Protocol.Encoders
55

66
import Data.Word (Word32)
77
import Data.Monoid ((<>))
8+
import Data.Char (ord)
89
import qualified Data.Vector as V
910
import qualified Data.ByteString as B
1011

@@ -21,8 +22,11 @@ encodeStartMessage (StartupMessage (Username uname) (DatabaseName dbname))
2122
where
2223
len = fromIntegral $ getEncodeLen payload
2324
payload = putWord32BE currentVersion <>
24-
putPgString "user" <> putPgString uname <>
25-
putPgString "database" <> putPgString dbname <> putWord8 0
25+
putByteStringNull "user" <>
26+
putByteStringNull uname <>
27+
putByteStringNull "database" <>
28+
putByteStringNull dbname <>
29+
putWord8 0
2630
encodeStartMessage SSLRequest
2731
-- Value hardcoded by PostgreSQL docs.
2832
= putWord32BE 8 <> putWord32BE 80877103
@@ -31,8 +35,8 @@ encodeClientMessage :: ClientMessage -> Encode
3135
encodeClientMessage (Bind (PortalName portalName) (StatementName stmtName)
3236
paramFormat values resultFormat)
3337
= prependHeader 'B' $
34-
putPgString portalName <>
35-
putPgString stmtName <>
38+
putByteStringNull portalName <>
39+
putByteStringNull stmtName <>
3640
-- `1` means that the specified format code is applied to all parameters
3741
putWord16BE 1 <>
3842
encodeFormat paramFormat <>
@@ -43,51 +47,57 @@ encodeClientMessage (Bind (PortalName portalName) (StatementName stmtName)
4347
putWord16BE 1 <>
4448
encodeFormat resultFormat
4549
encodeClientMessage (CloseStatement (StatementName stmtName))
46-
= prependHeader 'C' $ putChar8 'S' <> putPgString stmtName
50+
= prependHeader 'C' $ putChar8 'S' <> putByteStringNull stmtName
4751
encodeClientMessage (ClosePortal (PortalName portalName))
48-
= prependHeader 'C' $ putChar8 'P' <> putPgString portalName
52+
= prependHeader 'C' $ putChar8 'P' <> putByteStringNull portalName
4953
encodeClientMessage (DescribeStatement (StatementName stmtName))
50-
= prependHeader 'D' $ putChar8 'S' <> putPgString stmtName
54+
= prependHeader 'D' $ putChar8 'S' <> putByteStringNull stmtName
5155
encodeClientMessage (DescribePortal (PortalName portalName))
52-
= prependHeader 'D' $ putChar8 'P' <> putPgString portalName
56+
= prependHeader 'D' $ putChar8 'P' <> putByteStringNull portalName
5357
encodeClientMessage (Execute (PortalName portalName) (RowsToReceive rows))
5458
= prependHeader 'E' $
55-
putPgString portalName <>
59+
putByteStringNull portalName <>
5660
putWord32BE rows
5761
encodeClientMessage Flush
5862
= prependHeader 'H' mempty
5963
encodeClientMessage (Parse (StatementName stmtName) (StatementSQL stmt) oids)
6064
= prependHeader 'P' $
61-
putPgString stmtName <>
62-
putPgString stmt <>
65+
putByteStringNull stmtName <>
66+
putByteStringNull stmt <>
6367
putWord16BE (fromIntegral $ V.length oids) <>
6468
foldMap (putWord32BE . unOid) oids
6569
encodeClientMessage (PasswordMessage passtext)
66-
= prependHeader 'p' $ putPgString $ getPassword passtext
70+
= prependHeader 'p' $ putByteStringNull $ getPassword passtext
6771
where
6872
getPassword (PasswordPlain p) = p
6973
getPassword (PasswordMD5 p) = p
7074
encodeClientMessage (SimpleQuery (StatementSQL stmt))
71-
= prependHeader 'Q' $ putPgString stmt
75+
= prependHeader 'Q' $ putByteStringNull stmt
7276
encodeClientMessage Sync
7377
= prependHeader 'S' mempty
7478
encodeClientMessage Terminate
7579
= prependHeader 'X' mempty
7680

7781
-- | Encodes single data values. Length `-1` indicates a NULL parameter value.
7882
-- No value bytes follow in the NULL case.
83+
{-# INLINE encodeValue #-}
7984
encodeValue :: Maybe B.ByteString -> Encode
8085
encodeValue Nothing = putWord32BE (-1)
8186
encodeValue (Just v) = putWord32BE (fromIntegral $ B.length v)
8287
<> putByteString v
8388

89+
{-# INLINE encodeFormat #-}
8490
encodeFormat :: Format -> Encode
8591
encodeFormat Text = putWord16BE 0
8692
encodeFormat Binary = putWord16BE 1
8793

94+
{-# INLINE prependHeader #-}
8895
prependHeader :: Char -> Encode -> Encode
8996
prependHeader c payload =
9097
-- Length includes itself but not the first message-type byte
9198
let len = 4 + fromIntegral (getEncodeLen payload)
9299
in putChar8 c <> putWord32BE len <> payload
93100

101+
{-# INLINE putChar8 #-}
102+
putChar8 :: Char -> Encode
103+
putChar8 = putWord8 . fromIntegral . ord

‎src/Database/PostgreSQL/Protocol/Store/Encode.hs

Lines changed: 52 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,9 @@ module Database.PostgreSQL.Protocol.Store.Encode where
33
import Data.Monoid (Monoid(..), (<>))
44
import Foreign (poke, plusPtr, Ptr)
55
import Data.Int (Int16, Int32)
6-
import Data.Word (Word8, Word16, Word32)
7-
import Data.Char (ord)
8-
import Data.Bits (shiftR)
6+
import Data.Word
97

8+
import Foreign
109
import Data.ByteString (ByteString)
1110
import Data.ByteString.Internal as B(toForeignPtr)
1211
import Data.Store.Core (Poke(..), unsafeEncodeWith, pokeStatePtr,
@@ -15,65 +14,78 @@ import Data.Store.Core (Poke(..), unsafeEncodeWith, pokeStatePtr,
1514
data Encode = Encode {-# UNPACK #-} !Int !(Poke ())
1615

1716
instance Monoid Encode where
18-
mempty = Encode 0 . Poke $ \_ offset -> pure (offset, ())
1917
{-# INLINE mempty #-}
18+
mempty = Encode 0 . Poke $ \_ offset -> pure (offset, ())
2019

21-
(Encode len1 f1) `mappend` (Encode len2 f2) = Encode (len1 + len2) (f1 *> f2)
2220
{-# INLINE mappend #-}
21+
(Encode len1 f1) `mappend` (Encode len2 f2) = Encode (len1 + len2) (f1 *> f2)
2322

23+
{-# INLINE getEncodeLen #-}
2424
getEncodeLen :: Encode -> Int
2525
getEncodeLen (Encode len _) = len
26-
{-# INLINE getEncodeLen #-}
2726

27+
{-# INLINE runEncode #-}
2828
runEncode :: Encode -> ByteString
2929
runEncode (Encode len f) = unsafeEncodeWith f len
30-
{-# INLINE runEncode #-}
3130

32-
fixedPrim :: Int -> (Ptr Word8 -> IO ()) -> Encode
33-
fixedPrim len f = Encode len . Poke $ \state offset -> do
31+
{-# INLINE fixed #-}
32+
fixed :: Int -> (Ptr Word8 -> IO ()) -> Encode
33+
fixed len f = Encode len . Poke $ \state offset -> do
3434
f $ pokeStatePtr state `plusPtr` offset
3535
let !newOffset = offset + len
3636
return (newOffset, ())
37-
{-# INLINE fixedPrim #-}
3837

39-
putWord8 :: Word8 -> Encode
40-
putWord8 w = fixedPrim 1 $ \p -> poke p w
41-
{-# INLINE putWord8 #-}
38+
{-# INLINE putByteString #-}
39+
putByteString :: ByteString -> Encode
40+
putByteString bs =
41+
let (ptr, offset, len) = toForeignPtr bs
42+
in Encode len $ pokeFromForeignPtr ptr offset len
4243

43-
putChar8 :: Char -> Encode
44-
putChar8 = putWord8 . fromIntegral . ord
45-
{-# INLINE putChar8 #-}
44+
-- | C-like string
45+
{-# INLINE putByteStringNull #-}
46+
putByteStringNull :: ByteString -> Encode
47+
putByteStringNull bs = putByteString bs <> putWord8 0
48+
49+
{-# INLINE putWord8 #-}
50+
putWord8 :: Word8 -> Encode
51+
putWord8 w = fixed 1 $ \p -> poke p w
4652

47-
putWord16BE :: Word16 -> Encode
48-
putWord16BE w = fixedPrim 2 $ \p -> do
49-
poke p (fromIntegral (shiftR w 8) :: Word8)
50-
poke (p `plusPtr` 1) (fromIntegral w :: Word8)
5153
{-# INLINE putWord16BE #-}
54+
putWord16BE :: Word16 -> Encode
55+
putWord16BE w = fixed 2 $ \p -> poke (castPtr p) (byteSwap16 w)
5256

53-
putWord32BE :: Word32 -> Encode
54-
putWord32BE w = fixedPrim 4 $ \p -> do
55-
poke p (fromIntegral (shiftR w 24) :: Word8)
56-
poke (p `plusPtr` 1) (fromIntegral (shiftR w 16) :: Word8)
57-
poke (p `plusPtr` 2) (fromIntegral (shiftR w 8) :: Word8)
58-
poke (p `plusPtr` 3) (fromIntegral w :: Word8)
5957
{-# INLINE putWord32BE #-}
58+
putWord32BE :: Word32 -> Encode
59+
putWord32BE w = fixed 4 $ \p -> poke (castPtr p) (byteSwap32 w)
6060

61-
putInt32BE :: Int32 -> Encode
62-
putInt32BE = putWord32BE . fromIntegral
63-
{-# INLINE putInt32BE #-}
61+
{-# INLINE putWord64BE #-}
62+
putWord64BE :: Word64 -> Encode
63+
putWord64BE w = fixed 8 $ \p -> poke (castPtr p) (byteSwap64 w)
6464

65+
{-# INLINE putInt16BE #-}
6566
putInt16BE :: Int16 -> Encode
6667
putInt16BE = putWord16BE . fromIntegral
67-
{-# INLINE putInt16BE #-}
6868

69-
putByteString :: ByteString -> Encode
70-
putByteString bs =
71-
let (ptr, offset, len) = toForeignPtr bs
72-
in Encode len $ pokeFromForeignPtr ptr offset len
73-
{-# INLINE putByteString #-}
74-
75-
-- | C-like string
76-
putPgString :: ByteString -> Encode
77-
putPgString bs = putByteString bs <> putWord8 0
78-
{-# INLINE putPgString #-}
69+
{-# INLINE putInt32BE #-}
70+
putInt32BE :: Int32 -> Encode
71+
putInt32BE = putWord32BE . fromIntegral
7972

73+
{-# INLINE putInt64BE #-}
74+
putInt64BE :: Int64 -> Encode
75+
putInt64BE = putWord64BE . fromIntegral
76+
77+
{-# INLINE putFloat32BE #-}
78+
putFloat32BE :: Float -> Encode
79+
putFloat32BE float = fixed 4 $ \ptr -> byteSwap32 <$> floatToWord float
80+
>>= poke (castPtr ptr)
81+
82+
{-# INLINE putFloat64BE #-}
83+
putFloat64BE :: Double -> Encode
84+
putFloat64BE double = fixed 8 $ \ptr -> byteSwap64 <$> floatToWord double
85+
>>= poke (castPtr ptr)
86+
87+
{-# INLINE floatToWord #-}
88+
floatToWord :: (Storable word, Storable float) => float -> IO word
89+
floatToWord float = alloca $ \buf -> do
90+
poke (castPtr buf) float
91+
peek buf

0 commit comments

Comments
 (0)