module HaskellWorks.Data.Conduit.ByteString
( rechunk
) where
import Control.Monad
import qualified Data.ByteString as BS
import Data.Conduit
rechunk :: Monad m => Int -> Conduit BS.ByteString m BS.ByteString
rechunk = rechunk' BS.empty
rechunk' :: Monad m => BS.ByteString -> Int -> Conduit BS.ByteString m BS.ByteString
rechunk' as n | BS.length as >= n = do
yield (BS.take n as)
rechunk' (BS.drop n as) n
rechunk' as n = do
mbss <- slurp (n BS.length as)
case mbss of
Just bss -> do
let bs = BS.concat (as : bss)
yield (BS.take n bs)
rechunk' (BS.drop n bs) n
Nothing -> unless (BS.null as) $ yield as
slurp :: Monad m => Int -> ConduitM BS.ByteString BS.ByteString m (Maybe [BS.ByteString])
slurp = go []
where go rs n | n > 0 = do
mbs <- await
case mbs of
Just bs -> go (bs : rs) (n BS.length bs)
Nothing -> if null rs then return Nothing else return (Just (reverse rs))
go rs _ = if null rs then return Nothing else return (Just (reverse rs))