-- | Internal module. Not part of the public API.
module Data.PackStream.Put (
  putNull, putBoolean, putFloat,
  putInteger, putInteger64,
  putString, putBytes, putList, putList', putDictionary
  ) where

import           Compat.Prelude
import           Prelude

import qualified Data.ByteString          as S
import qualified Data.Text                as T
import qualified Data.Text.Encoding       as T
import qualified Data.Vector              as V
import qualified Data.HashMap.Lazy        as H

import           Compat.Binary
import           Data.PackStream.Integer
import           Data.PackStream.Tags

-- | Encode a PackStream null value.
putNull :: Put
putNull :: Put
putNull = Word8 -> Put
putWord8 Word8
forall a. (Eq a, Num a) => a
TAG_Null

-- | Encode a PackStream boolean value.
putBoolean :: Bool -> Put
putBoolean :: Bool -> Put
putBoolean Bool
False = Word8 -> Put
putWord8 Word8
forall a. (Eq a, Num a) => a
TAG_false
putBoolean Bool
True  = Word8 -> Put
putWord8 Word8
forall a. (Eq a, Num a) => a
TAG_true

-- | Encodes an 'Int' to PackStream
--
-- See also 'PSInteger' and its 'Binary' instance.
putInteger :: Int -> Put
putInteger :: Int -> Put
putInteger = PSInteger -> Put
forall t. Persist t => t -> Put
put (PSInteger -> Put) -> (Int -> PSInteger) -> Int -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> PSInteger
forall a. ToPSInteger a => a -> PSInteger
toPSInteger

-- putWord :: Word -> Put
-- putWord = put . toPSInteger

-- | Encode an 'Int64' to PackStream.
putInteger64 :: Int64 -> Put
putInteger64 :: Int64 -> Put
putInteger64 = PSInteger -> Put
forall t. Persist t => t -> Put
put (PSInteger -> Put) -> (Int64 -> PSInteger) -> Int64 -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> PSInteger
forall a. ToPSInteger a => a -> PSInteger
toPSInteger

-- putWord64 :: Word64 -> Put
-- putWord64 = put . toPSInteger

-- | Encode a 'Double' as a PackStream 64-bit float.
putFloat :: Double -> Put
putFloat :: Double -> Put
putFloat Double
d = Word8 -> Put
putWord8 Word8
forall a. (Eq a, Num a) => a
TAG_Float Put -> Put -> Put
forall a b. Put a -> Put b -> Put b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Double -> Put
putFloat64be Double
d

-- | Encode a 'T.Text' as a PackStream UTF-8 string.
putString :: T.Text -> Put
putString :: Text -> Put
putString Text
t = do
  let bs :: ByteString
bs = Text -> ByteString
T.encodeUtf8 Text
t
  String -> Int -> PutM Word32
toSizeM (String
"putString: data exceeds 2^32-1 byte limit of PackStream") (ByteString -> Int
S.length ByteString
bs) PutM Word32 -> (Word32 -> Put) -> Put
forall a b. Put a -> (a -> Put b) -> Put b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Word32
len | Word32
len Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word32
15     -> Word8 -> Put
putWord8 (Word8
forall a. (Eq a, Num a) => a
TAG_STRING_SHORT Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word32 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
len)
        | Word32
len Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
< Word32
0x100   -> Word8 -> Put
putWord8 Word8
forall a. (Eq a, Num a) => a
TAG_STRING_8  Put -> Put -> Put
forall a b. Put a -> Put b -> Put b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Word8 -> Put
putWord8    (Word32 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
len)
        | Word32
len Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
< Word32
0x10000 -> Word8 -> Put
putWord8 Word8
forall a. (Eq a, Num a) => a
TAG_STRING_16 Put -> Put -> Put
forall a b. Put a -> Put b -> Put b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Word16 -> Put
putWord16be (Word32 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
len)
        | Bool
otherwise     -> Word8 -> Put
putWord8 Word8
forall a. (Eq a, Num a) => a
TAG_STRING_32 Put -> Put -> Put
forall a b. Put a -> Put b -> Put b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Word32 -> Put
putWord32be Word32
len
  ByteString -> Put
putByteString ByteString
bs

-- | Encode a 'S.ByteString' as a PackStream byte array.
putBytes :: S.ByteString -> Put
putBytes :: ByteString -> Put
putBytes ByteString
bs = do
  String -> Int -> PutM Word32
toSizeM (String
"putBytes: data exceeds 2^32-1 byte limit of PackStream") (ByteString -> Int
S.length ByteString
bs) PutM Word32 -> (Word32 -> Put) -> Put
forall a b. Put a -> (a -> Put b) -> Put b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Word32
len | Word32
len Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
< Word32
0x100   -> Word8 -> Put
putWord8 Word8
forall a. (Eq a, Num a) => a
TAG_BYTE_ARRAY_8  Put -> Put -> Put
forall a b. Put a -> Put b -> Put b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Word8 -> Put
putWord8    (Word32 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
len)
        | Word32
len Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
< Word32
0x10000 -> Word8 -> Put
putWord8 Word8
forall a. (Eq a, Num a) => a
TAG_BYTE_ARRAY_16 Put -> Put -> Put
forall a b. Put a -> Put b -> Put b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Word16 -> Put
putWord16be (Word32 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
len)
        | Bool
