git.iko.soy/streamly
on cfcf26d help sign in
streamly/src/Streamly/Internal/Data/Stream/StreamK.hs
967 lines 30 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
86 -- ** Specialized Folds
87 , drain
88 , null
89 , head
90 , tail
91 , init
92 , elem
93 , notElem
94 , all
95 , any
96 , last
97 , minimum
98 , minimumBy
99 , maximum
100 , maximumBy
101 , findIndices
102 , lookup
103 , findM
104 , find
105 , (!!)
107 -- ** Map and Fold
108 , mapM_
110 -- ** Conversions
111 , toList
112 , hoist
114 -- * Transformation
115 -- ** By folding (scans)
116 , scanl'
117 , scanlx'
119 -- ** Filtering
120 , filter
121 , take
122 , takeWhile
123 , drop
124 , dropWhile
126 -- ** Mapping
127 , map
128 , mapM
129 , sequence
131 -- ** Inserting
132 , intersperseM
133 , intersperse
134 , insertBy
136 -- ** Deleting
137 , deleteBy
139 -- ** Reordering
140 , reverse
142 -- ** Map and Filter
143 , mapMaybe
145 -- ** Zipping
146 , zipWith
147 , zipWithM
149 -- ** Merging
150 , mergeBy
151 , mergeByM
153 -- ** Nesting
154 , concatMapWith
155 , concatMap
156 , bindWith
157 , concatPairsWith
158 , apWith
159 , apSerial
160 , apSerialDiscardFst
161 , apSerialDiscardSnd
163 -- ** Transformation comprehensions
164 , the
166 -- * Semigroup Style Composition
167 , serial
169 -- * Utilities
170 , consM
171 , withLocal
172 , mfix
173 )
174where
176import Control.Monad.Trans.Class (MonadTrans(lift))
177import Control.Monad (void, join)
178import Streamly.Internal.Data.SVar.Type (adaptState, defState)
180import qualified Streamly.Internal.Data.Fold.Type as FL
182import Prelude
183 hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,
184 take, filter, all, any, takeWhile, drop, dropWhile, minimum,
185 maximum, elem, notElem, null, head, tail, init, zipWith, lookup,
186 foldr1, (!!), replicate, reverse, concatMap, iterate)
188import Streamly.Internal.Data.Stream.StreamK.Type
190-- $setup
191-- >>> :m
193-------------------------------------------------------------------------------
194-- Generation
195-------------------------------------------------------------------------------
197{-# INLINE unfoldrM #-}
198unfoldrM :: Monad m => (b -> m (Maybe (a, b))) -> b -> Stream m a
199unfoldrM = unfoldrMWith consM
200{-
201-- Generalization of concurrent streams/SVar via unfoldr.
202--
203-- Unfold a value into monadic actions and then run the resulting monadic
204-- actions to generate a stream. Since the step of generating the monadic
205-- action and running them are decoupled we can run the monadic actions
206-- cooncurrently. For example, the seed could be a list of monadic actions or a
207-- pure stream of monadic actions.
208--
209-- We can have different flavors of this depending on the stream type t. The
210-- concurrent version could be async or ahead etc. Depending on how we queue
211-- back the feedback portion b, it could be DFS or BFS style.
212--
213unfoldrA :: (IsStream t, MonadAsync m) =>
214 (b -> Maybe (m a, b)) -> b -> Stream m a
215unfoldrA = undefined
216-}
218-------------------------------------------------------------------------------
219-- Special generation
220-------------------------------------------------------------------------------
222repeatM :: Monad m => m a -> Stream m a
223repeatM = repeatMWith consM
225{-# INLINE replicateM #-}
226replicateM :: Monad m => Int -> m a -> Stream m a
227replicateM = replicateMWith consM
228{-# INLINE replicate #-}
229replicate :: Int -> a -> Stream m a
230replicate n a = go n
231 where
232 go cnt = if cnt <= 0 then nil else a `cons` go (cnt - 1)
234{-# INLINE fromIndicesM #-}
235fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
236fromIndicesM = fromIndicesMWith consM
237{-# INLINE fromIndices #-}
238fromIndices :: (Int -> a) -> Stream m a
239fromIndices gen = go 0
240 where
241 go n = gen n `cons` go (n + 1)
243{-# INLINE iterate #-}
244iterate :: (a -> a) -> a -> Stream m a
245iterate step = go
246 where
247 go !s = cons s (go (step s))
249{-# INLINE iterateM #-}
250iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
251iterateM = iterateMWith consM
253-------------------------------------------------------------------------------
254-- Conversions
255-------------------------------------------------------------------------------
257{-# INLINE fromList #-}
258fromList :: [a] -> Stream m a
259fromList = fromFoldable
261-------------------------------------------------------------------------------
262-- Elimination by Folding
263-------------------------------------------------------------------------------
265-- | Lazy right associative fold.
266{-# INLINE foldr #-}
267foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
268foldr step acc = foldrM (\x xs -> xs >>= \b -> return (step x b)) (return acc)
270-- | Right associative fold to an arbitrary transformer monad.
271{-# INLINE foldrT #-}
272foldrT :: (Monad m, Monad (s m), MonadTrans s)
273 => (a -> s m b -> s m b) -> s m b -> Stream m a -> s m b
274foldrT step final = go
275 where
276 go m1 = do
277 res <- lift $ uncons m1
278 case res of
279 Just (h, t) -> step h (go t)
280 Nothing -> final
282{-# INLINE foldr1 #-}
283foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
284foldr1 step m = do
285 r <- uncons m
286 case r of
287 Nothing -> return Nothing
288 Just (h, t) -> fmap Just (go h t)
289 where
290 go p m1 =
291 let stp = return p
292 single a = return $ step a p
293 yieldk a r = fmap (step p) (go a r)
294 in foldStream defState yieldk single stp m1
296-- XXX replace the recursive "go" with explicit continuations.
297-- | Like 'foldx', but with a monadic step function.
298{-# INLINABLE foldlMx' #-}
299foldlMx' :: Monad m
300 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> m b
301foldlMx' step begin done = go begin
302 where
303 go !acc m1 =
304 let stop = acc >>= done
305 single a = acc >>= \b -> step b a >>= done
306 yieldk a r = acc >>= \b -> step b a >>= \x -> go (return x) r
307 in foldStream defState yieldk single stop m1
309{-# INLINABLE fold #-}
310fold :: Monad m => FL.Fold m a b -> Stream m a -> m b
311fold (FL.Fold step begin done) m = do
312 res <- begin
313 case res of
314 FL.Partial fs -> go fs m
315 FL.Done fb -> return fb
317 where
318 go !acc m1 =
319 let stop = done acc
320 single a = step acc a
321 >>= \case
322 FL.Partial s -> done s
323 FL.Done b1 -> return b1
324 yieldk a r = step acc a
325 >>= \case
326 FL.Partial s -> go s r
327 FL.Done b1 -> return b1
328 in foldStream defState yieldk single stop m1
330-- | Like 'foldl'' but with a monadic step function.
331{-# INLINE foldlM' #-}
332foldlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> m b
333foldlM' step begin = foldlMx' step begin return
335-- | Lazy left fold to an arbitrary transformer monad.
336{-# INLINE foldlT #-}
337foldlT :: (Monad m, Monad (s m), MonadTrans s)
338 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
339foldlT step = go
340 where
341 go acc m1 = do
342 res <- lift $ uncons m1
343 case res of
344 Just (h, t) -> go (step acc h) t
345 Nothing -> acc
347------------------------------------------------------------------------------
348-- Specialized folds
349------------------------------------------------------------------------------
351{-# INLINE head #-}
352head :: Monad m => Stream m a -> m (Maybe a)
353-- head = foldrM (\x _ -> return $ Just x) (return Nothing)
354head m =
355 let stop = return Nothing
356 single a = return (Just a)
357 yieldk a _ = return (Just a)
358 in foldStream defState yieldk single stop m
360{-# INLINE elem #-}
361elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
362elem e = go
363 where
364 go m1 =
365 let stop = return False
366 single a = return (a == e)
367 yieldk a r = if a == e then return True else go r
368 in foldStream defState yieldk single stop m1
370{-# INLINE notElem #-}
371notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
372notElem e = go
373 where
374 go m1 =
375 let stop = return True
376 single a = return (a /= e)
377 yieldk a r = if a == e then return False else go r
378 in foldStream defState yieldk single stop m1
380{-# INLINABLE all #-}
381all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
382all p = go
383 where
384 go m1 =
385 let single a | p a = return True
386 | otherwise = return False
387 yieldk a r | p a = go r
388 | otherwise = return False
389 in foldStream defState yieldk single (return True) m1
391{-# INLINABLE any #-}
392any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
393any p = go
394 where
395 go m1 =
396 let single a | p a = return True
397 | otherwise = return False
398 yieldk a r | p a = return True
399 | otherwise = go r
400 in foldStream defState yieldk single (return False) m1
402-- | Extract the last element of the stream, if any.
403{-# INLINE last #-}
404last :: Monad m => Stream m a -> m (Maybe a)
405last = foldlx' (\_ y -> Just y) Nothing id
407{-# INLINE minimum #-}
408minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
409minimum = go Nothing
410 where
411 go Nothing m1 =
412 let stop = return Nothing
413 single a = return (Just a)
414 yieldk a r = go (Just a) r
415 in foldStream defState yieldk single stop m1
417 go (Just res) m1 =
418 let stop = return (Just res)
419 single a =
420 if res <= a
421 then return (Just res)
422 else return (Just a)
423 yieldk a r =
424 if res <= a
425 then go (Just res) r
426 else go (Just a) r
427 in foldStream defState yieldk single stop m1
429{-# INLINE minimumBy #-}
430minimumBy
431 :: (Monad m)
432 => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
433minimumBy cmp = go Nothing
434 where
435 go Nothing m1 =
436 let stop = return Nothing
437 single a = return (Just a)
438 yieldk a r = go (Just a) r
439 in foldStream defState yieldk single stop m1
441 go (Just res) m1 =
442 let stop = return (Just res)
443 single a = case cmp res a of
444 GT -> return (Just a)
445 _ -> return (Just res)
446 yieldk a r = case cmp res a of
447 GT -> go (Just a) r
448 _ -> go (Just res) r
449 in foldStream defState yieldk single stop m1
451{-# INLINE maximum #-}
452maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
453maximum = go Nothing
454 where
455 go Nothing m1 =
456 let stop = return Nothing
457 single a = return (Just a)
458 yieldk a r = go (Just a) r
459 in foldStream defState yieldk single stop m1
461 go (Just res) m1 =
462 let stop = return (Just res)
463 single a =
464 if res <= a
465 then return (Just a)
466 else return (Just res)
467 yieldk a r =
468 if res <= a
469 then go (Just a) r
470 else go (Just res) r
471 in foldStream defState yieldk single stop m1
473{-# INLINE maximumBy #-}
474maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
475maximumBy cmp = go Nothing
476 where
477 go Nothing m1 =
478 let stop = return Nothing
479 single a = return (Just a)
480 yieldk a r = go (Just a) r
481 in foldStream defState yieldk single stop m1
483 go (Just res) m1 =
484 let stop = return (Just res)
485 single a = case cmp res a of
486 GT -> return (Just res)
487 _ -> return (Just a)
488 yieldk a r = case cmp res a of
489 GT -> go (Just res) r
490 _ -> go (Just a) r
491 in foldStream defState yieldk single stop m1
493{-# INLINE (!!) #-}
494(!!) :: Monad m => Stream m a -> Int -> m (Maybe a)
495m !! i = go i m
496 where
497 go n m1 =
498 let single a | n == 0 = return $ Just a
499 | otherwise = return Nothing
500 yieldk a x | n < 0 = return Nothing
501 | n == 0 = return $ Just a
502 | otherwise = go (n - 1) x
503 in foldStream defState yieldk single (return Nothing) m1
505{-# INLINE lookup #-}
506lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
507lookup e = go
508 where
509 go m1 =
510 let single (a, b) | a == e = return $ Just b
511 | otherwise = return Nothing
512 yieldk (a, b) x | a == e = return $ Just b
513 | otherwise = go x
514 in foldStream defState yieldk single (return Nothing) m1
516{-# INLINE findM #-}
517findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
518findM p = go
519 where
520 go m1 =
521 let single a = do
522 b <- p a
523 if b then return $ Just a else return Nothing
524 yieldk a x = do
525 b <- p a
526 if b then return $ Just a else go x
527 in foldStream defState yieldk single (return Nothing) m1
529{-# INLINE find #-}
530find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
531find p = findM (return . p)
533{-# INLINE findIndices #-}
534findIndices :: (a -> Bool) -> Stream m a -> Stream m Int
535findIndices p = go 0
536 where
537 go offset m1 = mkStream $ \st yld sng stp ->
538 let single a | p a = sng offset
539 | otherwise = stp
540 yieldk a x | p a = yld offset $ go (offset + 1) x
541 | otherwise = foldStream (adaptState st) yld sng stp $
542 go (offset + 1) x
543 in foldStream (adaptState st) yieldk single stp m1
545------------------------------------------------------------------------------
546-- Map and Fold
547------------------------------------------------------------------------------
549-- | Apply a monadic action to each element of the stream and discard the
550-- output of the action.
551{-# INLINE mapM_ #-}
552mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
553mapM_ f = go
554 where
555 go m1 =
556 let stop = return ()
557 single a = void (f a)
558 yieldk a r = f a >> go r
559 in foldStream defState yieldk single stop m1
561{-# INLINE mapM #-}
562mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
563mapM = mapMWith consM
565------------------------------------------------------------------------------
566-- Converting folds
567------------------------------------------------------------------------------
569{-# INLINABLE toList #-}
570toList :: Monad m => Stream m a -> m [a]
571toList = foldr (:) []
573-- Based on suggestions by David Feuer and Pranay Sashank
574{-# INLINE hoist #-}
575hoist :: (Monad m, Monad n)
576 => (forall x. m x -> n x) -> Stream m a -> Stream n a
577hoist f str =
578 mkStream $ \st yld sng stp ->
579 let single = return . sng
580 yieldk a s = return $ yld a (hoist f s)
581 stop = return stp
582 state = adaptState st
583 in join . f $ foldStreamShared state yieldk single stop str
585-------------------------------------------------------------------------------
586-- Transformation by folding (Scans)
587-------------------------------------------------------------------------------
589{-# INLINE scanlx' #-}
590scanlx' :: (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
591scanlx' step begin done m =
592 cons (done begin) $ go m begin
593 where
594 go m1 !acc = mkStream $ \st yld sng stp ->
595 let single a = sng (done $ step acc a)
596 yieldk a r =
597 let s = step acc a
598 in yld (done s) (go r s)
599 in foldStream (adaptState st) yieldk single stp m1
601{-# INLINE scanl' #-}
602scanl' :: (b -> a -> b) -> b -> Stream m a -> Stream m b
603scanl' step begin = scanlx' step begin id
605-------------------------------------------------------------------------------
606-- Filtering
607-------------------------------------------------------------------------------
609{-# INLINE filter #-}
610filter :: (a -> Bool) -> Stream m a -> Stream m a
611filter p = go
612 where
613 go m1 = mkStream $ \st yld sng stp ->
614 let single a | p a = sng a
615 | otherwise = stp
616 yieldk a r | p a = yld a (go r)
617 | otherwise = foldStream st yieldk single stp r
618 in foldStream st yieldk single stp m1
620{-# INLINE take #-}
621take :: Int -> Stream m a -> Stream m a
622take = go
623 where
624 go n1 m1 = mkStream $ \st yld sng stp ->
625 let yieldk a r = yld a (go (n1 - 1) r)
626 in if n1 <= 0
627 then stp
628 else foldStream st yieldk sng stp m1
630{-# INLINE takeWhile #-}
631takeWhile :: (a -> Bool) -> Stream m a -> Stream m a
632takeWhile p = go
633 where
634 go m1 = mkStream $ \st yld sng stp ->
635 let single a | p a = sng a
636 | otherwise = stp
637 yieldk a r | p a = yld a (go r)
638 | otherwise = stp
639 in foldStream st yieldk single stp m1
641{-# INLINE drop #-}
642drop :: Int -> Stream m a -> Stream m a
643drop n m = unShare (go n m)
644 where
645 go n1 m1 = mkStream $ \st yld sng stp ->
646 let single _ = stp
647 yieldk _ r = foldStreamShared st yld sng stp $ go (n1 - 1) r
648 -- Somehow "<=" check performs better than a ">"
649 in if n1 <= 0
650 then foldStreamShared st yld sng stp m1
651 else foldStreamShared st yieldk single stp m1
653{-# INLINE dropWhile #-}
654dropWhile :: (a -> Bool) -> Stream m a -> Stream m a
655dropWhile p = go
656 where
657 go m1 = mkStream $ \st yld sng stp ->
658 let single a | p a = stp
659 | otherwise = sng a
660 yieldk a r | p a = foldStream st yieldk single stp r
661 | otherwise = yld a r
662 in foldStream st yieldk single stp m1
664-------------------------------------------------------------------------------
665-- Mapping
666-------------------------------------------------------------------------------
668-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
669-- other stream types to overload it.
670{-# INLINE sequence #-}
671sequence :: Monad m => Stream m (m a) -> Stream m a
672sequence = go
673 where
674 go m1 = mkStream $ \st yld sng stp ->
675 let single ma = ma >>= sng
676 yieldk ma r = foldStreamShared st yld sng stp $ ma `consM` go r
677 in foldStream (adaptState st) yieldk single stp m1
679-------------------------------------------------------------------------------
680-- Inserting
681-------------------------------------------------------------------------------
683{-# INLINE intersperseM #-}
684intersperseM :: Monad m => m a -> Stream m a -> Stream m a
685intersperseM a = prependingStart
686 where
687 prependingStart m1 = mkStream $ \st yld sng stp ->
688 let yieldk i x =
689 foldStreamShared st yld sng stp $ return i `consM` go x
690 in foldStream st yieldk sng stp m1
691 go m2 = mkStream $ \st yld sng stp ->
692 let single i = foldStreamShared st yld sng stp $ a `consM` fromPure i
693 yieldk i x =
694 foldStreamShared
695 st yld sng stp $ a `consM` return i `consM` go x
696 in foldStream st yieldk single stp m2
698{-# INLINE intersperse #-}
699intersperse :: Monad m => a -> Stream m a -> Stream m a
700intersperse a = intersperseM (return a)
702{-# INLINE insertBy #-}
703insertBy :: (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
704insertBy cmp x = go
705 where
706 go m1 = mkStream $ \st yld _ _ ->
707 let single a = case cmp x a of
708 GT -> yld a (fromPure x)
709 _ -> yld x (fromPure a)
710 stop = yld x nil
711 yieldk a r = case cmp x a of
712 GT -> yld a (go r)
713 _ -> yld x (a `cons` r)
714 in foldStream st yieldk single stop m1
716------------------------------------------------------------------------------
717-- Deleting
718------------------------------------------------------------------------------
720{-# INLINE deleteBy #-}
721deleteBy :: (a -> a -> Bool) -> a -> Stream m a -> Stream m a
722deleteBy eq x = go
723 where
724 go m1 = mkStream $ \st yld sng stp ->
725 let single a = if eq x a then stp else sng a
726 yieldk a r = if eq x a
727 then foldStream st yld sng stp r
728 else yld a (go r)
729 in foldStream st yieldk single stp m1
731-------------------------------------------------------------------------------
732-- Map and Filter
733-------------------------------------------------------------------------------
735{-# INLINE mapMaybe #-}
736mapMaybe :: (a -> Maybe b) -> Stream m a -> Stream m b
737mapMaybe f = go
738 where
739 go m1 = mkStream $ \st yld sng stp ->
740 let single a = maybe stp sng (f a)
741 yieldk a r = case f a of
742 Just b -> yld b $ go r
743 Nothing -> foldStream (adaptState st) yieldk single stp r
744 in foldStream (adaptState st) yieldk single stp m1
746------------------------------------------------------------------------------
747-- Serial Zipping
748------------------------------------------------------------------------------
750-- | Zip two streams serially using a pure zipping function.
751--
752-- @since 0.1.0
753{-# INLINE zipWith #-}
754zipWith :: (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
755zipWith f = go
757 where
759 go mx my = mkStream $ \st yld sng stp -> do
760 let merge a ra =
761 let single2 b = sng (f a b)
762 yield2 b rb = yld (f a b) (go ra rb)
763 in foldStream (adaptState st) yield2 single2 stp my
764 let single1 a = merge a nil
765 yield1 = merge
766 foldStream (adaptState st) yield1 single1 stp mx
768-- | Zip two streams serially using a monadic zipping function.
769--
770-- @since 0.1.0
771{-# INLINE zipWithM #-}
772zipWithM :: Monad m =>
773 (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
774zipWithM f = go
776 where
778 go mx my = mkStream $ \st yld sng stp -> do
779 let merge a ra =
780 let single2 b = f a b >>= sng
781 yield2 b rb = f a b >>= \x -> yld x (go ra rb)
782 in foldStream (adaptState st) yield2 single2 stp my
783 let single1 a = merge a nil
784 yield1 = merge
785 foldStream (adaptState st) yield1 single1 stp mx
787------------------------------------------------------------------------------
788-- Merging
789------------------------------------------------------------------------------
791{-# INLINE mergeByM #-}
792mergeByM :: Monad m =>
793 (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
794mergeByM cmp = go
796 where
798 go mx my = mkStream $ \st yld sng stp -> do
799 let stop = foldStream st yld sng stp my
800 single x = foldStream st yld sng stp (goX0 x my)
801 yield x rx = foldStream st yld sng stp (goX x rx my)
802 foldStream st yield single stop mx
804 goX0 x my = mkStream $ \st yld sng _ -> do
805 let stop = sng x
806 single y = do
807 r <- cmp x y
808 case r of
809 GT -> yld y (fromPure x)
810 _ -> yld x (fromPure y)
811 yield y ry = do
812 r <- cmp x y
813 case r of
814 GT -> yld y (goX0 x ry)
815 _ -> yld x (y `cons` ry)
816 in foldStream st yield single stop my
818 goX x mx my = mkStream $ \st yld _ _ -> do
819 let stop = yld x mx
820 single y = do
821 r <- cmp x y
822 case r of
823 GT -> yld y (x `cons` mx)
824 _ -> yld x (goY0 mx y)
825 yield y ry = do
826 r <- cmp x y
827 case r of
828 GT -> yld y (goX x mx ry)
829 _ -> yld x (goY mx y ry)
830 in foldStream st yield single stop my
832 goY0 mx y = mkStream $ \st yld sng _ -> do
833 let stop = sng y
834 single x = do
835 r <- cmp x y
836 case r of
837 GT -> yld y (fromPure x)
838 _ -> yld x (fromPure y)
839 yield x rx = do
840 r <- cmp x y
841 case r of
842 GT -> yld y (x `cons` rx)
843 _ -> yld x (goY0 rx y)
844 in foldStream st yield single stop mx
846 goY mx y my = mkStream $ \st yld _ _ -> do
847 let stop = yld y my
848 single x = do
849 r <- cmp x y
850 case r of
851 GT -> yld y (goX0 x my)
852 _ -> yld x (y `cons` my)
853 yield x rx = do
854 r <- cmp x y
855 case r of
856 GT -> yld y (goX x rx my)
857 _ -> yld x (goY rx y my)
858 in foldStream st yield single stop mx
860{-# INLINE mergeBy #-}
861mergeBy :: (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
862-- XXX GHC: This has slightly worse performance than replacing "r <- cmp x y"
863-- with "let r = cmp x y" in the monadic version. The definition below is
864-- exactly the same as mergeByM except this change.
865-- mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
866mergeBy cmp = go
868 where
870 go mx my = mkStream $ \st yld sng stp -> do
871 let stop = foldStream st yld sng stp my
872 single x = foldStream st yld sng stp (goX0 x my)
873 yield x rx = foldStream st yld sng stp (goX x rx my)
874 foldStream st yield single stop mx
876 goX0 x my = mkStream $ \st yld sng _ -> do
877 let stop = sng x
878 single y = do
879 case cmp x y of
880 GT -> yld y (fromPure x)
881 _ -> yld x (fromPure y)
882 yield y ry = do
883 case cmp x y of
884 GT -> yld y (goX0 x ry)
885 _ -> yld x (y `cons` ry)
886 in foldStream st yield single stop my
888 goX x mx my = mkStream $ \st yld _ _ -> do
889 let stop = yld x mx
890 single y = do
891 case cmp x y of
892 GT -> yld y (x `cons` mx)
893 _ -> yld x (goY0 mx y)
894 yield y ry = do
895 case cmp x y of
896 GT -> yld y (goX x mx ry)
897 _ -> yld x (goY mx y ry)
898 in foldStream st yield single stop my
900 goY0 mx y = mkStream $ \st yld sng _ -> do
901 let stop = sng y
902 single x = do
903 case cmp x y of
904 GT -> yld y (fromPure x)
905 _ -> yld x (fromPure y)
906 yield x rx = do
907 case cmp x y of
908 GT -> yld y (x `cons` rx)
909 _ -> yld x (goY0 rx y)
910 in foldStream st yield single stop mx
912 goY mx y my = mkStream $ \st yld _ _ -> do
913 let stop = yld y my
914 single x = do
915 case cmp x y of
916 GT -> yld y (goX0 x my)
917 _ -> yld x (y `cons` my)
918 yield x rx = do
919 case cmp x y of
920 GT -> yld y (goX x rx my)
921 _ -> yld x (goY rx y my)
922 in foldStream st yield single stop mx
924------------------------------------------------------------------------------
925-- Transformation comprehensions
926------------------------------------------------------------------------------
928{-# INLINE the #-}
929the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
930the m = do
931 r <- uncons m
932 case r of
933 Nothing -> return Nothing
934 Just (h, t) -> go h t
935 where
936 go h m1 =
937 let single a | h == a = return $ Just h
938 | otherwise = return Nothing
939 yieldk a r | h == a = go h r
940 | otherwise = return Nothing
941 in foldStream defState yieldk single (return $ Just h) m1
943------------------------------------------------------------------------------
944-- Alternative & MonadPlus
945------------------------------------------------------------------------------
947_alt :: Stream m a -> Stream m a -> Stream m a
948_alt m1 m2 = mkStream $ \st yld sng stp ->
949 let stop = foldStream st yld sng stp m2
950 in foldStream st yld sng stop m1
952------------------------------------------------------------------------------
953-- MonadError
954------------------------------------------------------------------------------
956{-
957-- XXX handle and test cross thread state transfer
958withCatchError
959 :: MonadError e m
960 => Stream m a -> (e -> Stream m a) -> Stream m a
961withCatchError m h =
962 mkStream $ \_ stp sng yld ->
963 let run x = unStream x Nothing stp sng yieldk
964 handle r = r `catchError` \e -> run $ h e
965 yieldk a r = yld a (withCatchError r h)
966 in handle $ run m
967-}