|
| 1 | +module Database.PostgreSQL.Protocol.Store where |
| 2 | + |
| 3 | +import qualified Data.ByteString as B |
| 4 | +import qualified Data.ByteString.Internal as B(toForeignPtr) |
| 5 | +import qualified Data.Vector as V |
| 6 | +import Data.Store.Core |
| 7 | +import Data.Int (Int16, Int32) |
| 8 | +import Data.Word (Word8) |
| 9 | +import Foreign |
| 10 | +import Data.Monoid |
| 11 | +import Data.Foldable |
| 12 | + |
| 13 | + |
| 14 | +data Encode = Encode Int (Poke ()) |
| 15 | + |
| 16 | +instance Monoid Encode where |
| 17 | + mempty = Encode 0 . Poke $ \_ offset -> pure (offset, ()) |
| 18 | + (Encode len1 f1) `mappend` (Encode len2 f2) = Encode (len1 + len2) (f1 *> f2) |
| 19 | + |
| 20 | + |
| 21 | +runEncode :: Encode -> B.ByteString |
| 22 | +runEncode (Encode len f) = unsafeEncodeWith f len |
| 23 | + |
| 24 | +fixedPrim :: Int -> (Ptr Word8 -> IO ()) -> Encode |
| 25 | +fixedPrim len f = Encode len . Poke $ \state offset -> do |
| 26 | + f $ pokeStatePtr state `plusPtr` offset |
| 27 | + let !newOffset = offset + len |
| 28 | + return (newOffset, ()) |
| 29 | + |
| 30 | +putWord8 :: Word8 -> Encode |
| 31 | +putWord8 w = fixedPrim 1 $ \p -> poke p w |
| 32 | + |
| 33 | +putWord16BE :: Word16 -> Encode |
| 34 | +putWord16BE w = fixedPrim 2 $ \p -> do |
| 35 | + poke p (fromIntegral (shiftR w 8) :: Word8) |
| 36 | + poke (p `plusPtr` 1) (fromIntegral w :: Word8) |
| 37 | + |
| 38 | +putWord32BE :: Word32 -> Encode |
| 39 | +putWord32BE w = fixedPrim 4 $ \p -> do |
| 40 | + poke p (fromIntegral (shiftR w 24) :: Word8) |
| 41 | + poke (p `plusPtr` 1) (fromIntegral (shiftR w 16) :: Word8) |
| 42 | + poke (p `plusPtr` 2) (fromIntegral (shiftR w 8) :: Word8) |
| 43 | + poke (p `plusPtr` 3) (fromIntegral w :: Word8) |
| 44 | + |
| 45 | +putInt32BE :: Int32 -> Encode |
| 46 | +putInt32BE = putWord32BE . fromIntegral |
| 47 | + |
| 48 | +putInt16BE :: Int16 -> Encode |
| 49 | +putInt16BE = putWord16BE . fromIntegral |
| 50 | + |
| 51 | +putByteString :: B.ByteString ->
BDBE
Encode |
| 52 | +putByteString bs = |
| 53 | + let (ptr, offset, len) = B.toForeignPtr bs |
| 54 | + in Encode len $ pokeFromForeignPtr ptr offset len |
| 55 | + |
| 56 | +-- | C-like string |
| 57 | +putPgString :: B.ByteString -> Encode |
| 58 | +putPgString bs = putByteString bs <> putWord8 0 |
| 59 | + |
| 60 | +-- | List with prepended length |
| 61 | +putPgList :: V.Vector Encode -> Encode |
| 62 | +putPgList v = putInt16BE (fromIntegral $ V.length v) <> fold v |
| 63 | + |
0 commit comments