Generalized Conduit re-chunking

rechunk.hs

{-# LANGUAGE DefaultSignatures, LambdaCase #-}

import Conduit
import qualified Data.ByteString as BS
import qualified Data.DList as DL
import Data.Monoid
import Data.Sequences as S
import qualified Data.Vector as V
import qualified Data.Vector.Storable as SV
import qualified Data.Vector.Unboxed as UV

class FastLength a where
  flength :: a -> Index a

class Monoid a => Rechunkable a where
  strictSplitAt :: Index a -> a -> Either a ( a, a )
  default strictSplitAt :: (FastLength a, IsSequence a) => Index a -> a -> Either a ( a, a )
  strictSplitAt n xs
    | flength xs < n = Left xs
    | otherwise = Right (S.splitAt n xs)
  
instance Rechunkable [ a ] where
  strictSplitAt n = strictSplitAt' n DL.empty
    where
      strictSplitAt' 0 l xs = Right ( DL.toList l, xs ) 
      strictSplitAt' n l [] = Left (DL.toList l)
      strictSplitAt' n l (x : xs) = strictSplitAt' (n - 1) (l `DL.snoc` x) xs
      
instance FastLength BS.ByteString where
  flength = BS.length
instance Rechunkable BS.ByteString

instance FastLength (V.Vector a) where
  flength = V.length
instance Rechunkable (V.Vector a)

instance SV.Storable a => FastLength (SV.Vector a) where
  flength = SV.length
instance SV.Storable a => Rechunkable (SV.Vector a)

instance UV.Unbox a => FastLength (UV.Vector a) where
  flength = UV.length
instance UV.Unbox a => Rechunkable (UV.Vector a)

rechunk :: (Monad m, Rechunkable a) => Index a -> Conduit a m a
rechunk len = rechunk' mempty
  where
    rechunk' lo = await >>= \case
      Nothing -> return ()
      Just xs -> split (lo <> xs)
    split xs = case strictSplitAt len xs of
      Left lo -> rechunk' lo
      Right ( chunk, rest ) -> yield chunk >> split rest