{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ViewPatterns #-}

module U.Util.Serialization where

import Control.Applicative (liftA3)
import Control.Monad (foldM, replicateM, replicateM_, when)
import Data.Bits (Bits, clearBit, setBit, shiftL, shiftR, testBit, (.|.))
import Data.ByteString (ByteString, readFile, writeFile)
import qualified Data.ByteString as BS
import Data.ByteString.Short (ShortByteString)
import qualified Data.ByteString.Short as BSS
import Data.Bytes.Get (MonadGet, getByteString, getBytes, getWord8, remaining, runGetS, skip)
import Data.Bytes.Put (MonadPut, putByteString, putWord8, runPutS)
import Data.Bytes.VarInt (VarInt (VarInt))
import Data.Foldable (Foldable (toList), traverse_)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Sequence (Seq)
import qualified Data.Sequence as Seq
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Text (Text)
import Data.Text.Encoding (decodeUtf8, encodeUtf8)
import Data.Text.Short (ShortText)
import qualified Data.Text.Short as TS
import qualified Data.Text.Short.Unsafe as TSU
import Data.Vector (Vector)
import qualified Data.Vector as Vector
import Data.Word (Word8)
import Debug.Trace (traceM)
import GHC.Word (Word64)
import System.FilePath (takeDirectory)
import UnliftIO (MonadIO, liftIO)
import UnliftIO.Directory (createDirectoryIfMissing, doesFileExist)
import Prelude hiding (readFile, writeFile)

type Get a = forall m. (MonadGet m) => m a

type Put a = forall m. (MonadPut m) => a -> m ()

-- todo: do we use this?
data Format a = Format
  { forall a. Format a -> Get a
get :: Get a,
    forall a. Format a -> Put a
put :: Put a
  }

debug :: Bool
debug :: Bool
debug = Bool
False

getFromBytes :: Get a -> ByteString -> Maybe a
getFromBytes :: forall a. Get a -> ByteString -> Maybe a
getFromBytes Get a
getA ByteString
bytes =
  case Get a -> ByteString -> Either String a
forall a. Get a -> ByteString -> Either String a
runGetS Get a
Get a
getA ByteString
bytes of Left String
_ -> Maybe a
forall a. Maybe a
Nothing; Right a
a -> a -> Maybe a
forall a. a -> Maybe a
Just a
a

getFromFile :: (MonadIO m) => Get a -> FilePath -> m (Maybe a)
getFromFile :: forall (m :: * -> *) a. MonadIO m => Get a -> String -> m (Maybe a)
getFromFile Get a
getA String
file = do
  Bool
b <- String -> m Bool
forall (m :: * -> *). MonadIO m => String -> m Bool
doesFileExist String
file
  if Bool
b then Get a -> ByteString -> Maybe a
forall a. Get a -> ByteString -> Maybe a
getFromBytes m a
Get a
getA (ByteString -> Maybe a) -> m ByteString -> m (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO ByteString -> m ByteString
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (String -> IO ByteString
readFile String
file) else Maybe a -> m (Maybe a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing

getFromFile' :: (MonadIO m) => Get a -> FilePath -> m (Either String a)
getFromFile' :: forall (m :: * -> *) a.
MonadIO m =>
Get a -> String -> m (Either String a)
getFromFile' Get a
getA String
file = do
  Bool
b <- String -> m Bool
forall (m :: * -> *). MonadIO m => String -> m Bool
doesFileExist String
file
  if Bool
b
    then Get a -> ByteString -> Either String a
forall a. Get a -> ByteString -> Either String a
runGetS Get a
Get a
getA (ByteString -> Either String a)
-> m ByteString -> m (Either String a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO ByteString -> m ByteString
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (String -> IO ByteString
readFile String
file)
    else Either String a -> m (Either String a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either String a -> m (Either String a))
-> (String -> Either String a) -> String -> m (Either String a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Either String a
forall a b. a -> Either a b
Left (String -> m (Either String a)) -> String -> m (Either String a)
forall a b. (a -> b) -> a -> b
$ String
"No such file: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
file

putBytes :: Put a -> a -> ByteString
putBytes :: forall a. Put a -> a -> ByteString
putBytes Put a
put a
a = Put -> ByteString
runPutS (a -> Put
Put a
put a
a)

putWithParentDirs :: (MonadIO m) => Put a -> FilePath -> a -> m ()
putWithParentDirs :: forall (m :: * -> *) a. MonadIO m => Put a -> String -> a -> m ()
putWithParentDirs Put a
putA String
file a
a = do
  Bool -> String -> m ()
forall (m :: * -> *). MonadIO m => Bool -> String -> m ()
createDirectoryIfMissing Bool
True (String -> String
takeDirectory String
file)
  IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> (ByteString -> IO ()) -> ByteString -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ByteString -> IO ()
writeFile String
file (ByteString -> m ()) -> ByteString -> m ()
forall a b. (a -> b) -> a -> b
$ Put a -> a -> ByteString
forall a. Put a -> a -> ByteString
putBytes a -> m ()
Put a
putA a
a

putVarInt :: (MonadPut m, Integral a, Bits a) => a -> m ()
putVarInt :: forall (m :: * -> *) a.
(MonadPut m, Integral a, Bits a) =>
a -> m ()
putVarInt a
n
  | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
0x80 = Word8 -> m ()
forall (m :: * -> *). MonadPut m => Word8 -> m ()
putWord8 (Word8 -> m ()) -> Word8 -> m ()
forall a b. (a -> b) -> a -> b
$ a -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n
  | Bool
otherwise = do
      Word8 -> m ()
forall (m :: * -> *). MonadPut m => Word8 -> m ()
putWord8 (Word8 -> m ()) -> Word8 -> m ()
forall a b. (a -> b) -> a -> b
$ Word8 -> Int -> Word8
forall a. Bits a => a -> Int -> a
setBit (a -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
n) Int
7
      a -> m ()
forall (m :: * -> *) a.
(MonadPut m, Integral a, Bits a) =>
a -> m ()
putVarInt (a -> m ()) -> a -> m ()
forall a b. (a -> b) -> a -> b
$ a -> Int -> a
forall a. Bits a => a -> Int -> a
shiftR a
n Int
7
{-# INLINE putVarInt #-}

getVarInt :: (MonadGet m, Num b, Bits b) => m b
getVarInt :: forall (m :: * -> *) b. (MonadGet m, Num b, Bits b) => m b
getVarInt = m Word8
forall (m :: * -> *). MonadGet m => m Word8
getWord8 m Word8 -> (Word8 -> m b) -> m b
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Word8 -> m b
forall (m :: * -> *) b. (MonadGet m, Num b, Bits b) => Word8 -> m b
getVarInt
  where
    getVarInt :: (MonadGet m, Num b, Bits b) => Word8 -> m b
    getVarInt :: forall (m :: * -> *) b. (MonadGet m, Num b, Bits b) => Word8 -> m b
getVarInt Word8
n
      | Word8 -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit Word8
n Int
7 = do
          VarInt b
m <- m Word8
forall (m :: * -> *). MonadGet m => m Word8
getWord8 m Word8 -> (Word8 -> m (VarInt b)) -> m (VarInt b)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Word8 -> m (VarInt b)
forall (m :: * -> *) b. (MonadGet m, Num b, Bits b) => Word8 -> m b
getVarInt
          b -> m b
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (b -> m b) -> b -> m b
forall a b. (a -> b) -> a -> b
$ b -> Int -> b
forall a. Bits a => a -> Int -> a
shiftL b
m Int
7 b -> b -> b
forall a. Bits a => a -> a -> a
.|. b -> Int -> b
forall a. Bits a => a -> Int -> a
clearBit (Word8 -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
n) Int
7
      | Bool
otherwise = b -> m b
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (b -> m b) -> b -> m b
forall a b. (a -> b) -> a -> b
$ Word8 -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word8
n
    {-# INLINE getVarInt #-}
{-# INLINE getVarInt #-}

putText :: (MonadPut m) => Text -> m ()
putText :: forall (m :: * -> *). MonadPut m => Text -> m ()
putText Text
text = do
  let bs :: ByteString
bs = Text -> ByteString
encodeUtf8 Text
text
  Int -> m ()
forall (m :: * -> *) a.
(MonadPut m, Integral a, Bits a) =>
a -> m ()
putVarInt (Int -> m ()) -> Int -> m ()
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
BS.length ByteString
bs
  ByteString -> m ()
forall (m :: * -> *). MonadPut m => ByteString -> m ()
putByteString ByteString
bs

getText :: (MonadGet m) => m Text
getText :: forall (m :: * -> *). MonadGet m => m Text
getText = do
  Int
len <- m Int
forall (m :: * -> *) b. (MonadGet m, Num b, Bits b) => m b
getVarInt
  ByteString
bs <- ByteString -> ByteString
BS.copy (ByteString -> ByteString) -> m ByteString -> m ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m ByteString
forall (m :: * -> *). MonadGet m => Int -> m ByteString
getBytes Int
len
  Text -> m Text
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> m Text) -> Text -> m Text
forall a b. (a -> b) -> a -> b
$ ByteString -> Text
decodeUtf8 ByteString
bs

skipText :: (MonadGet m) => m ()
skipText :: forall (m :: * -> *). MonadGet m => m ()
skipText = Int -> m ()
forall (m :: * -> *). MonadGet m => Int -> m ()
skip (Int -> m ()) -> m Int -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m Int
forall (m :: * -> *) b. (MonadGet m, Num b, Bits b) => m b
getVarInt

putShortText :: (MonadPut m) => ShortText -> m ()
putShortText :: forall (m :: * -> *). MonadPut m => ShortText -> m ()
putShortText ShortText
text = do
  let sbs :: ShortByteString
sbs = ShortText -> ShortByteString
TS.toShortByteString ShortText
text
  Int -> m ()
forall (m :: * -> *) a.
(MonadPut m, Integral a, Bits a) =>
a -> m ()
putVarInt (Int -> m ()) -> Int -> m ()
forall a b. (a -> b) -> a -> b
$ ShortByteString -> Int
BSS.length ShortByteString
sbs
  ShortByteString -> m ()
forall (m :: * -> *). MonadPut m => ShortByteString -> m ()
putShortByteString ShortByteString
sbs

getShortText :: (MonadGet m) => m ShortText
getShortText :: forall (m :: * -> *). MonadGet m => m ShortText
getShortText = do
  Int
len <- m Int
forall (m :: * -> *) b. (MonadGet m, Num b, Bits b) => m b
getVarInt
  ShortByteString
sbs <- Int -> m ShortByteString
forall (m :: * -> *). MonadGet m => Int -> m ShortByteString
getShortByteString Int
len
  ShortText -> m ShortText
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ShortText -> m ShortText) -> ShortText -> m ShortText
forall a b. (a -> b) -> a -> b
$ ShortByteString -> ShortText
TSU.fromShortByteStringUnsafe ShortByteString
sbs

-- | the `binary` package has a native version of this,
--  which may be more efficient by a constant factor
putShortByteString :: (MonadPut m) => ShortByteString -> m ()
putShortByteString :: forall (m :: * -> *). MonadPut m => ShortByteString -> m ()
putShortByteString = ByteString -> m ()
forall (m :: * -> *). MonadPut m => ByteString -> m ()
putByteString (ByteString -> m ())
-> (ShortByteString -> ByteString) -> ShortByteString -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShortByteString -> ByteString
BSS.fromShort

-- | the `binary` package has a native version of this,
--  which may be more efficient by a constant factor
getShortByteString :: (MonadGet m) => Int -> m ShortByteString
getShortByteString :: forall (m :: * -> *). MonadGet m => Int -> m ShortByteString
getShortByteString Int
len = ByteString -> ShortByteString
BSS.toShort (ByteString -> ShortByteString)
-> m ByteString -> m ShortByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m ByteString
forall (m :: * -> *). MonadGet m => Int -> m ByteString
getByteString Int
len

putFoldable ::
  (Foldable f, MonadPut m) => (a -> m ()) -> f a -> m ()
putFoldable :: forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, MonadPut m) =>
(a -> m ()) -> f a -> m ()
putFoldable a -> m ()
putA f a
as = do
  Int -> m ()
forall (m :: * -> *) a.
(MonadPut m, Integral a, Bits a) =>
a -> m ()
putVarInt (f a -> Int
forall a. f a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length f a
as)
  (a -> m ()) -> f a -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ a -> m ()
putA f a
as

getList :: (MonadGet m) => m a -> m [a]
getList :: forall (m :: * -> *) a. MonadGet m => m a -> m [a]
getList m a
getA = do
  Int
length <- m Int
forall (m :: * -> *). MonadGet m => m Int
getListLength
  Int -> m a -> m [a]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
length m a
getA

getListLength :: (MonadGet m) => m Int
getListLength :: forall (m :: * -> *). MonadGet m => m Int
getListLength =
  m Int
forall (m :: * -> *) b. (MonadGet m, Num b, Bits b) => m b
getVarInt

getVector :: (MonadGet m) => m a -> m (Vector a)
getVector :: forall (m :: * -> *) a. MonadGet m => m a -> m (Vector a)
getVector m a
getA = do
  Int
length <- m Int
forall (m :: * -> *) b. (MonadGet m, Num b, Bits b) => m b
getVarInt
  Int -> m a -> m (Vector a)
forall (m :: * -> *) a. Monad m => Int -> m a -> m (Vector a)
Vector.replicateM Int
length m a
getA

skipVector :: (MonadGet m) => m a -> m ()
skipVector :: forall (m :: * -> *) a. MonadGet m => m a -> m ()
skipVector m a
getA = do
  Int
length <- m Int
forall (m :: * -> *) b. (MonadGet m, Num b, Bits b) => m b
getVarInt
  Int -> m a -> m ()
forall (m :: * -> *) a. Applicative m => Int -> m a -> m ()
replicateM_ Int
length m a
getA

getSequence :: (MonadGet m) => m a -> m (Seq a)
getSequence :: forall (m :: * -> *) a. MonadGet m => m a -> m (Seq a)
getSequence m a
getA = do
  Int
length <- m Int
forall (m :: * -> *) b. (MonadGet m, Num b, Bits b) => m b
getVarInt
  Int -> m a -> m (Seq a)
forall (m :: * -> *) a. Applicative m => Int -> m a -> m (Seq a)
Seq.replicateM Int
length m a
getA

getSet :: (MonadGet m, Ord a) => m a -> m (Set a)
getSet :: forall (m :: * -> *) a. (MonadGet m, Ord a) => m a -> m (Set a)
getSet m a
getA = do
  Int
length <- m Int
forall (m :: * -> *) b. (MonadGet m, Num b, Bits b) => m b
getVarInt
  -- avoid materializing intermediate list
  (Set a -> m a -> m (Set a)) -> Set a -> [m a] -> m (Set a)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\Set a
s m a
ma -> a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.insert (a -> Set a -> Set a) -> m a -> m (Set a -> Set a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m a
ma m (Set a -> Set a) -> m (Set a) -> m (Set a)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Set a -> m (Set a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Set a
s) Set a
forall a. Monoid a => a
mempty (Int -> m a -> [m a]
forall a. Int -> a -> [a]
replicate Int
length m a
getA)

putMap :: (MonadPut m) => (a -> m ()) -> (b -> m ()) -> Map a b -> m ()
putMap :: forall (m :: * -> *) a b.
MonadPut m =>
(a -> m ()) -> (b -> m ()) -> Map a b -> m ()
putMap a -> m ()
putA b -> m ()
putB Map a b
m = ((a, b) -> m ()) -> [(a, b)] -> m ()
forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, MonadPut m) =>
(a -> m ()) -> f a -> m ()
putFoldable ((a -> m ()) -> (b -> m ()) -> (a, b) -> m ()
forall (m :: * -> *) a b.
MonadPut m =>
(a -> m ()) -> (b -> m ()) -> (a, b) -> m ()
putPair a -> m ()
putA b -> m ()
putB) (Map a b -> [(a, b)]
forall k a. Map k a -> [(k, a)]
Map.toList Map a b
m)

addToExistingMap :: (MonadGet m, Ord a) => m a -> m b -> Map a b -> m (Map a b)
addToExistingMap :: forall (m :: * -> *) a b.
(MonadGet m, Ord a) =>
m a -> m b -> Map a b -> m (Map a b)
addToExistingMap m a
getA m b
getB Map a b
map = do
  Int
length <- m Int
forall (m :: * -> *) b. (MonadGet m, Num b, Bits b) => m b
getVarInt
  -- avoid materializing intermediate list
  (Map a b -> (m a, m b) -> m (Map a b))
-> Map a b -> [(m a, m b)] -> m (Map a b)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM
    (\Map a b
s (m a
ma, m b
mb) -> a -> b -> Map a b -> Map a b
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (a -> b -> Map a b -> Map a b)
-> m a -> m (b -> Map a b -> Map a b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m a
ma m (b -> Map a b -> Map a b) -> m b -> m (Map a b -> Map a b)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> m b
mb m (Map a b -> Map a b) -> m (Map a b) -> m (Map a b)
forall a b. m (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Map a b -> m (Map a b)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Map a b
s)
    Map a b
map
    (Int -> (m a, m b) -> [(m a, m b)]
forall a. Int -> a -> [a]
replicate Int
length (m a
getA, m b
getB))

getMap :: (MonadGet m, Ord a) => m a -> m b -> m (Map a b)
getMap :: forall (m :: * -> *) a b.
(MonadGet m, Ord a) =>
m a -> m b -> m (Map a b)
getMap m a
getA m b
getB = m a -> m b -> Map a b -> m (Map a b)
forall (m :: * -> *) a b.
(MonadGet m, Ord a) =>
m a -> m b -> Map a b -> m (Map a b)
addToExistingMap m a
getA m b
getB Map a b
forall a. Monoid a => a
mempty

getFramedByteString :: (MonadGet m) => m ByteString
getFramedByteString :: forall (m :: * -> *). MonadGet m => m ByteString
getFramedByteString = m Int
forall (m :: * -> *) b. (MonadGet m, Num b, Bits b) => m b
getVarInt m Int -> (Int -> m ByteString) -> m ByteString
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> m ByteString
forall (m :: * -> *). MonadGet m => Int -> m ByteString
getByteString

getRemainingByteString :: (MonadGet m) => m ByteString
getRemainingByteString :: forall (m :: * -> *). MonadGet m => m ByteString
getRemainingByteString = Remaining m -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Remaining m -> Int) -> m (Remaining m) -> m Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m (Remaining m)
forall (m :: * -> *). MonadGet m => m (Remaining m)
remaining m Int -> (Int -> m ByteString) -> m ByteString
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Int -> m ByteString
forall (m :: * -> *). MonadGet m => Int -> m ByteString
getByteString

getFramed :: (MonadGet m) => Get a -> m a
getFramed :: forall (m :: * -> *) a. MonadGet m => Get a -> m a
getFramed Get a
get =
  m ByteString
forall (m :: * -> *). MonadGet m => m ByteString
getFramedByteString m ByteString -> (ByteString -> m a) -> m a
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (String -> m a) -> (a -> m a) -> Either String a -> m a
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either String -> m a
forall a. String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail a -> m a
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either String a -> m a)
-> (ByteString -> Either String a) -> ByteString -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Get a -> ByteString -> Either String a
forall a. Get a -> ByteString -> Either String a
runGetS Get a
Get a
get

putFramedByteString :: (MonadPut m) => ByteString -> m ()
putFramedByteString :: forall (m :: * -> *). MonadPut m => ByteString -> m ()
putFramedByteString ByteString
bs = do
  Int -> m ()
forall (m :: * -> *) a.
(MonadPut m, Integral a, Bits a) =>
a -> m ()
putVarInt (ByteString -> Int
BS.length ByteString
bs)
  ByteString -> m ()
forall (m :: * -> *). MonadPut m => ByteString -> m ()
putByteString ByteString
bs

putFramed :: (MonadPut m) => Put a -> a -> m ()
putFramed :: forall (m :: * -> *) a. MonadPut m => Put a -> a -> m ()
putFramed Put a
put a
a = do
  -- 1. figure out the length `len` of serialized `a`
  -- 2. Put the length `len`
  -- 3. Put `a`
  let bs :: ByteString
bs = Put a -> a -> ByteString
forall a. Put a -> a -> ByteString
putBytes a -> m ()
Put a
put a
a
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ String -> m ()
forall (f :: * -> *). Applicative f => String -> f ()
traceM (String -> m ()) -> String -> m ()
forall a b. (a -> b) -> a -> b
$ String
"putFramed " String -> String -> String
forall a. [a] -> [a] -> [a]
++ (Int -> String
forall a. Show a => a -> String
show (Int -> String) -> Int -> String
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
BS.length ByteString
bs) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" bytes: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ ByteString -> String
forall a. Show a => a -> String
show ByteString
bs
  Int -> m ()
forall (m :: * -> *) a.
(MonadPut m, Integral a, Bits a) =>
a -> m ()
putVarInt (ByteString -> Int
BS.length ByteString
bs)
  ByteString -> m ()
forall (m :: * -> *). MonadPut m => ByteString -> m ()
putByteString ByteString
bs

skipFramed :: (MonadGet m) => m ()
skipFramed :: forall (m :: * -> *). MonadGet m => m ()
skipFramed = do
  Int
len <- m Int
forall (m :: * -> *) b. (MonadGet m, Num b, Bits b) => m b
getVarInt
  Int -> m ()
forall (m :: * -> *). MonadGet m => Int -> m ()
skip Int
len

putFramedArray :: (MonadPut m, Foldable f) => Put a -> f a -> m ()
putFramedArray :: forall (m :: * -> *) (f :: * -> *) a.
(MonadPut m, Foldable f) =>
Put a -> f a -> m ()
putFramedArray Put a
put (f a -> [a]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList -> [a]
as) = do
  let bss :: [ByteString]
bss = (a -> ByteString) -> [a] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Put a -> a -> ByteString
forall a. Put a -> a -> ByteString
putBytes a -> m ()
Put a
put) [a]
as
  let lengths :: [Int]
lengths = (ByteString -> Int) -> [ByteString] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ByteString -> Int
BS.length [ByteString]
bss
  let offsets :: [Int]
offsets = (Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) Int
0 [Int]
lengths
  (Int -> m ()) -> [Int] -> m ()
forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, MonadPut m) =>
(a -> m ()) -> f a -> m ()
putFoldable Int -> m ()
forall (m :: * -> *) a.
(MonadPut m, Integral a, Bits a) =>
a -> m ()
putVarInt [Int]
offsets
  (ByteString -> m ()) -> [ByteString] -> m ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ ByteString -> m ()
forall (m :: * -> *). MonadPut m => ByteString -> m ()
putByteString [ByteString]
bss

getFramedArray :: (MonadGet m) => m a -> m (Vector a)
getFramedArray :: forall (m :: * -> *) a. MonadGet m => m a -> m (Vector a)
getFramedArray m a
getA = do
  [Int]
offsets :: [Int] <- m Int -> m [Int]
forall (m :: * -> *) a. MonadGet m => m a -> m [a]
getList m Int
forall (m :: * -> *) b. (MonadGet m, Num b, Bits b) => m b
getVarInt
  let count :: Int
count = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
offsets Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
  Int -> m a -> m (Vector a)
forall (m :: * -> *) a. Monad m => Int -> m a -> m (Vector a)
Vector.replicateM Int
count m a
getA

-- | Look up a 0-based index in a framed array, O(num array elements),
--  because it reads the start indices for all elements first.
--  This could be skipped if the indices had a fixed size instead of varint
lookupFramedArray :: (MonadGet m) => m a -> Int -> m (Maybe a)
lookupFramedArray :: forall (m :: * -> *) a. MonadGet m => m a -> Int -> m (Maybe a)
lookupFramedArray m a
getA Int
index = do
  Vector Int
offsets <- m Int -> m (Vector Int)
forall (m :: * -> *) a. MonadGet m => m a -> m (Vector a)
getVector m Int
forall (m :: * -> *) b. (MonadGet m, Num b, Bits b) => m b
getVarInt
  if Int
index Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Vector Int -> Int
forall a. Vector a -> Int
Vector.length Vector Int
offsets Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
    then Maybe a -> m (Maybe a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing
    else do
      Int -> m ()
forall (m :: * -> *). MonadGet m => Int -> m ()
skip (Vector Int -> Int -> Int
forall a. Vector a -> Int -> a
Vector.unsafeIndex Vector Int
offsets Int
index)
      a -> Maybe a
forall a. a -> Maybe a
Just (a -> Maybe a) -> m a -> m (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m a
getA

lengthFramedArray :: (MonadGet m) => m Word64
lengthFramedArray :: forall (m :: * -> *). MonadGet m => m Word64
lengthFramedArray = (\Word64
offsetsLen -> Word64
offsetsLen Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
1) (Word64 -> Word64) -> m Word64 -> m Word64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m Word64
forall (m :: * -> *) b. (MonadGet m, Num b, Bits b) => m b
getVarInt

unsafeFramedArrayLookup :: (MonadGet m) => m a -> Int -> m a
unsafeFramedArrayLookup :: forall (m :: * -> *) a. MonadGet m => m a -> Int -> m a
unsafeFramedArrayLookup m a
getA Int
index = do
  Vector Int
offsets <- m Int -> m (Vector Int)
forall (m :: * -> *) a. MonadGet m => m a -> m (Vector a)
getVector m Int
forall (m :: * -> *) b. (MonadGet m, Num b, Bits b) => m b
getVarInt
  Int -> m ()
forall (m :: * -> *). MonadGet m => Int -> m ()
skip (Vector Int -> Int -> Int
forall a. Vector a -> Int -> a
Vector.unsafeIndex Vector Int
offsets Int
index)
  m a
getA

putPair :: (MonadPut m) => (a -> m ()) -> (b -> m ()) -> (a, b) -> m ()
putPair :: forall (m :: * -> *) a b.
MonadPut m =>
(a -> m ()) -> (b -> m ()) -> (a, b) -> m ()
putPair a -> m ()
putA b -> m ()
putB (a
a, b
b) = a -> m ()
putA a
a m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> b -> m ()
putB b
b

getPair :: (MonadGet m) => m a -> m b -> m (a, b)
getPair :: forall (m :: * -> *) a b. MonadGet m => m a -> m b -> m (a, b)
getPair = (a -> b -> (a, b)) -> m a -> m b -> m (a, b)
forall a b c. (a -> b -> c) -> m a -> m b -> m c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (,)

getTuple3 :: (MonadGet m) => m a -> m b -> m c -> m (a, b, c)
getTuple3 :: forall (m :: * -> *) a b c.
MonadGet m =>
m a -> m b -> m c -> m (a, b, c)
getTuple3 = (a -> b -> c -> (a, b, c)) -> m a -> m b -> m c -> m (a, b, c)
forall (f :: * -> *) a b c d.
Applicative f =>
(a -> b -> c -> d) -> f a -> f b -> f c -> f d
liftA3 (,,)