{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
module Data.ZMQParse (
    getInt8
  , getInt16
  , getInt32
  , parseString
  , parseStrings
  , parseLongString
  , parseKV
  , parseMap
  , putByteStringLen
  , putByteStrings
  , putLongByteStringLen
  , putKV
  , putMap
  , Get.Get()
  , runGet
  , Get.getByteString
  , Put.Put
  , Put.PutM
  , runPut
  , Put.putInt8
  , Put.putWord8
  , Put.putByteString
  , Put.putWord16be
  , Put.putWord32be
  , Put.putInt16be
  , Put.putInt32be
  )
  where

import Prelude hiding (putStrLn, take)
import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as B
import qualified Data.ByteString.Lazy as BL

import Data.Binary.Get hiding (getInt8, runGet)
import Data.Binary.Put hiding (runPut)

import qualified Data.Map as M

import qualified Data.Binary.Get as Get
import qualified Data.Binary.Put as Put

getInt8 :: (Integral a) => Get a
getInt8 :: Get a
getInt8  = Word8 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word8 -> a) -> Get Word8 -> Get a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word8
getWord8
getInt16 :: (Integral a) => Get a
getInt16 :: Get a
getInt16 = Word16 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word16 -> a) -> Get Word16 -> Get a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word16
getWord16be
getInt32 :: (Integral a) => Get a
getInt32 :: Get a
getInt32 = Word32 -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32 -> a) -> Get Word32 -> Get a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Word32
getWord32be

parseString :: Get ByteString
parseString :: Get ByteString
parseString = do
  Int
len <- Get Int
forall a. Integral a => Get a
getInt8
  ByteString
st <- Int -> Get ByteString
getByteString Int
len
  ByteString -> Get ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
st

parseLongString :: Get ByteString
parseLongString :: Get ByteString
parseLongString = do
  Int
len <- Get Int
forall a. Integral a => Get a
getInt32
  ByteString
st <- Int -> Get ByteString
getByteString Int
len
  ByteString -> Get ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
st

parseStrings :: Get [ByteString]
parseStrings :: Get [ByteString]
parseStrings = do
  Int
count <- Get Int
forall a. Integral a => Get a
getInt32
  [ByteString]
res <- [Get ByteString] -> Get [ByteString]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence ([Get ByteString] -> Get [ByteString])
-> [Get ByteString] -> Get [ByteString]
forall a b. (a -> b) -> a -> b
$ Int -> Get ByteString -> [Get ByteString]
forall a. Int -> a -> [a]
replicate Int
count Get ByteString
parseLongString
  [ByteString] -> Get [ByteString]
forall (m :: * -> *) a. Monad m => a -> m a
return [ByteString]
res

parseKV :: Get (ByteString, ByteString)
parseKV :: Get (ByteString, ByteString)
parseKV = do
  ByteString
key <- Get ByteString
parseString
  ByteString
value <- Get ByteString
parseLongString
  (ByteString, ByteString) -> Get (ByteString, ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
key, ByteString
value)

parseMap :: Get (M.Map ByteString ByteString)
parseMap :: Get (Map ByteString ByteString)
parseMap = do
  Int
count <- Get Int
forall a. Integral a => Get a
getInt32
  [(ByteString, ByteString)]