otherwise     -> Word8 -> Put
putWord8 Word8
forall a. (Eq a, Num a) => a
TAG_BYTE_ARRAY_32 Put -> Put -> Put
forall a b. Put a -> Put b -> Put b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Word32 -> Put
putWord32be Word32
len
  ByteString -> Put
putByteString ByteString
bs

-- | Encode a 'V.Vector' as a PackStream list using the given element encoder.
putList :: (a -> Put) -> V.Vector a -> Put
putList :: forall a. (a -> Put) -> Vector a -> Put
putList a -> Put
p Vector a
xs = do
  Word32
len <- String -> Int -> PutM Word32
toSizeM (String
"putList: data exceeds 2^32-1 element limit of PackStream") (Vector a -> Int
forall a. Vector a -> Int
V.length Vector a
xs)
  Word32 -> Put -> Put
putList' Word32
len ((a -> Put) -> Vector a -> Put
forall (m :: * -> *) a b. Monad m => (a -> m b) -> Vector a -> m ()
V.mapM_ a -> Put
p Vector a
xs)

putList' :: Word32 -- ^ number of array elements
          -> Put    -- ^ 'Put' action emitting array elements (__NOTE__: it's the responsibility of the caller to ensure that the declared array length matches exactly the data generated by the 'Put' action)
          -> Put
putList' :: Word32 -> Put -> Put
putList' Word32
len Put
putter = do
  case () of
    ()
_   | Word32
len Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word32
15     -> Word8 -> Put
putWord8 (Word8
forall a. (Eq a, Num a) => a
TAG_LIST_SHORT Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Word32 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
len)
        | Word32
len Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
< Word32
0x100   -> Word8 -> Put
putWord8 Word8
forall a. (Eq a, Num a) => a
TAG_LIST_8  Put -> Put -> Put
forall a b. Put a -> Put b -> Put b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Word8 -> Put
putWord8    (Word32 -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
len)
        | Word32
len Word32 -> Word32 -> Bool
forall a. Ord a => a -> a -> Bool
< Word32
0x10000 -> Word8 -> Put
putWord8 Word8
forall a. (Eq a, Num a) => a
TAG_LIST_16 Put -> Put -> Put
forall a b. Put a -> Put b -> Put b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Word16 -> Put
putWord16be (Word32 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
len)
        | Bool
otherwise     -> Word8 -> Put
putWord8 Word8
forall a. (Eq a, Num a) => a
TAG_LIST_32 Put -> Put -> Put
forall a b. Put a -> Put b -> Put b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Word32 -> Put
putWord32be Word32
len
  Put
putter

-- | Encode a 'H.HashMap' as a PackStream dictionary using the given key and value encoders.
putDictionary :: (a -> Put) -> (b -> Put) -> H.HashMap a b -> Put
putDictionary :: forall a b. (a -> Put) -> (b -> Put) -> HashMap a b -> Put
putDictionary a -> Put
p b -> Put
q HashMap a b
xs = do
  let len' :: Int
len' = HashMap a b -> Int
forall k v. HashMap k v -> Int
H.size HashMap a b
xs
  if Int
len' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
2_147_483_647 then
    String -> Put
forall a. HasCallStack => String -> a
error String
"putDictionary: data exceeds 2^32-1 element limit of PackStream"
  else case Int
len' of
    Int
len | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
15     -> Word8 -> Put
putWord8 (Word8
forall a. (Eq a, Num a) => a
TAG_DICTIONARY_SHORT Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.|. Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)
        | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0x100   -> Word8 -> Put
putWord8 Word8
forall a. (Eq a, Num a) => a
TAG_DICTIONARY_8  Put -> Put -> Put
forall a b. Put a -> Put b -> Put b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Word8 -> Put
putWord8    (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)
        | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0x10000 -> Word8 -> Put
putWord8 Word8
forall a. (Eq a, Num a) => a
TAG_DICTIONARY_16 Put -> Put -> Put
forall a b. Put a -> Put b -> Put b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Word16 -> Put
putWord16be (Int -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)
        | Bool
otherwise     -> Word8 -> Put
putWord8 Word8
forall a. (Eq a, Num a) => a
TAG_DICTIONARY_32 Put -> Put -> Put
forall a b. Put a -> Put b -> Put b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Word32 -> Put
putWord32be (Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len)
  ((a, b) -> Put) -> [(a, b)] -> Put
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\(a
a, b
b) -> a -> Put
p a
a Put -> Put -> Put
forall a b. Put a -> Put b -> Put b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> b -> Put
q b
b) ([(a, b)] -> Put) -> [(a, b)] -> Put
forall a b. (a -> b) -> a -> b
$ HashMap a b -> [(a, b)]
forall k v. HashMap k v -> [(k, v)]
H.toList HashMap a b
xs


toSizeM :: String -> Int -> PutM Word32
toSizeM :: String -> Int -> PutM Word32
toSizeM String
label Int
len0 = PutM Word32
-> (Word32 -> PutM Word32) -> Maybe Word32 -> PutM Word32
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (String -> PutM Word32
forall a. HasCallStack => String -> a
error String
label) Word32 -> PutM Word32
forall a. a -> Put a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> Maybe Word32
forall a b.
(Integral a, Integral b, Bits a, Bits b) =>
a -> Maybe b
intCastMaybe Int
len0)