{-# LANGUAGE DeriveDataTypeable #-}
-- | Turn a 'Get' into a 'Sink' and a 'Put' into a 'Source'
module Data.Conduit.Cereal where

import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as LBS
import Data.Typeable (Typeable)
import Control.Exception (Exception)
import Control.Monad.Trans
import qualified Data.Conduit as DC
import Data.Conduit.List (sourceList)
import Data.Serialize.Get
import Data.Serialize.Put
import Control.Exception (throw)

data GetException = GetException String
                  | GetDoesntConsumeInput
  deriving (Show, Typeable)

instance Exception GetException

-- | Convert a 'Get' into a 'Sink'. The 'Get' will be streamed bytes until it returns 'Done' or 'Fail'.
-- If the 'Get' fails, a GetException will be thrown with 'resourceThrow'. This function itself can also throw a GetException.
sinkGet :: DC.ResourceThrow m => Get output -> DC.Sink BS.ByteString m output
sinkGet get = case runGetPartial get BS.empty of
                Fail s -> throw $ GetException s
                Partial f -> DC.SinkData { DC.sinkPush = push f
                                         , DC.sinkClose = close f
                Done _ _ -> throw GetDoesntConsumeInput
  where push f input = case f input of
                         Fail s -> lift $ DC.resourceThrow $ GetException s
                         Partial f' -> return $ DC.Processing (push f') (close f')
                         Done r rest -> return $ DC.Done (if BS.null rest
                                                            then Nothing
                                                            else Just rest
                                                         ) r
        close f = let Fail s = f BS.empty in lift $ DC.resourceThrow $ GetException s

-- | Convert a 'Put' into a 'Source'. Runs in constant memory.
sourcePut :: DC.Resource m => Put -> DC.Source m BS.ByteString
sourcePut put = sourceList $ LBS.toChunks $ runPutLazy put