git.iko.soy/streamly
on 03218ad help sign in
streamly/src/Streamly/Internal/Data/Stream/StreamK.hs
1127 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 Prelude
194 hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,
195 take, filter, all, any, takeWhile, drop, dropWhile, minimum,
196 maximum, elem, notElem, null, head, tail, init, zipWith, lookup,
197 foldr1, (!!), replicate, reverse, concatMap, iterate)
198import qualified Prelude
200import Streamly.Internal.Data.SVar
201import Streamly.Internal.Data.Stream.StreamK.Type
203import qualified Streamly.Internal.Data.Fold.Type as FL
205-- $setup
206-- >>> :m
208-------------------------------------------------------------------------------
209-- Deconstruction
210-------------------------------------------------------------------------------
212{-# INLINE uncons #-}
213uncons :: (IsStream t, Monad m) => t m a -> m (Maybe (a, t m a))
214uncons m =
215 let stop = return Nothing
216 single a = return (Just (a, nil))
217 yieldk a r = return (Just (a, r))
218 in foldStream defState yieldk single stop m
220-------------------------------------------------------------------------------
221-- Generation
222-------------------------------------------------------------------------------
224{-# INLINE unfoldr #-}
225unfoldr :: IsStream t => (b -> Maybe (a, b)) -> b -> t m a
226unfoldr next s0 = build $ \yld stp ->
227 let go s =
228 case next s of
229 Just (a, b) -> yld a (go b)
230 Nothing -> stp
231 in go s0
233{-# INLINE unfoldrM #-}
234unfoldrM :: (IsStream t, MonadAsync m) => (b -> m (Maybe (a, b))) -> b -> t m a
235unfoldrM step = go
236 where
237 go s = sharedM $ \yld _ stp -> do
238 r <- step s
239 case r of
240 Just (a, b) -> yld a (go b)
241 Nothing -> stp
243{-
244-- Generalization of concurrent streams/SVar via unfoldr.
245--
246-- Unfold a value into monadic actions and then run the resulting monadic
247-- actions to generate a stream. Since the step of generating the monadic
248-- action and running them are decoupled we can run the monadic actions
249-- cooncurrently. For example, the seed could be a list of monadic actions or a
250-- pure stream of monadic actions.
251--
252-- We can have different flavors of this depending on the stream type t. The
253-- concurrent version could be async or ahead etc. Depending on how we queue
254-- back the feedback portion b, it could be DFS or BFS style.
255--
256unfoldrA :: (IsStream t, MonadAsync m) => (b -> Maybe (m a, b)) -> b -> t m a
257unfoldrA = undefined
258-}
260-------------------------------------------------------------------------------
261-- Special generation
262-------------------------------------------------------------------------------
264-- | Same as fromEffect
265--
266-- @since 0.2.0
267{-# DEPRECATED once "Please use fromEffect instead." #-}
268{-# INLINE once #-}
269once :: (Monad m, IsStream t) => m a -> t m a
270once = fromEffect
272-- |
273-- @
274-- repeatM = fix . cons
275-- repeatM = cycle1 . fromPure
276-- @
277--
278-- Generate an infinite stream by repeating a monadic value.
279--
280-- /Pre-release/
281repeatM :: (IsStream t, MonadAsync m) => m a -> t m a
282repeatM = go
283 where go m = m |: go m
285-- Generate an infinite stream by repeating a pure value.
286--
287-- /Pre-release/
288{-# INLINE repeat #-}
289repeat :: IsStream t => a -> t m a
290repeat a = let x = cons a x in x
292{-# INLINE replicateM #-}
293replicateM :: (IsStream t, MonadAsync m) => Int -> m a -> t m a
294replicateM n m = go n
295 where
296 go cnt = if cnt <= 0 then nil else m |: go (cnt - 1)
298{-# INLINE replicate #-}
299replicate :: IsStream t => Int -> a -> t m a
300replicate n a = go n
301 where
302 go cnt = if cnt <= 0 then nil else a `cons` go (cnt - 1)
304{-# INLINE fromIndicesM #-}
305fromIndicesM :: (IsStream t, MonadAsync m) => (Int -> m a) -> t m a
306fromIndicesM gen = go 0
307 where
308 go i = mkStream $ \st stp sng yld -> do
309 foldStreamShared st stp sng yld (gen i |: go (i + 1))
311{-# INLINE fromIndices #-}
312fromIndices :: IsStream t => (Int -> a) -> t m a
313fromIndices gen = go 0
314 where
315 go n = (gen n) `cons` go (n + 1)
317{-# INLINE iterate #-}
318iterate :: IsStream t => (a -> a) -> a -> t m a
319iterate step = go
320 where
321 go !s = cons s (go (step s))
323{-# INLINE iterateM #-}
324iterateM :: (IsStream t, MonadAsync m) => (a -> m a) -> m a -> t m a
325iterateM step = go
326 where
327 go s = mkStream $ \st stp sng yld -> do
328 !next <- s
329 foldStreamShared st stp sng yld (return next |: go (step next))
331-------------------------------------------------------------------------------
332-- Conversions
333-------------------------------------------------------------------------------
335-- |
336-- @
337-- fromFoldable = 'Prelude.foldr' 'cons' 'nil'
338-- @
339--
340-- Construct a stream from a 'Foldable' containing pure values:
341--
342-- @since 0.2.0
343{-# INLINE fromFoldable #-}
344fromFoldable :: (IsStream t, Foldable f) => f a -> t m a
345fromFoldable = Prelude.foldr cons nil
347{-# INLINE fromList #-}
348fromList :: IsStream t => [a] -> t m a
349fromList = fromFoldable
351{-# INLINE fromStreamK #-}
352fromStreamK :: IsStream t => Stream m a -> t m a
353fromStreamK = fromStream
355-------------------------------------------------------------------------------
356-- Elimination by Folding
357-------------------------------------------------------------------------------
359-- | Lazy right associative fold.
360{-# INLINE foldr #-}
361foldr :: (IsStream t, Monad m) => (a -> b -> b) -> b -> t m a -> m b
362foldr step acc = foldrM (\x xs -> xs >>= \b -> return (step x b)) (return acc)
364-- | Right associative fold to an arbitrary transformer monad.
365{-# INLINE foldrT #-}
366foldrT :: (IsStream t, Monad m, Monad (s m), MonadTrans s)
367 => (a -> s m b -> s m b) -> s m b -> t m a -> s m b
368foldrT step final m = go m
369 where
370 go m1 = do
371 res <- lift $ uncons m1
372 case res of
373 Just (h, t) -> step h (go t)
374 Nothing -> final
376{-# INLINE foldr1 #-}
377foldr1 :: (IsStream t, Monad m) => (a -> a -> a) -> t m a -> m (Maybe a)
378foldr1 step m = do
379 r <- uncons m
380 case r of
381 Nothing -> return Nothing
382 Just (h, t) -> fmap Just (go h t)
383 where
384 go p m1 =
385 let stp = return p
386 single a = return $ step a p
387 yieldk a r = fmap (step p) (go a r)
388 in foldStream defState yieldk single stp m1
390-- XXX replace the recursive "go" with explicit continuations.
391-- | Like 'foldx', but with a monadic step function.
392{-# INLINABLE foldlMx' #-}
393foldlMx' :: (IsStream t, Monad m)
394 => (x -> a -> m x) -> m x -> (x -> m b) -> t m a -> m b
395foldlMx' step begin done m = go begin m
396 where
397 go !acc m1 =
398 let stop = acc >>= done
399 single a = acc >>= \b -> step b a >>= done
400 yieldk a r = acc >>= \b -> step b a >>= \x -> go (return x) r
401 in foldStream defState yieldk single stop m1
403{-# INLINABLE fold #-}
404fold :: (IsStream t, Monad m) => FL.Fold m a b -> t m a -> m b
405fold (FL.Fold step begin done) m = do
406 res <- begin
407 case res of
408 FL.Partial fs -> go fs m
409 FL.Done fb -> return fb
411 where
412 go !acc m1 =
413 let stop = done acc
414 single a = step acc a
415 >>= \case
416 FL.Partial s -> done s
417 FL.Done b1 -> return b1
418 yieldk a r = step acc a
419 >>= \case
420 FL.Partial s -> go s r
421 FL.Done b1 -> return b1
422 in foldStream defState yieldk single stop m1
424-- | Like 'foldl'' but with a monadic step function.
425{-# INLINE foldlM' #-}
426foldlM' :: (IsStream t, Monad m) => (b -> a -> m b) -> m b -> t m a -> m b
427foldlM' step begin = foldlMx' step begin return
429-- | Lazy left fold to a stream.
430{-# INLINE foldlS #-}
431foldlS :: IsStream t => (t m b -> a -> t m b) -> t m b -> t m a -> t m b
432foldlS step begin m = go begin m
433 where
434 go acc rest = mkStream $ \st yld sng stp ->
435 let run x = foldStream st yld sng stp x
436 stop = run acc
437 single a = run $ step acc a
438 yieldk a r = run $ go (step acc a) r
439 in foldStream (adaptState st) yieldk single stop rest
441-- | Lazy left fold to an arbitrary transformer monad.
442{-# INLINE foldlT #-}
443foldlT :: (IsStream t, Monad m, Monad (s m), MonadTrans s)
444 => (s m b -> a -> s m b) -> s m b -> t m a -> s m b
445foldlT step begin m = go begin m
446 where
447 go acc m1 = do
448 res <- lift $ uncons m1
449 case res of
450 Just (h, t) -> go (step acc h) t
451 Nothing -> acc
453------------------------------------------------------------------------------
454-- Specialized folds
455------------------------------------------------------------------------------
457-- XXX use foldrM to implement folds where possible
458-- XXX This (commented) definition of drain and mapM_ perform much better on
459-- some benchmarks but worse on others. Need to investigate why, may there is
460-- an optimization opportunity that we can exploit.
461-- drain = foldrM (\_ xs -> return () >> xs) (return ())
463-- |
464-- > drain = foldl' (\_ _ -> ()) ()
465-- > drain = mapM_ (\_ -> return ())
466{-# INLINE drain #-}
467drain :: (Monad m, IsStream t) => t m a -> m ()
468drain = foldrM (\_ xs -> xs) (return ())
469{-
470drain = go
471 where
472 go m1 =
473 let stop = return ()
474 single _ = return ()
475 yieldk _ r = go r
476 in foldStream defState yieldk single stop m1
477-}
479{-# INLINE null #-}
480null :: (IsStream t, Monad m) => t m a -> m Bool
481-- null = foldrM (\_ _ -> return True) (return False)
482null m =
483 let stop = return True
484 single _ = return False
485 yieldk _ _ = return False
486 in foldStream defState yieldk single stop m
488{-# INLINE head #-}
489head :: (IsStream t, Monad m) => t m a -> m (Maybe a)
490-- head = foldrM (\x _ -> return $ Just x) (return Nothing)
491head m =
492 let stop = return Nothing
493 single a = return (Just a)
494 yieldk a _ = return (Just a)
495 in foldStream defState yieldk single stop m
497{-# INLINE tail #-}
498tail :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
499tail m =
500 let stop = return Nothing
501 single _ = return $ Just nil
502 yieldk _ r = return $ Just r
503 in foldStream defState yieldk single stop m
505{-# INLINE headPartial #-}
506headPartial :: (IsStream t, Monad m) => t m a -> m a
507headPartial = foldrM (\x _ -> return x) (error "head of nil")
509{-# INLINE tailPartial #-}
510tailPartial :: IsStream t => t m a -> t m a
511tailPartial m = mkStream $ \st yld sng stp ->
512 let stop = error "tail of nil"
513 single _ = stp
514 yieldk _ r = foldStream st yld sng stp r
515 in foldStream st yieldk single stop m
517-- | We can define cyclic structures using @let@:
518--
519-- >>> let (a, b) = ([1, b], head a) in (a, b)
520-- ([1,1],1)
521--
522-- The function @fix@ defined as:
523--
524-- > fix f = let x = f x in x
525--
526-- ensures that the argument of a function and its output refer to the same
527-- lazy value @x@ i.e. the same location in memory. Thus @x@ can be defined
528-- in terms of itself, creating structures with cyclic references.
529--
530-- >>> import Data.Function (fix)
531-- >>> f ~(a, b) = ([1, b], head a)
532-- >>> fix f
533-- ([1,1],1)
534--
535-- 'Control.Monad.mfix' is essentially the same as @fix@ but for monadic
536-- values.
537--
538-- Using 'mfix' for streams we can construct a stream in which each element of
539-- the stream is defined in a cyclic fashion. The argument of the function
540-- being fixed represents the current element of the stream which is being
541-- returned by the stream monad. Thus, we can use the argument to construct
542-- itself.
543--
544-- In the following example, the argument @action@ of the function @f@
545-- represents the tuple @(x,y)@ returned by it in a given iteration. We define
546-- the first element of the tuple in terms of the second.
547--
548-- @
549-- import Streamly.Internal.Data.Stream.IsStream as Stream
550-- import System.IO.Unsafe (unsafeInterleaveIO)
551--
552-- main = do
553-- Stream.mapM_ print $ Stream.mfix f
554--
555-- where
556--
557-- f action = do
558-- let incr n act = fmap ((+n) . snd) $ unsafeInterleaveIO act
559-- x <- Stream.fromListM [incr 1 action, incr 2 action]
560-- y <- Stream.fromList [4,5]
561-- return (x, y)
562-- @
563--
564-- Note: you cannot achieve this by just changing the order of the monad
565-- statements because that would change the order in which the stream elements
566-- are generated.
567--
568-- Note that the function @f@ must be lazy in its argument, that's why we use
569-- 'unsafeInterleaveIO' on @action@ because IO monad is strict.
570--
571-- /Pre-release/
573mfix :: (IsStream t, Monad m) => (m a -> t m a) -> t m a
574mfix f = mkStream $ \st yld sng stp ->
575 let single a = foldStream st yld sng stp $ a `cons` ys
576 yieldk a _ = foldStream st yld sng stp $ a `cons` ys
577 in foldStream st yieldk single stp xs
579 where
581 -- fix the head element of the stream
582 xs = fix (f . headPartial)
584 -- now fix the tail recursively
585 ys = mfix (tailPartial . f)
587{-# INLINE init #-}
588init :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
589init m = go1 m
590 where
591 go1 m1 = do
592 r <- uncons m1
593 case r of
594 Nothing -> return Nothing
595 Just (h, t) -> return . Just $ go h t
596 go p m1 = mkStream $ \_ yld sng stp ->
597 let single _ = sng p
598 yieldk a x = yld p $ go a x
599 in foldStream defState yieldk single stp m1
601{-# INLINE elem #-}
602elem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
603elem e m = go m
604 where
605 go m1 =
606 let stop = return False
607 single a = return (a == e)
608 yieldk a r = if a == e then return True else go r
609 in foldStream defState yieldk single stop m1
611{-# INLINE notElem #-}
612notElem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
613notElem e m = go m
614 where
615 go m1 =
616 let stop = return True
617 single a = return (a /= e)
618 yieldk a r = if a == e then return False else go r
619 in foldStream defState yieldk single stop m1
621{-# INLINABLE all #-}
622all :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
623all p m = go m
624 where
625 go m1 =
626 let single a | p a = return True
627 | otherwise = return False
628 yieldk a r | p a = go r
629 | otherwise = return False
630 in foldStream defState yieldk single (return True) m1
632{-# INLINABLE any #-}
633any :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
634any p m = go m
635 where
636 go m1 =
637 let single a | p a = return True
638 | otherwise = return False
639 yieldk a r | p a = return True
640 | otherwise = go r
641 in foldStream defState yieldk single (return False) m1
643-- | Extract the last element of the stream, if any.
644{-# INLINE last #-}
645last :: (IsStream t, Monad m) => t m a -> m (Maybe a)
646last = foldlx' (\_ y -> Just y) Nothing id
648{-# INLINE minimum #-}
649minimum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
650minimum m = go Nothing m
651 where
652 go Nothing m1 =
653 let stop = return Nothing
654 single a = return (Just a)
655 yieldk a r = go (Just a) r
656 in foldStream defState yieldk single stop m1
658 go (Just res) m1 =
659 let stop = return (Just res)
660 single a =
661 if res <= a
662 then return (Just res)
663 else return (Just a)
664 yieldk a r =
665 if res <= a
666 then go (Just res) r
667 else go (Just a) r
668 in foldStream defState yieldk single stop m1
670{-# INLINE minimumBy #-}
671minimumBy
672 :: (IsStream t, Monad m)
673 => (a -> a -> Ordering) -> t m a -> m (Maybe a)
674minimumBy cmp m = go Nothing m
675 where
676 go Nothing m1 =
677 let stop = return Nothing
678 single a = return (Just a)
679 yieldk a r = go (Just a) r
680 in foldStream defState yieldk single stop m1
682 go (Just res) m1 =
683 let stop = return (Just res)
684 single a = case cmp res a of
685 GT -> return (Just a)
686 _ -> return (Just res)
687 yieldk a r = case cmp res a of
688 GT -> go (Just a) r
689 _ -> go (Just res) r
690 in foldStream defState yieldk single stop m1
692{-# INLINE maximum #-}
693maximum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
694maximum m = go Nothing m
695 where
696 go Nothing m1 =
697 let stop = return Nothing
698 single a = return (Just a)
699 yieldk a r = go (Just a) r
700 in foldStream defState yieldk single stop m1
702 go (Just res) m1 =
703 let stop = return (Just res)
704 single a =
705 if res <= a
706 then return (Just a)
707 else return (Just res)
708 yieldk a r =
709 if res <= a
710 then go (Just a) r
711 else go (Just res) r
712 in foldStream defState yieldk single stop m1
714{-# INLINE maximumBy #-}
715maximumBy :: (IsStream t, Monad m) => (a -> a -> Ordering) -> t m a -> m (Maybe a)
716maximumBy cmp m = go Nothing m
717 where
718 go Nothing m1 =
719 let stop = return Nothing
720 single a = return (Just a)
721 yieldk a r = go (Just a) r
722 in foldStream defState yieldk single stop m1
724 go (Just res) m1 =
725 let stop = return (Just res)
726 single a = case cmp res a of
727 GT -> return (Just res)
728 _ -> return (Just a)
729 yieldk a r = case cmp res a of
730 GT -> go (Just res) r
731 _ -> go (Just a) r
732 in foldStream defState yieldk single stop m1
734{-# INLINE (!!) #-}
735(!!) :: (IsStream t, Monad m) => t m a -> Int -> m (Maybe a)
736m !! i = go i m
737 where
738 go n m1 =
739 let single a | n == 0 = return $ Just a
740 | otherwise = return Nothing
741 yieldk a x | n < 0 = return Nothing
742 | n == 0 = return $ Just a
743 | otherwise = go (n - 1) x
744 in foldStream defState yieldk single (return Nothing) m1
746{-# INLINE lookup #-}
747lookup :: (IsStream t, Monad m, Eq a) => a -> t m (a, b) -> m (Maybe b)
748lookup e m = go m
749 where
750 go m1 =
751 let single (a, b) | a == e = return $ Just b
752 | otherwise = return Nothing
753 yieldk (a, b) x | a == e = return $ Just b
754 | otherwise = go x
755 in foldStream defState yieldk single (return Nothing) m1
757{-# INLINE findM #-}
758findM :: (IsStream t, Monad m) => (a -> m Bool) -> t m a -> m (Maybe a)
759findM p m = go m
760 where
761 go m1 =
762 let single a = do
763 b <- p a
764 if b then return $ Just a else return Nothing
765 yieldk a x = do
766 b <- p a
767 if b then return $ Just a else go x
768 in foldStream defState yieldk single (return Nothing) m1
770{-# INLINE find #-}
771find :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m (Maybe a)
772find p = findM (return . p)
774{-# INLINE findIndices #-}
775findIndices :: IsStream t => (a -> Bool) -> t m a -> t m Int
776findIndices p = go 0
777 where
778 go offset m1 = mkStream $ \st yld sng stp ->
779 let single a | p a = sng offset
780 | otherwise = stp
781 yieldk a x | p a = yld offset $ go (offset + 1) x
782 | otherwise = foldStream (adaptState st) yld sng stp $
783 go (offset + 1) x
784 in foldStream (adaptState st) yieldk single stp m1
786------------------------------------------------------------------------------
787-- Map and Fold
788------------------------------------------------------------------------------
790-- | Apply a monadic action to each element of the stream and discard the
791-- output of the action.
792{-# INLINE mapM_ #-}
793mapM_ :: (IsStream t, Monad m) => (a -> m b) -> t m a -> m ()
794mapM_ f m = go m
795 where
796 go m1 =
797 let stop = return ()
798 single a = void (f a)
799 yieldk a r = f a >> go r
800 in foldStream defState yieldk single stop m1
802------------------------------------------------------------------------------
803-- Converting folds
804------------------------------------------------------------------------------
806{-# INLINABLE toList #-}
807toList :: (IsStream t, Monad m) => t m a -> m [a]
808toList = foldr (:) []
810{-# INLINE toStreamK #-}
811toStreamK :: Stream m a -> Stream m a
812toStreamK = id
814-- Based on suggestions by David Feuer and Pranay Sashank
815{-# INLINE hoist #-}
816hoist :: (IsStream t, Monad m, Monad n)
817 => (forall x. m x -> n x) -> t m a -> t n a
818hoist f str =
819 mkStream $ \st yld sng stp ->
820 let single = return . sng
821 yieldk a s = return $ yld a (hoist f s)
822 stop = return stp
823 state = adaptState st
824 in join . f $ foldStreamShared state yieldk single stop str
826-------------------------------------------------------------------------------
827-- Transformation by folding (Scans)
828-------------------------------------------------------------------------------
830{-# INLINE scanlx' #-}
831scanlx' :: IsStream t => (x -> a -> x) -> x -> (x -> b) -> t m a -> t m b
832scanlx' step begin done m =
833 cons (done begin) $ go m begin
834 where
835 go m1 !acc = mkStream $ \st yld sng stp ->
836 let single a = sng (done $ step acc a)
837 yieldk a r =
838 let s = step acc a
839 in yld (done s) (go r s)
840 in foldStream (adaptState st) yieldk single stp m1
842{-# INLINE scanl' #-}
843scanl' :: IsStream t => (b -> a -> b) -> b -> t m a -> t m b
844scanl' step begin = scanlx' step begin id
846-------------------------------------------------------------------------------
847-- Filtering
848-------------------------------------------------------------------------------
850{-# INLINE filter #-}
851filter :: IsStream t => (a -> Bool) -> t m a -> t m a
852filter p m = go m
853 where
854 go m1 = mkStream $ \st yld sng stp ->
855 let single a | p a = sng a
856 | otherwise = stp
857 yieldk a r | p a = yld a (go r)
858 | otherwise = foldStream st yieldk single stp r
859 in foldStream st yieldk single stp m1
861{-# INLINE take #-}
862take :: IsStream t => Int -> t m a -> t m a
863take n m = go n m
864 where
865 go n1 m1 = mkStream $ \st yld sng stp ->
866 let yieldk a r = yld a (go (n1 - 1) r)
867 in if n1 <= 0
868 then stp
869 else foldStream st yieldk sng stp m1
871{-# INLINE takeWhile #-}
872takeWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
873takeWhile p m = go m
874 where
875 go m1 = mkStream $ \st yld sng stp ->
876 let single a | p a = sng a
877 | otherwise = stp
878 yieldk a r | p a = yld a (go r)
879 | otherwise = stp
880 in foldStream st yieldk single stp m1
882{-# INLINE drop #-}
883drop :: IsStream t => Int -> t m a -> t m a
884drop n m = fromStream $ unShare (go n (toStream m))
885 where
886 go n1 m1 = mkStream $ \st yld sng stp ->
887 let single _ = stp
888 yieldk _ r = foldStreamShared st yld sng stp $ go (n1 - 1) r
889 -- Somehow "<=" check performs better than a ">"
890 in if n1 <= 0
891 then foldStreamShared st yld sng stp m1
892 else foldStreamShared st yieldk single stp m1
894{-# INLINE dropWhile #-}
895dropWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
896dropWhile p m = go m
897 where
898 go m1 = mkStream $ \st yld sng stp ->
899 let single a | p a = stp
900 | otherwise = sng a
901 yieldk a r | p a = foldStream st yieldk single stp r
902 | otherwise = yld a r
903 in foldStream st yieldk single stp m1
905-------------------------------------------------------------------------------
906-- Mapping
907-------------------------------------------------------------------------------
909-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
910-- other stream types to overload it.
911{-# INLINE sequence #-}
912sequence :: (IsStream t, MonadAsync m) => t m (m a) -> t m a
913sequence m = go m
914 where
915 go m1 = mkStream $ \st yld sng stp ->
916 let single ma = ma >>= sng
917 yieldk ma r = foldStreamShared st yld sng stp $ ma |: go r
918 in foldStream (adaptState st) yieldk single stp m1
920-------------------------------------------------------------------------------
921-- Inserting
922-------------------------------------------------------------------------------
924{-# INLINE intersperseM #-}
925intersperseM :: (IsStream t, MonadAsync m) => m a -> t m a -> t m a
926intersperseM a m = prependingStart m
927 where
928 prependingStart m1 = mkStream $ \st yld sng stp ->
929 let yieldk i x = foldStreamShared st yld sng stp $ return i |: go x
930 in foldStream st yieldk sng stp m1
931 go m2 = mkStream $ \st yld sng stp ->
932 let single i = foldStreamShared st yld sng stp $ a |: fromPure i
933 yieldk i x = foldStreamShared st yld sng stp $ a |: return i |: go x
934 in foldStream st yieldk single stp m2
936{-# INLINE intersperse #-}
937intersperse :: (IsStream t, MonadAsync m) => a -> t m a -> t m a
938intersperse a = intersperseM (return a)
940{-# INLINE insertBy #-}
941insertBy :: IsStream t => (a -> a -> Ordering) -> a -> t m a -> t m a
942insertBy cmp x m = go m
943 where
944 go m1 = mkStream $ \st yld _ _ ->
945 let single a = case cmp x a of
946 GT -> yld a (fromPure x)
947 _ -> yld x (fromPure a)
948 stop = yld x nil
949 yieldk a r = case cmp x a of
950 GT -> yld a (go r)
951 _ -> yld x (a `cons` r)
952 in foldStream st yieldk single stop m1
954------------------------------------------------------------------------------
955-- Deleting
956------------------------------------------------------------------------------
958{-# INLINE deleteBy #-}
959deleteBy :: IsStream t => (a -> a -> Bool) -> a -> t m a -> t m a
960deleteBy eq x m = go m
961 where
962 go m1 = mkStream $ \st yld sng stp ->
963 let single a = if eq x a then stp else sng a
964 yieldk a r = if eq x a
965 then foldStream st yld sng stp r
966 else yld a (go r)
967 in foldStream st yieldk single stp m1
969------------------------------------------------------------------------------
970-- Reordering
971------------------------------------------------------------------------------
973{-# INLINE reverse #-}
974reverse :: IsStream t => t m a -> t m a
975reverse = foldlS (flip cons) nil
977-------------------------------------------------------------------------------
978-- Map and Filter
979-------------------------------------------------------------------------------
981{-# INLINE mapMaybe #-}
982mapMaybe :: IsStream t => (a -> Maybe b) -> t m a -> t m b
983mapMaybe f m = go m
984 where
985 go m1 = mkStream $ \st yld sng stp ->
986 let single a = case f a of
987 Just b -> sng b
988 Nothing -> stp
989 yieldk a r = case f a of
990 Just b -> yld b $ go r
991 Nothing -> foldStream (adaptState st) yieldk single stp r
992 in foldStream (adaptState st) yieldk single stp m1
994------------------------------------------------------------------------------
995-- Serial Zipping
996------------------------------------------------------------------------------
998-- XXX We can probably implement zipWith in terms of zipWithM
999-- | Zip two streams serially using a pure zipping function. The zipping
1000-- function is applied concurrently for concurrent streams.
1002-- @since 0.1.0
1003{-# INLINABLE zipWith #-}
1004zipWith ::
1005 (IsStream t, MonadAsync m) => (a -> b -> c) -> t m a -> t m b -> t m c
1006zipWith f = go
1007 where
1008 go mx my = mkStream $ \st yld sng stp -> do
1009 let merge a ra =
1010 let single2 b = sng (f a b)
1011 runIt x = foldStream st yld sng stp x
1012 yield2 b rb = runIt (return (f a b) `consM` go ra rb)
1013 in foldStream (adaptState st) yield2 single2 stp my
1014 let single1 a = merge a nil
1015 yield1 = merge
1016 foldStream (adaptState st) yield1 single1 stp mx
1018-- | Zip two streams serially using a monadic zipping function. The zipping
1019-- function is applied concurrently for concurrent streams.
1021-- @since 0.1.0
1022{-# INLINABLE zipWithM #-}
1023zipWithM ::
1024 (IsStream t, MonadAsync m) => (a -> b -> m c) -> t m a -> t m b -> t m c
1025zipWithM f m1 m2 = go m1 m2
1026 where
1027 go mx my = mkStream $ \st yld sng stp -> do
1028 let merge a ra =
1029 let runIt x = foldStream st yld sng stp x
1030 single2 b = f a b >>= sng
1031 yield2 b rb = runIt (f a b `consM` go ra rb)
1032 in foldStream (adaptState st) yield2 single2 stp my
1033 let single1 a = merge a nil
1034 yield1 = merge
1035 foldStream (adaptState st) yield1 single1 stp mx
1037------------------------------------------------------------------------------
1038-- Merging
1039------------------------------------------------------------------------------
1041{-# INLINE mergeByM #-}
1042mergeByM
1043 :: (IsStream t, Monad m)
1044 => (a -> a -> m Ordering) -> t m a -> t m a -> t m a
1045mergeByM cmp = go
1046 where
1047 go mx my = mkStream $ \st yld sng stp -> do
1048 let mergeWithY a ra =
1049 let stop2 = foldStream st yld sng stp mx
1050 single2 b = do
1051 r <- cmp a b
1052 case r of
1053 GT -> yld b (go (a `cons` ra) nil)
1054 _ -> yld a (go ra (b `cons` nil))
1055 yield2 b rb = do
1056 r <- cmp a b
1057 case r of
1058 GT -> yld b (go (a `cons` ra) rb)
1059 _ -> yld a (go ra (b `cons` rb))
1060 in foldStream st yield2 single2 stop2 my
1061 let stopX = foldStream st yld sng stp my
1062 singleX a = mergeWithY a nil
1063 yieldX = mergeWithY
1064 foldStream st yieldX singleX stopX mx
1066{-# INLINABLE mergeBy #-}
1067mergeBy
1068 :: (IsStream t, Monad m)
1069 => (a -> a -> Ordering) -> t m a -> t m a -> t m a
1070mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
1072------------------------------------------------------------------------------
1073-- Transformation comprehensions
1074------------------------------------------------------------------------------
1076{-# INLINE the #-}
1077the :: (Eq a, IsStream t, Monad m) => t m a -> m (Maybe a)
1078the m = do
1079 r <- uncons m
1080 case r of
1081 Nothing -> return Nothing
1082 Just (h, t) -> go h t
1083 where
1084 go h m1 =
1085 let single a | h == a = return $ Just h
1086 | otherwise = return Nothing
1087 yieldk a r | h == a = go h r
1088 | otherwise = return Nothing
1089 in foldStream defState yieldk single (return $ Just h) m1
1091------------------------------------------------------------------------------
1092-- Alternative & MonadPlus
1093------------------------------------------------------------------------------
1095_alt :: Stream m a -> Stream m a -> Stream m a
1096_alt m1 m2 = mkStream $ \st yld sng stp ->
1097 let stop = foldStream st yld sng stp m2
1098 in foldStream st yld sng stop m1
1100------------------------------------------------------------------------------
1101-- MonadReader
1102------------------------------------------------------------------------------
1104{-# INLINABLE withLocal #-}
1105withLocal :: MonadReader r m => (r -> r) -> Stream m a -> Stream m a
1106withLocal f m =
1107 mkStream $ \st yld sng stp ->
1108 let single = local f . sng
1109 yieldk a r = local f $ yld a (withLocal f r)
1110 in foldStream st yieldk single (local f stp) m
1112------------------------------------------------------------------------------
1113-- MonadError
1114------------------------------------------------------------------------------
1117-- XXX handle and test cross thread state transfer
1118withCatchError
1119 :: MonadError e m
1120 => Stream m a -> (e -> Stream m a) -> Stream m a
1121withCatchError m h =
1122 mkStream $ \_ stp sng yld ->
1123 let run x = unStream x Nothing stp sng yieldk
1124 handle r = r `catchError` \e -> run $ h e
1125 yieldk a r = yld a (withCatchError r h)
1126 in handle $ run m