{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

module Data.Conduit.Serialise where

import Codec.Serialise as S
import Conduit hiding (leftover)
import Control.Monad.ST
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BL
import Data.Function (fix)

serialiseC :: forall a m. (Serialise a, Monad m) => ConduitM a ByteString m ()
serialiseC :: ConduitM a ByteString m ()
serialiseC = (a -> ByteString) -> ConduitM a ByteString m ()
forall (m :: * -> *) a b. Monad m => (a -> b) -> ConduitT a b m ()
mapC ((a -> ByteString) -> ConduitM a ByteString m ())
-> (a -> ByteString) -> ConduitM a ByteString m ()
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
BL.toStrict (ByteString -> ByteString) -> (a -> ByteString) -> a -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ByteString
forall a. Serialise a => a -> ByteString
S.serialise

serialiseWithLocC :: forall a k m. (Eq k, Serialise a, Monad m) => ConduitM (k, a) ByteString m [(k, (Integer, Integer))]
serialiseWithLocC :: ConduitM (k, a) ByteString m [(k, (Integer, Integer))]
serialiseWithLocC =
  [(k, (Integer, Integer))]
-> ConduitM (k, a) ByteString m [(k, (Integer, Integer))]
forall (m :: * -> *) b a b.
(Monad m, Eq a, Serialise b, Num b) =>
[(a, (b, b))] -> ConduitT (a, b) ByteString m [(a, (b, b))]
go []
  where
    go :: [(a, (b, b))] -> ConduitT (a, b) ByteString m [(a, (b, b))]
go [(a, (b, b))]
acc =
      ConduitT (a, b) ByteString m (Maybe (a, b))
forall (m :: * -> *) i. Monad m => Consumer i m (Maybe i)
await ConduitT (a, b) ByteString m (Maybe (a, b))
-> (Maybe (a, b) -> ConduitT (a, b) ByteString m [(a, (b, b))])
-> ConduitT (a, b) ByteString m [(a, (b, b))]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Maybe (a, b)
Nothing -> [(a, (b, b))] -> ConduitT (a, b) ByteString m [(a, (b, b))]
forall (m :: * -> *) a. Monad m => a -> m a
return ([(a, (b, b))] -> ConduitT (a, b) ByteString m [(a, (b, b))])
-> [(a, (b, b))] -> ConduitT (a, b) ByteString m [(a, (b, b))]
forall a b. (a -> b) -> a -> b
$ [(a, (b, b))] -> [(a, (b, b))]
forall a. [a] -> [a]
reverse [(a, (b, b))]
acc
        Just (a
k, b
a) -> do
          let ser :: ByteString
ser = ByteString -> ByteString
BL.toStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ b -> ByteString
forall a. Serialise a => a -> ByteString
S.serialise b
a
              len :: b
len = Int -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> b) -> Int -> b
forall a b. (a -> b) -> a -> b
$ ByteString -> Int
BS.length ByteString
ser
          ByteString -> ConduitT (a, b) ByteString m ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
yield ByteString
ser
          case [(a, (b, b))]
acc of
            [] -> [(a, (b, b))] -> ConduitT (a, b) ByteString m [(a, (b, b))]
go [(a
k, (b
0, b
len b -> b -> b
forall a. Num a => a -> a -> a
-b
1))]
            (a
k', (b
beg, b
end)) : [(a, (b, b))]
rest ->
              if a
k' a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
k
                then [(a, (b, b))] -> ConduitT (a, b) ByteString m [(a, (b, b))]
go ((a
k', (b
beg, b
end b -> b -> b
forall a. Num a => a -> a -> a
+ b
len)) (a, (b, b)) -> [(a, (b, b))] -> [(a, (b, b))]
forall a. a -> [a] -> [a]
: [(a, (b, b))]
rest)
                else [(a, (b, b))] -> ConduitT (a, b) ByteString m [(a, (b, b))]
go ((a
k, (b
end b -> b -> b
forall a. Num a => a -> a -> a
+ b
1, b
end b -> b -> b
forall a. Num a => a -> a -> a
+ b
len)) (a, (b, b)) -> [(a, (b, b))] -> [(a, (b, b))]
forall a. a -> [a] -> [a]
: [(a, (b, b))]
acc)

deserialiseC :: forall a m. (Serialise a, MonadIO m) => ConduitM ByteString a m ()
deserialiseC :: ConduitM ByteString a m ()
deserialiseC = (ByteString -> Bool) -> ConduitT ByteString ByteString m ()
forall (m :: * -> *) a. Monad m => (a -> Bool) -> ConduitT a a m ()
filterC (Bool -> Bool
not (Bool -> Bool) -> (ByteString -> Bool) -> ByteString -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Bool
BS.null) ConduitT ByteString ByteString m ()
-> ConduitM ByteString a m () -> ConduitM ByteString a m ()
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| ConduitM ByteString a m ()
pipe
  where
    pipe :: ConduitM ByteString a m ()
