module Data.PackStream
( PackStream(..)
, ToPackStream(..)
, FromPackStream(..)
, Parser
, parse
, parsefail
, parseEither
, parseMaybe
, pack
, unpack
, pretty
, prettyStruct
, genericStructName
, (.=)
, (.:)
, (.:?)
, (.!=)
)
where
import Control.Monad
import Data.Bifunctor
import Data.Bits
import qualified Data.ByteString as BS
import Data.Hashable
import qualified Data.HashMap.Strict as HM
import Data.Int
import qualified Data.Map.Strict as M
import Data.Monoid
import Data.Scientific
import Data.Serialize.Get
import Data.Serialize.IEEE754
import Data.Serialize.Put
import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Data.Vector as V
import Data.Word
import GHC.Generics (Generic)
import Text.Printf
instance (Hashable a) => Hashable (V.Vector a) where
hashWithSalt salt = hashWithSalt salt . V.toList
data PackStream = Null
| Bool !Bool
| Int !Int64
| Float !Double
| String !Text
| List !(V.Vector PackStream)
| Map !(HM.HashMap PackStream PackStream)
| Struct !Word8 ![PackStream]
deriving (Show, Eq, Generic, Hashable)
type Parser = Either String
parse :: (a -> Parser b) -> a -> Parser b
parse = parseEither
parseEither :: (a -> Parser b) -> a -> Either String b
parseEither f = f
parseMaybe :: (a -> Parser b) -> a -> Maybe b
parseMaybe f = either (const Nothing) Just . parseEither f
parsefail :: a -> Either a b
parsefail = Left
class FromPackStream a where
parsePackStream :: PackStream -> Parser a
instance FromPackStream PackStream where
parsePackStream = return
instance FromPackStream Bool where
parsePackStream (Bool b) = return b
parsePackStream _ = parsefail "Expecting Bool"
instance FromPackStream Int64 where
parsePackStream (Int i) = return i
parsePackStream _ = parsefail "Expecting Int"
instance FromPackStream Double where
parsePackStream (Float d) = return d
parsePackStream _ = parsefail "Expecting Float"
instance FromPackStream Text where
parsePackStream (String t) = return t
parsePackStream _ = parsefail "Expecting String"
instance FromPackStream a => FromPackStream [a] where
parsePackStream (List l) = mapM parsePackStream (V.toList l)
parsePackStream _ = parsefail "Expecting List"
instance FromPackStream a => FromPackStream (V.Vector a) where
parsePackStream (List l) = mapM parsePackStream l
parsePackStream _ = parsefail "Expecting List"
instance (Eq a, Hashable a, FromPackStream a, FromPackStream b) => FromPackStream (HM.HashMap a b) where
parsePackStream (Map m) = do
kvs <- mapM parseAssoc (HM.toList m)
return $ HM.fromList kvs
where
parseAssoc (k, v) = (,) <$> parsePackStream k <*> parsePackStream v
parsePackStream _ = parsefail "Expecting Map"
instance (Ord a, FromPackStream a, FromPackStream b) => FromPackStream (M.Map a b) where
parsePackStream (Map m) = do
kvs <- mapM parseAssoc (HM.toList m)
return $ M.fromList kvs
where
parseAssoc (k, v) = (,) <$> parsePackStream k <*> parsePackStream v
parsePackStream _ = parsefail "Expecting Map"
instance FromPackStream Scientific where
parsePackStream (Int i) = return $ fromIntegral i
parsePackStream (Float f) = return $ fromFloatDigits f
parsePackStream _ = parsefail "Expecting Int or Float"
class ToPackStream a where
toPackStream :: a -> PackStream
instance ToPackStream PackStream where
toPackStream = id
instance ToPackStream Bool where
toPackStream = Bool
instance ToPackStream Int64 where
toPackStream = Int
instance ToPackStream Double where
toPackStream = Float
instance ToPackStream Text where
toPackStream = String
instance ToPackStream a => ToPackStream [a] where
toPackStream = List . V.fromList . fmap toPackStream
instance ToPackStream a => ToPackStream (V.Vector a) where
toPackStream = List . fmap toPackStream
instance (ToPackStream a, ToPackStream b) => ToPackStream (M.Map a b) where
toPackStream = Map . HM.fromList . fmap (bimap toPackStream toPackStream) . M.toList
instance (ToPackStream a, ToPackStream b) => ToPackStream (HM.HashMap a b) where
toPackStream = Map . HM.fromList . fmap (bimap toPackStream toPackStream) . HM.toList
instance ToPackStream Scientific where
toPackStream s =
case Data.Scientific.floatingOrInteger s of
Left rf -> Float rf
Right i -> Int i
pack :: ToPackStream a => Putter a
pack = putPackStream . toPackStream
unpack :: FromPackStream a => Get (Parser a)
unpack = parsePackStream <$> getPackStream
getPackStream :: Get PackStream
getPackStream = do
marker <- getWord8
if | marker == 0xc0 -> return Null
| marker == 0xc1 -> Float <$> getFloat64be
| marker == 0xc2 -> return $ Bool False
| marker == 0xc3 -> return $ Bool True
| marker < 0x80 -> return $ Int (fromIntegral marker)
| marker >= 0xf0 -> return $ Int (fromIntegral marker 0x100)
| marker == 0xc8 -> Int . fromIntegral <$> getWord8
| marker == 0xc9 -> Int . fromIntegral <$> getWord16be
| marker == 0xca -> Int . fromIntegral <$> getWord32be
| marker == 0xcb -> Int . fromIntegral <$> getWord64be
| 0x80 <= marker && marker < 0x90
-> getString (fromIntegral marker .&. 0x0f)
| marker == 0xd0 -> fromIntegral <$> getWord8 >>= getString
| marker == 0xd1 -> fromIntegral <$> getWord16be >>= getString
| marker == 0xd2 -> fromIntegral <$> getWord32be >>= getString
| 0x90 <= marker && marker < 0xa0
-> getList (fromIntegral marker .&. 0x0f)
| marker == 0xd4 -> fromIntegral <$> getWord8 >>= getList
| marker == 0xd5 -> fromIntegral <$> getWord16be >>= getList
| marker == 0xd6 -> fromIntegral <$> getWord32be >>= getList
| 0xa0 <= marker && marker < 0xb0
-> getMap (fromIntegral marker .&. 0x0f)
| marker == 0xd8 -> fromIntegral <$> getWord8 >>= getMap
| marker == 0xd9 -> fromIntegral <$> getWord16be >>= getMap
| marker == 0xda -> fromIntegral <$> getWord32be >>= getMap
| 0xb0 <= marker && marker < 0xc0
-> getStruct (fromIntegral marker .&. 0x0f)
| marker == 0xdc -> fromIntegral <$> getWord8 >>= getStruct
| marker == 0xdd -> fromIntegral <$> getWord16be >>= getStruct
| otherwise -> fail $ "Unknown marker " ++ printf "0x%02x" marker
getString :: Int -> Get PackStream
getString n = String . T.decodeUtf8 <$> getByteString n
getList :: Int -> Get PackStream
getList n = List . V.fromList <$> replicateM n getPackStream
getMap :: Int -> Get PackStream
getMap n = Map . HM.fromList <$> replicateM n getPair
where
getPair = (,) <$> getPackStream <*> getPackStream
getStruct :: Int -> Get PackStream
getStruct n = Struct <$> getWord8 <*> replicateM n getPackStream
putPackStream :: Putter PackStream
putPackStream Null = putWord8 0xc0
putPackStream (Float d) = putWord8 0xc1 >> putFloat64be d
putPackStream (Bool False) = putWord8 0xc2
putPackStream (Bool True) = putWord8 0xc3
putPackStream (Int i)
| 0x10 <= i && i < 0x80 = putWord8 (fromIntegral i)
| 0x80 <= i && i < 0x80 = putWord8 0xc8 >> putWord8 (fromIntegral i)
| 0x8000 <= i && i < 0x8000 = putWord8 0xc9 >> putWord16be (fromIntegral i)
| 0x80000000 <= i && i < 0x80000000 = putWord8 0xca >> putWord32be (fromIntegral i)
| otherwise = putWord8 0xcb >> putWord64be (fromIntegral i)
putPackStream (String t) = do
let bstr = T.encodeUtf8 t
size = BS.length bstr
if | size < 0x10 -> putWord8 (0x80 + fromIntegral size)
| size < 0x100 -> putWord8 0xd0 >> putWord8 (fromIntegral size)
| size < 0x10000 -> putWord8 0xd1 >> putWord16be (fromIntegral size)
| size < 0x100000000 -> putWord8 0xd2 >> putWord16be (fromIntegral size)
| otherwise -> fail "String too long"
putByteString $ T.encodeUtf8 t
putPackStream (List xs) = do
let size = V.length xs
if | size < 0x10 -> putWord8 (0x90 + fromIntegral size)
| size < 0x100 -> putWord8 0xd4 >> putWord8 (fromIntegral size)
| size < 0x10000 -> putWord8 0xd5 >> putWord16be (fromIntegral size)
| size < 0x100000000 -> putWord8 0xd6 >> putWord16be (fromIntegral size)
| otherwise -> fail "List too long"
mapM_ putPackStream xs
putPackStream (Map m) = do
let size = HM.size m
if | size < 0x10 -> putWord8 (0xa0 + fromIntegral size)
| size < 0x100 -> putWord8 0xd8 >> putWord8 (fromIntegral size)
| size < 0x10000 -> putWord8 0xd9 >> putWord16be (fromIntegral size)
| size < 0x100000000 -> putWord8 0xda >> putWord16be (fromIntegral size)
| otherwise -> fail "Map too large"
mapM_ (uncurry putPair) (HM.toList m)
where
putPair k v = putPackStream k >> putPackStream v
putPackStream (Struct sig fs) = do
let size = length fs
if | size < 0x10 -> putWord8 (0xb0 + fromIntegral size)
| size < 0x100 -> putWord8 0xdc >> putWord8 (fromIntegral size)
| size < 0x10000 -> putWord8 0xdd >> putWord16be (fromIntegral size)
| otherwise -> fail "Structure too big"
putWord8 sig
mapM_ putPackStream fs
pretty :: PackStream -> Text
pretty = prettyStruct genericStructName
prettyStruct :: (Word8 -> Text) -> PackStream -> Text
prettyStruct _ Null = "null"
prettyStruct _ (Bool True) = "true"
prettyStruct _ (Bool False) = "false"
prettyStruct _ (Int i) = T.pack (show i)
prettyStruct _ (Float d) = T.pack (show d)
prettyStruct _ (String t) = "\"" <> t <> "\""
prettyStruct _ (List xs) = "[" <> T.intercalate ", " (pretty <$> V.toList xs) <> "]"
prettyStruct _ (Map ps) = "{" <> T.intercalate ", " (fmap (\(k, v) -> pretty k <> ": " <> pretty v) (HM.toList ps)) <> "}"
prettyStruct sn (Struct s fs) = sn s <> "{" <> T.intercalate ", " (pretty <$> fs) <> "}"
genericStructName :: Word8 -> Text
genericStructName n = "Struct(signature=" <> T.pack (printf "0x%02x" n) <> ")"
(.=) :: ToPackStream a => Text -> a -> (PackStream, PackStream)
k .= v = (String k, toPackStream v)
(.:) :: FromPackStream a => HM.HashMap PackStream PackStream -> Text -> Parser a
m .: k = maybe (parsefail "Expected Key missing in map") parsePackStream (HM.lookup (String k) m)
(.:?) :: FromPackStream a => HM.HashMap PackStream PackStream -> Text -> Parser (Maybe a)
m .:? k = maybe (return Nothing) (fmap Just . parsePackStream) (HM.lookup (String k) m)
(.!=) :: Parser (Maybe a) -> a -> Parser a
p .!= d = do
ma <- p
maybe (return d) return ma