git.iko.soy/streamly
on a296fac help sign in
streamly/src/Streamly/Internal/Data/Stream/StreamK.hs
1123 lines 35 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 -- * A class for streams
25 IsStream (..)
26 , adapt
28 -- * The stream type
29 , Stream(..)
31 -- * Construction Primitives
32 , mkStream
33 , nil
34 , nilM
35 , cons
36 , (.:)
38 -- * Elimination Primitives
39 , foldStream
40 , foldStreamShared
42 -- * Transformation Primitives
43 , unShare
45 -- * Deconstruction
46 , uncons
48 -- * Generation
49 -- ** Unfolds
50 , unfoldr
51 , unfoldrM
53 -- ** Specialized Generation
54 , repeat
55 , repeatM
56 , replicate
57 , replicateM
58 , fromIndices
59 , fromIndicesM
60 , iterate
61 , iterateM
63 -- ** Conversions
64 , fromPure
65 , fromEffect
66 , fromFoldable
67 , fromList
68 , fromStreamK
70 -- * foldr/build
71 , foldrS
72 , foldrSM
73 , buildS
74 , buildM
75 , augmentS
76 , augmentSM
78 -- * Elimination
79 -- ** General Folds
80 , foldr
81 , foldr1
82 , foldrM
83 , foldrT
85 , foldl'
86 , foldlM'
87 , foldlS
88 , foldlT
89 , foldlx'
90 , foldlMx'
91 , fold
93 -- ** Specialized Folds
94 , drain
95 , null
96 , head
97 , tail
98 , init
99 , elem
100 , notElem
101 , all
102 , any
103 , last
104 , minimum
105 , minimumBy
106 , maximum
107 , maximumBy
108 , findIndices
109 , lookup
110 , findM
111 , find
112 , (!!)
114 -- ** Map and Fold
115 , mapM_
117 -- ** Conversions
118 , toList
119 , toStreamK
120 , hoist
122 -- * Transformation
123 -- ** By folding (scans)
124 , scanl'
125 , scanlx'
127 -- ** Filtering
128 , filter
129 , take
130 , takeWhile
131 , drop
132 , dropWhile
134 -- ** Mapping
135 , map
136 , mapM
137 , mapMSerial
138 , sequence
140 -- ** Inserting
141 , intersperseM
142 , intersperse
143 , insertBy
145 -- ** Deleting
146 , deleteBy
148 -- ** Reordering
149 , reverse
151 -- ** Map and Filter
152 , mapMaybe
154 -- ** Zipping
155 , zipWith
156 , zipWithM
158 -- ** Merging
159 , mergeBy
160 , mergeByM
162 -- ** Nesting
163 , concatMapBy
164 , concatMap
165 , bindWith
166 , concatPairsWith
167 , apWith
168 , apSerial
169 , apSerialDiscardFst
170 , apSerialDiscardSnd
172 -- ** Transformation comprehensions
173 , the
175 -- * Semigroup Style Composition
176 , serial
178 -- * Utilities
179 , consMStream
180 , withLocal
181 , mfix
183 -- * Deprecated
184 , Streaming -- deprecated
185 , once -- deprecated
186 )
187where
189import Control.Monad.Trans.Class (MonadTrans(lift))
190import Control.Monad (void, join)
191import Control.Monad.Reader.Class (MonadReader(..))
192import Data.Function (fix)
193import Streamly.Internal.Control.Concurrent (MonadAsync)
194import Streamly.Internal.Data.SVar.Type (adaptState, defState)
196import qualified Prelude
197import qualified Streamly.Internal.Data.Fold.Type as FL
199import Prelude
200 hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,
201 take, filter, all, any, takeWhile, drop, dropWhile, minimum,
202 maximum, elem, notElem, null, head, tail, init, zipWith, lookup,
203 foldr1, (!!), replicate, reverse, concatMap, iterate)
205import Streamly.Internal.Data.Stream.StreamK.Type
207-- $setup
208-- >>> :m
210-------------------------------------------------------------------------------
211-- Deconstruction
212-------------------------------------------------------------------------------
214{-# INLINE uncons #-}
215uncons :: (IsStream t, Monad m) => t m a -> m (Maybe (a, t m a))
216uncons m =
217 let stop = return Nothing
218 single a = return (Just (a, nil))
219 yieldk a r = return (Just (a, r))
220 in foldStream defState yieldk single stop m
222-------------------------------------------------------------------------------
223-- Generation
224-------------------------------------------------------------------------------
226{-# INLINE unfoldr #-}
227unfoldr :: IsStream t => (b -> Maybe (a, b)) -> b -> t m a
228unfoldr next s0 = build $ \yld stp ->
229 let go s =
230 case next s of
231 Just (a, b) -> yld a (go b)
232 Nothing -> stp
233 in go s0
235{-# INLINE unfoldrM #-}
236unfoldrM :: (IsStream t, MonadAsync m) => (b -> m (Maybe (a, b))) -> b -> t m a
237unfoldrM step = go
238 where
239 go s = sharedM $ \yld _ stp -> do
240 r <- step s
241 case r of
242 Just (a, b) -> yld a (go b)
243 Nothing -> stp
245{-
246-- Generalization of concurrent streams/SVar via unfoldr.
247--
248-- Unfold a value into monadic actions and then run the resulting monadic
249-- actions to generate a stream. Since the step of generating the monadic
250-- action and running them are decoupled we can run the monadic actions
251-- cooncurrently. For example, the seed could be a list of monadic actions or a
252-- pure stream of monadic actions.
253--
254-- We can have different flavors of this depending on the stream type t. The
255-- concurrent version could be async or ahead etc. Depending on how we queue
256-- back the feedback portion b, it could be DFS or BFS style.
257--
258unfoldrA :: (IsStream t, MonadAsync m) => (b -> Maybe (m a, b)) -> b -> t m a
259unfoldrA = undefined
260-}
262-------------------------------------------------------------------------------
263-- Special generation
264-------------------------------------------------------------------------------
266-- | Same as fromEffect
267--
268-- @since 0.2.0
269{-# DEPRECATED once "Please use fromEffect instead." #-}
270{-# INLINE once #-}
271once :: (Monad m, IsStream t) => m a -> t m a
272once = fromEffect
274-- |
275-- @
276-- repeatM = fix . cons
277-- repeatM = cycle1 . fromPure
278-- @
279--
280-- Generate an infinite stream by repeating a monadic value.
281--
282-- /Pre-release/
283repeatM :: (IsStream t, MonadAsync m) => m a -> t m a
284repeatM = go
285 where go m = m |: go m
287-- Generate an infinite stream by repeating a pure value.
288--
289-- /Pre-release/
290{-# INLINE repeat #-}
291repeat :: IsStream t => a -> t m a
292repeat a = let x = cons a x in x
294{-# INLINE replicateM #-}
295replicateM :: (IsStream t, MonadAsync m) => Int -> m a -> t m a
296replicateM n m = go n
297 where
298 go cnt = if cnt <= 0 then nil else m |: go (cnt - 1)
300{-# INLINE replicate #-}
301replicate :: IsStream t => Int -> a -> t m a
302replicate n a = go n
303 where
304 go cnt = if cnt <= 0 then nil else a `cons` go (cnt - 1)
306{-# INLINE fromIndicesM #-}
307fromIndicesM :: (IsStream t, MonadAsync m) => (Int -> m a) -> t m a
308fromIndicesM gen = go 0
309 where
310 go i = mkStream $ \st stp sng yld -> do
311 foldStreamShared st stp sng yld (gen i |: go (i + 1))
313{-# INLINE fromIndices #-}
314fromIndices :: IsStream t => (Int -> a) -> t m a
315fromIndices gen = go 0
316 where
317 go n = (gen n) `cons` go (n + 1)
319{-# INLINE iterate #-}
320iterate :: IsStream t => (a -> a) -> a -> t m a
321iterate step = go
322 where
323 go !s = cons s (go (step s))
325{-# INLINE iterateM #-}
326iterateM :: (IsStream t, MonadAsync m) => (a -> m a) -> m a -> t m a
327iterateM step = go
328 where
329 go s = mkStream $ \st stp sng yld -> do
330 !next <- s
331 foldStreamShared st stp sng yld (return next |: go (step next))
333-------------------------------------------------------------------------------
334-- Conversions
335-------------------------------------------------------------------------------
337-- |
338-- @
339-- fromFoldable = 'Prelude.foldr' 'cons' 'nil'
340-- @
341--
342-- Construct a stream from a 'Foldable' containing pure values:
343--
344-- @since 0.2.0
345{-# INLINE fromFoldable #-}
346fromFoldable :: (IsStream t, Foldable f) => f a -> t m a
347fromFoldable = Prelude.foldr cons nil
349{-# INLINE fromList #-}
350fromList :: IsStream t => [a] -> t m a
351fromList = fromFoldable
353{-# INLINE fromStreamK #-}
354fromStreamK :: IsStream t => Stream m a -> t m a
355fromStreamK = fromStream
357-------------------------------------------------------------------------------
358-- Elimination by Folding
359-------------------------------------------------------------------------------
361-- | Lazy right associative fold.
362{-# INLINE foldr #-}
363foldr :: (IsStream t, Monad m) => (a -> b -> b) -> b -> t m a -> m b
364foldr step acc = foldrM (\x xs -> xs >>= \b -> return (step x b)) (return acc)
366-- | Right associative fold to an arbitrary transformer monad.
367{-# INLINE foldrT #-}
368foldrT :: (IsStream t, Monad m, Monad (s m), MonadTrans s)
369 => (a -> s m b -> s m b) -> s m b -> t m a -> s m b
370foldrT step final m = go m
371 where
372 go m1 = do
373 res <- lift $ uncons m1
374 case res of
375 Just (h, t) -> step h (go t)
376 Nothing -> final
378{-# INLINE foldr1 #-}
379foldr1 :: (IsStream t, Monad m) => (a -> a -> a) -> t m a -> m (Maybe a)
380foldr1 step m = do
381 r <- uncons m
382 case r of
383 Nothing -> return Nothing
384 Just (h, t) -> fmap Just (go h t)
385 where
386 go p m1 =
387 let stp = return p
388 single a = return $ step a p
389 yieldk a r = fmap (step p) (go a r)
390 in foldStream defState yieldk single stp m1
392-- XXX replace the recursive "go" with explicit continuations.
393-- | Like 'foldx', but with a monadic step function.
394{-# INLINABLE foldlMx' #-}
395foldlMx' :: (IsStream t, Monad m)
396 => (x -> a -> m x) -> m x -> (x -> m b) -> t m a -> m b
397foldlMx' step begin done m = go begin m
398 where
399 go !acc m1 =
400 let stop = acc >>= done
401 single a = acc >>= \b -> step b a >>= done
402 yieldk a r = acc >>= \b -> step b a >>= \x -> go (return x) r
403 in foldStream defState yieldk single stop m1
405{-# INLINABLE fold #-}
406fold :: (IsStream t, Monad m) => FL.Fold m a b -> t m a -> m b
407fold (FL.Fold step begin done) m = do
408 res <- begin
409 case res of
410 FL.Partial fs -> go fs m
411 FL.Done fb -> return fb
413 where
414 go !acc m1 =
415 let stop = done acc
416 single a = step acc a
417 >>= \case
418 FL.Partial s -> done s
419 FL.Done b1 -> return b1
420 yieldk a r = step acc a
421 >>= \case
422 FL.Partial s -> go s r
423 FL.Done b1 -> return b1
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) -> m b -> t m a -> m b
429foldlM' step begin = foldlMx' step begin return
431-- | Lazy left fold to a stream.
432{-# INLINE foldlS #-}
433foldlS :: IsStream t => (t m b -> a -> t m b) -> t m b -> t m a -> t m b
434foldlS step begin m = go begin m
435 where
436 go acc rest = mkStream $ \st yld sng stp ->
437 let run x = foldStream st yld sng stp x
438 stop = run acc
439 single a = run $ step acc a
440 yieldk a r = run $ go (step acc a) r
441 in foldStream (adaptState st) yieldk single stop rest
443-- | Lazy left fold to an arbitrary transformer monad.
444{-# INLINE foldlT #-}
445foldlT :: (IsStream t, Monad m, Monad (s m), MonadTrans s)
446 => (s m b -> a -> s m b) -> s m b -> t m a -> s m b
447foldlT step begin m = go begin m
448 where
449 go acc m1 = do
450 res <- lift $ uncons m1
451 case res of
452 Just (h, t) -> go (step acc h) t
453 Nothing -> acc
455------------------------------------------------------------------------------
456-- Specialized folds
457------------------------------------------------------------------------------
459-- XXX use foldrM to implement folds where possible
460-- XXX This (commented) definition of drain and mapM_ perform much better on
461-- some benchmarks but worse on others. Need to investigate why, may there is
462-- an optimization opportunity that we can exploit.
463-- drain = foldrM (\_ xs -> return () >> xs) (return ())
465-- |
466-- > drain = foldl' (\_ _ -> ()) ()
467-- > drain = mapM_ (\_ -> return ())
468{-# INLINE drain #-}
469drain :: (Monad m, IsStream t) => t m a -> m ()
470drain = foldrM (\_ xs -> xs) (return ())
471{-
472drain = go
473 where
474 go m1 =
475 let stop = return ()
476 single _ = return ()
477 yieldk _ r = go r
478 in foldStream defState yieldk single stop m1
479-}
481{-# INLINE null #-}
482null :: (IsStream t, Monad m) => t m a -> m Bool
483-- null = foldrM (\_ _ -> return True) (return False)
484null m =
485 let stop = return True
486 single _ = return False
487 yieldk _ _ = return False
488 in foldStream defState yieldk single stop m
490{-# INLINE head #-}
491head :: (IsStream t, Monad m) => t m a -> m (Maybe a)
492-- head = foldrM (\x _ -> return $ Just x) (return Nothing)
493head m =
494 let stop = return Nothing
495 single a = return (Just a)
496 yieldk a _ = return (Just a)
497 in foldStream defState yieldk single stop m
499{-# INLINE tail #-}
500tail :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
501tail m =
502 let stop = return Nothing
503 single _ = return $ Just nil
504 yieldk _ r = return $ Just r
505 in foldStream defState yieldk single stop m
507{-# INLINE headPartial #-}
508headPartial :: (IsStream t, Monad m) => t m a -> m a
509headPartial = foldrM (\x _ -> return x) (error "head of nil")
511{-# INLINE tailPartial #-}
512tailPartial :: IsStream t => t m a -> t m a
513tailPartial m = mkStream $ \st yld sng stp ->
514 let stop = error "tail of nil"
515 single _ = stp
516 yieldk _ r = foldStream st yld sng stp r
517 in foldStream st yieldk single stop m
519-- | We can define cyclic structures using @let@:
520--
521-- >>> let (a, b) = ([1, b], head a) in (a, b)
522-- ([1,1],1)
523--
524-- The function @fix@ defined as:
525--
526-- > fix f = let x = f x in x
527--
528-- ensures that the argument of a function and its output refer to the same
529-- lazy value @x@ i.e. the same location in memory. Thus @x@ can be defined
530-- in terms of itself, creating structures with cyclic references.
531--
532-- >>> import Data.Function (fix)
533-- >>> f ~(a, b) = ([1, b], head a)
534-- >>> fix f
535-- ([1,1],1)
536--
537-- 'Control.Monad.mfix' is essentially the same as @fix@ but for monadic
538-- values.
539--
540-- Using 'mfix' for streams we can construct a stream in which each element of
541-- the stream is defined in a cyclic fashion. The argument of the function
542-- being fixed represents the current element of the stream which is being
543-- returned by the stream monad. Thus, we can use the argument to construct
544-- itself.
545--
546-- In the following example, the argument @action@ of the function @f@
547-- represents the tuple @(x,y)@ returned by it in a given iteration. We define
548-- the first element of the tuple in terms of the second.
549--
550-- @
551-- import Streamly.Internal.Data.Stream.IsStream as Stream
552-- import System.IO.Unsafe (unsafeInterleaveIO)
553--
554-- main = do
555-- Stream.mapM_ print $ Stream.mfix f
556--
557-- where
558--
559-- f action = do
560-- let incr n act = fmap ((+n) . snd) $ unsafeInterleaveIO act
561-- x <- Stream.fromListM [incr 1 action, incr 2 action]
562-- y <- Stream.fromList [4,5]
563-- return (x, y)
564-- @
565--
566-- Note: you cannot achieve this by just changing the order of the monad
567-- statements because that would change the order in which the stream elements
568-- are generated.
569--
570-- Note that the function @f@ must be lazy in its argument, that's why we use
571-- 'unsafeInterleaveIO' on @action@ because IO monad is strict.
572--
573-- /Pre-release/
575mfix :: (IsStream t, Monad m) => (m a -> t m a) -> t m a
576mfix f = mkStream $ \st yld sng stp ->
577 let single a = foldStream st yld sng stp $ a `cons` ys
578 yieldk a _ = foldStream st yld sng stp $ a `cons` ys
579 in foldStream st yieldk single stp xs
581 where
583 -- fix the head element of the stream
584 xs = fix (f . headPartial)
586 -- now fix the tail recursively
587 ys = mfix (tailPartial . f)
589{-# INLINE init #-}
590init :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
591init m = go1 m
592 where
593 go1 m1 = do
594 r <- uncons m1
595 case r of
596 Nothing -> return Nothing
597 Just (h, t) -> return . Just $ go h t
598 go p m1 = mkStream $ \_ yld sng stp ->
599 let single _ = sng p
600 yieldk a x = yld p $ go a x
601 in foldStream defState yieldk single stp m1
603{-# INLINE elem #-}
604elem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
605elem e m = go m
606 where
607 go m1 =
608 let stop = return False
609 single a = return (a == e)
610 yieldk a r = if a == e then return True else go r
611 in foldStream defState yieldk single stop m1
613{-# INLINE notElem #-}
614notElem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
615notElem e m = go m
616 where
617 go m1 =
618 let stop = return True
619 single a = return (a /= e)
620 yieldk a r = if a == e then return False else go r
621 in foldStream defState yieldk single stop m1
623{-# INLINABLE all #-}
624all :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
625all p m = go m
626 where
627 go m1 =
628 let single a | p a = return True
629 | otherwise = return False
630 yieldk a r | p a = go r
631 | otherwise = return False
632 in foldStream defState yieldk single (return True) m1
634{-# INLINABLE any #-}
635any :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
636any p m = go m
637 where
638 go m1 =
639 let single a | p a = return True
640 | otherwise = return False
641 yieldk a r | p a = return True
642 | otherwise = go r
643 in foldStream defState yieldk single (return False) m1
645-- | Extract the last element of the stream, if any.
646{-# INLINE last #-}
647last :: (IsStream t, Monad m) => t m a -> m (Maybe a)
648last = foldlx' (\_ y -> Just y) Nothing id
650{-# INLINE minimum #-}
651minimum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
652minimum m = go Nothing m
653 where
654 go Nothing m1 =
655 let stop = return Nothing
656 single a = return (Just a)
657 yieldk a r = go (Just a) r
658 in foldStream defState yieldk single stop m1
660 go (Just res) m1 =
661 let stop = return (Just res)
662 single a =
663 if res <= a
664 then return (Just res)
665 else return (Just a)
666 yieldk a r =
667 if res <= a
668 then go (Just res) r
669 else go (Just a) r
670 in foldStream defState yieldk single stop m1
672{-# INLINE minimumBy #-}
673minimumBy
674 :: (IsStream t, Monad m)
675 => (a -> a -> Ordering) -> t m a -> m (Maybe a)
676minimumBy cmp m = go Nothing m
677 where
678 go Nothing m1 =
679 let stop = return Nothing
680 single a = return (Just a)
681 yieldk a r = go (Just a) r
682 in foldStream defState yieldk single stop m1
684 go (Just res) m1 =
685 let stop = return (Just res)
686 single a = case cmp res a of
687 GT -> return (Just a)
688 _ -> return (Just res)
689 yieldk a r = case cmp res a of
690 GT -> go (Just a) r
691 _ -> go (Just res) r
692 in foldStream defState yieldk single stop m1
694{-# INLINE maximum #-}
695maximum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
696maximum m = go Nothing m
697 where
698 go Nothing m1 =
699 let stop = return Nothing
700 single a = return (Just a)
701 yieldk a r = go (Just a) r
702 in foldStream defState yieldk single stop m1
704 go (Just res) m1 =
705 let stop = return (Just res)
706 single a =
707 if res <= a
708 then return (Just a)
709 else return (Just res)
710 yieldk a r =
711 if res <= a
712 then go (Just a) r
713 else go (Just res) r
714 in foldStream defState yieldk single stop m1
716{-# INLINE maximumBy #-}
717maximumBy :: (IsStream t, Monad m) => (a -> a -> Ordering) -> t m a -> m (Maybe a)
718maximumBy cmp m = go Nothing m
719 where
720 go Nothing m1 =
721 let stop = return Nothing
722 single a = return (Just a)
723 yieldk a r = go (Just a) r
724 in foldStream defState yieldk single stop m1
726 go (Just res) m1 =
727 let stop = return (Just res)
728 single a = case cmp res a of
729 GT -> return (Just res)
730 _ -> return (Just a)
731 yieldk a r = case cmp res a of
732 GT -> go (Just res) r
733 _ -> go (Just a) r
734 in foldStream defState yieldk single stop m1
736{-# INLINE (!!) #-}
737(!!) :: (IsStream t, Monad m) => t m a -> Int -> m (Maybe a)
738m !! i = go i m
739 where
740 go n m1 =
741 let single a | n == 0 = return $ Just a
742 | otherwise = return Nothing
743 yieldk a x | n < 0 = return Nothing
744 | n == 0 = return $ Just a
745 | otherwise = go (n - 1) x
746 in foldStream defState yieldk single (return Nothing) m1
748{-# INLINE lookup #-}
749lookup :: (IsStream t, Monad m, Eq a) => a -> t m (a, b) -> m (Maybe b)
750lookup e m = go m
751 where
752 go m1 =
753 let single (a, b) | a == e = return $ Just b
754 | otherwise = return Nothing
755 yieldk (a, b) x | a == e = return $ Just b
756 | otherwise = go x
757 in foldStream defState yieldk single (return Nothing) m1
759{-# INLINE findM #-}
760findM :: (IsStream t, Monad m) => (a -> m Bool) -> t m a -> m (Maybe a)
761findM p m = go m
762 where
763 go m1 =
764 let single a = do
765 b <- p a
766 if b then return $ Just a else return Nothing
767 yieldk a x = do
768 b <- p a
769 if b then return $ Just a else go x
770 in foldStream defState yieldk single (return Nothing) m1
772{-# INLINE find #-}
773find :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m (Maybe a)
774find p = findM (return . p)
776{-# INLINE findIndices #-}
777findIndices :: IsStream t => (a -> Bool) -> t m a -> t m Int
778findIndices p = go 0
779 where
780 go offset m1 = mkStream $ \st yld sng stp ->
781 let single a | p a = sng offset
782 | otherwise = stp
783 yieldk a x | p a = yld offset $ go (offset + 1) x
784 | otherwise = foldStream (adaptState st) yld sng stp $
785 go (offset + 1) x
786 in foldStream (adaptState st) yieldk single stp m1
788------------------------------------------------------------------------------
789-- Map and Fold
790------------------------------------------------------------------------------
792-- | Apply a monadic action to each element of the stream and discard the
793-- output of the action.
794{-# INLINE mapM_ #-}
795mapM_ :: (IsStream t, Monad m) => (a -> m b) -> t m a -> m ()
796mapM_ f m = go m
797 where
798 go m1 =
799 let stop = return ()
800 single a = void (f a)
801 yieldk a r = f a >> go r
802 in foldStream defState yieldk single stop m1
804------------------------------------------------------------------------------
805-- Converting folds
806------------------------------------------------------------------------------
808{-# INLINABLE toList #-}
809toList :: (IsStream t, Monad m) => t m a -> m [a]
810toList = foldr (:) []
812{-# INLINE toStreamK #-}
813toStreamK :: Stream m a -> Stream m a
814toStreamK = id
816-- Based on suggestions by David Feuer and Pranay Sashank
817{-# INLINE hoist #-}
818hoist :: (IsStream t, Monad m, Monad n)
819 => (forall x. m x -> n x) -> t m a -> t n a
820hoist f str =
821 mkStream $ \st yld sng stp ->
822 let single = return . sng
823 yieldk a s = return $ yld a (hoist f s)
824 stop = return stp
825 state = adaptState st
826 in join . f $ foldStreamShared state yieldk single stop str
828-------------------------------------------------------------------------------
829-- Transformation by folding (Scans)
830-------------------------------------------------------------------------------
832{-# INLINE scanlx' #-}
833scanlx' :: IsStream t => (x -> a -> x) -> x -> (x -> b) -> t m a -> t m b
834scanlx' step begin done m =
835 cons (done begin) $ go m begin
836 where
837 go m1 !acc = mkStream $ \st yld sng stp ->
838 let single a = sng (done $ step acc a)
839 yieldk a r =
840 let s = step acc a
841 in yld (done s) (go r s)
842 in foldStream (adaptState st) yieldk single stp m1
844{-# INLINE scanl' #-}
845scanl' :: IsStream t => (b -> a -> b) -> b -> t m a -> t m b
846scanl' step begin = scanlx' step begin id
848-------------------------------------------------------------------------------
849-- Filtering
850-------------------------------------------------------------------------------
852{-# INLINE filter #-}
853filter :: IsStream t => (a -> Bool) -> t m a -> t m a
854filter p m = go m
855 where
856 go m1 = mkStream $ \st yld sng stp ->
857 let single a | p a = sng a
858 | otherwise = stp
859 yieldk a r | p a = yld a (go r)
860 | otherwise = foldStream st yieldk single stp r
861 in foldStream st yieldk single stp m1
863{-# INLINE take #-}
864take :: IsStream t => Int -> t m a -> t m a
865take n m = go n m
866 where
867 go n1 m1 = mkStream $ \st yld sng stp ->
868 let yieldk a r = yld a (go (n1 - 1) r)
869 in if n1 <= 0
870 then stp
871 else foldStream st yieldk sng stp m1
873{-# INLINE takeWhile #-}
874takeWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
875takeWhile p m = go m
876 where
877 go m1 = mkStream $ \st yld sng stp ->
878 let single a | p a = sng a
879 | otherwise = stp
880 yieldk a r | p a = yld a (go r)
881 | otherwise = stp
882 in foldStream st yieldk single stp m1
884{-# INLINE drop #-}
885drop :: IsStream t => Int -> t m a -> t m a
886drop n m = fromStream $ unShare (go n (toStream m))
887 where
888 go n1 m1 = mkStream $ \st yld sng stp ->
889 let single _ = stp
890 yieldk _ r = foldStreamShared st yld sng stp $ go (n1 - 1) r
891 -- Somehow "<=" check performs better than a ">"
892 in if n1 <= 0
893 then foldStreamShared st yld sng stp m1
894 else foldStreamShared st yieldk single stp m1
896{-# INLINE dropWhile #-}
897dropWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
898dropWhile p m = go m
899 where
900 go m1 = mkStream $ \st yld sng stp ->
901 let single a | p a = stp
902 | otherwise = sng a
903 yieldk a r | p a = foldStream st yieldk single stp r
904 | otherwise = yld a r
905 in foldStream st yieldk single stp m1
907-------------------------------------------------------------------------------
908-- Mapping
909-------------------------------------------------------------------------------
911-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
912-- other stream types to overload it.
913{-# INLINE sequence #-}
914sequence :: (IsStream t, MonadAsync m) => t m (m a) -> t m a
915sequence m = go m
916 where
917 go m1 = mkStream $ \st yld sng stp ->
918 let single ma = ma >>= sng
919 yieldk ma r = foldStreamShared st yld sng stp $ ma |: go r
920 in foldStream (adaptState st) yieldk single stp m1
922-------------------------------------------------------------------------------
923-- Inserting
924-------------------------------------------------------------------------------
926{-# INLINE intersperseM #-}
927intersperseM :: (IsStream t, MonadAsync m) => m a -> t m a -> t m a
928intersperseM a m = prependingStart m
929 where
930 prependingStart m1 = mkStream $ \st yld sng stp ->
931 let yieldk i x = foldStreamShared st yld sng stp $ return i |: go x
932 in foldStream st yieldk sng stp m1
933 go m2 = mkStream $ \st yld sng stp ->
934 let single i = foldStreamShared st yld sng stp $ a |: fromPure i
935 yieldk i x = foldStreamShared st yld sng stp $ a |: return i |: go x
936 in foldStream st yieldk single stp m2
938{-# INLINE intersperse #-}
939intersperse :: (IsStream t, MonadAsync m) => a -> t m a -> t m a
940intersperse a = intersperseM (return a)
942{-# INLINE insertBy #-}
943insertBy :: IsStream t => (a -> a -> Ordering) -> a -> t m a -> t m a
944insertBy cmp x m = go m
945 where
946 go m1 = mkStream $ \st yld _ _ ->
947 let single a = case cmp x a of
948 GT -> yld a (fromPure x)
949 _ -> yld x (fromPure a)
950 stop = yld x nil
951 yieldk a r = case cmp x a of
952 GT -> yld a (go r)
953 _ -> yld x (a `cons` r)
954 in foldStream st yieldk single stop m1
956------------------------------------------------------------------------------
957-- Deleting
958------------------------------------------------------------------------------
960{-# INLINE deleteBy #-}
961deleteBy :: IsStream t => (a -> a -> Bool) -> a -> t m a -> t m a
962deleteBy eq x m = go m
963 where
964 go m1 = mkStream $ \st yld sng stp ->
965 let single a = if eq x a then stp else sng a
966 yieldk a r = if eq x a
967 then foldStream st yld sng stp r
968 else yld a (go r)
969 in foldStream st yieldk single stp m1
971------------------------------------------------------------------------------
972-- Reordering
973------------------------------------------------------------------------------
975{-# INLINE reverse #-}
976reverse :: IsStream t => t m a -> t m a
977reverse = foldlS (flip cons) nil
979-------------------------------------------------------------------------------
980-- Map and Filter
981-------------------------------------------------------------------------------
983{-# INLINE mapMaybe #-}
984mapMaybe :: IsStream t => (a -> Maybe b) -> t m a -> t m b
985mapMaybe f m = go m
986 where
987 go m1 = mkStream $ \st yld sng stp ->
988 let single a = case f a of
989 Just b -> sng b
990 Nothing -> stp
991 yieldk a r = case f a of
992 Just b -> yld b $ go r
993 Nothing -> foldStream (adaptState st) yieldk single stp r
994 in foldStream (adaptState st) yieldk single stp m1
996------------------------------------------------------------------------------
997-- Serial Zipping
998------------------------------------------------------------------------------
1000-- | Zip two streams serially using a pure zipping function.
1002-- @since 0.1.0
1003{-# INLINABLE zipWith #-}
1004zipWith :: IsStream t => (a -> b -> c) -> t m a -> t m b -> t m c
1005zipWith f = go
1006 where
1007 go mx my = mkStream $ \st yld sng stp -> do
1008 let merge a ra =
1009 let single2 b = sng (f a b)
1010 yield2 b rb = yld (f a b) (go ra rb)
1011 in foldStream (adaptState st) yield2 single2 stp my
1012 let single1 a = merge a nil
1013 yield1 = merge
1014 foldStream (adaptState st) yield1 single1 stp mx
1016-- | Zip two streams serially using a monadic zipping function.
1018-- @since 0.1.0
1019{-# INLINABLE zipWithM #-}
1020zipWithM :: (IsStream t, Monad m) => (a -> b -> m c) -> t m a -> t m b -> t m c
1021zipWithM f m1 m2 = go m1 m2
1022 where
1023 go mx my = mkStream $ \st yld sng stp -> do
1024 let merge a ra =
1025 let runIt x = foldStream st yld sng stp x
1026 single2 b = f a b >>= sng
1027 yield2 b rb = f a b >>= \x -> runIt (x `cons` go ra rb)
1028 in foldStream (adaptState st) yield2 single2 stp my
1029 let single1 a = merge a nil
1030 yield1 = merge
1031 foldStream (adaptState st) yield1 single1 stp mx
1033------------------------------------------------------------------------------
1034-- Merging
1035------------------------------------------------------------------------------
1037{-# INLINE mergeByM #-}
1038mergeByM
1039 :: (IsStream t, Monad m)
1040 => (a -> a -> m Ordering) -> t m a -> t m a -> t m a
1041mergeByM cmp = go
1042 where
1043 go mx my = mkStream $ \st yld sng stp -> do
1044 let mergeWithY a ra =
1045 let stop2 = foldStream st yld sng stp mx
1046 single2 b = do
1047 r <- cmp a b
1048 case r of
1049 GT -> yld b (go (a `cons` ra) nil)
1050 _ -> yld a (go ra (b `cons` nil))
1051 yield2 b rb = do
1052 r <- cmp a b
1053 case r of
1054 GT -> yld b (go (a `cons` ra) rb)
1055 _ -> yld a (go ra (b `cons` rb))
1056 in foldStream st yield2 single2 stop2 my
1057 let stopX = foldStream st yld sng stp my
1058 singleX a = mergeWithY a nil
1059 yieldX = mergeWithY
1060 foldStream st yieldX singleX stopX mx
1062{-# INLINABLE mergeBy #-}
1063mergeBy
1064 :: (IsStream t, Monad m)
1065 => (a -> a -> Ordering) -> t m a -> t m a -> t m a
1066mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
1068------------------------------------------------------------------------------
1069-- Transformation comprehensions
1070------------------------------------------------------------------------------
1072{-# INLINE the #-}
1073the :: (Eq a, IsStream t, Monad m) => t m a -> m (Maybe a)
1074the m = do
1075 r <- uncons m
1076 case r of
1077 Nothing -> return Nothing
1078 Just (h, t) -> go h t
1079 where
1080 go h m1 =
1081 let single a | h == a = return $ Just h
1082 | otherwise = return Nothing
1083 yieldk a r | h == a = go h r
1084 | otherwise = return Nothing
1085 in foldStream defState yieldk single (return $ Just h) m1
1087------------------------------------------------------------------------------
1088-- Alternative & MonadPlus
1089------------------------------------------------------------------------------
1091_alt :: Stream m a -> Stream m a -> Stream m a
1092_alt m1 m2 = mkStream $ \st yld sng stp ->
1093 let stop = foldStream st yld sng stp m2
1094 in foldStream st yld sng stop m1
1096------------------------------------------------------------------------------
1097-- MonadReader
1098------------------------------------------------------------------------------
1100{-# INLINABLE withLocal #-}
1101withLocal :: MonadReader r m => (r -> r) -> Stream m a -> Stream m a
1102withLocal f m =
1103 mkStream $ \st yld sng stp ->
1104 let single = local f . sng
1105 yieldk a r = local f $ yld a (withLocal f r)
1106 in foldStream st yieldk single (local f stp) m
1108------------------------------------------------------------------------------
1109-- MonadError
1110------------------------------------------------------------------------------
1113-- XXX handle and test cross thread state transfer
1114withCatchError
1115 :: MonadError e m
1116 => Stream m a -> (e -> Stream m a) -> Stream m a
1117withCatchError m h =
1118 mkStream $ \_ stp sng yld ->
1119 let run x = unStream x Nothing stp sng yieldk
1120 handle r = r `catchError` \e -> run $ h e
1121 yieldk a r = yld a (withCatchError r h)
1122 in handle $ run m