ConduitCerealExt.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE BangPatterns #-}

module ConduitCerealExt where

import           Control.Exception.Base
import           Control.Monad.Trans.Resource (MonadThrow, throwM)
import qualified Data.ByteString as BS
import           Data.Conduit (ConduitT, await, yield)
import           Data.Serialize
import           Data.Typeable


newtype GetException2 = GetException2 String
  deriving (Show, Typeable)

instance Exception GetException2


conduitGet3 :: MonadThrow m => Get o -> ConduitT BS.ByteString (o, Int) m ()
conduitGet3 g =
    awaitNE >>= start 0
  where
    -- Get the next chunk of data, only returning an empty ByteString at the
    -- end of the stream.
    awaitNE =
        loop
      where
        loop = await >>= maybe (return BS.empty) check
        check bs
            | BS.null bs = loop
            | otherwise = return bs

    start !p bs
        | BS.null bs = return ()
        | otherwise = result p (runGetPartial ((,) <$> g <*> bytesRead) bs)

    result _ (Fail msg _) = throwM (GetException2 msg)
    -- This will feed an empty ByteString into f at end of stream, which is how
    -- we indicate to cereal that there is no data left. If we wanted to be
    -- more pedantic, we could ensure that cereal only ever consumes a single
    -- ByteString to avoid a loop, but that is the contract that cereal is
    -- giving us anyway.
    result p (Partial f) = awaitNE >>= result p . f
    result p (Done (x, l) rest) = do
        yield (x, p)
        let next = start (p + l)
        if BS.null rest
            then awaitNE >>= next
            else next rest