module Lmdb.Codec where
import Data.Word
import Control.Monad
import Lmdb.Types
import Foreign.Storable (peek,poke,sizeOf)
import Foreign.Ptr (Ptr,castPtr,plusPtr)
import Foreign.C.Types (CSize(..))
import Data.Text.Internal (Text(..))
import GHC.Int (Int64(I64#))
import GHC.Ptr (Ptr(Ptr))
import GHC.Types (IO(IO),Int(I#))
import Data.Bits (unsafeShiftR,unsafeShiftL)
import qualified Data.Text.Array as TextArray
import GHC.Prim (newByteArray#,copyAddrToByteArray#,unsafeFreezeByteArray#,copyByteArrayToAddr#)
import Data.ByteString.Internal (ByteString(PS))
import Foreign.Marshal.Utils (copyBytes)
import Foreign.ForeignPtr (withForeignPtr,mallocForeignPtrBytes)
import qualified Data.Vector.Unboxed as UVector
import qualified Data.Vector.Unboxed.Mutable as MUVector
import qualified Data.Vector.Primitive as PVector
import Data.Primitive.Types (Prim)
import Data.Vector.Unboxed (Unbox)
import Foreign.Storable (Storable,sizeOf,peekElemOff)
import Control.Monad.Trans.State.Strict (StateT(runStateT),get,put)
import Control.Monad.Trans.Class
text :: Codec 'Variable Text
text = Codec encodeText (Decoding decodeText)
word :: Codec 'MachineWord Word
word = Codec unsafeEncodePaddedIntegral unsafeDecodePaddedIntegral
int :: Codec 'MachineWord Int
int = Codec unsafeEncodePaddedIntegral unsafeDecodePaddedIntegral
word32Padded :: Codec 'MachineWord Word32
word32Padded = Codec unsafeEncodePaddedIntegral unsafeDecodePaddedIntegral
word64 :: Codec 'Fixed Word64
word64 = Codec encodeStorable (decodeStorable "Word64")
word16 :: Codec 'Fixed Word16
word16 = Codec encodeStorable (decodeStorable "Word16")
plusPtrTyped :: Ptr a -> Int -> Ptr a
plusPtrTyped = plusPtr
byteString :: Codec 'Variable ByteString
byteString = Codec encodeByteString (Decoding decodeByteString)
unit :: Codec 'Fixed ()
unit = Codec encodeEmpty (decodeConst ())
encodeEmpty :: Encoding 'Fixed a
encodeEmpty = EncodingFixed 0 $ \_ -> FixedPoke $ \_ -> return ()
decodeConst :: a -> Decoding a
decodeConst a = Decoding $ \sz _ -> if sz == 0
then return a
else fail "decodeConst: encountered a non-zero size"
throughByteString :: (a -> ByteString) -> (ByteString -> Maybe a) -> Codec 'Variable a
throughByteString encode decode = Codec
(encodeThroughByteString encode)
(Decoding
(\sz ptr -> do
bs <- decodeByteString sz ptr
case decode bs of
Just a -> return a
Nothing -> fail "throughByteString: failed while decoding LMDB data"
)
)
encodeByteString :: Encoding 'Variable ByteString
encodeByteString = encodeThroughByteString id
encodeThroughByteString :: (a -> ByteString) -> Encoding 'Variable a
encodeThroughByteString f =
EncodingVariable $ \a -> let (PS fptr off len) = f a in SizedPoke
(fromIntegral len)
(\targetPtr -> withForeignPtr fptr $ \sourcePtr ->
fastMemcpyBytePtr sourcePtr targetPtr off len
)
decodeByteString :: CSize -> Ptr Word8 -> IO ByteString
decodeByteString sz source = do
let szInt = fromIntegral sz
fptr <- mallocForeignPtrBytes szInt
withForeignPtr fptr $ \target -> do
fastMemcpyBytePtr source target 0 szInt
return (PS fptr 0 szInt)
decodeText :: CSize -> Ptr Word8 -> IO Text
decodeText (CSize szWord64) (Ptr addr) =
let !szInt@(I# sz) = fromIntegral szWord64
in IO (\ s1 ->
case newByteArray# sz s1 of
(# s2, mutByteArr #) -> case copyAddrToByteArray# addr mutByteArr 0# sz s2 of
s3 -> case unsafeFreezeByteArray# mutByteArr s3 of
(# s4, byteArr #) -> (# s4, Text (TextArray.Array byteArr) 0 (unsafeShiftR szInt 1) #)
)
encodeText :: Encoding 'Variable Text
encodeText = EncodingVariable $ \(Text arr off len) -> SizedPoke
(CSize $ fromIntegral $ unsafeShiftL len 1)
(\ptr -> fastMemcpyTextArray arr (castPtr ptr) off len)
fastMemcpyBytePtr :: Ptr Word8 -> Ptr Word8 -> Int -> Int -> IO ()
fastMemcpyBytePtr source target off len =
copyBytes target (plusPtr source off :: Ptr Word8) len
fastMemcpyTextArray :: TextArray.Array -> Ptr Word8 -> Int -> Int -> IO ()
fastMemcpyTextArray (TextArray.Array byteArr) (Ptr addr) off len =
let !(I# offWord8) = unsafeShiftL off 1
!(I# lenWord8) = unsafeShiftL len 1
in IO (\ s1 -> case copyByteArrayToAddr# byteArr offWord8 addr lenWord8 s1 of
s2 -> (# s2, () #)
)
slowMemcpyTextArray :: TextArray.Array -> Ptr Word16 -> Int -> Int -> IO ()
slowMemcpyTextArray arr ptr off len = go off
where
end = off + len
go !ix
| ix >= end = return ()
| otherwise = do
let w16 = TextArray.unsafeIndex arr ix
poke (plusPtr ptr (unsafeShiftL ix 1) :: Ptr Word16) w16
go (ix + 1)
unsafeEncodePaddedIntegral :: Integral a => Encoding 'MachineWord a
unsafeEncodePaddedIntegral = EncodingMachineWord $ \w -> FixedPoke $ \ptr ->
poke (castPtr ptr :: Ptr Word) (fromIntegral w :: Word)
unsafeDecodePaddedIntegral :: Integral a => Decoding a
unsafeDecodePaddedIntegral = Decoding $ \sz ptr -> if sz == sizeOfMachineWord
then fmap fromIntegral (peek (castPtr ptr :: Ptr Word))
else fail "lmdb failure decoding machine sized word or integral"
decodeIntegral :: Integral a => CSize -> Ptr Word8 -> IO a
decodeIntegral sz = fmap fromIntegral . decodeWord64 sz
encodePaddedWord32 :: Encoding 'MachineWord Word32
encodePaddedWord32 = EncodingMachineWord $ \w -> FixedPoke $ \ptr ->
poke (castPtr ptr :: Ptr Word) (fromIntegral w :: Word)
encodeWord :: Encoding 'MachineWord Word
encodeWord = EncodingMachineWord $ \w -> FixedPoke $ \ptr ->
poke (castPtr ptr :: Ptr Word) w
encodeStorable :: forall a. Storable a => Encoding 'Fixed a
encodeStorable = EncodingFixed (fromIntegral $ sizeOf (undefined :: a)) $ \a -> FixedPoke $ \ptr ->
poke (castPtr ptr :: Ptr a) a
encodeWord64 :: Encoding 'Fixed Word64
encodeWord64 = EncodingFixed 8 $ \w -> FixedPoke $ \ptr ->
poke (castPtr ptr :: Ptr Word64) w
decodeWord64 :: CSize -> Ptr Word8 -> IO Word64
decodeWord64 sz ptr = if sz == 8
then peek (castPtr ptr :: Ptr Word64)
else fail "lmdb failure decoding 64-bit integral type"
decodeStorable :: forall a. Storable a => String -> Decoding a
decodeStorable descr = Decoding $ \sz ptr -> if fromIntegral sz == sizeOf (undefined :: a)
then peek (castPtr ptr :: Ptr a)
else fail $ "lmdb failure decoding " ++ descr ++ " storable type due to mismatched size"
sizeOfMachineWord :: CSize
sizeOfMachineWord = fromIntegral (sizeOf (undefined :: Word))
whileM_ :: (Monad m) => m Bool -> m a -> m ()
whileM_ p f = go
where go = do
x <- p
if x
then f >> go
else return ()