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