pipe =
      ((Maybe (IDecode RealWorld a) -> ConduitM ByteString a m ())
 -> Maybe (IDecode RealWorld a) -> ConduitM ByteString a m ())
-> Maybe (IDecode RealWorld a) -> ConduitM ByteString a m ()
forall a. (a -> a) -> a
fix
        ( \Maybe (IDecode RealWorld a) -> ConduitM ByteString a m ()
rec' -> \case
            Maybe (IDecode RealWorld a)
Nothing ->
              ConduitT ByteString a m (Maybe ByteString)
forall (m :: * -> *) i. Monad m => Consumer i m (Maybe i)
await ConduitT ByteString a m (Maybe ByteString)
-> (Maybe ByteString -> ConduitM ByteString a m ())
-> ConduitM ByteString a m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
                Maybe ByteString
Nothing -> () -> ConduitM ByteString a m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
                Just ByteString
bs ->
                  IO (IDecode RealWorld a)
-> ConduitT ByteString a m (IDecode RealWorld a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (ST RealWorld (IDecode RealWorld a) -> IO (IDecode RealWorld a)
forall a. ST RealWorld a -> IO a
stToIO (ST RealWorld (IDecode RealWorld a) -> IO (IDecode RealWorld a))
-> ST RealWorld (IDecode RealWorld a) -> IO (IDecode RealWorld a)
forall a b. (a -> b) -> a -> b
$ forall s. Serialise a => ST s (IDecode s a)
forall a s. Serialise a => ST s (IDecode s a)
deserialiseIncremental @a) ConduitT ByteString a m (IDecode RealWorld a)
-> (IDecode RealWorld a -> ConduitM ByteString a m ())
-> ConduitM ByteString a m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
                    Fail {} ->
                      [Char] -> ConduitM ByteString a m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"couldn't init parser"
                    Done {} ->
                      [Char] -> ConduitM ByteString a m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"couldn't init parser"
                    Partial Maybe ByteString -> ST RealWorld (IDecode RealWorld a)
cont ->
                      IO (IDecode RealWorld a)
-> ConduitT ByteString a m (IDecode RealWorld a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (ST RealWorld (IDecode RealWorld a) -> IO (IDecode RealWorld a)
forall a. ST RealWorld a -> IO a
stToIO (ST RealWorld (IDecode RealWorld a) -> IO (IDecode RealWorld a))
-> ST RealWorld (IDecode RealWorld a) -> IO (IDecode RealWorld a)
forall a b. (a -> b) -> a -> b
$ Maybe ByteString -> ST RealWorld (IDecode RealWorld a)
cont (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
bs)) ConduitT ByteString a m (IDecode RealWorld a)
-> (IDecode RealWorld a -> ConduitM ByteString a m ())
-> ConduitM ByteString a m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Maybe (IDecode RealWorld a) -> ConduitM ByteString a m ()
rec' (Maybe (IDecode RealWorld a) -> ConduitM ByteString a m ())
-> (IDecode RealWorld a -> Maybe (IDecode RealWorld a))
-> IDecode RealWorld a
-> ConduitM ByteString a m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IDecode RealWorld a -> Maybe (IDecode RealWorld a)
forall a. a -> Maybe a
Just
            Just Fail {} ->
              m () -> ConduitM ByteString a m ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m () -> ConduitM ByteString a m ())
-> m () -> ConduitM ByteString a m ()
forall a b. (a -> b) -> a -> b
$ [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"failed"
            Just (Done ByteString
leftover ByteOffset
_ a
ret) -> do
              a -> ConduitM ByteString a m ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
yield a
ret
              if ByteString -> Bool
BS.null ByteString
leftover
                then Maybe (IDecode RealWorld a) -> ConduitM ByteString a m ()
rec' Maybe (IDecode RealWorld a)
forall a. Maybe a
Nothing
                else IO (IDecode RealWorld a)
-> ConduitT ByteString a m (IDecode RealWorld a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (ST RealWorld (IDecode RealWorld a) -> IO (IDecode RealWorld a)
forall a. ST RealWorld a -> IO a
stToIO (ST RealWorld (IDecode RealWorld a) -> IO (IDecode RealWorld a))
-> ST RealWorld (IDecode RealWorld a) -> IO (IDecode RealWorld a)
forall a b. (a -> b) -> a -> b
$ forall s. Serialise a => ST s (IDecode s a)
forall a s. Serialise a => ST s (IDecode s a)
deserialiseIncremental @a) ConduitT ByteString a m (IDecode RealWorld a)
-> (IDecode RealWorld a -> ConduitM ByteString a m ())
-> ConduitM ByteString a m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
                  Fail {} ->
                    [Char] -> ConduitM ByteString a m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"couldn't init parser"
                  Done {} ->
                    [Char] -> ConduitM ByteString a m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"couldn't init parser"
                  Partial Maybe ByteString -> ST RealWorld (IDecode RealWorld a)
cont ->
                    IO (IDecode RealWorld a)
-> ConduitT ByteString a m (IDecode RealWorld a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (ST RealWorld (IDecode RealWorld a) -> IO (IDecode RealWorld a)
forall a. ST RealWorld a -> IO a
stToIO (ST RealWorld (IDecode RealWorld a) -> IO (IDecode RealWorld a))
-> ST RealWorld (IDecode RealWorld a) -> IO (IDecode RealWorld a)
forall a b. (a -> b) -> a -> b
$ Maybe ByteString -> ST RealWorld (IDecode RealWorld a)
cont (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
leftover)) ConduitT ByteString a m (IDecode RealWorld a)
-> (IDecode RealWorld a -> ConduitM ByteString a m ())
-> ConduitM ByteString a m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Maybe (IDecode RealWorld a) -> ConduitM ByteString a m ()
rec' (Maybe (IDecode RealWorld a) -> ConduitM ByteString a m ())
-> (IDecode RealWorld a -> Maybe (IDecode RealWorld a))
-> IDecode RealWorld a
-> ConduitM ByteString a m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IDecode RealWorld a -> Maybe (IDecode RealWorld a)
forall a. a -> Maybe a
Just
            Just (Partial Maybe ByteString -> ST RealWorld (IDecode RealWorld a)
cont) ->
              ConduitT ByteString a m (Maybe ByteString)
forall (m :: * -> *) i. Monad m => Consumer i m (Maybe i)
await ConduitT ByteString a m (Maybe ByteString)
-> (Maybe ByteString -> ConduitM ByteString a m ())
-> ConduitM ByteString a m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
                Maybe ByteString
Nothing ->
                  IO (IDecode RealWorld a)
-> ConduitT ByteString a m (IDecode RealWorld a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (ST RealWorld (IDecode RealWorld a) -> IO (IDecode RealWorld a)
forall a. ST RealWorld a -> IO a
stToIO (ST RealWorld (IDecode RealWorld a) -> IO (IDecode RealWorld a))
-> ST RealWorld (IDecode RealWorld a) -> IO (IDecode RealWorld a)
forall a b. (a -> b) -> a -> b
$ Maybe ByteString -> ST RealWorld (IDecode RealWorld a)
cont Maybe ByteString
forall a. Maybe a
Nothing) ConduitT ByteString a m (IDecode RealWorld a)
-> (IDecode RealWorld a -> ConduitM ByteString a m ())
-> ConduitM ByteString a m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
                    Fail {} ->
                      [Char] -> ConduitM ByteString a m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"parse failed"
                    Done ByteString
_ ByteOffset
_ a
ret ->
                      a -> ConduitM ByteString a m ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
yield a
ret
                    Partial Maybe ByteString -> ST RealWorld (IDecode RealWorld a)
_ ->
                      [Char] -> ConduitM ByteString a m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"parse failed. abrubt message?"
                Just ByteString
bs ->
                  IO (IDecode RealWorld a)
-> ConduitT ByteString a m (IDecode RealWorld a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (ST RealWorld (IDecode RealWorld a) -> IO (IDecode RealWorld a)
forall a. ST RealWorld a -> IO a
stToIO (ST RealWorld (IDecode RealWorld a) -> IO (IDecode RealWorld a))
-> ST RealWorld (IDecode RealWorld a) -> IO (IDecode RealWorld a)
forall a b. (a -> b) -> a -> b
$ Maybe ByteString -> ST RealWorld (IDecode RealWorld a)
cont (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
bs)) ConduitT ByteString a m (IDecode RealWorld a)
-> (IDecode RealWorld a -> ConduitM ByteString a m ())
-> ConduitM ByteString a m ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Maybe (IDecode RealWorld a) -> ConduitM ByteString a m ()
rec' (Maybe (IDecode RealWorld a) -> ConduitM ByteString a m ())
-> (IDecode RealWorld a -> Maybe (IDecode RealWorld a))
-> IDecode RealWorld a
-> ConduitM ByteString a m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IDecode RealWorld a -> Maybe (IDecode RealWorld a)
forall a. a -> Maybe a
Just
        )
        Maybe (IDecode RealWorld a)
forall a. Maybe a
Nothing