git.iko.soy/streamly
on cbccb77 help sign in
streamly/core/src/Streamly/Internal/Data/Stream/StreamK.hs
1078 lines 34 kB history raw
1{-# LANGUAGE UndecidableInstances #-}
3#include "inline.hs"
5-- |
6-- Module : Streamly.Internal.Data.Stream.StreamK
7-- Copyright : (c) 2017 Composewell Technologies
8--
9-- License : BSD3
10-- Maintainer : streamly@composewell.com
11-- Stability : experimental
12-- Portability : GHC
13--
14--
15-- Continuation passing style (CPS) stream implementation. The symbol 'K' below
16-- denotes a function as well as a Kontinuation.
17--
18-- @
19-- import qualified Streamly.Internal.Data.Stream.StreamK as K
20-- @
21--
22module Streamly.Internal.Data.Stream.StreamK
23 (
24 -- * The stream type
25 Stream(..)
27 -- * Construction Primitives
28 , mkStream
29 , nil
30 , nilM
31 , cons
32 , (.:)
34 -- * Elimination Primitives
35 , foldStream
36 , foldStreamShared
38 -- * Transformation Primitives
39 , unShare
41 -- * Deconstruction
42 , uncons
44 -- * Generation
45 -- ** Unfolds
46 , unfoldr
47 , unfoldrM
49 -- ** Specialized Generation
50 , repeat
51 , repeatM
52 , replicate
53 , replicateM
54 , fromIndices
55 , fromIndicesM
56 , iterate
57 , iterateM
59 -- ** Conversions
60 , fromPure
61 , fromEffect
62 , fromFoldable
63 , fromList
65 -- * foldr/build
66 , foldrS
67 , foldrSM
68 , buildS
69 , augmentS
71 -- * Elimination
72 -- ** General Folds
73 , foldr
74 , foldr1
75 , foldrM
76 , foldrT
78 , foldl'
79 , foldlM'
80 , foldlS
81 , foldlT
82 , foldlx'
83 , foldlMx'
84 , fold
85 , foldBreak
86 , parseBreak
88 -- ** Specialized Folds
89 , drain
90 , null
91 , head
92 , tail
93 , init
94 , elem
95 , notElem
96 , all
97 , any
98 , last
99 , minimum
100 , minimumBy
101 , maximum
102 , maximumBy
103 , findIndices
104 , lookup
105 , findM
106 , find
107 , (!!)
109 -- ** Map and Fold
110 , mapM_
112 -- ** Conversions
113 , toList
114 , hoist
116 -- * Transformation
117 -- ** By folding (scans)
118 , scanl'
119 , scanlx'
121 -- ** Filtering
122 , filter
123 , take
124 , takeWhile
125 , drop
126 , dropWhile
128 -- ** Mapping
129 , map
130 , mapM
131 , sequence
133 -- ** Inserting
134 , intersperseM
135 , intersperse
136 , insertBy
138 -- ** Deleting
139 , deleteBy
141 -- ** Reordering
142 , reverse
144 -- ** Map and Filter
145 , mapMaybe
147 -- ** Zipping
148 , zipWith
149 , zipWithM
151 -- ** Merging
152 , mergeBy
153 , mergeByM
155 -- ** Nesting
156 , concatMapWith
157 , concatMap
158 , bindWith
159 , concatPairsWith
160 , apWith
161 , apSerial
162 , apSerialDiscardFst
163 , apSerialDiscardSnd
165 -- ** Transformation comprehensions
166 , the
168 -- * Semigroup Style Composition
169 , serial
171 -- * Utilities
172 , consM
173 , withLocal
174 , mfix
175 )
176where
178import Control.Exception (assert)
179import Control.Monad.Catch (MonadThrow, throwM)
180import Control.Monad.Trans.Class (MonadTrans(lift))
181import Control.Monad (void, join)
182import Streamly.Internal.Data.SVar.Type (adaptState, defState)
184import qualified Streamly.Internal.Data.Fold.Type as FL
185import qualified Streamly.Internal.Data.Parser.ParserD.Type as PR
186import qualified Prelude
188import Prelude
189 hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,
190 take, filter, all, any, takeWhile, drop, dropWhile, minimum,
191 maximum, elem, notElem, null, head, tail, init, zipWith, lookup,
192 foldr1, (!!), replicate, reverse, concatMap, iterate, splitAt)
194import Streamly.Internal.Data.Stream.StreamK.Type
196-- $setup
197-- >>> :m
199-------------------------------------------------------------------------------
200-- Generation
201-------------------------------------------------------------------------------
203{-# INLINE unfoldrM #-}
204unfoldrM :: Monad m => (b -> m (Maybe (a, b))) -> b -> Stream m a
205unfoldrM = unfoldrMWith consM
206{-
207-- Generalization of concurrent streams/SVar via unfoldr.
208--
209-- Unfold a value into monadic actions and then run the resulting monadic
210-- actions to generate a stream. Since the step of generating the monadic
211-- action and running them are decoupled we can run the monadic actions
212-- cooncurrently. For example, the seed could be a list of monadic actions or a
213-- pure stream of monadic actions.
214--
215-- We can have different flavors of this depending on the stream type t. The
216-- concurrent version could be async or ahead etc. Depending on how we queue
217-- back the feedback portion b, it could be DFS or BFS style.
218--
219unfoldrA :: (IsStream t, MonadAsync m) =>
220 (b -> Maybe (m a, b)) -> b -> Stream m a
221unfoldrA = undefined
222-}
224-------------------------------------------------------------------------------
225-- Special generation
226-------------------------------------------------------------------------------
228repeatM :: Monad m => m a -> Stream m a
229repeatM = repeatMWith consM
231{-# INLINE replicateM #-}
232replicateM :: Monad m => Int -> m a -> Stream m a
233replicateM = replicateMWith consM
234{-# INLINE replicate #-}
235replicate :: Int -> a -> Stream m a
236replicate n a = go n
237 where
238 go cnt = if cnt <= 0 then nil else a `cons` go (cnt - 1)
240{-# INLINE fromIndicesM #-}
241fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
242fromIndicesM = fromIndicesMWith consM
243{-# INLINE fromIndices #-}
244fromIndices :: (Int -> a) -> Stream m a
245fromIndices gen = go 0
246 where
247 go n = gen n `cons` go (n + 1)
249{-# INLINE iterate #-}
250iterate :: (a -> a) -> a -> Stream m a
251iterate step = go
252 where
253 go !s = cons s (go (step s))
255{-# INLINE iterateM #-}
256iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
257iterateM = iterateMWith consM
259-------------------------------------------------------------------------------
260-- Conversions
261-------------------------------------------------------------------------------
263{-# INLINE fromList #-}
264fromList :: [a] -> Stream m a
265fromList = fromFoldable
267-------------------------------------------------------------------------------
268-- Elimination by Folding
269-------------------------------------------------------------------------------
271-- | Lazy right associative fold.
272{-# INLINE foldr #-}
273foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
274foldr step acc = foldrM (\x xs -> xs >>= \b -> return (step x b)) (return acc)
276-- | Right associative fold to an arbitrary transformer monad.
277{-# INLINE foldrT #-}
278foldrT :: (Monad m, Monad (s m), MonadTrans s)
279 => (a -> s m b -> s m b) -> s m b -> Stream m a -> s m b
280foldrT step final = go
281 where
282 go m1 = do
283 res <- lift $ uncons m1
284 case res of
285 Just (h, t) -> step h (go t)
286 Nothing -> final
288{-# INLINE foldr1 #-}
289foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
290foldr1 step m = do
291 r <- uncons m
292 case r of
293 Nothing -> return Nothing
294 Just (h, t) -> fmap Just (go h t)
295 where
296 go p m1 =
297 let stp = return p
298 single a = return $ step a p
299 yieldk a r = fmap (step p) (go a r)
300 in foldStream defState yieldk single stp m1
302-- XXX replace the recursive "go" with explicit continuations.
303-- | Like 'foldx', but with a monadic step function.
304{-# INLINABLE foldlMx' #-}
305foldlMx' :: Monad m
306 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> m b
307foldlMx' step begin done = go begin
308 where
309 go !acc m1 =
310 let stop = acc >>= done
311 single a = acc >>= \b -> step b a >>= done
312 yieldk a r = acc >>= \b -> step b a >>= \x -> go (return x) r
313 in foldStream defState yieldk single stop m1
315{-# INLINABLE fold #-}
316fold :: Monad m => FL.Fold m a b -> Stream m a -> m b
317fold (FL.Fold step begin done) m = do
318 res <- begin
319 case res of
320 FL.Partial fs -> go fs m
321 FL.Done fb -> return fb
323 where
324 go !acc m1 =
325 let stop = done acc
326 single a = step acc a
327 >>= \case
328 FL.Partial s -> done s
329 FL.Done b1 -> return b1
330 yieldk a r = step acc a
331 >>= \case
332 FL.Partial s -> go s r
333 FL.Done b1 -> return b1
334 in foldStream defState yieldk single stop m1
336{-# INLINE foldBreak #-}
337foldBreak :: Monad m => FL.Fold m a b -> Stream m a -> m (b, Stream m a)
338foldBreak (FL.Fold step begin done) m = do
339 res <- begin
340 case res of
341 FL.Partial fs -> go fs m
342 FL.Done fb -> return (fb, m)
344 where
346 go !acc m1 =
347 let stop = (, nil) <$> done acc
348 single a =
349 step acc a
350 >>= \case
351 FL.Partial s -> (, nil) <$> done s
352 FL.Done b1 -> return (b1, nil)
353 yieldk a r =
354 step acc a
355 >>= \case
356 FL.Partial s -> go s r
357 FL.Done b1 -> return (b1, r)
358 in foldStream defState yieldk single stop m1
360-- | Like 'foldl'' but with a monadic step function.
361{-# INLINE foldlM' #-}
362foldlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> m b
363foldlM' step begin = foldlMx' step begin return
365-- | Lazy left fold to an arbitrary transformer monad.
366{-# INLINE foldlT #-}
367foldlT :: (Monad m, Monad (s m), MonadTrans s)
368 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
369foldlT step = go
370 where
371 go acc m1 = do
372 res <- lift $ uncons m1
373 case res of
374 Just (h, t) -> go (step acc h) t
375 Nothing -> acc
377------------------------------------------------------------------------------
378-- Specialized folds
379------------------------------------------------------------------------------
381{-# INLINE head #-}
382head :: Monad m => Stream m a -> m (Maybe a)
383-- head = foldrM (\x _ -> return $ Just x) (return Nothing)
384head m =
385 let stop = return Nothing
386 single a = return (Just a)
387 yieldk a _ = return (Just a)
388 in foldStream defState yieldk single stop m
390{-# INLINE elem #-}
391elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
392elem e = go
393 where
394 go m1 =
395 let stop = return False
396 single a = return (a == e)
397 yieldk a r = if a == e then return True else go r
398 in foldStream defState yieldk single stop m1
400{-# INLINE notElem #-}
401notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
402notElem e = go
403 where
404 go m1 =
405 let stop = return True
406 single a = return (a /= e)
407 yieldk a r = if a == e then return False else go r
408 in foldStream defState yieldk single stop m1
410{-# INLINABLE all #-}
411all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
412all p = go
413 where
414 go m1 =
415 let single a | p a = return True
416 | otherwise = return False
417 yieldk a r | p a = go r
418 | otherwise = return False
419 in foldStream defState yieldk single (return True) m1
421{-# INLINABLE any #-}
422any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
423any p = go
424 where
425 go m1 =
426 let single a | p a = return True
427 | otherwise = return False
428 yieldk a r | p a = return True
429 | otherwise = go r
430 in foldStream defState yieldk single (return False) m1
432-- | Extract the last element of the stream, if any.
433{-# INLINE last #-}
434last :: Monad m => Stream m a -> m (Maybe a)
435last = foldlx' (\_ y -> Just y) Nothing id
437{-# INLINE minimum #-}
438minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
439minimum = go Nothing
440 where
441 go Nothing m1 =
442 let stop = return Nothing
443 single a = return (Just a)
444 yieldk a r = go (Just a) r
445 in foldStream defState yieldk single stop m1
447 go (Just res) m1 =
448 let stop = return (Just res)
449 single a =
450 if res <= a
451 then return (Just res)
452 else return (Just a)
453 yieldk a r =
454 if res <= a
455 then go (Just res) r
456 else go (Just a) r
457 in foldStream defState yieldk single stop m1
459{-# INLINE minimumBy #-}
460minimumBy
461 :: (Monad m)
462 => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
463minimumBy cmp = go Nothing
464 where
465 go Nothing m1 =
466 let stop = return Nothing
467 single a = return (Just a)
468 yieldk a r = go (Just a) r
469 in foldStream defState yieldk single stop m1
471 go (Just res) m1 =
472 let stop = return (Just res)
473 single a = case cmp res a of
474 GT -> return (Just a)
475 _ -> return (Just res)
476 yieldk a r = case cmp res a of
477 GT -> go (Just a) r
478 _ -> go (Just res) r
479 in foldStream defState yieldk single stop m1
481{-# INLINE maximum #-}
482maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
483maximum = go Nothing
484 where
485 go Nothing m1 =
486 let stop = return Nothing
487 single a = return (Just a)
488 yieldk a r = go (Just a) r
489 in foldStream defState yieldk single stop m1
491 go (Just res) m1 =
492 let stop = return (Just res)
493 single a =
494 if res <= a
495 then return (Just a)
496 else return (Just res)
497 yieldk a r =
498 if res <= a
499 then go (Just a) r
500 else go (Just res) r
501 in foldStream defState yieldk single stop m1
503{-# INLINE maximumBy #-}
504maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
505maximumBy cmp = go Nothing
506 where
507 go Nothing m1 =
508 let stop = return Nothing
509 single a = return (Just a)
510 yieldk a r = go (Just a) r
511 in foldStream defState yieldk single stop m1
513 go (Just res) m1 =
514 let stop = return (Just res)
515 single a = case cmp res a of
516 GT -> return (Just res)
517 _ -> return (Just a)
518 yieldk a r = case cmp res a of
519 GT -> go (Just res) r
520 _ -> go (Just a) r
521 in foldStream defState yieldk single stop m1
523{-# INLINE (!!) #-}
524(!!) :: Monad m => Stream m a -> Int -> m (Maybe a)
525m !! i = go i m
526 where
527 go n m1 =
528 let single a | n == 0 = return $ Just a
529 | otherwise = return Nothing
530 yieldk a x | n < 0 = return Nothing
531 | n == 0 = return $ Just a
532 | otherwise = go (n - 1) x
533 in foldStream defState yieldk single (return Nothing) m1
535{-# INLINE lookup #-}
536lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
537lookup e = go
538 where
539 go m1 =
540 let single (a, b) | a == e = return $ Just b
541 | otherwise = return Nothing
542 yieldk (a, b) x | a == e = return $ Just b
543 | otherwise = go x
544 in foldStream defState yieldk single (return Nothing) m1
546{-# INLINE findM #-}
547findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
548findM p = go
549 where
550 go m1 =
551 let single a = do
552 b <- p a
553 if b then return $ Just a else return Nothing
554 yieldk a x = do
555 b <- p a
556 if b then return $ Just a else go x
557 in foldStream defState yieldk single (return Nothing) m1
559{-# INLINE find #-}
560find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
561find p = findM (return . p)
563{-# INLINE findIndices #-}
564findIndices :: (a -> Bool) -> Stream m a -> Stream m Int
565findIndices p = go 0
566 where
567 go offset m1 = mkStream $ \st yld sng stp ->
568 let single a | p a = sng offset
569 | otherwise = stp
570 yieldk a x | p a = yld offset $ go (offset + 1) x
571 | otherwise = foldStream (adaptState st) yld sng stp $
572 go (offset + 1) x
573 in foldStream (adaptState st) yieldk single stp m1
575------------------------------------------------------------------------------
576-- Map and Fold
577------------------------------------------------------------------------------
579-- | Apply a monadic action to each element of the stream and discard the
580-- output of the action.
581{-# INLINE mapM_ #-}
582mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
583mapM_ f = go
584 where
585 go m1 =
586 let stop = return ()
587 single a = void (f a)
588 yieldk a r = f a >> go r
589 in foldStream defState yieldk single stop m1
591{-# INLINE mapM #-}
592mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
593mapM = mapMWith consM
595------------------------------------------------------------------------------
596-- Converting folds
597------------------------------------------------------------------------------
599{-# INLINABLE toList #-}
600toList :: Monad m => Stream m a -> m [a]
601toList = foldr (:) []
603-- Based on suggestions by David Feuer and Pranay Sashank
604{-# INLINE hoist #-}
605hoist :: (Monad m, Monad n)
606 => (forall x. m x -> n x) -> Stream m a -> Stream n a
607hoist f str =
608 mkStream $ \st yld sng stp ->
609 let single = return . sng
610 yieldk a s = return $ yld a (hoist f s)
611 stop = return stp
612 state = adaptState st
613 in join . f $ foldStreamShared state yieldk single stop str
615-------------------------------------------------------------------------------
616-- Transformation by folding (Scans)
617-------------------------------------------------------------------------------
619{-# INLINE scanlx' #-}
620scanlx' :: (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
621scanlx' step begin done m =
622 cons (done begin) $ go m begin
623 where
624 go m1 !acc = mkStream $ \st yld sng stp ->
625 let single a = sng (done $ step acc a)
626 yieldk a r =
627 let s = step acc a
628 in yld (done s) (go r s)
629 in foldStream (adaptState st) yieldk single stp m1
631{-# INLINE scanl' #-}
632scanl' :: (b -> a -> b) -> b -> Stream m a -> Stream m b
633scanl' step begin = scanlx' step begin id
635-------------------------------------------------------------------------------
636-- Filtering
637-------------------------------------------------------------------------------
639{-# INLINE filter #-}
640filter :: (a -> Bool) -> Stream m a -> Stream m a
641filter p = go
642 where
643 go m1 = mkStream $ \st yld sng stp ->
644 let single a | p a = sng a
645 | otherwise = stp
646 yieldk a r | p a = yld a (go r)
647 | otherwise = foldStream st yieldk single stp r
648 in foldStream st yieldk single stp m1
650{-# INLINE take #-}
651take :: Int -> Stream m a -> Stream m a
652take = go
653 where
654 go n1 m1 = mkStream $ \st yld sng stp ->
655 let yieldk a r = yld a (go (n1 - 1) r)
656 in if n1 <= 0
657 then stp
658 else foldStream st yieldk sng stp m1
660{-# INLINE takeWhile #-}
661takeWhile :: (a -> Bool) -> Stream m a -> Stream m a
662takeWhile p = go
663 where
664 go m1 = mkStream $ \st yld sng stp ->
665 let single a | p a = sng a
666 | otherwise = stp
667 yieldk a r | p a = yld a (go r)
668 | otherwise = stp
669 in foldStream st yieldk single stp m1
671{-# INLINE drop #-}
672drop :: Int -> Stream m a -> Stream m a
673drop n m = unShare (go n m)
674 where
675 go n1 m1 = mkStream $ \st yld sng stp ->
676 let single _ = stp
677 yieldk _ r = foldStreamShared st yld sng stp $ go (n1 - 1) r
678 -- Somehow "<=" check performs better than a ">"
679 in if n1 <= 0
680 then foldStreamShared st yld sng stp m1
681 else foldStreamShared st yieldk single stp m1
683{-# INLINE dropWhile #-}
684dropWhile :: (a -> Bool) -> Stream m a -> Stream m a
685dropWhile p = go
686 where
687 go m1 = mkStream $ \st yld sng stp ->
688 let single a | p a = stp
689 | otherwise = sng a
690 yieldk a r | p a = foldStream st yieldk single stp r
691 | otherwise = yld a r
692 in foldStream st yieldk single stp m1
694-------------------------------------------------------------------------------
695-- Mapping
696-------------------------------------------------------------------------------
698-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
699-- other stream types to overload it.
700{-# INLINE sequence #-}
701sequence :: Monad m => Stream m (m a) -> Stream m a
702sequence = go
703 where
704 go m1 = mkStream $ \st yld sng stp ->
705 let single ma = ma >>= sng
706 yieldk ma r = foldStreamShared st yld sng stp $ ma `consM` go r
707 in foldStream (adaptState st) yieldk single stp m1
709-------------------------------------------------------------------------------
710-- Inserting
711-------------------------------------------------------------------------------
713{-# INLINE intersperseM #-}
714intersperseM :: Monad m => m a -> Stream m a -> Stream m a
715intersperseM a = prependingStart
716 where
717 prependingStart m1 = mkStream $ \st yld sng stp ->
718 let yieldk i x =
719 foldStreamShared st yld sng stp $ return i `consM` go x
720 in foldStream st yieldk sng stp m1
721 go m2 = mkStream $ \st yld sng stp ->
722 let single i = foldStreamShared st yld sng stp $ a `consM` fromPure i
723 yieldk i x =
724 foldStreamShared
725 st yld sng stp $ a `consM` return i `consM` go x
726 in foldStream st yieldk single stp m2
728{-# INLINE intersperse #-}
729intersperse :: Monad m => a -> Stream m a -> Stream m a
730intersperse a = intersperseM (return a)
732{-# INLINE insertBy #-}
733insertBy :: (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
734insertBy cmp x = go
735 where
736 go m1 = mkStream $ \st yld _ _ ->
737 let single a = case cmp x a of
738 GT -> yld a (fromPure x)
739 _ -> yld x (fromPure a)
740 stop = yld x nil
741 yieldk a r = case cmp x a of
742 GT -> yld a (go r)
743 _ -> yld x (a `cons` r)
744 in foldStream st yieldk single stop m1
746------------------------------------------------------------------------------
747-- Deleting
748------------------------------------------------------------------------------
750{-# INLINE deleteBy #-}
751deleteBy :: (a -> a -> Bool) -> a -> Stream m a -> Stream m a
752deleteBy eq x = go
753 where
754 go m1 = mkStream $ \st yld sng stp ->
755 let single a = if eq x a then stp else sng a
756 yieldk a r = if eq x a
757 then foldStream st yld sng stp r
758 else yld a (go r)
759 in foldStream st yieldk single stp m1
761-------------------------------------------------------------------------------
762-- Map and Filter
763-------------------------------------------------------------------------------
765{-# INLINE mapMaybe #-}
766mapMaybe :: (a -> Maybe b) -> Stream m a -> Stream m b
767mapMaybe f = go
768 where
769 go m1 = mkStream $ \st yld sng stp ->
770 let single a = maybe stp sng (f a)
771 yieldk a r = case f a of
772 Just b -> yld b $ go r
773 Nothing -> foldStream (adaptState st) yieldk single stp r
774 in foldStream (adaptState st) yieldk single stp m1
776------------------------------------------------------------------------------
777-- Serial Zipping
778------------------------------------------------------------------------------
780-- | Zip two streams serially using a pure zipping function.
781--
782-- @since 0.1.0
783{-# INLINE zipWith #-}
784zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
785zipWith f = zipWithM (\a b -> return (f a b))
787-- | Zip two streams serially using a monadic zipping function.
788--
789-- @since 0.1.0
790{-# INLINE zipWithM #-}
791zipWithM :: Monad m =>
792 (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
793zipWithM f = go
795 where
797 go mx my = mkStream $ \st yld sng stp -> do
798 let merge a ra =
799 let single2 b = f a b >>= sng
800 yield2 b rb = f a b >>= \x -> yld x (go ra rb)
801 in foldStream (adaptState st) yield2 single2 stp my
802 let single1 a = merge a nil
803 yield1 = merge
804 foldStream (adaptState st) yield1 single1 stp mx
806------------------------------------------------------------------------------
807-- Merging
808------------------------------------------------------------------------------
810{-# INLINE mergeByM #-}
811mergeByM :: Monad m =>
812 (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
813mergeByM cmp = go
815 where
817 go mx my = mkStream $ \st yld sng stp -> do
818 let stop = foldStream st yld sng stp my
819 single x = foldStream st yld sng stp (goX0 x my)
820 yield x rx = foldStream st yld sng stp (goX x rx my)
821 foldStream st yield single stop mx
823 goX0 x my = mkStream $ \st yld sng _ -> do
824 let stop = sng x
825 single y = do
826 r <- cmp x y
827 case r of
828 GT -> yld y (fromPure x)
829 _ -> yld x (fromPure y)
830 yield y ry = do
831 r <- cmp x y
832 case r of
833 GT -> yld y (goX0 x ry)
834 _ -> yld x (y `cons` ry)
835 in foldStream st yield single stop my
837 goX x mx my = mkStream $ \st yld _ _ -> do
838 let stop = yld x mx
839 single y = do
840 r <- cmp x y
841 case r of
842 GT -> yld y (x `cons` mx)
843 _ -> yld x (goY0 mx y)
844 yield y ry = do
845 r <- cmp x y
846 case r of
847 GT -> yld y (goX x mx ry)
848 _ -> yld x (goY mx y ry)
849 in foldStream st yield single stop my
851 goY0 mx y = mkStream $ \st yld sng _ -> do
852 let stop = sng y
853 single x = do
854 r <- cmp x y
855 case r of
856 GT -> yld y (fromPure x)
857 _ -> yld x (fromPure y)
858 yield x rx = do
859 r <- cmp x y
860 case r of
861 GT -> yld y (x `cons` rx)
862 _ -> yld x (goY0 rx y)
863 in foldStream st yield single stop mx
865 goY mx y my = mkStream $ \st yld _ _ -> do
866 let stop = yld y my
867 single x = do
868 r <- cmp x y
869 case r of
870 GT -> yld y (goX0 x my)
871 _ -> yld x (y `cons` my)
872 yield x rx = do
873 r <- cmp x y
874 case r of
875 GT -> yld y (goX x rx my)
876 _ -> yld x (goY rx y my)
877 in foldStream st yield single stop mx
879{-# INLINE mergeBy #-}
880mergeBy :: (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
881-- XXX GHC: This has slightly worse performance than replacing "r <- cmp x y"
882-- with "let r = cmp x y" in the monadic version. The definition below is
883-- exactly the same as mergeByM except this change.
884-- mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
885mergeBy cmp = go
887 where
889 go mx my = mkStream $ \st yld sng stp -> do
890 let stop = foldStream st yld sng stp my
891 single x = foldStream st yld sng stp (goX0 x my)
892 yield x rx = foldStream st yld sng stp (goX x rx my)
893 foldStream st yield single stop mx
895 goX0 x my = mkStream $ \st yld sng _ -> do
896 let stop = sng x
897 single y = do
898 case cmp x y of
899 GT -> yld y (fromPure x)
900 _ -> yld x (fromPure y)
901 yield y ry = do
902 case cmp x y of
903 GT -> yld y (goX0 x ry)
904 _ -> yld x (y `cons` ry)
905 in foldStream st yield single stop my
907 goX x mx my = mkStream $ \st yld _ _ -> do
908 let stop = yld x mx
909 single y = do
910 case cmp x y of
911 GT -> yld y (x `cons` mx)
912 _ -> yld x (goY0 mx y)
913 yield y ry = do
914 case cmp x y of
915 GT -> yld y (goX x mx ry)
916 _ -> yld x (goY mx y ry)
917 in foldStream st yield single stop my
919 goY0 mx y = mkStream $ \st yld sng _ -> do
920 let stop = sng y
921 single x = do
922 case cmp x y of
923 GT -> yld y (fromPure x)
924 _ -> yld x (fromPure y)
925 yield x rx = do
926 case cmp x y of
927 GT -> yld y (x `cons` rx)
928 _ -> yld x (goY0 rx y)
929 in foldStream st yield single stop mx
931 goY mx y my = mkStream $ \st yld _ _ -> do
932 let stop = yld y my
933 single x = do
934 case cmp x y of
935 GT -> yld y (goX0 x my)
936 _ -> yld x (y `cons` my)
937 yield x rx = do
938 case cmp x y of
939 GT -> yld y (goX x rx my)
940 _ -> yld x (goY rx y my)
941 in foldStream st yield single stop mx
943------------------------------------------------------------------------------
944-- Transformation comprehensions
945------------------------------------------------------------------------------
947{-# INLINE the #-}
948the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
949the m = do
950 r <- uncons m
951 case r of
952 Nothing -> return Nothing
953 Just (h, t) -> go h t
954 where
955 go h m1 =
956 let single a | h == a = return $ Just h
957 | otherwise = return Nothing
958 yieldk a r | h == a = go h r
959 | otherwise = return Nothing
960 in foldStream defState yieldk single (return $ Just h) m1
962------------------------------------------------------------------------------
963-- Alternative & MonadPlus
964------------------------------------------------------------------------------
966_alt :: Stream m a -> Stream m a -> Stream m a
967_alt m1 m2 = mkStream $ \st yld sng stp ->
968 let stop = foldStream st yld sng stp m2
969 in foldStream st yld sng stop m1
971------------------------------------------------------------------------------
972-- MonadError
973------------------------------------------------------------------------------
975{-
976-- XXX handle and test cross thread state transfer
977withCatchError
978 :: MonadError e m
979 => Stream m a -> (e -> Stream m a) -> Stream m a
980withCatchError m h =
981 mkStream $ \_ stp sng yld ->
982 let run x = unStream x Nothing stp sng yieldk
983 handle r = r `catchError` \e -> run $ h e
984 yieldk a r = yld a (withCatchError r h)
985 in handle $ run m
986-}
988-------------------------------------------------------------------------------
989-- Parsing
990-------------------------------------------------------------------------------
992-- Inlined definition.
993{-# INLINE splitAt #-}
994splitAt :: Int -> [a] -> ([a],[a])
995splitAt n ls
996 | n <= 0 = ([], ls)
997 | otherwise = splitAt' n ls
998 where
999 splitAt' :: Int -> [a] -> ([a], [a])
1000 splitAt' _ [] = ([], [])
1001 splitAt' 1 (x:xs) = ([x], xs)
1002 splitAt' m (x:xs) = (x:xs', xs'')
1003 where
1004 (xs', xs'') = splitAt' (m - 1) xs
1006-- | Run a 'Parser' over a stream and return rest of the Stream.
1007{-# INLINE_NORMAL parseBreak #-}
1008parseBreak
1009 :: MonadThrow m
1010 => PR.Parser m a b
1011 -> Stream m a
1012 -> m (b, Stream m a)
1013parseBreak (PR.Parser pstep initial extract) stream = do
1014 res <- initial
1015 case res of
1016 PR.IPartial s -> goStream stream [] s
1017 PR.IDone b -> return (b, stream)
1018 PR.IError err -> throwM $ PR.ParseError err
1020 where
1022 -- "buf" contains last few items in the stream that we may have to
1023 -- backtrack to.
1024 --
1025 -- XXX currently we are using a dumb list based approach for backtracking
1026 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
1027 -- That will allow us more efficient random back and forth movement.
1028 goStream st buf !pst =
1029 let stop = do
1030 b <- extract pst
1031 return (b, fromList buf)
1032 single x = yieldk x nil
1033 yieldk x r = do
1034 res <- pstep pst x
1035 case res of
1036 PR.Partial 0 s -> goStream r [] s
1037 PR.Partial n s -> do
1038 assert (n <= length (x:buf)) (return ())
1039 let src0 = Prelude.take n (x:buf)
1040 src = Prelude.reverse src0
1041 goBuf r [] src s
1042 PR.Continue 0 s -> goStream r (x:buf) s
1043 PR.Continue n s -> do
1044 assert (n <= length (x:buf)) (return ())
1045 let (src0, buf1) = splitAt n (x:buf)
1046 src = Prelude.reverse src0
1047 goBuf r buf1 src s
1048 PR.Done 0 b -> return (b, r)
1049 PR.Done n b -> do
1050 assert (n <= length (x:buf)) (return ())
1051 let src0 = Prelude.take n (x:buf)
1052 src = Prelude.reverse src0
1053 return (b, serial (fromList src) r)
1054 PR.Error err -> throwM $ PR.ParseError err
1055 in foldStream defState yieldk single stop st
1057 goBuf st buf [] !pst = goStream st buf pst
1058 goBuf st buf (x:xs) !pst = do
1059 pRes <- pstep pst x
1060 case pRes of
1061 PR.Partial 0 s -> goBuf st [] xs s
1062 PR.Partial n s -> do
1063 assert (n <= length (x:buf)) (return ())
1064 let src0 = Prelude.take n (x:buf)
1065 src = Prelude.reverse src0 ++ xs
1066 goBuf st [] src s
1067 PR.Continue 0 s -> goBuf st (x:buf) xs s
1068 PR.Continue n s -> do
1069 assert (n <= length (x:buf)) (return ())
1070 let (src0, buf1) = splitAt n (x:buf)
1071 src = Prelude.reverse src0 ++ xs
1072 goBuf st buf1 src s
1073 PR.Done n b -> do
1074 assert (n <= length (x:buf)) (return ())
1075 let src0 = Prelude.take n (x:buf)
1076 src = Prelude.reverse src0
1077 return (b, serial (fromList src) st)
1078 PR.Error err -> throwM $ PR.ParseError err