module HaskellWorks.Data.Conduit.ByteString
  ( rechunk
  ) where

import Control.Monad
import Data.Conduit

import qualified Data.ByteString as BS

rechunk :: Monad m
  => Int
  -> ConduitT BS.ByteString BS.ByteString m ()
rechunk :: Int -> ConduitT ByteString ByteString m ()
rechunk = ByteString -> Int -> ConduitT ByteString ByteString m ()
forall (m :: * -> *).
Monad m =>
ByteString -> Int -> ConduitT ByteString ByteString m ()
rechunk' ByteString
BS.empty

rechunk' :: Monad m
  => BS.ByteString
  -> Int
  -> ConduitT BS.ByteString BS.ByteString m ()
rechunk' :: ByteString -> Int -> ConduitT ByteString ByteString m ()
rechunk' ByteString
as Int
n | ByteString -> Int
BS.length ByteString
as Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = do
  ByteString -> ConduitT ByteString ByteString m ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
yield (Int -> ByteString -> ByteString
BS.take Int
n ByteString
as)
  ByteString -> Int -> ConduitT ByteString ByteString m ()
forall (m :: * -> *).
Monad m =>
ByteString -> Int -> ConduitT ByteString ByteString m ()
rechunk' (Int -> ByteString -> ByteString
BS.drop Int
n ByteString
as) Int
n
rechunk' ByteString
as Int
n = do
  Maybe [ByteString]
mbss <- Int -> ConduitM ByteString ByteString m (Maybe [ByteString])
forall (m :: * -> *).
Monad m =>
Int -> ConduitM ByteString ByteString m (Maybe [ByteString])
slurp (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- ByteString -> Int
BS.length ByteString
as)
  case Maybe [ByteString]
mbss of
    Just [ByteString]
bss -> do
      let bs :: ByteString
bs = [ByteString] -> ByteString
BS.concat (ByteString
as ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [ByteString]
bss)
      ByteString -> ConduitT ByteString ByteString m ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
yield (Int -> ByteString -> ByteString
BS.take Int
n ByteString
bs)
      ByteString -> Int -> ConduitT ByteString ByteString m ()
forall (m :: * -> *).
Monad m =>
ByteString -> Int -> ConduitT ByteString ByteString m ()
rechunk' (Int -> ByteString -> ByteString
BS.drop Int
n ByteString
bs) Int
n
    Maybe [ByteString]
Nothing -> Bool
-> ConduitT ByteString ByteString m ()
-> ConduitT ByteString ByteString m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ByteString -> Bool
BS.null ByteString
as) (ConduitT ByteString ByteString m ()
 -> ConduitT ByteString ByteString m ())
-> ConduitT ByteString ByteString m ()
-> ConduitT ByteString ByteString m ()
forall a b. (a -> b) -> a -> b
$ ByteString -> ConduitT ByteString ByteString m ()
forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
yield ByteString
as

slurp :: Monad m
  => Int
  -> ConduitM BS.ByteString BS.ByteString m (Maybe [BS.ByteString])
slurp :: Int -> ConduitM ByteString ByteString m (Maybe [ByteString])
slurp = [ByteString]
-> Int -> ConduitM ByteString ByteString m (Maybe [ByteString])
forall (m :: * -> *) o.
Monad m =>
[ByteString] -> Int -> ConduitT ByteString o m (Maybe [ByteString])
go []
  where go :: [ByteString] -> Int -> ConduitT ByteString o m (Maybe [ByteString])
go [ByteString]
rs Int
n | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 = do
          Maybe ByteString
mbs <- ConduitT ByteString o m (Maybe ByteString)
forall (m :: * -> *) i. Monad m => Consumer i m (Maybe i)
await
          case Maybe ByteString
mbs of
            Just ByteString
bs -> [ByteString] -> Int -> ConduitT ByteString o m (Maybe [ByteString])
go (ByteString
bs ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [ByteString]
rs) (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- ByteString -> Int
BS.length ByteString
bs)
            Maybe ByteString
Nothing -> if [ByteString] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [ByteString]
rs then Maybe [ByteString] -> ConduitT ByteString o m (Maybe [ByteString])
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe [ByteString]
forall a. Maybe a
Nothing else Maybe [ByteString] -> ConduitT ByteString o m (Maybe [ByteString])
forall (m :: * -> *) a. Monad m => a -> m a
return ([ByteString] -> Maybe [ByteString]
forall a. a -> Maybe a
Just ([ByteString] -> [ByteString]
forall a. [a] -> [a]
reverse [ByteString]
rs))
        go [ByteString]
rs Int
_ = if [ByteString] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [ByteString]
rs then Maybe [ByteString] -> ConduitT ByteString o m (Maybe [ByteString])
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe [ByteString]
forall a. Maybe a
Nothing else Maybe [ByteString] -> ConduitT ByteString o m (Maybe [ByteString])
forall (m :: * -> *) a. Monad m => a -> m a
return ([ByteString] -> Maybe [ByteString]
forall a. a -> Maybe a
Just ([ByteString] -> [ByteString]
forall a. [a] -> [a]
reverse [ByteString]
rs))