Put.hs

-- | 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 = putWord8 TAG_Null

-- | Encode a PackStream boolean value.
putBoolean :: Bool -> Put
putBoolean False = putWord8 TAG_false
putBoolean True  = putWord8 TAG_true

-- | Encodes an 'Int' to PackStream
--
-- See also 'PSInteger' and its 'Binary' instance.
putInteger :: Int -> Put
putInteger = put . toPSInteger

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

-- | Encode an 'Int64' to PackStream.
putInteger64 :: Int64 -> Put
putInteger64 = put . toPSInteger

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

-- | Encode a 'Double' as a PackStream 64-bit float.
putFloat :: Double -> Put
putFloat d = putWord8 TAG_Float >> putFloat64be d

-- | Encode a 'T.Text' as a PackStream UTF-8 string.
putString :: T.Text -> Put
putString t = do
  let bs = T.encodeUtf8 t
  toSizeM ("putString: data exceeds 2^32-1 byte limit of PackStream") (S.length bs) >>= \case
    len | len <= 15     -> putWord8 (TAG_STRING_SHORT .|. fromIntegral len)
        | len < 0x100   -> putWord8 TAG_STRING_8  >> putWord8    (fromIntegral len)
        | len < 0x10000 -> putWord8 TAG_STRING_16 >> putWord16be (fromIntegral len)
        | otherwise     -> putWord8 TAG_STRING_32 >> putWord32be len
  putByteString bs

-- | Encode a 'S.ByteString' as a PackStream byte array.
putBytes :: S.ByteString -> Put
putBytes bs = do
  toSizeM ("putBytes: data exceeds 2^32-1 byte limit of PackStream") (S.length bs) >>= \case
    len | len < 0x100   -> putWord8 TAG_BYTE_ARRAY_8  >> putWord8    (fromIntegral len)
        | len < 0x10000 -> putWord8 TAG_BYTE_ARRAY_16 >> putWord16be (fromIntegral len)
        | otherwise     -> putWord8 TAG_BYTE_ARRAY_32 >> putWord32be len
  putByteString bs

-- | Encode a 'V.Vector' as a PackStream list using the given element encoder.
putList :: (a -> Put) -> V.Vector a -> Put
putList p xs = do
  len <- toSizeM ("putList: data exceeds 2^32-1 element limit of PackStream") (V.length xs)
  putList' len (V.mapM_ p 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' len putter = do
  case () of
    _   | len <= 15     -> putWord8 (TAG_LIST_SHORT .|. fromIntegral len)
        | len < 0x100   -> putWord8 TAG_LIST_8  >> putWord8    (fromIntegral len)
        | len < 0x10000 -> putWord8 TAG_LIST_16 >> putWord16be (fromIntegral len)
        | otherwise     -> putWord8 TAG_LIST_32 >> putWord32be len
  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 p q xs = do
  let len' = H.size xs
  if len' > 2_147_483_647 then
    error "putDictionary: data exceeds 2^32-1 element limit of PackStream"
  else case len' of
    len | len <= 15     -> putWord8 (TAG_DICTIONARY_SHORT .|. fromIntegral len)
        | len < 0x100   -> putWord8 TAG_DICTIONARY_8  >> putWord8    (fromIntegral len)
        | len < 0x10000 -> putWord8 TAG_DICTIONARY_16 >> putWord16be (fromIntegral len)
        | otherwise     -> putWord8 TAG_DICTIONARY_32 >> putWord32be (fromIntegral len)
  mapM_ (\(a, b) -> p a >> q b) $ H.toList xs


toSizeM :: String -> Int -> PutM Word32
toSizeM label len0 = maybe (error label) pure (intCastMaybe len0)