{-# LANGUAGE CPP #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE TypeFamilies #-}
module Data.Conduit.Combinators.Stream
( yieldManyS
, repeatMS
, repeatWhileMS
, foldl1S
, allS
, anyS
, sinkLazyS
, sinkVectorS
, sinkVectorNS
, sinkLazyBuilderS
, lastS
, lastES
, findS
, concatMapS
, concatMapMS
, concatS
, scanlS
, scanlMS
, mapAccumWhileS
, mapAccumWhileMS
, intersperseS
, slidingWindowS
, filterMS
, splitOnUnboundedES
, initReplicateS
, initRepeatS
)
where
import Control.Monad (liftM)
import Control.Monad.Primitive (PrimMonad)
import qualified Data.ByteString.Lazy as BL
import Data.ByteString.Builder (Builder, toLazyByteString)
import Data.Conduit.Internal.Fusion
import Data.Conduit.Internal.List.Stream (foldS)
import Data.Maybe (isNothing, isJust)
import Data.MonoTraversable
#if ! MIN_VERSION_base(4,8,0)
import Data.Monoid (Monoid (..))
#endif
import qualified Data.NonNull as NonNull
import qualified Data.Sequences as Seq
import qualified Data.Vector.Generic as V
import qualified Data.Vector.Generic.Mutable as VM
import Prelude
#if MIN_VERSION_mono_traversable(1,0,0)
import Data.Sequences (LazySequence (..))
#else
import Data.Sequences.Lazy
#endif
yieldManyS :: (Monad m, MonoFoldable mono)
=> mono
-> StreamProducer m (Element mono)
yieldManyS mono _ =
Stream (return . step) (return (otoList mono))
where
step [] = Stop ()
step (x:xs) = Emit xs x
{-# INLINE yieldManyS #-}
repeatMS :: Monad m
=> m a
-> StreamProducer m a
repeatMS m _ =
Stream step (return ())
where
step _ = liftM (Emit ()) m
{-# INLINE repeatMS #-}
repeatWhileMS :: Monad m
=> m a
-> (a -> Bool)
-> StreamProducer m a
repeatWhileMS m f _ =
Stream step (return ())
where
step _ = do
x <- m
return $ if f x
then Emit () x
else Stop ()
{-# INLINE repeatWhileMS #-}
foldl1S :: Monad m
=> (a -> a -> a)
-> StreamConsumer a m (Maybe a)
foldl1S f (Stream step ms0) =
Stream step' (liftM (Nothing, ) ms0)
where
step' (mprev, s) = do
res <- step s
return $ case res of
Stop () -> Stop mprev
Skip s' -> Skip (mprev, s')
Emit s' a -> Skip (Just $ maybe a (`f` a) mprev, s')
{-# INLINE foldl1S #-}
allS :: Monad m
=> (a -> Bool)
-> StreamConsumer a m Bool
allS f = fmapS isNothing (findS (Prelude.not . f))
{-# INLINE allS #-}
anyS :: Monad m
=> (a -> Bool)
-> StreamConsumer a m Bool
anyS f = fmapS isJust (findS f)
{-# INLINE anyS #-}
sinkLazyS :: (Monad m, LazySequence lazy strict)
=> StreamConsumer strict m lazy
sinkLazyS = fmapS (fromChunks . ($ [])) $ foldS (\front next -> front . (next:)) id
{-# INLINE sinkLazyS #-}
sinkVectorS :: (V.Vector v a, PrimMonad m)
=> StreamConsumer a m (v a)
sinkVectorS (Stream step ms0) = do
Stream step' $ do
s0 <- ms0
mv0 <- VM.new initSize
return (initSize, 0, mv0, s0)
where
initSize = 10
step' (maxSize, i, mv, s) = do
res <- step s
case res of
Stop () -> liftM (Stop . V.slice 0 i) $ V.unsafeFreeze mv
Skip s' -> return $ Skip (maxSize, i, mv, s')
Emit s' x -> do
VM.write mv i x
let i' = i + 1
if i' >= maxSize
then do
let newMax = maxSize * 2
mv' <- VM.grow mv maxSize
return $ Skip (newMax, i', mv', s')
else return $ Skip (maxSize, i', mv, s')
{-# INLINE sinkVectorS #-}
sinkVectorNS :: (V.Vector v a, PrimMonad m)
=> Int
-> StreamConsumer a m (v a)
sinkVectorNS maxSize (Stream step ms0) = do
Stream step' $ do
s0 <- ms0
mv0 <- VM.new maxSize
return (0, mv0, s0)
where
step' (i, mv, _) | i >= maxSize = liftM Stop $ V.unsafeFreeze mv
step' (i, mv, s) = do
res <- step s
case res of
Stop () -> liftM (Stop . V.slice 0 i) $ V.unsafeFreeze mv
Skip s' -> return $ Skip (i, mv, s')
Emit s' x -> do
VM.write mv i x
let i' = i + 1
return $ Skip (i', mv, s')
{-# INLINE sinkVectorNS #-}
sinkLazyBuilderS :: Monad m => StreamConsumer Builder m BL.ByteString
sinkLazyBuilderS = fmapS toLazyByteString (foldS mappend mempty)
{-# INLINE sinkLazyBuilderS #-}
lastS :: Monad m
=> StreamConsumer a m (Maybe a)
lastS (Stream step ms0) =
Stream step' (liftM (Nothing,) ms0)
where
step' (mlast, s) = do
res <- step s
return $ case res of
Stop () -> Stop mlast
Skip s' -> Skip (mlast, s')
Emit s' x -> Skip (Just x, s')
{-# INLINE lastS #-}
lastES :: (Monad m, Seq.IsSequence seq)
=> StreamConsumer seq m (Maybe (Element seq))
lastES (Stream step ms0) =
Stream step' (liftM (Nothing, ) ms0)
where
step' (mlast, s) = do
res <- step s
return $ case res of
Stop () -> Stop (fmap NonNull.last mlast)
Skip s' -> Skip (mlast, s')
Emit s' (NonNull.fromNullable -> mlast'@(Just _)) -> Skip (mlast', s')
Emit s' _ -> Skip (mlast, s')
{-# INLINE lastES #-}
findS :: Monad m
=> (a -> Bool) -> StreamConsumer a m (Maybe a)
findS f (Stream step ms0) =
Stream step' ms0
where
step' s = do
res <- step s
return $ case res of
Stop () -> Stop Nothing
Skip s' -> Skip s'
Emit s' x ->
if f x
then Stop (Just x)
else Skip s'
{-# INLINE findS #-}
concatMapS :: (Monad m, MonoFoldable mono)
=> (a -> mono)
-> StreamConduit a m (Element mono)
concatMapS f (Stream step ms0) =
Stream step' (liftM ([], ) ms0)
where
step' ([], s) = do
res <- step s
return $ case res of
Stop () -> Stop ()
Skip s' -> Skip ([], s')
Emit s' x -> Skip (otoList (f x), s')
step' ((x:xs), s) = return (Emit (xs, s) x)
{-# INLINE concatMapS #-}
concatMapMS :: (Monad m, MonoFoldable mono)
=> (a -> m mono)
-> StreamConduit a m (Element mono)
concatMapMS f (Stream step ms0) =
Stream step' (liftM ([], ) ms0)
where
step' ([], s) = do
res <- step s
case res of
Stop () -> return $ Stop ()
Skip s' -> return $ Skip ([], s')
Emit s' x -> do
o <- f x
return $ Skip (otoList o, s')
step' ((x:xs), s) = return (Emit (xs, s) x)
{-# INLINE concatMapMS #-}
concatS :: (Monad m, MonoFoldable mono)
=> StreamConduit mono m (Element mono)
concatS = concatMapS id
{-# INLINE concatS #-}
data ScanState a s
= ScanEnded
| ScanContinues a s
scanlS :: Monad m => (a -> b -> a) -> a -> StreamConduit b m a
scanlS f seed0 (Stream step ms0) =
Stream step' (liftM (ScanContinues seed0) ms0)
where
step' ScanEnded = return $ Stop ()
step' (ScanContinues seed s) = do
res <- step s
return $ case res of
Stop () -> Emit ScanEnded seed
Skip s' -> Skip (ScanContinues seed s')
Emit s' x -> Emit (ScanContinues seed' s') seed
where
!seed' = f seed x
{-# INLINE scanlS #-}
scanlMS :: Monad m => (a -> b -> m a) -> a -> StreamConduit b m a
scanlMS f seed0 (Stream step ms0) =
Stream step' (liftM (ScanContinues seed0) ms0)
where
step' ScanEnded = return $ Stop ()
step' (ScanContinues seed s) = do
res <- step s
case res of
Stop () -> return $ Emit ScanEnded seed
Skip s' -> return $ Skip (ScanContinues seed s')
Emit s' x -> do
!seed' <- f seed x
return $ Emit (ScanContinues seed' s') seed
{-# INLINE scanlMS #-}
mapAccumWhileS :: Monad m =>
(a -> s -> Either s (s, b)) -> s -> StreamConduitT a b m s
mapAccumWhileS f initial (Stream step ms0) =
Stream step' (liftM (initial, ) ms0)
where
step' (!accum, s) = do
res <- step s
return $ case res of
Stop () -> Stop accum
Skip s' -> Skip (accum, s')
Emit s' x -> case f x accum of
Right (!accum', r) -> Emit (accum', s') r
Left !accum' -> Stop accum'
{-# INLINE mapAccumWhileS #-}
mapAccumWhileMS :: Monad m =>
(a -> s -> m (Either s (s, b))) -> s -> StreamConduitT a b m s
mapAccumWhileMS f initial (Stream step ms0) =
Stream step' (liftM (initial, ) ms0)
where
step' (!accum, s) = do
res <- step s
case res of
Stop () -> return $ Stop accum
Skip s' -> return $ Skip (accum, s')
Emit s' x -> do
lr <- f x accum
return $ case lr of
Right (!accum', r) -> Emit (accum', s') r
Left !accum' -> Stop accum'
{-# INLINE mapAccumWhileMS #-}
data IntersperseState a s
= IFirstValue s
| IGotValue s a
| IEmitValue s a
intersperseS :: Monad m => a -> StreamConduit a m a
intersperseS sep (Stream step ms0) =
Stream step' (liftM IFirstValue ms0)
where
step' (IFirstValue s) = do
res <- step s
return $ case res of
Stop () -> Stop ()
Skip s' -> Skip (IFirstValue s')
Emit s' x -> Emit (IGotValue s' x) x
step' (IGotValue s x) = do
res <- step s
return $ case res of
Stop () -> Stop ()
Skip s' -> Skip (IGotValue s' x)
Emit s' x' -> Emit (IEmitValue s' x') sep
step' (IEmitValue s x) = return $ Emit (IGotValue s x) x
{-# INLINE intersperseS #-}
data SlidingWindowState seq s
= SWInitial Int seq s
| SWSliding seq s
| SWEarlyExit
slidingWindowS :: (Monad m, Seq.IsSequence seq, Element seq ~ a) => Int -> StreamConduit a m seq
slidingWindowS sz (Stream step ms0) =
Stream step' (liftM (SWInitial (max 1 sz) mempty) ms0)
where
step' (SWInitial n st s) = do
res <- step s
return $ case res of
Stop () -> Emit SWEarlyExit st
Skip s' -> Skip (SWInitial n st s')
Emit s' x ->
if n == 1
then Emit (SWSliding (Seq.unsafeTail st') s') st'
else Skip (SWInitial (n - 1) st' s')
where
st' = Seq.snoc st x
step' (SWSliding st s) = do
res <- step s
return $ case res of
Stop () -> Stop ()
Skip s' -> Skip (SWSliding st s')
Emit s' x -> Emit (SWSliding (Seq.unsafeTail st') s') st'
where
st' = Seq.snoc st x
step' SWEarlyExit = return $ Stop ()
{-# INLINE slidingWindowS #-}
filterMS :: Monad m
=> (a -> m Bool)
-> StreamConduit a m a
filterMS f (Stream step ms0) = do
Stream step' ms0
where
step' s = do
res <- step s
case res of
Stop () -> return $ Stop ()
Skip s' -> return $ Skip s'
Emit s' x -> do
r <- f x
return $
if r
then Emit s' x
else Skip s'
{-# INLINE filterMS #-}
data SplitState seq s
= SplitDone
| SplitNoSep seq s
| SplitState seq s
splitOnUnboundedES :: (Monad m, Seq.IsSequence seq)
=> (Element seq -> Bool) -> StreamConduit seq m seq
splitOnUnboundedES f (Stream step ms0) =
Stream step' (liftM (SplitState mempty) ms0)
where
step' SplitDone = return $ Stop ()
step' (SplitNoSep t s) = do
res <- step s
return $ case res of
Stop () | not (onull t) -> Emit SplitDone t
| otherwise -> Stop ()
Skip s' -> Skip (SplitNoSep t s')
Emit s' t' -> Skip (SplitState (t `mappend` t') s')
step' (SplitState t s) = do
if onull y
then do
res <- step s
return $ case res of
Stop () | not (onull t) -> Emit SplitDone t
| otherwise -> Stop ()
Skip s' -> Skip (SplitNoSep t s')
Emit s' t' -> Skip (SplitState (t `mappend` t') s')
else return $ Emit (SplitState (Seq.drop 1 y) s) x
where
(x, y) = Seq.break f t
{-# INLINE splitOnUnboundedES #-}
initReplicateS :: Monad m => m seed -> (seed -> m a) -> Int -> StreamProducer m a
initReplicateS mseed f cnt _ =
Stream step (liftM (cnt, ) mseed)
where
step (ix, _) | ix <= 0 = return $ Stop ()
step (ix, seed) = do
x <- f seed
return $ Emit (ix - 1, seed) x
{-# INLINE initReplicateS #-}
initRepeatS :: Monad m => m seed -> (seed -> m a) -> StreamProducer m a
initRepeatS mseed f _ =
Stream step mseed
where
step seed = do
x <- f seed
return $ Emit seed x
{-# INLINE initRepeatS #-}
fmapS :: Monad m
=> (a -> b)
-> StreamConduitT i o m a
-> StreamConduitT i o m b
fmapS f s inp =
case s inp of
Stream step ms0 -> Stream (fmap (liftM (fmap f)) step) ms0
{-# INLINE fmapS #-}