module Data.Flat.Encoder.Prim (
eBitsF,
eFloatF,
eDoubleF,
eUTF16F,
eUTF8F,
eCharF,
eNaturalF,
eIntegerF,
eInt64F,
eInt32F,
eIntF,
eInt16F,
eInt8F,
eWordF,
eWord64F,
eWord32F,
eWord16F,
eBytesF,
eLazyBytesF,
eShortBytesF,
eWord8F,
eFillerF,
eBoolF,
eTrueF,
eFalseF,
varWordF,
w7l,
) where
import Control.Monad
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as L
import qualified Data.ByteString.Lazy.Internal as L
import qualified Data.ByteString.Short.Internal as SBS
import Data.Char
import Data.Flat.Encoder.Types
import Data.Flat.Memory
import Data.Flat.Types
import Data.FloatCast
import Data.Primitive.ByteArray
import qualified Data.Text as T
import qualified Data.Text.Array as TA
import qualified Data.Text.Encoding as T
import qualified Data.Text.Internal as T
import Data.ZigZag
import Foreign
import System.Endian
#include "MachDeps.h"
eFloatF :: Float -> Prim
eFloatF = eWord32BEF . floatToWord
eDoubleF :: Double -> Prim
eDoubleF = eWord64BEF . doubleToWord
eWord64BEF :: Word64 -> Prim
eWord64BEF = eWord64E toBE64
eWord32BEF :: Word32 -> Prim
eWord32BEF = eWord32E toBE32
eCharF :: Char -> Prim
eCharF = eWord32F . fromIntegral . ord
eWordF :: Word -> Prim
eIntF :: Int -> Prim
#if WORD_SIZE_IN_BITS == 64
eWordF = eWord64F . (fromIntegral :: Word -> Word64)
eIntF = eInt64F . (fromIntegral :: Int -> Int64)
#elif WORD_SIZE_IN_BITS == 32
eWordF = eWord32F . (fromIntegral :: Word -> Word32)
eIntF = eInt32F . (fromIntegral :: Int -> Int32)
#else
#error expected WORD_SIZE_IN_BITS to be 32 or 64
#endif
eInt8F :: Int8 -> Prim
eInt8F = eWord8F . zzEncode
eInt16F :: Int16 -> Prim
eInt16F = eWord16F . zzEncode
eInt32F :: Int32 -> Prim
eInt32F = eWord32F . zzEncode
eInt64F :: Int64 -> Prim
eInt64F = eWord64F . zzEncode
eIntegerF :: Integer -> Prim
eIntegerF = eIntegralF . zzEncodeInteger
eNaturalF :: Natural -> Prim
eNaturalF = eIntegralF . toInteger
eIntegralF :: (Bits t, Integral t) => t -> Prim
eIntegralF t = let vs = w7l t
in eIntegralW vs
w7l :: (Bits t, Integral t) => t -> [Word8]
w7l t = let l = low7 t
t' = t `shiftR` 7
in if t' == 0
then [l]
else w7 l : w7l t'
where
w7 :: Word8 -> Word8
w7 l = l .|. 0x80
eIntegralW :: [Word8] -> Prim
eIntegralW vs s@(S op _ o) | o == 0 = foldM pokeWord' op vs >>= \op' -> return (S op' 0 0)
| otherwise = foldM (flip eWord8F) s vs
eWord8F :: Word8 -> Prim
eWord8F t s@(S op _ o) | o==0 = pokeWord op t
| otherwise = pokeByteUnaligned t s
eWord32E :: (Word32 -> Word32) -> Word32 -> Prim
eWord32E conv t (S op w o) | o==0 = pokeW conv op t >> skipBytes op 4
| otherwise = pokeW conv op (fromIntegral w `shiftL` 24 .|. t `shiftR` o) >> return (S (plusPtr op 4) (fromIntegral t `shiftL` (8o)) o)
eWord64E :: (Word64 -> Word64) -> Word64 -> Prim
eWord64E conv t (S op w o) | o==0 = pokeW conv op t >> skipBytes op 8
| otherwise = pokeW conv op (fromIntegral w `shiftL` 56 .|. t `shiftR` o) >> return (S (plusPtr op 8) (fromIntegral t `shiftL` (8o)) o)
eWord16F :: Word16 -> Prim
eWord16F = varWordF
eWord32F :: Word32 -> Prim
eWord32F = varWordF
eWord64F :: Word64 -> Prim
eWord64F = varWordF
varWordF :: (Bits t, Integral t) => t -> Prim
varWordF t s@(S _ _ o) | o == 0 = varWord pokeByteAligned t s
| otherwise = varWord pokeByteUnaligned t s
varWord :: (Bits t, Integral t) => (Word8 -> Prim) -> t -> Prim
varWord writeByte t s
| t < 128 = writeByte (fromIntegral t) s
| t < 16384 = varWord2_ writeByte t s
| t < 2097152 = varWord3_ writeByte t s
| otherwise = varWordN_ writeByte t s
where
varWord2_ writeByte t s = writeByte (fromIntegral t .|. 0x80) s >>= writeByte (fromIntegral (t `shiftR` 7) .&. 0x7F)
varWord3_ writeByte t s = writeByte (fromIntegral t .|. 0x80) s >>= writeByte (fromIntegral (t `shiftR` 7) .|. 0x80) >>= writeByte (fromIntegral (t `shiftR` 14) .&. 0x7F)
varWordN_ :: (Bits t, Integral t) => (Word8 -> Prim) -> t -> Prim
varWordN_ writeByte = go
where
go !v !st =
let !l = low7 v
!v' = v `shiftR` 7
in if v' == 0
then writeByte l st
else writeByte (l .|. 0x80) st >>= go v'
low7 :: (Integral a) => a -> Word8
low7 t = fromIntegral t .&. 0x7F
eUTF8F :: T.Text -> Prim
eUTF8F = eBytesF . T.encodeUtf8
eUTF16F :: T.Text -> Prim
eUTF16F t = eFillerF >=> eUTF16F_ t
where
eUTF16F_ !(T.Text (TA.Array array) w16Off w16Len) s = writeArray array (2 * w16Off) (2 * w16Len) (nextPtr s)
eLazyBytesF :: L.ByteString -> Prim
eLazyBytesF bs = eFillerF >=> \s -> write bs (nextPtr s)
where
write lbs op = do
case lbs of
L.Chunk h t -> writeBS h op >>= write t
L.Empty -> pokeWord op 0
eShortBytesF :: SBS.ShortByteString -> Prim
eShortBytesF bs = eFillerF >=> eShortBytesF_ bs
eShortBytesF_ :: SBS.ShortByteString -> Prim
eShortBytesF_ bs@(SBS.SBS arr) = \(S op _ 0) -> writeArray arr 0 (SBS.length bs) op
writeArray :: ByteArray# -> Int -> Int -> Ptr Word8 -> IO S
writeArray arr soff slen sop = do
op' <- go soff slen sop
pokeWord op' 0
where
go !off !len !op
| len == 0 = return op
| otherwise =
let l = min 255 len
in pokeWord' op (fromIntegral l) >>=
pokeByteArray arr off l >>=
go (off+l) (len l)
eBytesF :: B.ByteString -> Prim
eBytesF bs = eFillerF >=> eBytesF_
where
eBytesF_ s = do
op' <- writeBS bs (nextPtr s)
pokeWord op' 0
eBitsF :: NumBits -> Word8 -> Prim
eBitsF 1 0 = eFalseF
eBitsF 1 1 = eTrueF
eBitsF 2 0 = eFalseF >=> eFalseF
eBitsF 2 1 = eFalseF >=> eTrueF
eBitsF 2 2 = eTrueF >=> eFalseF
eBitsF 2 3 = eTrueF >=> eTrueF
eBitsF n t = \(S op w o) ->
let o' = o + n
f = 8 o'
in if | f > 0 -> return $ S op (w .|. (t `shiftL` f)) o'
| f == 0 -> pokeWord op (w .|. t)
| otherwise -> let o'' = f
in poke op (w .|. (t `shiftR` o'')) >> return (S (plusPtr op 1) (t `shiftL` (8o'')) o'')
eBoolF :: Bool -> Prim
eBoolF False = eFalseF
eBoolF True = eTrueF
eTrueF :: Prim
eTrueF (S op w o) | o == 7 = pokeWord op (w .|. 1)
| otherwise = return (S op (setBit w (7o)) (o+1))
eFalseF :: Prim
eFalseF (S op w o) | o == 7 = pokeWord op w
| otherwise = return (S op w (o+1))
eFillerF :: Prim
eFillerF (S op w _) = pokeWord op (w .|. 1)
pokeByteUnaligned :: Word8 -> Prim
pokeByteUnaligned t (S op w o) = poke op (w .|. (t `shiftR` o)) >> return (S (plusPtr op 1) (t `shiftL` (8o)) o)
pokeByteAligned :: Word8 -> Prim
pokeByteAligned t (S op _ _) = pokeWord op t
pokeWord :: Storable a => Ptr a -> a -> IO S
pokeWord op w = poke op w >> skipByte op
pokeWord' :: Storable a => Ptr a -> a -> IO (Ptr b)
pokeWord' op w = poke op w >> return (plusPtr op 1)
pokeW :: Storable a => (t -> a) -> Ptr a1 -> t -> IO ()
pokeW conv op t = poke (castPtr op) (conv t)
skipByte :: Monad m => Ptr a -> m S
skipByte op = return (S (plusPtr op 1) 0 0)
skipBytes :: Monad m => Ptr a -> Int -> m S
skipBytes op n = return (S (plusPtr op n) 0 0)
writeBS :: B.ByteString -> Ptr Word8 -> IO (Ptr Word8)
writeBS bs op
| B.length bs == 0 = return op
| otherwise =
let (h, t) = B.splitAt 255 bs
in pokeWord' op (fromIntegral $ B.length h :: Word8) >>= pokeByteString h >>= writeBS t