module Data.Conduit.Extra.Utils where
import Control.Applicative
import Control.Monad.Loops
import Control.Monad.Primitive
import Control.Monad.Trans.Class
import Control.Monad.Trans.State
import Data.Conduit
import Data.Conduit.List as CL
import Data.Foldable
import Data.Sequence as Seq
import Data.Vector as Boxed (Vector, freeze)
import Data.Vector.Mutable as Boxed hiding (length)
import qualified Data.Vector.Unboxed as Unboxed
import qualified Data.Vector.Unboxed.Mutable as Unboxed
takeWhile :: Monad m => (a -> Bool) -> Conduit a m a
takeWhile f = loop where
loop = await >>= maybe (return ()) go
go x | f x = yield x >> loop
| otherwise = leftover x
collect :: PrimMonad m => Int -> Sink a m (Vector a)
collect size = do
v <- lift $ unsafeNew size
forM_ [0..size1] $ \i -> do
me <- await
case me of
Nothing ->
error $ "Too many elements for a vector of size "
++ show size
Just e -> lift $ unsafeWrite v i e
lift $ freeze v
collectUnboxed :: (PrimMonad m, Unboxed.Unbox a)
=> Int -> Sink a m (Unboxed.Vector a)
collectUnboxed size = do
v <- lift $ Unboxed.unsafeNew size
forM_ [0..size1] $ \i -> do
me <- await
case me of
Nothing ->
error $ "Too many elements for an unboxed vector of size "
++ show size
Just e -> lift $ Unboxed.unsafeWrite v i e
lift $ Unboxed.freeze v
dropRight :: Monad m => Int -> Conduit a m a
dropRight size = do
xs <- Seq.fromList <$> CL.take size
flip evalStateT xs $ whileM_ ((== size) . Seq.length <$> get) $ do
xs' <- get
case viewl xs' of
EmptyL -> error "impossible"
y :< ys -> do
mz <- lift await
case mz of
Nothing -> put Seq.empty
Just z -> put (ys |> z) >> lift (yield y)