res <- [Get (ByteString, ByteString)] -> Get [(ByteString, ByteString)]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence ([Get (ByteString, ByteString)] -> Get [(ByteString, ByteString)])
-> [Get (ByteString, ByteString)] -> Get [(ByteString, ByteString)]
forall a b. (a -> b) -> a -> b
$ Int
-> Get (ByteString, ByteString) -> [Get (ByteString, ByteString)]
forall a. Int -> a -> [a]
replicate Int
count Get (ByteString, ByteString)
parseKV
  Map ByteString ByteString -> Get (Map ByteString ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (Map ByteString ByteString -> Get (Map ByteString ByteString))
-> Map ByteString ByteString -> Get (Map ByteString ByteString)
forall a b. (a -> b) -> a -> b
$ [(ByteString, ByteString)] -> Map ByteString ByteString
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList [(ByteString, ByteString)]
res

putByteStringLen :: ByteString -> PutM ()
putByteStringLen :: ByteString -> PutM ()
putByteStringLen ByteString
x = do
  Int8 -> PutM ()
putInt8 (Int8 -> PutM ()) -> Int8 -> PutM ()
forall a b. (a -> b) -> a -> b
$ Int -> Int8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int8) -> Int -> Int8
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
B.length ByteString
x
  ByteString -> PutM ()
putByteString ByteString
x

putLongByteStringLen :: ByteString -> PutM ()
putLongByteStringLen :: ByteString -> PutM ()
putLongByteStringLen ByteString
x = do
  Int32 -> PutM ()
putInt32be (Int32 -> PutM ()) -> Int32 -> PutM ()
forall a b. (a -> b) -> a -> b
$ Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int32) -> Int -> Int32
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
B.length ByteString
x
  ByteString -> PutM ()
putByteString ByteString
x

putByteStrings :: Foldable t => t ByteString -> PutM ()
putByteStrings :: t ByteString -> PutM ()
putByteStrings t ByteString
x = do
  Int32 -> PutM ()
putInt32be (Int32 -> PutM ()) -> Int32 -> PutM ()
forall a b. (a -> b) -> a -> b
$ Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int32) -> Int -> Int32
forall a b. (a -> b) -> a -> b
$ t ByteString -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t ByteString
x
  (ByteString -> PutM ()) -> t ByteString -> PutM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ByteString -> PutM ()
putLongByteStringLen t ByteString
x

putKV :: (ByteString, ByteString) -> PutM ()
putKV :: (ByteString, ByteString) -> PutM ()
putKV (ByteString
k, ByteString
v) = do
  ByteString -> PutM ()
putByteStringLen ByteString
k
  ByteString -> PutM ()
putLongByteStringLen ByteString
v

putMap :: M.Map ByteString ByteString -> PutM ()
putMap :: Map ByteString ByteString -> PutM ()
putMap Map ByteString ByteString
m = do
  Int32 -> PutM ()
putInt32be (Int32 -> PutM ()) -> Int32 -> PutM ()
forall a b. (a -> b) -> a -> b
$ Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int32) -> Int -> Int32
forall a b. (a -> b) -> a -> b
$ [(ByteString, ByteString)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(ByteString, ByteString)]
ml
  ((ByteString, ByteString) -> PutM ())
-> [(ByteString, ByteString)] -> PutM ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (ByteString, ByteString) -> PutM ()
putKV [(ByteString, ByteString)]
ml
  where ml :: [(ByteString, ByteString)]
ml = Map ByteString ByteString -> [(ByteString, ByteString)]
forall k a. Map k a -> [(k, a)]
M.toList Map ByteString ByteString
m

runGet :: Get a -> ByteString -> Either String a
runGet :: Get a -> ByteString -> Either String a
runGet Get a
g ByteString
b = case Get a
-> ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, a)
forall a.
Get a
-> ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, a)
Get.runGetOrFail Get a
g (ByteString -> ByteString
BL.fromStrict ByteString
b) of
  (Left (ByteString
_unconsumed, ByteOffset
_offset, String
err)) -> String -> Either String a
forall a b. a -> Either a b
Left String
err
  (Right (ByteString
_unconsumed, ByteOffset
_offset, a
res)) -> a -> Either String a
forall a b. b -> Either a b
Right a
res

runPut :: Put -> ByteString
runPut :: PutM () -> ByteString
runPut = ByteString -> ByteString
BL.toStrict (ByteString -> ByteString)
-> (PutM () -> ByteString) -> PutM () -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PutM () -> ByteString
Put.runPut