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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy