git.iko.soy/streamly
on ede29df help sign in
streamly/src/Streamly/Streams/StreamK.hs
1013 lines 32 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE CPP #-}
3{-# LANGUAGE ConstraintKinds #-}
4{-# LANGUAGE FlexibleContexts #-}
5{-# LANGUAGE FlexibleInstances #-}
6{-# LANGUAGE InstanceSigs #-}
7{-# LANGUAGE MultiParamTypeClasses #-}
8{-# LANGUAGE RankNTypes #-}
9{-# LANGUAGE ScopedTypeVariables #-}
10{-# LANGUAGE UndecidableInstances #-} -- XXX
12#include "inline.hs"
14-- |
15-- Module : Streamly.Streams.StreamK
16-- Copyright : (c) 2017 Harendra Kumar
17--
18-- License : BSD3
19-- Maintainer : harendra.kumar@gmail.com
20-- Stability : experimental
21-- Portability : GHC
22--
23--
24-- Continuation passing style (CPS) stream implementation. The symbol 'K' below
25-- denotes a function as well as a Kontinuation.
26--
27-- @
28-- import qualified Streamly.Streams.StreamK as K
29-- @
30--
31module Streamly.Streams.StreamK
32 (
33 -- * A class for streams
34 IsStream (..)
35 , adapt
37 -- * The stream type
38 , Stream
40 -- * Construction Primitives
41 , mkStream
42 , nil
43 , cons
44 , (.:)
46 -- * Elimination Primitives
47 , foldStream
48 , foldStreamShared
49 , foldStreamSVar
51 -- * Transformation Primitives
52 , unShare
54 -- * Deconstruction
55 , uncons
57 -- * Generation
58 -- ** Unfolds
59 , unfoldr
60 , unfoldrM
62 -- ** Specialized Generation
63 , repeat
64 , replicate
65 , replicateM
67 -- ** Conversions
68 , yield
69 , yieldM
70 , fromFoldable
71 , fromList
72 , fromStreamK
74 -- * Elimination
75 -- ** General Folds
76 , foldr
77 , foldrM
78 , foldr1
79 , foldl'
80 , foldlM'
81 , foldx
82 , foldxM
84 -- ** Specialized Folds
85 , runStream
86 , null
87 , head
88 , tail
89 , init
90 , elem
91 , notElem
92 , all
93 , any
94 , last
95 , minimum
96 , minimumBy
97 , maximum
98 , maximumBy
99 , findIndices
100 , lookup
101 , findM
102 , find
103 , (!!)
105 -- ** Map and Fold
106 , mapM_
108 -- ** Conversions
109 , toList
110 , toStreamK
112 -- * Transformation
113 -- ** By folding (scans)
114 , scanl'
115 , scanx
117 -- ** Filtering
118 , filter
119 , take
120 , takeWhile
121 , drop
122 , dropWhile
124 -- ** Mapping
125 , map
126 , mapM
127 , sequence
129 -- ** Inserting
130 , intersperseM
131 , insertBy
133 -- ** Deleting
134 , deleteBy
136 -- ** Map and Filter
137 , mapMaybe
139 -- ** Zipping
140 , zipWith
141 , zipWithM
143 -- ** Merging
144 , mergeBy
145 , mergeByM
147 -- ** Transformation comprehensions
148 , the
150 -- * Semigroup Style Composition
151 , serial
153 -- * Utilities
154 , consMStream
155 , bindWith
156 , withLocal
158 -- * Deprecated
159 , Streaming -- deprecated
160 , once -- deprecated
161 )
162where
164import Control.Monad (void)
165import Control.Monad.Reader.Class (MonadReader(..))
166import Prelude
167 hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,
168 take, filter, all, any, takeWhile, drop, dropWhile, minimum,
169 maximum, elem, notElem, null, head, tail, init, zipWith, lookup,
170 foldr1, (!!), replicate)
171import qualified Prelude
173import Streamly.SVar
174import Streamly.Streams.StreamK.Type
176-- | Detach a stream from an SVar
177{-# INLINE unShare #-}
178unShare :: IsStream t => t m a -> t m a
179unShare x = mkStream $ \st yld sng stp ->
180 foldStream st yld sng stp x
182------------------------------------------------------------------------------
183-- Construction
184------------------------------------------------------------------------------
186infixr 5 `cons`
188-- faster than consM because there is no bind.
189-- | Construct a stream by adding a pure value at the head of an existing
190-- stream. For serial streams this is the same as @(return a) \`consM` r@ but
191-- more efficient. For concurrent streams this is not concurrent whereas
192-- 'consM' is concurrent. For example:
193--
194-- @
195-- > toList $ 1 \`cons` 2 \`cons` 3 \`cons` nil
196-- [1,2,3]
197-- @
198--
199-- @since 0.1.0
200{-# INLINE cons #-}
201cons :: IsStream t => a -> t m a -> t m a
202cons a r = mkStream $ \_ yld _ _ -> yld a r
204infixr 5 .:
206-- | Operator equivalent of 'cons'.
207--
208-- @
209-- > toList $ 1 .: 2 .: 3 .: nil
210-- [1,2,3]
211-- @
212--
213-- @since 0.1.1
214{-# INLINE (.:) #-}
215(.:) :: IsStream t => a -> t m a -> t m a
216(.:) = cons
218-------------------------------------------------------------------------------
219-- Deconstruction
220-------------------------------------------------------------------------------
222{-# INLINE uncons #-}
223uncons :: (IsStream t, Monad m) => t m a -> m (Maybe (a, t m a))
224uncons m =
225 let stop = return Nothing
226 single a = return (Just (a, nil))
227 yieldk a r = return (Just (a, r))
228 in foldStream defState yieldk single stop m
230-------------------------------------------------------------------------------
231-- Generation
232-------------------------------------------------------------------------------
234{-# INLINE_NORMAL build #-}
235build :: IsStream t => forall a. (forall b. (a -> b -> b) -> b -> b) -> t m a
236build g = g cons nil
238{-# INLINE_NORMAL _augment #-}
239_augment
240 :: IsStream t
241 => forall a. (forall b. (a -> b -> b) -> b -> b) -> t m a -> t m a
242_augment g xs = g cons xs
244{-# INLINE_NORMAL _buildM #-}
245_buildM
246 :: (IsStream t, MonadAsync m)
247 => forall a. ((m a -> t m a -> t m a) -> t m a -> t m a) -> t m a
248_buildM g = g consM nil
250{-# INLINE unfoldr #-}
251unfoldr :: IsStream t => (b -> Maybe (a, b)) -> b -> t m a
252unfoldr step b0 = build $ \cns nl ->
253 let go s =
254 case step s of
255 Just (a, b) -> a `cns` go b
256 Nothing -> nl
257 in go b0
259{-# INLINE unfoldrM #-}
260unfoldrM :: (IsStream t, MonadAsync m) => (b -> m (Maybe (a, b))) -> b -> t m a
261unfoldrM step = go
262 where
263 go s = mkStream $ \st yld sng stp -> do
264 mayb <- step s
265 case mayb of
266 Nothing -> stp
267 Just (a, b) ->
268 foldStreamShared st yld sng stp $ return a |: go b
270-------------------------------------------------------------------------------
271-- Special generation
272-------------------------------------------------------------------------------
274{-# INLINE yield #-}
275yield :: IsStream t => a -> t m a
276yield a = mkStream $ \_ _ single _ -> single a
278-- | Same as yieldM
279--
280-- @since 0.2.0
281{-# DEPRECATED once "Please use yieldM instead." #-}
282{-# INLINE once #-}
283once :: (Monad m, IsStream t) => m a -> t m a
284once = yieldM
286-- |
287-- @
288-- repeatM = fix . cons
289-- repeatM = cycle1 . yield
290-- @
291--
292-- Generate an infinite stream by repeating a pure value.
293--
294-- @since 0.4.0
295{-# INLINE repeat #-}
296repeat :: IsStream t => a -> t m a
297repeat a = let x = cons a x in x
299{-# INLINE replicateM #-}
300replicateM :: (IsStream t, MonadAsync m) => Int -> m a -> t m a
301replicateM n m = go n
302 where
303 go cnt = if cnt <= 0 then nil else m |: go (cnt - 1)
305{-# INLINE replicate #-}
306replicate :: IsStream t => Int -> a -> t m a
307replicate n a = go n
308 where
309 go cnt = if cnt <= 0 then nil else a `cons` go (cnt - 1)
311-------------------------------------------------------------------------------
312-- Conversions
313-------------------------------------------------------------------------------
315-- |
316-- @
317-- fromFoldable = 'Prelude.foldr' 'cons' 'nil'
318-- @
319--
320-- Construct a stream from a 'Foldable' containing pure values:
321--
322-- @since 0.2.0
323{-# INLINE fromFoldable #-}
324fromFoldable :: (IsStream t, Foldable f) => f a -> t m a
325fromFoldable = Prelude.foldr cons nil
327{-# INLINE fromList #-}
328fromList :: IsStream t => [a] -> t m a
329fromList = fromFoldable
331{-# INLINE fromStreamK #-}
332fromStreamK :: IsStream t => Stream m a -> t m a
333fromStreamK = fromStream
335-------------------------------------------------------------------------------
336-- Elimination by Folding
337-------------------------------------------------------------------------------
339-- | Lazy right associative fold.
340{-# INLINE foldr #-}
341foldr :: (IsStream t, Monad m) => (a -> b -> b) -> b -> t m a -> m b
342foldr step acc m = go m
343 where
344 go m1 =
345 let stop = return acc
346 single a = return (step a acc)
347 yieldk a r = go r >>= \b -> return (step a b)
348 in foldStream defState yieldk single stop m1
350-- | Lazy right fold with a monadic step function.
351{-# INLINE foldrM #-}
352foldrM :: (IsStream t, Monad m) => (a -> b -> m b) -> b -> t m a -> m b
353foldrM step acc m = go m
354 where
355 go m1 =
356 let stop = return acc
357 single a = step a acc
358 yieldk a r = go r >>= step a
359 in foldStream defState yieldk single stop m1
361{-# INLINE foldr1 #-}
362foldr1 :: (IsStream t, Monad m) => (a -> a -> a) -> t m a -> m (Maybe a)
363foldr1 step m = do
364 r <- uncons m
365 case r of
366 Nothing -> return Nothing
367 Just (h, t) -> fmap Just (go h t)
368 where
369 go p m1 =
370 let stp = return p
371 single a = return $ step a p
372 yieldk a r = fmap (step p) (go a r)
373 in foldStream defState yieldk single stp m1
375-- | Strict left fold with an extraction function. Like the standard strict
376-- left fold, but applies a user supplied extraction function (the third
377-- argument) to the folded value at the end. This is designed to work with the
378-- @foldl@ library. The suffix @x@ is a mnemonic for extraction.
379--
380-- Note that the accumulator is always evaluated including the initial value.
381{-# INLINE foldx #-}
382foldx :: forall t m a b x. (IsStream t, Monad m)
383 => (x -> a -> x) -> x -> (x -> b) -> t m a -> m b
384foldx step begin done m = get $ go m begin
385 where
386 {-# NOINLINE get #-}
387 get :: t m x -> m b
388 get m1 =
389 -- XXX we are not strictly evaluating the accumulator here. Is this
390 -- okay?
391 let single = return . done
392 -- XXX this is foldSingleton. why foldStreamShared?
393 in foldStreamShared undefined undefined single undefined m1
395 -- Note, this can be implemented by making a recursive call to "go",
396 -- however that is more expensive because of unnecessary recursion
397 -- that cannot be tail call optimized. Unfolding recursion explicitly via
398 -- continuations is much more efficient.
399 go :: t m a -> x -> t m x
400 go m1 !acc = mkStream $ \_ yld sng _ ->
401 let stop = sng acc
402 single a = sng $ step acc a
403 -- XXX this is foldNonEmptyStream
404 yieldk a r = foldStream defState yld sng undefined $
405 go r (step acc a)
406 in foldStream defState yieldk single stop m1
408-- | Strict left associative fold.
409{-# INLINE foldl' #-}
410foldl' :: (IsStream t, Monad m) => (b -> a -> b) -> b -> t m a -> m b
411foldl' step begin = foldx step begin id
413-- XXX replace the recursive "go" with explicit continuations.
414-- | Like 'foldx', but with a monadic step function.
415{-# INLINABLE foldxM #-}
416foldxM :: (IsStream t, Monad m)
417 => (x -> a -> m x) -> m x -> (x -> m b) -> t m a -> m b
418foldxM step begin done m = go begin m
419 where
420 go !acc m1 =
421 let stop = acc >>= done
422 single a = acc >>= \b -> step b a >>= done
423 yieldk a r = acc >>= \b -> step b a >>= \x -> go (return x) r
424 in foldStream defState yieldk single stop m1
426-- | Like 'foldl'' but with a monadic step function.
427{-# INLINE foldlM' #-}
428foldlM' :: (IsStream t, Monad m) => (b -> a -> m b) -> b -> t m a -> m b
429foldlM' step begin = foldxM step (return begin) return
431------------------------------------------------------------------------------
432-- Specialized folds
433------------------------------------------------------------------------------
435-- |
436-- > runStream = foldl' (\_ _ -> ()) ()
437-- > runStream = mapM_ (\_ -> return ())
438{-# INLINE runStream #-}
439runStream :: (Monad m, IsStream t) => t m a -> m ()
440runStream = go
441 where
442 go m1 =
443 let stop = return ()
444 single _ = return ()
445 yieldk _ r = go r
446 in foldStream defState yieldk single stop m1
448{-# INLINE null #-}
449null :: (IsStream t, Monad m) => t m a -> m Bool
450null m =
451 let stop = return True
452 single _ = return False
453 yieldk _ _ = return False
454 in foldStream defState yieldk single stop m
456{-# INLINE head #-}
457head :: (IsStream t, Monad m) => t m a -> m (Maybe a)
458head m =
459 let stop = return Nothing
460 single a = return (Just a)
461 yieldk a _ = return (Just a)
462 in foldStream defState yieldk single stop m
464{-# INLINE tail #-}
465tail :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
466tail m =
467 let stop = return Nothing
468 single _ = return $ Just nil
469 yieldk _ r = return $ Just r
470 in foldStream defState yieldk single stop m
472{-# INLINE init #-}
473init :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
474init m = go1 m
475 where
476 go1 m1 = do
477 r <- uncons m1
478 case r of
479 Nothing -> return Nothing
480 Just (h, t) -> return . Just $ go h t
481 go p m1 = mkStream $ \_ yld sng stp ->
482 let single _ = sng p
483 yieldk a x = yld p $ go a x
484 in foldStream defState yieldk single stp m1
486{-# INLINE elem #-}
487elem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
488elem e m = go m
489 where
490 go m1 =
491 let stop = return False
492 single a = return (a == e)
493 yieldk a r = if a == e then return True else go r
494 in foldStream defState yieldk single stop m1
496{-# INLINE notElem #-}
497notElem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
498notElem e m = go m
499 where
500 go m1 =
501 let stop = return True
502 single a = return (a /= e)
503 yieldk a r = if a == e then return False else go r
504 in foldStream defState yieldk single stop m1
506{-# INLINABLE all #-}
507all :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
508all p m = go m
509 where
510 go m1 =
511 let single a | p a = return True
512 | otherwise = return False
513 yieldk a r | p a = go r
514 | otherwise = return False
515 in foldStream defState yieldk single (return True) m1
517{-# INLINABLE any #-}
518any :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
519any p m = go m
520 where
521 go m1 =
522 let single a | p a = return True
523 | otherwise = return False
524 yieldk a r | p a = return True
525 | otherwise = go r
526 in foldStream defState yieldk single (return False) m1
528-- | Extract the last element of the stream, if any.
529{-# INLINE last #-}
530last :: (IsStream t, Monad m) => t m a -> m (Maybe a)
531last = foldx (\_ y -> Just y) Nothing id
533{-# INLINE minimum #-}
534minimum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
535minimum m = go Nothing m
536 where
537 go Nothing m1 =
538 let stop = return Nothing
539 single a = return (Just a)
540 yieldk a r = go (Just a) r
541 in foldStream defState yieldk single stop m1
543 go (Just res) m1 =
544 let stop = return (Just res)
545 single a =
546 if res <= a
547 then return (Just res)
548 else return (Just a)
549 yieldk a r =
550 if res <= a
551 then go (Just res) r
552 else go (Just a) r
553 in foldStream defState yieldk single stop m1
555{-# INLINE minimumBy #-}
556minimumBy
557 :: (IsStream t, Monad m)
558 => (a -> a -> Ordering) -> t m a -> m (Maybe a)
559minimumBy cmp m = go Nothing m
560 where
561 go Nothing m1 =
562 let stop = return Nothing
563 single a = return (Just a)
564 yieldk a r = go (Just a) r
565 in foldStream defState yieldk single stop m1
567 go (Just res) m1 =
568 let stop = return (Just res)
569 single a = case cmp res a of
570 GT -> return (Just a)
571 _ -> return (Just res)
572 yieldk a r = case cmp res a of
573 GT -> go (Just a) r
574 _ -> go (Just res) r
575 in foldStream defState yieldk single stop m1
577{-# INLINE maximum #-}
578maximum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
579maximum m = go Nothing m
580 where
581 go Nothing m1 =
582 let stop = return Nothing
583 single a = return (Just a)
584 yieldk a r = go (Just a) r
585 in foldStream defState yieldk single stop m1
587 go (Just res) m1 =
588 let stop = return (Just res)
589 single a =
590 if res <= a
591 then return (Just a)
592 else return (Just res)
593 yieldk a r =
594 if res <= a
595 then go (Just a) r
596 else go (Just res) r
597 in foldStream defState yieldk single stop m1
599{-# INLINE maximumBy #-}
600maximumBy :: (IsStream t, Monad m) => (a -> a -> Ordering) -> t m a -> m (Maybe a)
601maximumBy cmp m = go Nothing m
602 where
603 go Nothing m1 =
604 let stop = return Nothing
605 single a = return (Just a)
606 yieldk a r = go (Just a) r
607 in foldStream defState yieldk single stop m1
609 go (Just res) m1 =
610 let stop = return (Just res)
611 single a = case cmp res a of
612 GT -> return (Just res)
613 _ -> return (Just a)
614 yieldk a r = case cmp res a of
615 GT -> go (Just res) r
616 _ -> go (Just a) r
617 in foldStream defState yieldk single stop m1
619{-# INLINE (!!) #-}
620(!!) :: (IsStream t, Monad m) => t m a -> Int -> m (Maybe a)
621m !! i = go i m
622 where
623 go n m1 =
624 let single a | n == 0 = return $ Just a
625 | otherwise = return Nothing
626 yieldk a x | n < 0 = return Nothing
627 | n == 0 = return $ Just a
628 | otherwise = go (n - 1) x
629 in foldStream defState yieldk single (return Nothing) m1
631{-# INLINE lookup #-}
632lookup :: (IsStream t, Monad m, Eq a) => a -> t m (a, b) -> m (Maybe b)
633lookup e m = go m
634 where
635 go m1 =
636 let single (a, b) | a == e = return $ Just b
637 | otherwise = return Nothing
638 yieldk (a, b) x | a == e = return $ Just b
639 | otherwise = go x
640 in foldStream defState yieldk single (return Nothing) m1
642{-# INLINE findM #-}
643findM :: (IsStream t, Monad m) => (a -> m Bool) -> t m a -> m (Maybe a)
644findM p m = go m
645 where
646 go m1 =
647 let single a = do
648 b <- p a
649 if b then return $ Just a else return Nothing
650 yieldk a x = do
651 b <- p a
652 if b then return $ Just a else go x
653 in foldStream defState yieldk single (return Nothing) m1
655{-# INLINE find #-}
656find :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m (Maybe a)
657find p = findM (return . p)
659{-# INLINE findIndices #-}
660findIndices :: IsStream t => (a -> Bool) -> t m a -> t m Int
661findIndices p = go 0
662 where
663 go offset m1 = mkStream $ \st yld sng stp ->
664 let single a | p a = sng offset
665 | otherwise = stp
666 yieldk a x | p a = yld offset $ go (offset + 1) x
667 | otherwise = foldStream (adaptState st) yld sng stp $
668 go (offset + 1) x
669 in foldStream (adaptState st) yieldk single stp m1
671------------------------------------------------------------------------------
672-- Map and Fold
673------------------------------------------------------------------------------
675-- | Apply a monadic action to each element of the stream and discard the
676-- output of the action.
677{-# INLINE mapM_ #-}
678mapM_ :: (IsStream t, Monad m) => (a -> m b) -> t m a -> m ()
679mapM_ f m = go m
680 where
681 go m1 =
682 let stop = return ()
683 single a = void (f a)
684 yieldk a r = f a >> go r
685 in foldStream defState yieldk single stop m1
687------------------------------------------------------------------------------
688-- Converting folds
689------------------------------------------------------------------------------
691{-# INLINABLE toList #-}
692toList :: (IsStream t, Monad m) => t m a -> m [a]
693toList = foldr (:) []
695{-# INLINE toStreamK #-}
696toStreamK :: Stream m a -> Stream m a
697toStreamK = id
699-------------------------------------------------------------------------------
700-- Transformation by folding (Scans)
701-------------------------------------------------------------------------------
703{-# INLINE scanx #-}
704scanx :: IsStream t => (x -> a -> x) -> x -> (x -> b) -> t m a -> t m b
705scanx step begin done m =
706 cons (done begin) $ go m begin
707 where
708 go m1 !acc = mkStream $ \st yld sng stp ->
709 let single a = sng (done $ step acc a)
710 yieldk a r =
711 let s = step acc a
712 in yld (done s) (go r s)
713 in foldStream (adaptState st) yieldk single stp m1
715{-# INLINE scanl' #-}
716scanl' :: IsStream t => (b -> a -> b) -> b -> t m a -> t m b
717scanl' step begin = scanx step begin id
719-------------------------------------------------------------------------------
720-- Filtering
721-------------------------------------------------------------------------------
723{-# INLINE filter #-}
724filter :: IsStream t => (a -> Bool) -> t m a -> t m a
725filter p m = go m
726 where
727 go m1 = mkStream $ \st yld sng stp ->
728 let single a | p a = sng a
729 | otherwise = stp
730 yieldk a r | p a = yld a (go r)
731 | otherwise = foldStream st yieldk single stp r
732 in foldStream st yieldk single stp m1
734{-# INLINE take #-}
735take :: IsStream t => Int -> t m a -> t m a
736take n m = go n m
737 where
738 go n1 m1 = mkStream $ \st yld sng stp ->
739 let yieldk a r = yld a (go (n1 - 1) r)
740 in if n1 <= 0
741 then stp
742 else foldStream st yieldk sng stp m1
744{-# INLINE takeWhile #-}
745takeWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
746takeWhile p m = go m
747 where
748 go m1 = mkStream $ \st yld sng stp ->
749 let single a | p a = sng a
750 | otherwise = stp
751 yieldk a r | p a = yld a (go r)
752 | otherwise = stp
753 in foldStream st yieldk single stp m1
755{-# INLINE drop #-}
756drop :: IsStream t => Int -> t m a -> t m a
757drop n m = go n m
758 where
759 go n1 m1 = mkStream $ \st yld sng stp ->
760 let single _ = stp
761 yieldk _ r = foldStream st yld sng stp $ go (n1 - 1) r
762 -- Somehow "<=" check performs better than a ">"
763 in if n1 <= 0
764 then foldStream st yld sng stp m1
765 else foldStream st yieldk single stp m1
767{-# INLINE dropWhile #-}
768dropWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
769dropWhile p m = go m
770 where
771 go m1 = mkStream $ \st yld sng stp ->
772 let single a | p a = stp
773 | otherwise = sng a
774 yieldk a r | p a = foldStream st yieldk single stp r
775 | otherwise = yld a r
776 in foldStream st yieldk single stp m1
778-------------------------------------------------------------------------------
779-- Mapping
780-------------------------------------------------------------------------------
782-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
783-- other stream types to overload it.
784{-# INLINE mapM #-}
785mapM :: (IsStream t, MonadAsync m) => (a -> m b) -> t m a -> t m b
786mapM f m = go m
787 where
788 go m1 = mkStream $ \st yld sng stp ->
789 let single a = f a >>= sng
790 yieldk a r = foldStreamShared st yld sng stp $ f a |: go r
791 in foldStream (adaptState st) yieldk single stp m1
793-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
794-- other stream types to overload it.
795{-# INLINE sequence #-}
796sequence :: (IsStream t, MonadAsync m) => t m (m a) -> t m a
797sequence m = go m
798 where
799 go m1 = mkStream $ \st yld sng stp ->
800 let single ma = ma >>= sng
801 yieldk ma r = foldStreamShared st yld sng stp $ ma |: go r
802 in foldStream (adaptState st) yieldk single stp m1
804-------------------------------------------------------------------------------
805-- Inserting
806-------------------------------------------------------------------------------
808{-# INLINE intersperseM #-}
809intersperseM :: (IsStream t, MonadAsync m) => m a -> t m a -> t m a
810intersperseM a m = prependingStart m
811 where
812 prependingStart m1 = mkStream $ \st yld sng stp ->
813 let yieldk i x = foldStreamShared st yld sng stp $ return i |: go x
814 in foldStream st yieldk sng stp m1
815 go m2 = mkStream $ \st yld sng stp ->
816 let single i = foldStreamShared st yld sng stp $ a |: yield i
817 yieldk i x = foldStreamShared st yld sng stp $ a |: return i |: go x
818 in foldStream st yieldk single stp m2
820{-# INLINE insertBy #-}
821insertBy :: IsStream t => (a -> a -> Ordering) -> a -> t m a -> t m a
822insertBy cmp x m = go m
823 where
824 go m1 = mkStream $ \st yld _ _ ->
825 let single a = case cmp x a of
826 GT -> yld a (yield x)
827 _ -> yld x (yield a)
828 stop = yld x nil
829 yieldk a r = case cmp x a of
830 GT -> yld a (go r)
831 _ -> yld x (a `cons` r)
832 in foldStream st yieldk single stop m1
834------------------------------------------------------------------------------
835-- Deleting
836------------------------------------------------------------------------------
838{-# INLINE deleteBy #-}
839deleteBy :: IsStream t => (a -> a -> Bool) -> a -> t m a -> t m a
840deleteBy eq x m = go m
841 where
842 go m1 = mkStream $ \st yld sng stp ->
843 let single a = if eq x a then stp else sng a
844 yieldk a r = if eq x a
845 then foldStream st yld sng stp r
846 else yld a (go r)
847 in foldStream st yieldk single stp m1
849-------------------------------------------------------------------------------
850-- Map and Filter
851-------------------------------------------------------------------------------
853{-# INLINE mapMaybe #-}
854mapMaybe :: IsStream t => (a -> Maybe b) -> t m a -> t m b
855mapMaybe f m = go m
856 where
857 go m1 = mkStream $ \st yld sng stp ->
858 let single a = case f a of
859 Just b -> sng b
860 Nothing -> stp
861 yieldk a r = case f a of
862 Just b -> yld b $ go r
863 Nothing -> foldStream (adaptState st) yieldk single stp r
864 in foldStream (adaptState st) yieldk single stp m1
866------------------------------------------------------------------------------
867-- Serial Zipping
868------------------------------------------------------------------------------
870-- | Zip two streams serially using a pure zipping function.
871--
872-- @since 0.1.0
873{-# INLINABLE zipWith #-}
874zipWith :: IsStream t => (a -> b -> c) -> t m a -> t m b -> t m c
875zipWith f = go
876 where
877 go mx my = mkStream $ \st yld sng stp -> do
878 let merge a ra =
879 let single2 b = sng (f a b)
880 yield2 b rb = yld (f a b) (go ra rb)
881 in foldStream (adaptState st) yield2 single2 stp my
882 let single1 a = merge a nil
883 yield1 = merge
884 foldStream (adaptState st) yield1 single1 stp mx
886-- | Zip two streams serially using a monadic zipping function.
887--
888-- @since 0.1.0
889{-# INLINABLE zipWithM #-}
890zipWithM :: (IsStream t, Monad m) => (a -> b -> m c) -> t m a -> t m b -> t m c
891zipWithM f m1 m2 = go m1 m2
892 where
893 go mx my = mkStream $ \st yld sng stp -> do
894 let merge a ra =
895 let runIt x = foldStream st yld sng stp x
896 single2 b = f a b >>= sng
897 yield2 b rb = f a b >>= \x -> runIt (x `cons` go ra rb)
898 in foldStream (adaptState st) yield2 single2 stp my
899 let single1 a = merge a nil
900 yield1 = merge
901 foldStream (adaptState st) yield1 single1 stp mx
903------------------------------------------------------------------------------
904-- Merging
905------------------------------------------------------------------------------
907{-# INLINE mergeByM #-}
908mergeByM
909 :: (IsStream t, Monad m)
910 => (a -> a -> m Ordering) -> t m a -> t m a -> t m a
911mergeByM cmp = go
912 where
913 go mx my = mkStream $ \st yld sng stp -> do
914 let mergeWithY a ra =
915 let stop2 = foldStream st yld sng stp mx
916 single2 b = do
917 r <- cmp a b
918 case r of
919 GT -> yld b (go (a `cons` ra) nil)
920 _ -> yld a (go ra (b `cons` nil))
921 yield2 b rb = do
922 r <- cmp a b
923 case r of
924 GT -> yld b (go (a `cons` ra) rb)
925 _ -> yld a (go ra (b `cons` rb))
926 in foldStream st yield2 single2 stop2 my
927 let stopX = foldStream st yld sng stp my
928 singleX a = mergeWithY a nil
929 yieldX = mergeWithY
930 foldStream st yieldX singleX stopX mx
932{-# INLINABLE mergeBy #-}
933mergeBy
934 :: (IsStream t, Monad m)
935 => (a -> a -> Ordering) -> t m a -> t m a -> t m a
936mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
938------------------------------------------------------------------------------
939-- Transformation comprehensions
940------------------------------------------------------------------------------
942{-# INLINE the #-}
943the :: (Eq a, IsStream t, Monad m) => t m a -> m (Maybe a)
944the m = do
945 r <- uncons m
946 case r of
947 Nothing -> return Nothing
948 Just (h, t) -> go h t
949 where
950 go h m1 =
951 let single a | h == a = return $ Just h
952 | otherwise = return Nothing
953 yieldk a r | h == a = go h r
954 | otherwise = return Nothing
955 in foldStream defState yieldk single (return $ Just h) m1
957-------------------------------------------------------------------------------
958-- Bind utility
959-------------------------------------------------------------------------------
961{-# INLINE bindWith #-}
962bindWith
963 :: IsStream t
964 => (forall c. t m c -> t m c -> t m c)
965 -> t m a
966 -> (a -> t m b)
967 -> t m b
968bindWith par m1 f = go m1
969 where
970 go m =
971 mkStream $ \st yld sng stp ->
972 let foldShared = foldStreamShared st yld sng stp
973 single a = foldShared $ unShare (f a)
974 yieldk a r = foldShared $ unShare (f a) `par` go r
975 in foldStream (adaptState st) yieldk single stp m
977------------------------------------------------------------------------------
978-- Alternative & MonadPlus
979------------------------------------------------------------------------------
981_alt :: Stream m a -> Stream m a -> Stream m a
982_alt m1 m2 = mkStream $ \st yld sng stp ->
983 let stop = foldStream st yld sng stp m2
984 in foldStream st yld sng stop m1
986------------------------------------------------------------------------------
987-- MonadReader
988------------------------------------------------------------------------------
990{-# INLINABLE withLocal #-}
991withLocal :: MonadReader r m => (r -> r) -> Stream m a -> Stream m a
992withLocal f m =
993 mkStream $ \st yld sng stp ->
994 let single = local f . sng
995 yieldk a r = local f $ yld a (withLocal f r)
996 in foldStream st yieldk single (local f stp) m
998------------------------------------------------------------------------------
999-- MonadError
1000------------------------------------------------------------------------------
1003-- XXX handle and test cross thread state transfer
1004withCatchError
1005 :: MonadError e m
1006 => Stream m a -> (e -> Stream m a) -> Stream m a
1007withCatchError m h =
1008 mkStream $ \_ stp sng yld ->
1009 let run x = unStream x Nothing stp sng yieldk
1010 handle r = r `catchError` \e -> run $ h e
1011 yieldk a r = yld a (withCatchError r h)
1012 in handle $ run m