git.iko.soy/streamly
on master help sign in
Move transform combinators from Stream.Eliminate to Transform.
97d5dc6Pranay Sashank
src/Streamly/Internal/Data/Stream/StreamD/Eliminate.hs+5−673
compare with master
unifiedsplit
@@ -14,36 +14,23 @@ module Streamly.Internal.Data.Stream.StreamD.Eliminate
14 (
14 (
15 -- * Elimination
15 -- * Elimination
16 -- ** General Folds
16 -- ** General Folds
17− foldrS
18− , foldrT
19− , foldrM
17+ foldrM
20 , foldrMx
18 , foldrMx
21 , foldr
19 , foldr
22 , foldr1
20 , foldr1
23
21
24 , foldl'
22 , foldl'
25 , foldlM'
23 , foldlM'
26− , foldlS
27− , foldlT
28 , toList
24 , toList
29 , toListRev
25 , toListRev
30− , reverse
31− , reverse'
32
26
33 , foldlx'
27 , foldlx'
34 , foldlMx'
28 , foldlMx'
35 , foldOnce
29 , foldOnce
36
30
37 , parselMx'
31 , parselMx'
38− , parseMany
39− , parseIterate
40
32
41 -- ** Specialized Folds
33 -- ** Specialized Folds
42− , tap
43− , tapOffsetEvery
44− , tapAsync
45− , tapRate
46− , pollCounts
47 , drain
34 , drain
48 , null
35 , null
49 , head
36 , head
@@ -79,55 +66,20 @@ module Streamly.Internal.Data.Stream.StreamD.Eliminate
79
66
80 -- ** Map and Fold
67 -- ** Map and Fold
81 , mapM_
68 , mapM_
82−
83− -- ** Conversions
84− -- | Transform a stream into another type.
85− , hoist
86− , generally
87−
88− , liftInner
89− , runReaderT
90− , evalStateT
91− , runStateT
92 )
69 )
93where
70where
94
71
95−import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)
72+import Control.Exception (assert)
96−import Control.Exception
97− (assert, AsyncException)
98−import Control.Monad (void, when)
99−import Control.Monad.Catch (MonadCatch, MonadThrow, throwM)
73+import Control.Monad.Catch (MonadThrow, throwM)
100−import Control.Monad.IO.Class (MonadIO(..))
101−import Control.Monad.Trans.Class (MonadTrans(lift))
102−import Data.IORef (newIORef, mkWeakIORef)
103−import Foreign.Storable (Storable(..))
104import GHC.Exts (SpecConstrAnnotation(..))
74import GHC.Exts (SpecConstrAnnotation(..))
105import GHC.Types (SPEC(..))
75import GHC.Types (SPEC(..))
106−import Fusion.Plugin.Types (Fuse(..))
107−import Streamly.Internal.Data.Fold.Types (Fold(..))
108import Streamly.Internal.Data.Parser (ParseError(..))
76import Streamly.Internal.Data.Parser (ParseError(..))
109−import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
110−import Control.Monad.Trans.Reader (ReaderT)
111−import Control.Monad.Trans.State.Strict (StateT)
112−import Data.Functor.Identity (Identity(..))
113
77
114−import qualified Streamly.Internal.Data.IORef.Prim as Prim
115−import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A
116−import qualified Streamly.Internal.Data.Fold as FL
117−import qualified Streamly.Internal.Data.Stream.StreamK as K
118import qualified Streamly.Internal.Data.Parser as PR
78import qualified Streamly.Internal.Data.Parser as PR
119−import qualified Streamly.Internal.Data.Parser.ParserD as PRD
120−import qualified Control.Monad.Catch as MC
121−import qualified Control.Monad.Trans.Reader as Reader
122−import qualified Control.Monad.Trans.State.Strict as State
123−import qualified Prelude
124
79
125import Prelude hiding
80import Prelude hiding
126− ( map, mapM, mapM_, repeat, foldr, last, take, filter
81+ ( all, any, elem, foldr, foldr1, head, last, lookup, mapM, mapM_
127− , takeWhile, drop, dropWhile, all, any, maximum, minimum, elem
82+ , maximum, minimum, notElem, null, splitAt, tail, (!!))
128− , notElem, null, head, tail, zipWith, lookup, foldr1, sequence
129− , (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat
130− , reverse, iterate, splitAt)
131import Streamly.Internal.Data.Stream.StreamD.Type
83import Streamly.Internal.Data.Stream.StreamD.Type
132import Streamly.Internal.Data.SVar
84import Streamly.Internal.Data.SVar
133
85
@@ -147,44 +99,6 @@ foldr1 f m = do
147 Nothing -> return Nothing
99 Nothing -> return Nothing
148 Just (h, t) -> fmap Just (foldr f h t)
100 Just (h, t) -> fmap Just (foldr f h t)
149
101
150−------------------------------------------------------------------------------
151−-- Left Folds
152−------------------------------------------------------------------------------
153−
154−{-# INLINE_NORMAL foldlT #-}
155−foldlT :: (Monad m, Monad (s m), MonadTrans s)
156− => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
157−foldlT fstep begin (Stream step state) = go SPEC begin state
158− where
159− go !_ acc st = do
160− r <- lift $ step defState st
161− case r of
162− Yield x s -> go SPEC (fstep acc x) s
163− Skip s -> go SPEC acc s
164− Stop -> acc
165−
166−-- Note, this is going to have horrible performance, because of the nature of
167−-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
168−-- likely be practically unusable.
169−{-# INLINE_NORMAL foldlS #-}
170−foldlS :: Monad m
171− => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
172−foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
173− where
174− step' gst (Left (st, acc)) = do
175− r <- step (adaptState gst) st
176− return $ case r of
177− Yield x s -> Skip (Left (s, fstep acc x))
178− Skip s -> Skip (Left (s, acc))
179− Stop -> Skip (Right acc)
180−
181− step' gst (Right (Stream stp stt)) = do
182− r <- stp (adaptState gst) stt
183− return $ case r of
184− Yield x s -> Yield x (Right (Stream stp s))
185− Skip s -> Skip (Right (Stream stp s))
186− Stop -> Stop
187−
188------------------------------------------------------------------------------
102------------------------------------------------------------------------------
189-- Parses
103-- Parses
190------------------------------------------------------------------------------
104------------------------------------------------------------------------------
@@ -272,227 +186,6 @@ parselMx' pstep initial extract (Stream step state) = do
272 PR.Done _ b -> return b
186 PR.Done _ b -> return b
273 PR.Error err -> throwM $ ParseError err
187 PR.Error err -> throwM $ ParseError err
274
188
275−------------------------------------------------------------------------------
276−-- Repeated parsing
277−------------------------------------------------------------------------------
278−
279−{-# ANN type ParseChunksState Fuse #-}
280−data ParseChunksState x inpBuf st pst =
281− ParseChunksInit inpBuf st
282− | ParseChunksInitLeftOver inpBuf
283− | ParseChunksStream st inpBuf !pst
284− | ParseChunksBuf inpBuf st inpBuf !pst
285− | ParseChunksYield x (ParseChunksState x inpBuf st pst)
286−
287−{-# INLINE_NORMAL parseMany #-}
288−parseMany
289− :: MonadThrow m
290− => PRD.Parser m a b
291− -> Stream m a
292− -> Stream m b
293−parseMany (PRD.Parser pstep initial extract) (Stream step state) =
294− Stream stepOuter (ParseChunksInit [] state)
295−
296− where
297−
298− {-# INLINE_LATE stepOuter #-}
299− -- Buffer is empty, get the first element from the stream, initialize the
300− -- fold and then go to stream processing loop.
301− stepOuter gst (ParseChunksInit [] st) = do
302− r <- step (adaptState gst) st
303− case r of
304− Yield x s -> initial >>= return . Skip . ParseChunksBuf [x] s []
305− Skip s -> return $ Skip $ ParseChunksInit [] s
306− Stop -> return Stop
307−
308− -- Buffer is not empty, go to buffered processing loop
309− stepOuter _ (ParseChunksInit src st) = do
310− initial >>= return . Skip . ParseChunksBuf src st []
311−
312− -- XXX we just discard any leftover input at the end
313− stepOuter _ (ParseChunksInitLeftOver _) = return Stop
314−
315− -- Buffer is empty, process elements from the stream
316− stepOuter gst (ParseChunksStream st buf pst) = do
317− r <- step (adaptState gst) st
318− case r of
319− Yield x s -> do
320− pRes <- pstep pst x
321− case pRes of
322− PR.Partial 0 pst1 ->
323− return $ Skip $ ParseChunksStream s [] pst1
324− PR.Partial n pst1 -> do
325− assert (n <= length (x:buf)) (return ())
326− let src0 = Prelude.take n (x:buf)
327− src = Prelude.reverse src0
328− return $ Skip $ ParseChunksBuf src s [] pst1
329− PR.Continue 0 pst1 ->
330− return $ Skip $ ParseChunksStream s (x:buf) pst1
331− PR.Continue n pst1 -> do
332− assert (n <= length (x:buf)) (return ())
333− let (src0, buf1) = splitAt n (x:buf)
334− src = Prelude.reverse src0
335− return $ Skip $ ParseChunksBuf src s buf1 pst1
336− PR.Done 0 b -> do
337− return $ Skip $
338− ParseChunksYield b (ParseChunksInit [] s)
339− PR.Done n b -> do
340− assert (n <= length (x:buf)) (return ())
341− let src = Prelude.reverse (Prelude.take n (x:buf))
342− return $ Skip $
343− ParseChunksYield b (ParseChunksInit src s)
344− PR.Error err -> throwM $ ParseError err
345− Skip s -> return $ Skip $ ParseChunksStream s buf pst
346− Stop -> do
347− b <- extract pst
348− let src = Prelude.reverse buf
349− return $ Skip $
350− ParseChunksYield b (ParseChunksInitLeftOver src)
351−
352− -- go back to stream processing mode
353− stepOuter _ (ParseChunksBuf [] s buf pst) =
354− return $ Skip $ ParseChunksStream s buf pst
355−
356− -- buffered processing loop
357− stepOuter _ (ParseChunksBuf (x:xs) s buf pst) = do
358− pRes <- pstep pst x
359− case pRes of
360− PR.Partial 0 pst1 ->
361− return $ Skip $ ParseChunksBuf xs s [] pst1
362− PR.Partial n pst1 -> do
363− assert (n <= length (x:buf)) (return ())
364− let src0 = Prelude.take n (x:buf)
365− src = Prelude.reverse src0 ++ xs
366− return $ Skip $ ParseChunksBuf src s [] pst1
367− PR.Continue 0 pst1 ->
368− return $ Skip $ ParseChunksBuf xs s (x:buf) pst1
369− PR.Continue n pst1 -> do
370− assert (n <= length (x:buf)) (return ())
371− let (src0, buf1) = splitAt n (x:buf)
372− src = Prelude.reverse src0 ++ xs
373− return $ Skip $ ParseChunksBuf src s buf1 pst1
374− PR.Done 0 b ->
375− return $ Skip $ ParseChunksYield b (ParseChunksInit xs s)
376− PR.Done n b -> do
377− assert (n <= length (x:buf)) (return ())
378− let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
379− return $ Skip $ ParseChunksYield b (ParseChunksInit src s)
380− PR.Error err -> throwM $ ParseError err
381−
382− stepOuter _ (ParseChunksYield a next) = return $ Yield a next
383−
384−{-# ANN type ConcatParseState Fuse #-}
385−data ConcatParseState x inpBuf st p =
386− ConcatParseInit inpBuf st p
387− | ConcatParseInitLeftOver inpBuf
388− | ConcatParseStream st inpBuf p
389− | ConcatParseBuf inpBuf st inpBuf p
390− | ConcatParseYield x (ConcatParseState x inpBuf st p)
391−
392−{-# INLINE_NORMAL parseIterate #-}
393−parseIterate
394− :: MonadThrow m
395− => (b -> PRD.Parser m a b)
396− -> b
397− -> Stream m a
398− -> Stream m b
399−parseIterate func seed (Stream step state) =
400− Stream stepOuter (ConcatParseInit [] state (func seed))
401−
402− where
403−
404− {-# INLINE_LATE stepOuter #-}
405− -- Buffer is empty, go to stream processing loop
406− stepOuter _ (ConcatParseInit [] st (PRD.Parser pstep initial extract)) = do
407− initial >>= \r -> return $ Skip $ ConcatParseStream st []
408− (PRD.Parser pstep (return r) extract)
409−
410− -- Buffer is not empty, go to buffered processing loop
411− stepOuter _ (ConcatParseInit src st
412− (PRD.Parser pstep initial extract)) = do
413− initial >>= \r -> return $ Skip $ ConcatParseBuf src st []
414− (PRD.Parser pstep (return r) extract)
415−
416− -- XXX we just discard any leftover input at the end
417− stepOuter _ (ConcatParseInitLeftOver _) = return Stop
418−
419− -- Buffer is empty process elements from the stream
420− stepOuter gst (ConcatParseStream st buf
421− p@(PRD.Parser pstep initial extract)) = do
422− r <- step (adaptState gst) st
423− case r of
424− Yield x s -> do
425− pst <- initial
426− pRes <- pstep pst x
427− case pRes of
428− PR.Partial 0 pst1 ->
429− return $ Skip $ ConcatParseStream s []
430− (PRD.Parser pstep (return pst1) extract)
431− PR.Partial n pst1 -> do
432− assert (n <= length (x:buf)) (return ())
433− let src0 = Prelude.take n (x:buf)
434− src = Prelude.reverse src0
435− return $ Skip $ ConcatParseBuf src s []
436− (PRD.Parser pstep (return pst1) extract)
437− -- PR.Continue 0 pst1 ->
438− -- return $ Skip $ ConcatParseStream s (x:buf) pst1
439− PR.Continue n pst1 -> do
440− assert (n <= length (x:buf)) (return ())
441− let (src0, buf1) = splitAt n (x:buf)
442− src = Prelude.reverse src0
443− return $ Skip $ ConcatParseBuf src s buf1
444− (PRD.Parser pstep (return pst1) extract)
445− -- XXX Specialize for Stop 0 common case?
446− PR.Done n b -> do
447− assert (n <= length (x:buf)) (return ())
448− let src = Prelude.reverse (Prelude.take n (x:buf))
449− return $ Skip $
450− ConcatParseYield b (ConcatParseInit src s (func b))
451− PR.Error err -> throwM $ ParseError err
452− Skip s -> return $ Skip $ ConcatParseStream s buf p
453− Stop -> do
454− pst <- initial
455− b <- extract pst
456− let src = Prelude.reverse buf
457− return $ Skip $
458− ConcatParseYield b (ConcatParseInitLeftOver src)
459−
460− -- go back to stream processing mode
461− stepOuter _ (ConcatParseBuf [] s buf p) =
462− return $ Skip $ ConcatParseStream s buf p
463−
464− -- buffered processing loop
465− stepOuter _ (ConcatParseBuf (x:xs) s buf
466− (PRD.Parser pstep initial extract)) = do
467− pst <- initial
468− pRes <- pstep pst x
469− case pRes of
470− PR.Partial 0 pst1 ->
471− return $ Skip $ ConcatParseBuf xs s []
472− (PRD.Parser pstep (return pst1) extract)
473− PR.Partial n pst1 -> do
474− assert (n <= length (x:buf)) (return ())
475− let src0 = Prelude.take n (x:buf)
476− src = Prelude.reverse src0 ++ xs
477− return $ Skip $ ConcatParseBuf src s []
478− (PRD.Parser pstep (return pst1) extract)
479− -- PR.Continue 0 pst1 -> return $ Skip $ ConcatParseBuf xs s (x:buf) pst1
480− PR.Continue n pst1 -> do
481− assert (n <= length (x:buf)) (return ())
482− let (src0, buf1) = splitAt n (x:buf)
483− src = Prelude.reverse src0 ++ xs
484− return $ Skip $ ConcatParseBuf src s buf1
485− (PRD.Parser pstep (return pst1) extract)
486− -- XXX Specialize for Stop 0 common case?
487− PR.Done n b -> do
488− assert (n <= length (x:buf)) (return ())
489− let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
490− return $ Skip $ ConcatParseYield b
491− (ConcatParseInit src s (func b))
492− PR.Error err -> throwM $ ParseError err
493−
494− stepOuter _ (ConcatParseYield a next) = return $ Yield a next
495−
496------------------------------------------------------------------------------
189------------------------------------------------------------------------------
497-- Specialized Folds
190-- Specialized Folds
498------------------------------------------------------------------------------
191------------------------------------------------------------------------------
@@ -704,367 +397,6 @@ findIndices p (Stream step state) = Stream step' (state, 0)
704toListRev :: Monad m => Stream m a -> m [a]
397toListRev :: Monad m => Stream m a -> m [a]
705toListRev = foldl' (flip (:)) []
398toListRev = foldl' (flip (:)) []
706
399
707−-- We can implement reverse as:
708−--
709−-- > reverse = foldlS (flip cons) nil
710−--
711−-- However, this implementation is unusable because of the horrible performance
712−-- of cons. So we just convert it to a list first and then stream from the
713−-- list.
714−--
715−-- XXX Maybe we can use an Array instead of a list here?
716−{-# INLINE_NORMAL reverse #-}
717−reverse :: Monad m => Stream m a -> Stream m a
718−reverse m = Stream step Nothing
719− where
720− {-# INLINE_LATE step #-}
721− step _ Nothing = do
722− xs <- toListRev m
723− return $ Skip (Just xs)
724− step _ (Just (x:xs)) = return $ Yield x (Just xs)
725− step _ (Just []) = return Stop
726−
727−-- Much faster reverse for Storables
728−{-# INLINE_NORMAL reverse' #-}
729−reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
730−{-
731−-- This commented implementation copies the whole stream into one single array
732−-- and then streams from that array, this is 3-4x faster than the chunked code
733−-- that follows. Though this could be problematic due to unbounded large
734−-- allocations. We need to figure out why the chunked code is slower and if we
735−-- can optimize the chunked code to work as fast as this one. It may be a
736−-- fusion issue?
737−import Foreign.ForeignPtr (touchForeignPtr)
738−import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
739−import Foreign.Ptr (Ptr, plusPtr)
740−reverse' m = Stream step Nothing
741− where
742− {-# INLINE_LATE step #-}
743− step _ Nothing = do
744− arr <- A.fromStreamD m
745− let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
746− return $ Skip $ Just (aStart arr, p)
747−
748− step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
749−
750− step _ (Just (start, p)) = do
751− let !x = A.unsafeInlineIO $ do
752− r <- peek p
753− touchForeignPtr start
754− return r
755− next = p `plusPtr` negate (sizeOf (undefined :: a))
756− return $ Yield x (Just (start, next))
757−-}
758−reverse' m =
759− A.flattenArraysRev
760− $ fromStreamK
761− $ K.reverse
762− $ toStreamK
763− $ A.fromStreamDArraysOf A.defaultChunkSize m
764−
765−------------------------------------------------------------------------------
766−-- Tapping/Distributing
767−------------------------------------------------------------------------------
768−
769−data TapState fs st a
770− = TapInit | Tapping !fs st | TapDone st
771−
772−-- XXX Multiple yield points
773−{-# INLINE tap #-}
774−tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
775−tap (Fold fstep initial extract) (Stream step state) = Stream step' TapInit
776−
777− where
778−
779− step' _ TapInit = do
780− r <- initial
781− return $ Skip (Tapping r state)
782− step' gst (Tapping acc st) = do
783− r <- step gst st
784− case r of
785− -- XXX Abstract Yield?
786− Yield x s -> do
787− res <- fstep acc x
788− return
789− $ Yield x
790− $ case res of
791− FL.Partial fs -> Tapping fs s
792− FL.Done _ -> TapDone s
793− Skip s -> return $ Skip (Tapping acc s)
794− Stop -> do
795− void $ extract acc
796− return Stop
797− step' gst (TapDone st) = do
798− r <- step gst st
799− return
800− $ case r of
801− Yield x s -> Yield x (TapDone s)
802− Skip s -> Skip (TapDone s)
803− Stop -> Stop
804−
805−data TapOffState fs s a
806− = TapOffInit
807− | TapOffTapping !fs s Int
808− | TapOffDone s
809−
810−-- XXX Multiple yield points
811−{-# INLINE_NORMAL tapOffsetEvery #-}
812−tapOffsetEvery :: Monad m
813− => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
814−tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
815− Stream step' TapOffInit
816−
817− where
818−
819− {-# INLINE_LATE step' #-}
820− step' _ TapOffInit = do
821− r <- initial
822− return $ Skip $ TapOffTapping r state (offset `mod` n)
823− step' gst (TapOffTapping acc st count) = do
824− r <- step gst st
825− case r of
826− Yield x s -> do
827− next <-
828− if count <= 0
829− then do
830− res <- fstep acc x
831− return
832− $ case res of
833− FL.Partial sres ->
834− TapOffTapping sres s (n - 1)
835− FL.Done _ -> TapOffDone s
836− else return $ TapOffTapping acc s (count - 1)
837− return $ Yield x next
838− Skip s -> return $ Skip (TapOffTapping acc s count)
839− Stop -> do
840− void $ extract acc
841− return Stop
842− step' gst (TapOffDone st) = do
843− r <- step gst st
844− return
845− $ case r of
846− Yield x s -> Yield x (TapOffDone s)
847− Skip s -> Skip (TapOffDone s)
848− Stop -> Stop
849−
850−{-# INLINE_NORMAL pollCounts #-}
851−pollCounts
852− :: MonadAsync m
853− => (a -> Bool)
854− -> (Stream m Int -> Stream m Int)
855− -> Fold m Int b
856− -> Stream m a
857− -> Stream m a
858−pollCounts predicate transf fld (Stream step state) = Stream step' Nothing
859− where
860−
861− {-# INLINE_LATE step' #-}
862− step' _ Nothing = do
863− -- As long as we are using an "Int" for counts lockfree reads from
864− -- Var should work correctly on both 32-bit and 64-bit machines.
865− -- However, an Int on a 32-bit machine may overflow quickly.
866− countVar <- liftIO $ Prim.newIORef (0 :: Int)
867− tid <- forkManaged
868− $ void $ foldOnce fld
869− $ transf $ fromPrimIORef countVar
870− return $ Skip (Just (countVar, tid, state))
871−
872− step' gst (Just (countVar, tid, st)) = do
873− r <- step gst st
874− case r of
875− Yield x s -> do
876− when (predicate x) $ liftIO $ Prim.modifyIORef' countVar (+ 1)
877− return $ Yield x (Just (countVar, tid, s))
878− Skip s -> return $ Skip (Just (countVar, tid, s))
879− Stop -> do
880− liftIO $ killThread tid
881− return Stop
882−
883−{-# INLINE_NORMAL tapRate #-}
884−tapRate ::
885− (MonadAsync m, MonadCatch m)
886− => Double
887− -> (Int -> m b)
888− -> Stream m a
889− -> Stream m a
890−tapRate samplingRate action (Stream step state) = Stream step' Nothing
891− where
892− {-# NOINLINE loop #-}
893− loop countVar prev = do
894− i <-
895− MC.catch
896− (do liftIO $ threadDelay (round $ samplingRate * 1000000)
897− i <- liftIO $ Prim.readIORef countVar
898− let !diff = i - prev
899− void $ action diff
900− return i)
901− (\(e :: AsyncException) -> do
902− i <- liftIO $ Prim.readIORef countVar
903− let !diff = i - prev
904− void $ action diff
905− throwM (MC.toException e))
906− loop countVar i
907−
908− {-# INLINE_LATE step' #-}
909− step' _ Nothing = do
910− countVar <- liftIO $ Prim.newIORef 0
911− tid <- fork $ loop countVar 0
912− ref <- liftIO $ newIORef ()
913− _ <- liftIO $ mkWeakIORef ref (killThread tid)
914− return $ Skip (Just (countVar, tid, state, ref))
915−
916− step' gst (Just (countVar, tid, st, ref)) = do
917− r <- step gst st
918− case r of
919− Yield x s -> do
920− liftIO $ Prim.modifyIORef' countVar (+ 1)
921− return $ Yield x (Just (countVar, tid, s, ref))
922− Skip s -> return $ Skip (Just (countVar, tid, s, ref))
923− Stop -> do
924− liftIO $ killThread tid
925− return Stop
926−
927−-------------------------------------------------------------------------------
928−-- Concurrent tap
929−-------------------------------------------------------------------------------
930−
931−-- | Create an SVar with a fold consumer that will fold any elements sent to it
932−-- using the supplied fold function.
933−{-# INLINE newFoldSVar #-}
934−newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
935−newFoldSVar stt f = do
936− -- Buffer size for the SVar is derived from the current state
937− sv <- newParallelVar StopAny (adaptState stt)
938− -- Add the producer thread-id to the SVar.
939− liftIO myThreadId >>= modifyThread sv
940− void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
941− return sv
942−
943− where
944−
945− {-# NOINLINE work #-}
946− work sv = void $ foldOnce f $ fromProducer sv
947−
948−{-# INLINE_NORMAL tapAsync #-}
949−tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
950−tapAsync f (Stream step1 state1) = Stream step TapInit
951− where
952−
953− drainFold svr = do
954− -- In general, a Stop event would come equipped with the result
955− -- of the fold. It is not used here but it would be useful in
956− -- applicative and distribute.
957− done <- fromConsumer svr
958− when (not done) $ do
959− liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
960− $ takeMVar (outputDoorBellFromConsumer svr)
961− drainFold svr
962−
963− stopFold svr = do
964− liftIO $ sendStop svr Nothing
965− -- drain/wait until a stop event arrives from the fold.
966− drainFold svr
967−
968− {-# INLINE_LATE step #-}
969− step gst TapInit = do
970− sv <- newFoldSVar gst f
971− return $ Skip (Tapping sv state1)
972−
973− step gst (Tapping sv st) = do
974− r <- step1 gst st
975− case r of
976− Yield a s -> do
977− done <- pushToFold sv a
978− if done
979− then do
980− -- XXX we do not need to wait synchronously here
981− stopFold sv
982− return $ Yield a (TapDone s)
983− else return $ Yield a (Tapping sv s)
984− Skip s -> return $ Skip (Tapping sv s)
985− Stop -> do
986− stopFold sv
987− return $ Stop
988−
989− step gst (TapDone st) = do
990− r <- step1 gst st
991− return $ case r of
992− Yield a s -> Yield a (TapDone s)
993− Skip s -> Skip (TapDone s)
994− Stop -> Stop
995−
996−-------------------------------------------------------------------------------
997−-- Hoisting the inner monad
998−-------------------------------------------------------------------------------
999−
1000−{-# INLINE_NORMAL hoist #-}
1001−hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
1002−hoist f (Stream step state) = (Stream step' state)
1003− where
1004− {-# INLINE_LATE step' #-}
1005− step' gst st = do
1006− r <- f $ step (adaptState gst) st
1007− return $ case r of
1008− Yield x s -> Yield x s
1009− Skip s -> Skip s
1010− Stop -> Stop
1011−
1012−{-# INLINE generally #-}
1013−generally :: Monad m => Stream Identity a -> Stream m a
1014−generally = hoist (return . runIdentity)
1015−
1016−{-# INLINE_NORMAL liftInner #-}
1017−liftInner :: (Monad m, MonadTrans t, Monad (t m))
1018− => Stream m a -> Stream (t m) a
1019−liftInner (Stream step state) = Stream step' state
1020− where
1021− {-# INLINE_LATE step' #-}
1022− step' gst st = do
1023− r <- lift $ step (adaptState gst) st
1024− return $ case r of
1025− Yield x s -> Yield x s
1026− Skip s -> Skip s
1027− Stop -> Stop
1028−
1029−{-# INLINE_NORMAL runReaderT #-}
1030−runReaderT :: Monad m => m s -> Stream (ReaderT s m) a -> Stream m a
1031−runReaderT env (Stream step state) = Stream step' (state, env)
1032− where
1033− {-# INLINE_LATE step' #-}
1034− step' gst (st, action) = do
1035− sv <- action
1036− r <- Reader.runReaderT (step (adaptState gst) st) sv
1037− return $ case r of
1038− Yield x s -> Yield x (s, return sv)
1039− Skip s -> Skip (s, return sv)
1040− Stop -> Stop
1041−
1042−{-# INLINE_NORMAL evalStateT #-}
1043−evalStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m a
1044−evalStateT initial (Stream step state) = Stream step' (state, initial)
1045− where
1046− {-# INLINE_LATE step' #-}
1047− step' gst (st, action) = do
1048− sv <- action
1049− (r, sv') <- State.runStateT (step (adaptState gst) st) sv
1050− return $ case r of
1051− Yield x s -> Yield x (s, return sv')
1052− Skip s -> Skip (s, return sv')
1053− Stop -> Stop
1054−
1055−{-# INLINE_NORMAL runStateT #-}
1056−runStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m (s, a)
1057−runStateT initial (Stream step state) = Stream step' (state, initial)
1058− where
1059− {-# INLINE_LATE step' #-}
1060− step' gst (st, action) = do
1061− sv <- action
1062− (r, sv') <- State.runStateT (step (adaptState gst) st) sv
1063− return $ case r of
1064− Yield x s -> Yield (sv', x) (s, return sv')
1065− Skip s -> Skip (s, return sv')
1066− Stop -> Stop
1067−
1068------------------------------------------------------------------------------
400------------------------------------------------------------------------------
1069-- Transformation comprehensions
401-- Transformation comprehensions
1070------------------------------------------------------------------------------
402------------------------------------------------------------------------------
src/Streamly/Internal/Data/Stream/StreamD/Transform.hs+684−7
@@ -13,7 +13,12 @@ module Streamly.Internal.Data.Stream.StreamD.Transform
13 -- * Transformation
13 -- * Transformation
14 transform
14 transform
15
15
16− -- ** By folding (scans)
16+ -- ** By folding
17+ , foldrS
18+ , foldrT
19+ , foldlS
20+ , foldlT
21+
17 , scanlM'
22 , scanlM'
18 , scanlMAfter'
23 , scanlMAfter'
19 , scanl'
24 , scanl'
@@ -40,6 +45,10 @@ module Streamly.Internal.Data.Stream.StreamD.Transform
40 , postscanOnce
45 , postscanOnce
41 , scanOnce
46 , scanOnce
42
47
48+ -- * Reordering
49+ , reverse
50+ , reverse'
51+
43 -- * Filtering
52 -- * Filtering
44 , filter
53 , filter
45 , filterM
54 , filterM
@@ -60,6 +69,13 @@ module Streamly.Internal.Data.Stream.StreamD.Transform
60 , rollingMap
69 , rollingMap
61 , rollingMapM
70 , rollingMapM
62
71
72+ -- ** Special Maps
73+ , tap
74+ , tapOffsetEvery
75+ , tapAsync
76+ , tapRate
77+ , pollCounts
78+
63 -- * Inserting
79 -- * Inserting
64 , intersperseM
80 , intersperseM
65 , intersperseM_
81 , intersperseM_
@@ -82,6 +98,10 @@ module Streamly.Internal.Data.Stream.StreamD.Transform
82 , zipWith
98 , zipWith
83 , zipWithM
99 , zipWithM
84
100
101+ -- * Parsers
102+ , parseMany
103+ , parseIterate
104+
85 -- * Merging
105 -- * Merging
86 , mergeBy
106 , mergeBy
87 , mergeByM
107 , mergeByM
@@ -89,25 +109,58 @@ module Streamly.Internal.Data.Stream.StreamD.Transform
89 -- * Concurrent Application
109 -- * Concurrent Application
90 , mkParallel
110 , mkParallel
91 , mkParallelD
111 , mkParallelD
112+
113+ -- ** Conversions
114+ -- | Transform a stream into another type.
115+ , hoist
116+ , generally
117+
118+ , liftInner
119+ , runReaderT
120+ , evalStateT
121+ , runStateT
92 )
122 )
93where
123where
94
124
95−import Control.Monad (void)
125+import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)
126+import Control.Exception (assert, AsyncException)
127+import Control.Monad (void, when)
128+import Control.Monad.Catch (MonadCatch, MonadThrow, throwM)
96import Control.Monad.IO.Class (MonadIO(..))
129import Control.Monad.IO.Class (MonadIO(..))
130+import Control.Monad.Trans.Class (MonadTrans(lift))
131+import Control.Monad.Trans.Reader (ReaderT)
132+import Control.Monad.Trans.State.Strict (StateT)
97import Data.Functor.Identity (Identity(..))
133import Data.Functor.Identity (Identity(..))
134+import Data.IORef (newIORef, mkWeakIORef)
98import Data.Maybe (fromJust, isJust)
135import Data.Maybe (fromJust, isJust)
136+import Foreign.Storable (Storable(..))
137+import Fusion.Plugin.Types (Fuse(..))
138+import GHC.Types (SPEC(..))
139+import qualified Control.Monad.Catch as MC
140+import qualified Control.Monad.Trans.Reader as Reader
141+import qualified Control.Monad.Trans.State.Strict as State
142+
99−import Streamly.Internal.Data.Time.Units
143+import Streamly.Internal.Data.Fold.Types (Fold(..))
100− (TimeUnit64, toRelTime64, diffAbsTime64)
144+import Streamly.Internal.Data.Parser (ParseError(..))
101import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
145import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
146+import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
102import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
147import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
148+import Streamly.Internal.Data.Time.Units
149+ (TimeUnit64, toRelTime64, diffAbsTime64)
103
150
104−import qualified Streamly.Internal.Data.Pipe.Types as Pipe
151+import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A
105import qualified Streamly.Internal.Data.Fold as FL
152import qualified Streamly.Internal.Data.Fold as FL
153+import qualified Streamly.Internal.Data.IORef.Prim as Prim
154+import qualified Streamly.Internal.Data.Parser as PR
155+import qualified Streamly.Internal.Data.Parser.ParserD as PRD
156+import qualified Streamly.Internal.Data.Pipe.Types as Pipe
106import qualified Streamly.Internal.Data.Stream.StreamK as K
157import qualified Streamly.Internal.Data.Stream.StreamK as K
107
158
159+import qualified Prelude
108import Prelude hiding
160import Prelude hiding
109− ( drop, dropWhile, filter, map, mapM
161+ ( drop, dropWhile, filter, map, mapM, reverse
110 , scanl, scanl1, sequence, take, takeWhile, zipWith)
162 , scanl, scanl1, sequence, take, takeWhile, zipWith)
163+
111import Streamly.Internal.Data.Stream.StreamD.Type
164import Streamly.Internal.Data.Stream.StreamD.Type
112import Streamly.Internal.Data.SVar
165import Streamly.Internal.Data.SVar
113
166
@@ -142,9 +195,47 @@ transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
142 Pipe.Continue pst' -> return $ Skip (pst', st)
195 Pipe.Continue pst' -> return $ Skip (pst', st)
143
196
144------------------------------------------------------------------------------
197------------------------------------------------------------------------------
145−-- Transformation by Folding (Scans)
198+-- Transformation by Folding
146------------------------------------------------------------------------------
199------------------------------------------------------------------------------
147
200
201+------------------------------------------------------------------------------
202+-- Left Folds
203+------------------------------------------------------------------------------
204+
205+{-# INLINE_NORMAL foldlT #-}
206+foldlT :: (Monad m, Monad (s m), MonadTrans s)
207+ => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
208+foldlT fstep begin (Stream step state) = go SPEC begin state
209+ where
210+ go !_ acc st = do
211+ r <- lift $ step defState st
212+ case r of
213+ Yield x s -> go SPEC (fstep acc x) s
214+ Skip s -> go SPEC acc s
215+ Stop -> acc
216+
217+-- Note, this is going to have horrible performance, because of the nature of
218+-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
219+-- likely be practically unusable.
220+{-# INLINE_NORMAL foldlS #-}
221+foldlS :: Monad m
222+ => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
223+foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
224+ where
225+ step' gst (Left (st, acc)) = do
226+ r <- step (adaptState gst) st
227+ return $ case r of
228+ Yield x s -> Skip (Left (s, fstep acc x))
229+ Skip s -> Skip (Left (s, acc))
230+ Stop -> Skip (Right acc)
231+
232+ step' gst (Right (Stream stp stt)) = do
233+ r <- stp (adaptState gst) stt
234+ return $ case r of
235+ Yield x s -> Yield x (Right (Stream stp s))
236+ Skip s -> Skip (Right (Stream stp s))
237+ Stop -> Stop
238+
148------------------------------------------------------------------------------
239------------------------------------------------------------------------------
149-- Prescans
240-- Prescans
150------------------------------------------------------------------------------
241------------------------------------------------------------------------------
@@ -456,6 +547,68 @@ rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
456rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
547rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
457rollingMap f = rollingMapM (\x y -> return $ f x y)
548rollingMap f = rollingMapM (\x y -> return $ f x y)
458
549
550+------------------------------------------------------------------------------
551+-- Reordering
552+------------------------------------------------------------------------------
553+
554+-- We can implement reverse as:
555+--
556+-- > reverse = foldlS (flip cons) nil
557+--
558+-- However, this implementation is unusable because of the horrible performance
559+-- of cons. So we just convert it to a list first and then stream from the
560+-- list.
561+--
562+-- XXX Maybe we can use an Array instead of a list here?
563+{-# INLINE_NORMAL reverse #-}
564+reverse :: Monad m => Stream m a -> Stream m a
565+reverse m = Stream step Nothing
566+ where
567+ {-# INLINE_LATE step #-}
568+ step _ Nothing = do
569+ xs <- foldl' (flip (:)) [] m
570+ return $ Skip (Just xs)
571+ step _ (Just (x:xs)) = return $ Yield x (Just xs)
572+ step _ (Just []) = return Stop
573+
574+-- Much faster reverse for Storables
575+{-# INLINE_NORMAL reverse' #-}
576+reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
577+{-
578+-- This commented implementation copies the whole stream into one single array
579+-- and then streams from that array, this is 3-4x faster than the chunked code
580+-- that follows. Though this could be problematic due to unbounded large
581+-- allocations. We need to figure out why the chunked code is slower and if we
582+-- can optimize the chunked code to work as fast as this one. It may be a
583+-- fusion issue?
584+import Foreign.ForeignPtr (touchForeignPtr)
585+import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
586+import Foreign.Ptr (Ptr, plusPtr)
587+reverse' m = Stream step Nothing
588+ where
589+ {-# INLINE_LATE step #-}
590+ step _ Nothing = do
591+ arr <- A.fromStreamD m
592+ let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
593+ return $ Skip $ Just (aStart arr, p)
594+
595+ step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
596+
597+ step _ (Just (start, p)) = do
598+ let !x = A.unsafeInlineIO $ do
599+ r <- peek p
600+ touchForeignPtr start
601+ return r
602+ next = p `plusPtr` negate (sizeOf (undefined :: a))
603+ return $ Yield x (Just (start, next))
604+-}
605+reverse' m =
606+ A.flattenArraysRev
607+ $ fromStreamK
608+ $ K.reverse
609+ $ toStreamK
610+ $ A.fromStreamDArraysOf A.defaultChunkSize m
611+
459-------------------------------------------------------------------------------
612-------------------------------------------------------------------------------
460-- Filtering
613-- Filtering
461-------------------------------------------------------------------------------
614-------------------------------------------------------------------------------
@@ -572,6 +725,237 @@ sequence (Stream step state) = Stream step' state
572 Skip s -> return $ Skip s
725 Skip s -> return $ Skip s
573 Stop -> return Stop
726 Stop -> return Stop
574
727
728+------------------------------------------------------------------------------
729+-- Tapping/Distributing
730+------------------------------------------------------------------------------
731+
732+data TapState fs st a
733+ = TapInit | Tapping !fs st | TapDone st
734+
735+-- XXX Multiple yield points
736+{-# INLINE tap #-}
737+tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
738+tap (Fold fstep initial extract) (Stream step state) = Stream step' TapInit
739+
740+ where
741+
742+ step' _ TapInit = do
743+ r <- initial
744+ return $ Skip (Tapping r state)
745+ step' gst (Tapping acc st) = do
746+ r <- step gst st
747+ case r of
748+ -- XXX Abstract Yield?
749+ Yield x s -> do
750+ res <- fstep acc x
751+ return
752+ $ Yield x
753+ $ case res of
754+ FL.Partial fs -> Tapping fs s
755+ FL.Done _ -> TapDone s
756+ Skip s -> return $ Skip (Tapping acc s)
757+ Stop -> do
758+ void $ extract acc
759+ return Stop
760+ step' gst (TapDone st) = do
761+ r <- step gst st
762+ return
763+ $ case r of
764+ Yield x s -> Yield x (TapDone s)
765+ Skip s -> Skip (TapDone s)
766+ Stop -> Stop
767+
768+data TapOffState fs s a
769+ = TapOffInit
770+ | TapOffTapping !fs s Int
771+ | TapOffDone s
772+
773+-- XXX Multiple yield points
774+{-# INLINE_NORMAL tapOffsetEvery #-}
775+tapOffsetEvery :: Monad m
776+ => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
777+tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
778+ Stream step' TapOffInit
779+
780+ where
781+
782+ {-# INLINE_LATE step' #-}
783+ step' _ TapOffInit = do
784+ r <- initial
785+ return $ Skip $ TapOffTapping r state (offset `mod` n)
786+ step' gst (TapOffTapping acc st count) = do
787+ r <- step gst st
788+ case r of
789+ Yield x s -> do
790+ next <-
791+ if count <= 0
792+ then do
793+ res <- fstep acc x
794+ return
795+ $ case res of
796+ FL.Partial sres ->
797+ TapOffTapping sres s (n - 1)
798+ FL.Done _ -> TapOffDone s
799+ else return $ TapOffTapping acc s (count - 1)
800+ return $ Yield x next
801+ Skip s -> return $ Skip (TapOffTapping acc s count)
802+ Stop -> do
803+ void $ extract acc
804+ return Stop
805+ step' gst (TapOffDone st) = do
806+ r <- step gst st
807+ return
808+ $ case r of
809+ Yield x s -> Yield x (TapOffDone s)
810+ Skip s -> Skip (TapOffDone s)
811+ Stop -> Stop
812+
813+{-# INLINE_NORMAL pollCounts #-}
814+pollCounts
815+ :: MonadAsync m
816+ => (a -> Bool)
817+ -> (Stream m Int -> Stream m Int)
818+ -> Fold m Int b
819+ -> Stream m a
820+ -> Stream m a
821+pollCounts predicate transf fld (Stream step state) = Stream step' Nothing
822+ where
823+
824+ {-# INLINE_LATE step' #-}
825+ step' _ Nothing = do
826+ -- As long as we are using an "Int" for counts lockfree reads from
827+ -- Var should work correctly on both 32-bit and 64-bit machines.
828+ -- However, an Int on a 32-bit machine may overflow quickly.
829+ countVar <- liftIO $ Prim.newIORef (0 :: Int)
830+ tid <- forkManaged
831+ $ void $ foldOnce fld
832+ $ transf $ fromPrimIORef countVar
833+ return $ Skip (Just (countVar, tid, state))
834+
835+ step' gst (Just (countVar, tid, st)) = do
836+ r <- step gst st
837+ case r of
838+ Yield x s -> do
839+ when (predicate x) $ liftIO $ Prim.modifyIORef' countVar (+ 1)
840+ return $ Yield x (Just (countVar, tid, s))
841+ Skip s -> return $ Skip (Just (countVar, tid, s))
842+ Stop -> do
843+ liftIO $ killThread tid
844+ return Stop
845+
846+{-# INLINE_NORMAL tapRate #-}
847+tapRate ::
848+ (MonadAsync m, MonadCatch m)
849+ => Double
850+ -> (Int -> m b)
851+ -> Stream m a
852+ -> Stream m a
853+tapRate samplingRate action (Stream step state) = Stream step' Nothing
854+ where
855+ {-# NOINLINE loop #-}
856+ loop countVar prev = do
857+ i <-
858+ MC.catch
859+ (do liftIO $ threadDelay (round $ samplingRate * 1000000)
860+ i <- liftIO $ Prim.readIORef countVar
861+ let !diff = i - prev
862+ void $ action diff
863+ return i)
864+ (\(e :: AsyncException) -> do
865+ i <- liftIO $ Prim.readIORef countVar
866+ let !diff = i - prev
867+ void $ action diff
868+ throwM (MC.toException e))
869+ loop countVar i
870+
871+ {-# INLINE_LATE step' #-}
872+ step' _ Nothing = do
873+ countVar <- liftIO $ Prim.newIORef 0
874+ tid <- fork $ loop countVar 0
875+ ref <- liftIO $ newIORef ()
876+ _ <- liftIO $ mkWeakIORef ref (killThread tid)
877+ return $ Skip (Just (countVar, tid, state, ref))
878+
879+ step' gst (Just (countVar, tid, st, ref)) = do
880+ r <- step gst st
881+ case r of
882+ Yield x s -> do
883+ liftIO $ Prim.modifyIORef' countVar (+ 1)
884+ return $ Yield x (Just (countVar, tid, s, ref))
885+ Skip s -> return $ Skip (Just (countVar, tid, s, ref))
886+ Stop -> do
887+ liftIO $ killThread tid
888+ return Stop
889+
890+-------------------------------------------------------------------------------
891+-- Concurrent tap
892+-------------------------------------------------------------------------------
893+
894+-- | Create an SVar with a fold consumer that will fold any elements sent to it
895+-- using the supplied fold function.
896+{-# INLINE newFoldSVar #-}
897+newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
898+newFoldSVar stt f = do
899+ -- Buffer size for the SVar is derived from the current state
900+ sv <- newParallelVar StopAny (adaptState stt)
901+ -- Add the producer thread-id to the SVar.
902+ liftIO myThreadId >>= modifyThread sv
903+ void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
904+ return sv
905+
906+ where
907+
908+ {-# NOINLINE work #-}
909+ work sv = void $ foldOnce f $ fromProducer sv
910+
911+{-# INLINE_NORMAL tapAsync #-}
912+tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
913+tapAsync f (Stream step1 state1) = Stream step TapInit
914+ where
915+
916+ drainFold svr = do
917+ -- In general, a Stop event would come equipped with the result
918+ -- of the fold. It is not used here but it would be useful in
919+ -- applicative and distribute.
920+ done <- fromConsumer svr
921+ when (not done) $ do
922+ liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
923+ $ takeMVar (outputDoorBellFromConsumer svr)
924+ drainFold svr
925+
926+ stopFold svr = do
927+ liftIO $ sendStop svr Nothing
928+ -- drain/wait until a stop event arrives from the fold.
929+ drainFold svr
930+
931+ {-# INLINE_LATE step #-}
932+ step gst TapInit = do
933+ sv <- newFoldSVar gst f
934+ return $ Skip (Tapping sv state1)
935+
936+ step gst (Tapping sv st) = do
937+ r <- step1 gst st
938+ case r of
939+ Yield a s -> do
940+ done <- pushToFold sv a
941+ if done
942+ then do
943+ -- XXX we do not need to wait synchronously here
944+ stopFold sv
945+ return $ Yield a (TapDone s)
946+ else return $ Yield a (Tapping sv s)
947+ Skip s -> return $ Skip (Tapping sv s)
948+ Stop -> do
949+ stopFold sv
950+ return $ Stop
951+
952+ step gst (TapDone st) = do
953+ r <- step1 gst st
954+ return $ case r of
955+ Yield a s -> Yield a (TapDone s)
956+ Skip s -> Skip (TapDone s)
957+ Stop -> Stop
958+
575------------------------------------------------------------------------------
959------------------------------------------------------------------------------
576-- Inserting
960-- Inserting
577------------------------------------------------------------------------------
961------------------------------------------------------------------------------
@@ -812,6 +1196,227 @@ zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
812zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
1196zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
813zipWith f = zipWithM (\a b -> return (f a b))
1197zipWith f = zipWithM (\a b -> return (f a b))
814
1198
1199+------------------------------------------------------------------------------
1200+-- Repeated parsing
1201+------------------------------------------------------------------------------
1202+
1203+{-# ANN type ParseChunksState Fuse #-}
1204+data ParseChunksState x inpBuf st pst =
1205+ ParseChunksInit inpBuf st
1206+ | ParseChunksInitLeftOver inpBuf
1207+ | ParseChunksStream st inpBuf !pst
1208+ | ParseChunksBuf inpBuf st inpBuf !pst
1209+ | ParseChunksYield x (ParseChunksState x inpBuf st pst)
1210+
1211+{-# INLINE_NORMAL parseMany #-}
1212+parseMany
1213+ :: MonadThrow m
1214+ => PRD.Parser m a b
1215+ -> Stream m a
1216+ -> Stream m b
1217+parseMany (PRD.Parser pstep initial extract) (Stream step state) =
1218+ Stream stepOuter (ParseChunksInit [] state)
1219+
1220+ where
1221+
1222+ {-# INLINE_LATE stepOuter #-}
1223+ -- Buffer is empty, get the first element from the stream, initialize the
1224+ -- fold and then go to stream processing loop.
1225+ stepOuter gst (ParseChunksInit [] st) = do
1226+ r <- step (adaptState gst) st
1227+ case r of
1228+ Yield x s -> initial >>= return . Skip . ParseChunksBuf [x] s []
1229+ Skip s -> return $ Skip $ ParseChunksInit [] s
1230+ Stop -> return Stop
1231+
1232+ -- Buffer is not empty, go to buffered processing loop
1233+ stepOuter _ (ParseChunksInit src st) = do
1234+ initial >>= return . Skip . ParseChunksBuf src st []
1235+
1236+ -- XXX we just discard any leftover input at the end
1237+ stepOuter _ (ParseChunksInitLeftOver _) = return Stop
1238+
1239+ -- Buffer is empty, process elements from the stream
1240+ stepOuter gst (ParseChunksStream st buf pst) = do
1241+ r <- step (adaptState gst) st
1242+ case r of
1243+ Yield x s -> do
1244+ pRes <- pstep pst x
1245+ case pRes of
1246+ PR.Partial 0 pst1 ->
1247+ return $ Skip $ ParseChunksStream s [] pst1
1248+ PR.Partial n pst1 -> do
1249+ assert (n <= length (x:buf)) (return ())
1250+ let src0 = Prelude.take n (x:buf)
1251+ src = Prelude.reverse src0
1252+ return $ Skip $ ParseChunksBuf src s [] pst1
1253+ PR.Continue 0 pst1 ->
1254+ return $ Skip $ ParseChunksStream s (x:buf) pst1
1255+ PR.Continue n pst1 -> do
1256+ assert (n <= length (x:buf)) (return ())
1257+ let (src0, buf1) = splitAt n (x:buf)
1258+ src = Prelude.reverse src0
1259+ return $ Skip $ ParseChunksBuf src s buf1 pst1
1260+ PR.Done 0 b -> do
1261+ return $ Skip $
1262+ ParseChunksYield b (ParseChunksInit [] s)
1263+ PR.Done n b -> do
1264+ assert (n <= length (x:buf)) (return ())
1265+ let src = Prelude.reverse (Prelude.take n (x:buf))
1266+ return $ Skip $
1267+ ParseChunksYield b (ParseChunksInit src s)
1268+ PR.Error err -> throwM $ ParseError err
1269+ Skip s -> return $ Skip $ ParseChunksStream s buf pst
1270+ Stop -> do
1271+ b <- extract pst
1272+ let src = Prelude.reverse buf
1273+ return $ Skip $
1274+ ParseChunksYield b (ParseChunksInitLeftOver src)
1275+
1276+ -- go back to stream processing mode
1277+ stepOuter _ (ParseChunksBuf [] s buf pst) =
1278+ return $ Skip $ ParseChunksStream s buf pst
1279+
1280+ -- buffered processing loop
1281+ stepOuter _ (ParseChunksBuf (x:xs) s buf pst) = do
1282+ pRes <- pstep pst x
1283+ case pRes of
1284+ PR.Partial 0 pst1 ->
1285+ return $ Skip $ ParseChunksBuf xs s [] pst1
1286+ PR.Partial n pst1 -> do
1287+ assert (n <= length (x:buf)) (return ())
1288+ let src0 = Prelude.take n (x:buf)
1289+ src = Prelude.reverse src0 ++ xs
1290+ return $ Skip $ ParseChunksBuf src s [] pst1
1291+ PR.Continue 0 pst1 ->
1292+ return $ Skip $ ParseChunksBuf xs s (x:buf) pst1
1293+ PR.Continue n pst1 -> do
1294+ assert (n <= length (x:buf)) (return ())
1295+ let (src0, buf1) = splitAt n (x:buf)
1296+ src = Prelude.reverse src0 ++ xs
1297+ return $ Skip $ ParseChunksBuf src s buf1 pst1
1298+ PR.Done 0 b ->
1299+ return $ Skip $ ParseChunksYield b (ParseChunksInit xs s)
1300+ PR.Done n b -> do
1301+ assert (n <= length (x:buf)) (return ())
1302+ let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1303+ return $ Skip $ ParseChunksYield b (ParseChunksInit src s)
1304+ PR.Error err -> throwM $ ParseError err
1305+
1306+ stepOuter _ (ParseChunksYield a next) = return $ Yield a next
1307+
1308+{-# ANN type ConcatParseState Fuse #-}
1309+data ConcatParseState x inpBuf st p =
1310+ ConcatParseInit inpBuf st p
1311+ | ConcatParseInitLeftOver inpBuf
1312+ | ConcatParseStream st inpBuf p
1313+ | ConcatParseBuf inpBuf st inpBuf p
1314+ | ConcatParseYield x (ConcatParseState x inpBuf st p)
1315+
1316+{-# INLINE_NORMAL parseIterate #-}
1317+parseIterate
1318+ :: MonadThrow m
1319+ => (b -> PRD.Parser m a b)
1320+ -> b
1321+ -> Stream m a
1322+ -> Stream m b
1323+parseIterate func seed (Stream step state) =
1324+ Stream stepOuter (ConcatParseInit [] state (func seed))
1325+
1326+ where
1327+
1328+ {-# INLINE_LATE stepOuter #-}
1329+ -- Buffer is empty, go to stream processing loop
1330+ stepOuter _ (ConcatParseInit [] st (PRD.Parser pstep initial extract)) = do
1331+ initial >>= \r -> return $ Skip $ ConcatParseStream st []
1332+ (PRD.Parser pstep (return r) extract)
1333+
1334+ -- Buffer is not empty, go to buffered processing loop
1335+ stepOuter _ (ConcatParseInit src st
1336+ (PRD.Parser pstep initial extract)) = do
1337+ initial >>= \r -> return $ Skip $ ConcatParseBuf src st []
1338+ (PRD.Parser pstep (return r) extract)
1339+
1340+ -- XXX we just discard any leftover input at the end
1341+ stepOuter _ (ConcatParseInitLeftOver _) = return Stop
1342+
1343+ -- Buffer is empty process elements from the stream
1344+ stepOuter gst (ConcatParseStream st buf
1345+ p@(PRD.Parser pstep initial extract)) = do
1346+ r <- step (adaptState gst) st
1347+ case r of
1348+ Yield x s -> do
1349+ pst <- initial
1350+ pRes <- pstep pst x
1351+ case pRes of
1352+ PR.Partial 0 pst1 ->
1353+ return $ Skip $ ConcatParseStream s []
1354+ (PRD.Parser pstep (return pst1) extract)
1355+ PR.Partial n pst1 -> do
1356+ assert (n <= length (x:buf)) (return ())
1357+ let src0 = Prelude.take n (x:buf)
1358+ src = Prelude.reverse src0
1359+ return $ Skip $ ConcatParseBuf src s []
1360+ (PRD.Parser pstep (return pst1) extract)
1361+ -- PR.Continue 0 pst1 ->
1362+ -- return $ Skip $ ConcatParseStream s (x:buf) pst1
1363+ PR.Continue n pst1 -> do
1364+ assert (n <= length (x:buf)) (return ())
1365+ let (src0, buf1) = splitAt n (x:buf)
1366+ src = Prelude.reverse src0
1367+ return $ Skip $ ConcatParseBuf src s buf1
1368+ (PRD.Parser pstep (return pst1) extract)
1369+ -- XXX Specialize for Stop 0 common case?
1370+ PR.Done n b -> do
1371+ assert (n <= length (x:buf)) (return ())
1372+ let src = Prelude.reverse (Prelude.take n (x:buf))
1373+ return $ Skip $
1374+ ConcatParseYield b (ConcatParseInit src s (func b))
1375+ PR.Error err -> throwM $ ParseError err
1376+ Skip s -> return $ Skip $ ConcatParseStream s buf p
1377+ Stop -> do
1378+ pst <- initial
1379+ b <- extract pst
1380+ let src = Prelude.reverse buf
1381+ return $ Skip $
1382+ ConcatParseYield b (ConcatParseInitLeftOver src)
1383+
1384+ -- go back to stream processing mode
1385+ stepOuter _ (ConcatParseBuf [] s buf p) =
1386+ return $ Skip $ ConcatParseStream s buf p
1387+
1388+ -- buffered processing loop
1389+ stepOuter _ (ConcatParseBuf (x:xs) s buf
1390+ (PRD.Parser pstep initial extract)) = do
1391+ pst <- initial
1392+ pRes <- pstep pst x
1393+ case pRes of
1394+ PR.Partial 0 pst1 ->
1395+ return $ Skip $ ConcatParseBuf xs s []
1396+ (PRD.Parser pstep (return pst1) extract)
1397+ PR.Partial n pst1 -> do
1398+ assert (n <= length (x:buf)) (return ())
1399+ let src0 = Prelude.take n (x:buf)
1400+ src = Prelude.reverse src0 ++ xs
1401+ return $ Skip $ ConcatParseBuf src s []
1402+ (PRD.Parser pstep (return pst1) extract)
1403+ -- PR.Continue 0 pst1 -> return $ Skip $ ConcatParseBuf xs s (x:buf) pst1
1404+ PR.Continue n pst1 -> do
1405+ assert (n <= length (x:buf)) (return ())
1406+ let (src0, buf1) = splitAt n (x:buf)
1407+ src = Prelude.reverse src0 ++ xs
1408+ return $ Skip $ ConcatParseBuf src s buf1
1409+ (PRD.Parser pstep (return pst1) extract)
1410+ -- XXX Specialize for Stop 0 common case?
1411+ PR.Done n b -> do
1412+ assert (n <= length (x:buf)) (return ())
1413+ let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1414+ return $ Skip $ ConcatParseYield b
1415+ (ConcatParseInit src s (func b))
1416+ PR.Error err -> throwM $ ParseError err
1417+
1418+ stepOuter _ (ConcatParseYield a next) = return $ Yield a next
1419+
815------------------------------------------------------------------------------
1420------------------------------------------------------------------------------
816-- Merging
1421-- Merging
817------------------------------------------------------------------------------
1422------------------------------------------------------------------------------
@@ -982,3 +1587,75 @@ dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
982 Yield x s -> Yield x (DropByTimeYield s)
1587 Yield x s -> Yield x (DropByTimeYield s)
983 Skip s -> Skip (DropByTimeYield s)
1588 Skip s -> Skip (DropByTimeYield s)
984 Stop -> Stop
1589 Stop -> Stop
1590+
1591+-------------------------------------------------------------------------------
1592+-- Hoisting the inner monad
1593+-------------------------------------------------------------------------------
1594+
1595+{-# INLINE_NORMAL hoist #-}
1596+hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
1597+hoist f (Stream step state) = (Stream step' state)
1598+ where
1599+ {-# INLINE_LATE step' #-}
1600+ step' gst st = do
1601+ r <- f $ step (adaptState gst) st
1602+ return $ case r of
1603+ Yield x s -> Yield x s
1604+ Skip s -> Skip s
1605+ Stop -> Stop
1606+
1607+{-# INLINE generally #-}
1608+generally :: Monad m => Stream Identity a -> Stream m a
1609+generally = hoist (return . runIdentity)
1610+
1611+{-# INLINE_NORMAL liftInner #-}
1612+liftInner :: (Monad m, MonadTrans t, Monad (t m))
1613+ => Stream m a -> Stream (t m) a
1614+liftInner (Stream step state) = Stream step' state
1615+ where
1616+ {-# INLINE_LATE step' #-}
1617+ step' gst st = do
1618+ r <- lift $ step (adaptState gst) st
1619+ return $ case r of
1620+ Yield x s -> Yield x s
1621+ Skip s -> Skip s
1622+ Stop -> Stop
1623+
1624+{-# INLINE_NORMAL runReaderT #-}
1625+runReaderT :: Monad m => m s -> Stream (ReaderT s m) a -> Stream m a
1626+runReaderT env (Stream step state) = Stream step' (state, env)
1627+ where
1628+ {-# INLINE_LATE step' #-}
1629+ step' gst (st, action) = do
1630+ sv <- action
1631+ r <- Reader.runReaderT (step (adaptState gst) st) sv
1632+ return $ case r of
1633+ Yield x s -> Yield x (s, return sv)
1634+ Skip s -> Skip (s, return sv)
1635+ Stop -> Stop
1636+
1637+{-# INLINE_NORMAL evalStateT #-}
1638+evalStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m a
1639+evalStateT initial (Stream step state) = Stream step' (state, initial)
1640+ where
1641+ {-# INLINE_LATE step' #-}
1642+ step' gst (st, action) = do
1643+ sv <- action
1644+ (r, sv') <- State.runStateT (step (adaptState gst) st) sv
1645+ return $ case r of
1646+ Yield x s -> Yield x (s, return sv')
1647+ Skip s -> Skip (s, return sv')
1648+ Stop -> Stop
1649+
1650+{-# INLINE_NORMAL runStateT #-}
1651+runStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m (s, a)
1652+runStateT initial (Stream step state) = Stream step' (state, initial)
1653+ where
1654+ {-# INLINE_LATE step' #-}
1655+ step' gst (st, action) = do
1656+ sv <- action
1657+ (r, sv') <- State.runStateT (step (adaptState gst) st) sv
1658+ return $ case r of
1659+ Yield x s -> Yield (sv', x) (s, return sv')
1660+ Skip s -> Skip (s, return sv')
1661+ Stop -> Stop