git.iko.soy/streamly
on 44e3cb3 help sign in
streamly/src/Streamly/Streams/StreamK.hs
1024 lines 32 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE CPP #-}
3{-# LANGUAGE ConstraintKinds #-}
4{-# LANGUAGE FlexibleContexts #-}
5{-# LANGUAGE FlexibleInstances #-}
6{-# LANGUAGE InstanceSigs #-}
7{-# LANGUAGE MultiParamTypeClasses #-}
8{-# LANGUAGE RankNTypes #-}
9{-# LANGUAGE ScopedTypeVariables #-}
10{-# LANGUAGE UndecidableInstances #-} -- XXX
12#include "inline.hs"
14-- |
15-- Module : Streamly.Streams.StreamK
16-- Copyright : (c) 2017 Harendra Kumar
17--
18-- License : BSD3
19-- Maintainer : harendra.kumar@gmail.com
20-- Stability : experimental
21-- Portability : GHC
22--
23--
24-- Continuation passing style (CPS) stream implementation. The symbol 'K' below
25-- denotes a function as well as a Kontinuation.
26--
27-- @
28-- import qualified Streamly.Streams.StreamK as K
29-- @
30--
31module Streamly.Streams.StreamK
32 (
33 -- * A class for streams
34 IsStream (..)
35 , adapt
37 -- * The stream type
38 , Stream
40 -- * Construction Primitives
41 , mkStream
42 , nil
43 , nilM
44 , cons
45 , (.:)
47 -- * Elimination Primitives
48 , foldStream
49 , foldStreamShared
50 , foldStreamSVar
52 -- * Transformation Primitives
53 , unShare
55 -- * Deconstruction
56 , uncons
58 -- * Generation
59 -- ** Unfolds
60 , unfoldr
61 , unfoldrM
63 -- ** Specialized Generation
64 , repeat
65 , repeatM
66 , replicate
67 , replicateM
68 , fromIndices
69 , fromIndicesM
71 -- ** Conversions
72 , yield
73 , yieldM
74 , fromFoldable
75 , fromList
76 , fromStreamK
78 -- * foldr/build
79 , foldrS
80 , foldrSM
81 , buildS
82 , buildM
83 , augmentS
84 , augmentSM
86 -- * Elimination
87 -- ** General Folds
88 , foldr
89 , foldr1
90 , foldrM
91 , foldrT
93 , foldl'
94 , foldlM'
95 , foldlS
96 , foldlT
97 , foldlx'
98 , foldlMx'
100 -- ** Specialized Folds
101 , drain
102 , null
103 , head
104 , tail
105 , init
106 , elem
107 , notElem
108 , all
109 , any
110 , last
111 , minimum
112 , minimumBy
113 , maximum
114 , maximumBy
115 , findIndices
116 , lookup
117 , findM
118 , find
119 , (!!)
121 -- ** Map and Fold
122 , mapM_
124 -- ** Conversions
125 , toList
126 , toStreamK
128 -- * Transformation
129 -- ** By folding (scans)
130 , scanl'
131 , scanlx'
133 -- ** Filtering
134 , filter
135 , take
136 , takeWhile
137 , drop
138 , dropWhile
140 -- ** Mapping
141 , map
142 , mapM
143 , mapMSerial
144 , sequence
146 -- ** Inserting
147 , intersperseM
148 , intersperse
149 , insertBy
151 -- ** Deleting
152 , deleteBy
154 -- ** Reordering
155 , reverse
157 -- ** Map and Filter
158 , mapMaybe
160 -- ** Zipping
161 , zipWith
162 , zipWithM
164 -- ** Merging
165 , mergeBy
166 , mergeByM
168 -- ** Nesting
169 , concatMapBy
170 , concatMap
171 , bindWith
173 -- ** Transformation comprehensions
174 , the
176 -- * Semigroup Style Composition
177 , serial
179 -- * Utilities
180 , consMStream
181 , withLocal
183 -- * Deprecated
184 , Streaming -- deprecated
185 , once -- deprecated
186 )
187where
189import Control.Monad.Trans (MonadTrans(lift))
190import Control.Monad (void)
191import Control.Monad.Reader.Class (MonadReader(..))
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)
197import qualified Prelude
199import Streamly.Internal.Data.SVar
200import Streamly.Streams.StreamK.Type
202-------------------------------------------------------------------------------
203-- Deconstruction
204-------------------------------------------------------------------------------
206{-# INLINE uncons #-}
207uncons :: (IsStream t, Monad m) => t m a -> m (Maybe (a, t m a))
208uncons m =
209 let stop = return Nothing
210 single a = return (Just (a, nil))
211 yieldk a r = return (Just (a, r))
212 in foldStream defState yieldk single stop m
214-------------------------------------------------------------------------------
215-- Generation
216-------------------------------------------------------------------------------
218{-# INLINE unfoldr #-}
219unfoldr :: IsStream t => (b -> Maybe (a, b)) -> b -> t m a
220unfoldr next s0 = build $ \yld stp ->
221 let go s =
222 case next s of
223 Just (a, b) -> yld a (go b)
224 Nothing -> stp
225 in go s0
227{-# INLINE unfoldrM #-}
228unfoldrM :: (IsStream t, MonadAsync m) => (b -> m (Maybe (a, b))) -> b -> t m a
229unfoldrM step = go
230 where
231 go s = sharedM $ \yld _ stp -> do
232 r <- step s
233 case r of
234 Just (a, b) -> yld a (go b)
235 Nothing -> stp
237{-
238-- Generalization of concurrent streams/SVar via unfoldr.
239--
240-- Unfold a value into monadic actions and then run the resulting monadic
241-- actions to generate a stream. Since the step of generating the monadic
242-- action and running them are decoupled we can run the monadic actions
243-- cooncurrently. For example, the seed could be a list of monadic actions or a
244-- pure stream of monadic actions.
245--
246-- We can have different flavors of this depending on the stream type t. The
247-- concurrent version could be async or ahead etc. Depending on how we queue
248-- back the feedback portion b, it could be DFS or BFS style.
249--
250unfoldrA :: (IsStream t, MonadAsync m) => (b -> Maybe (m a, b)) -> b -> t m a
251unfoldrA = undefined
252-}
254-------------------------------------------------------------------------------
255-- Special generation
256-------------------------------------------------------------------------------
258-- | Same as yieldM
259--
260-- @since 0.2.0
261{-# DEPRECATED once "Please use yieldM instead." #-}
262{-# INLINE once #-}
263once :: (Monad m, IsStream t) => m a -> t m a
264once = yieldM
266-- |
267-- @
268-- repeatM = fix . cons
269-- repeatM = cycle1 . yield
270-- @
271--
272-- Generate an infinite stream by repeating a monadic value.
273--
274-- /Internal/
275repeatM :: (IsStream t, MonadAsync m) => m a -> t m a
276repeatM = go
277 where go m = m |: go m
279-- Generate an infinite stream by repeating a pure value.
280--
281-- /Internal/
282{-# INLINE repeat #-}
283repeat :: IsStream t => a -> t m a
284repeat a = let x = cons a x in x
286{-# INLINE replicateM #-}
287replicateM :: (IsStream t, MonadAsync m) => Int -> m a -> t m a
288replicateM n m = go n
289 where
290 go cnt = if cnt <= 0 then nil else m |: go (cnt - 1)
292{-# INLINE replicate #-}
293replicate :: IsStream t => Int -> a -> t m a
294replicate n a = go n
295 where
296 go cnt = if cnt <= 0 then nil else a `cons` go (cnt - 1)
298{-# INLINE fromIndicesM #-}
299fromIndicesM :: (IsStream t, MonadAsync m) => (Int -> m a) -> t m a
300fromIndicesM gen = go 0
301 where
302 go i = mkStream $ \st stp sng yld -> do
303 foldStreamShared st stp sng yld (gen i |: go (i + 1))
305{-# INLINE fromIndices #-}
306fromIndices :: IsStream t => (Int -> a) -> t m a
307fromIndices gen = go 0
308 where
309 go n = (gen n) `cons` go (n + 1)
311-------------------------------------------------------------------------------
312-- Conversions
313-------------------------------------------------------------------------------
315-- |
316-- @
317-- fromFoldable = 'Prelude.foldr' 'cons' 'nil'
318-- @
319--
320-- Construct a stream from a 'Foldable' containing pure values:
321--
322-- @since 0.2.0
323{-# INLINE fromFoldable #-}
324fromFoldable :: (IsStream t, Foldable f) => f a -> t m a
325fromFoldable = Prelude.foldr cons nil
327{-# INLINE fromList #-}
328fromList :: IsStream t => [a] -> t m a
329fromList = fromFoldable
331{-# INLINE fromStreamK #-}
332fromStreamK :: IsStream t => Stream m a -> t m a
333fromStreamK = fromStream
335-------------------------------------------------------------------------------
336-- Elimination by Folding
337-------------------------------------------------------------------------------
339-- | Lazy right associative fold.
340{-# INLINE foldr #-}
341foldr :: (IsStream t, Monad m) => (a -> b -> b) -> b -> t m a -> m b
342foldr step acc = foldrM (\x xs -> xs >>= \b -> return (step x b)) (return acc)
344-- | Right associative fold to an arbitrary transformer monad.
345{-# INLINE foldrT #-}
346foldrT :: (IsStream t, Monad m, Monad (s m), MonadTrans s)
347 => (a -> s m b -> s m b) -> s m b -> t m a -> s m b
348foldrT step final m = go m
349 where
350 go m1 = do
351 res <- lift $ uncons m1
352 case res of
353 Just (h, t) -> step h (go t)
354 Nothing -> final
356{-# INLINE foldr1 #-}
357foldr1 :: (IsStream t, Monad m) => (a -> a -> a) -> t m a -> m (Maybe a)
358foldr1 step m = do
359 r <- uncons m
360 case r of
361 Nothing -> return Nothing
362 Just (h, t) -> fmap Just (go h t)
363 where
364 go p m1 =
365 let stp = return p
366 single a = return $ step a p
367 yieldk a r = fmap (step p) (go a r)
368 in foldStream defState yieldk single stp m1
370-- | Strict left fold with an extraction function. Like the standard strict
371-- left fold, but applies a user supplied extraction function (the third
372-- argument) to the folded value at the end. This is designed to work with the
373-- @foldl@ library. The suffix @x@ is a mnemonic for extraction.
374--
375-- Note that the accumulator is always evaluated including the initial value.
376{-# INLINE foldlx' #-}
377foldlx' :: forall t m a b x. (IsStream t, Monad m)
378 => (x -> a -> x) -> x -> (x -> b) -> t m a -> m b
379foldlx' step begin done m = get $ go m begin
380 where
381 {-# NOINLINE get #-}
382 get :: t m x -> m b
383 get m1 =
384 -- XXX we are not strictly evaluating the accumulator here. Is this
385 -- okay?
386 let single = return . done
387 -- XXX this is foldSingleton. why foldStreamShared?
388 in foldStreamShared undefined undefined single undefined m1
390 -- Note, this can be implemented by making a recursive call to "go",
391 -- however that is more expensive because of unnecessary recursion
392 -- that cannot be tail call optimized. Unfolding recursion explicitly via
393 -- continuations is much more efficient.
394 go :: t m a -> x -> t m x
395 go m1 !acc = mkStream $ \_ yld sng _ ->
396 let stop = sng acc
397 single a = sng $ step acc a
398 -- XXX this is foldNonEmptyStream
399 yieldk a r = foldStream defState yld sng undefined $
400 go r (step acc a)
401 in foldStream defState yieldk single stop m1
403-- | Strict left associative fold.
404{-# INLINE foldl' #-}
405foldl' :: (IsStream t, Monad m) => (b -> a -> b) -> b -> t m a -> m b
406foldl' step begin = foldlx' step begin id
408-- XXX replace the recursive "go" with explicit continuations.
409-- | Like 'foldx', but with a monadic step function.
410{-# INLINABLE foldlMx' #-}
411foldlMx' :: (IsStream t, Monad m)
412 => (x -> a -> m x) -> m x -> (x -> m b) -> t m a -> m b
413foldlMx' step begin done m = go begin m
414 where
415 go !acc m1 =
416 let stop = acc >>= done
417 single a = acc >>= \b -> step b a >>= done
418 yieldk a r = acc >>= \b -> step b a >>= \x -> go (return x) r
419 in foldStream defState yieldk single stop m1
421-- | Like 'foldl'' but with a monadic step function.
422{-# INLINE foldlM' #-}
423foldlM' :: (IsStream t, Monad m) => (b -> a -> m b) -> b -> t m a -> m b
424foldlM' step begin = foldlMx' step (return begin) return
426-- | Lazy left fold to a stream.
427{-# INLINE foldlS #-}
428foldlS :: IsStream t => (t m b -> a -> t m b) -> t m b -> t m a -> t m b
429foldlS step begin m = go begin m
430 where
431 go acc rest = mkStream $ \st yld sng stp ->
432 let run x = foldStream st yld sng stp x
433 stop = run acc
434 single a = run $ step acc a
435 yieldk a r = run $ go (step acc a) r
436 in foldStream (adaptState st) yieldk single stop rest
438-- | Lazy left fold to an arbitrary transformer monad.
439{-# INLINE foldlT #-}
440foldlT :: (IsStream t, Monad m, Monad (s m), MonadTrans s)
441 => (s m b -> a -> s m b) -> s m b -> t m a -> s m b
442foldlT step begin m = go begin m
443 where
444 go acc m1 = do
445 res <- lift $ uncons m1
446 case res of
447 Just (h, t) -> go (step acc h) t
448 Nothing -> acc
450------------------------------------------------------------------------------
451-- Specialized folds
452------------------------------------------------------------------------------
454-- XXX use foldrM to implement folds where possible
455-- XXX This (commented) definition of drain and mapM_ perform much better on
456-- some benchmarks but worse on others. Need to investigate why, may there is
457-- an optimization opportunity that we can exploit.
458-- drain = foldrM (\_ xs -> return () >> xs) (return ())
460-- |
461-- > drain = foldl' (\_ _ -> ()) ()
462-- > drain = mapM_ (\_ -> return ())
463{-# INLINE drain #-}
464drain :: (Monad m, IsStream t) => t m a -> m ()
465drain = foldrM (\_ xs -> xs) (return ())
466{-
467drain = go
468 where
469 go m1 =
470 let stop = return ()
471 single _ = return ()
472 yieldk _ r = go r
473 in foldStream defState yieldk single stop m1
474-}
476{-# INLINE null #-}
477null :: (IsStream t, Monad m) => t m a -> m Bool
478-- null = foldrM (\_ _ -> return True) (return False)
479null m =
480 let stop = return True
481 single _ = return False
482 yieldk _ _ = return False
483 in foldStream defState yieldk single stop m
485{-# INLINE head #-}
486head :: (IsStream t, Monad m) => t m a -> m (Maybe a)
487-- head = foldrM (\x _ -> return $ Just x) (return Nothing)
488head m =
489 let stop = return Nothing
490 single a = return (Just a)
491 yieldk a _ = return (Just a)
492 in foldStream defState yieldk single stop m
494{-# INLINE tail #-}
495tail :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
496tail m =
497 let stop = return Nothing
498 single _ = return $ Just nil
499 yieldk _ r = return $ Just r
500 in foldStream defState yieldk single stop m
502{-# INLINE init #-}
503init :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
504init m = go1 m
505 where
506 go1 m1 = do
507 r <- uncons m1
508 case r of
509 Nothing -> return Nothing
510 Just (h, t) -> return . Just $ go h t
511 go p m1 = mkStream $ \_ yld sng stp ->
512 let single _ = sng p
513 yieldk a x = yld p $ go a x
514 in foldStream defState yieldk single stp m1
516{-# INLINE elem #-}
517elem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
518elem e m = go m
519 where
520 go m1 =
521 let stop = return False
522 single a = return (a == e)
523 yieldk a r = if a == e then return True else go r
524 in foldStream defState yieldk single stop m1
526{-# INLINE notElem #-}
527notElem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
528notElem e m = go m
529 where
530 go m1 =
531 let stop = return True
532 single a = return (a /= e)
533 yieldk a r = if a == e then return False else go r
534 in foldStream defState yieldk single stop m1
536{-# INLINABLE all #-}
537all :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
538all p m = go m
539 where
540 go m1 =
541 let single a | p a = return True
542 | otherwise = return False
543 yieldk a r | p a = go r
544 | otherwise = return False
545 in foldStream defState yieldk single (return True) m1
547{-# INLINABLE any #-}
548any :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
549any p m = go m
550 where
551 go m1 =
552 let single a | p a = return True
553 | otherwise = return False
554 yieldk a r | p a = return True
555 | otherwise = go r
556 in foldStream defState yieldk single (return False) m1
558-- | Extract the last element of the stream, if any.
559{-# INLINE last #-}
560last :: (IsStream t, Monad m) => t m a -> m (Maybe a)
561last = foldlx' (\_ y -> Just y) Nothing id
563{-# INLINE minimum #-}
564minimum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
565minimum m = go Nothing m
566 where
567 go Nothing m1 =
568 let stop = return Nothing
569 single a = return (Just a)
570 yieldk a r = go (Just a) r
571 in foldStream defState yieldk single stop m1
573 go (Just res) m1 =
574 let stop = return (Just res)
575 single a =
576 if res <= a
577 then return (Just res)
578 else return (Just a)
579 yieldk a r =
580 if res <= a
581 then go (Just res) r
582 else go (Just a) r
583 in foldStream defState yieldk single stop m1
585{-# INLINE minimumBy #-}
586minimumBy
587 :: (IsStream t, Monad m)
588 => (a -> a -> Ordering) -> t m a -> m (Maybe a)
589minimumBy cmp m = go Nothing m
590 where
591 go Nothing m1 =
592 let stop = return Nothing
593 single a = return (Just a)
594 yieldk a r = go (Just a) r
595 in foldStream defState yieldk single stop m1
597 go (Just res) m1 =
598 let stop = return (Just res)
599 single a = case cmp res a of
600 GT -> return (Just a)
601 _ -> return (Just res)
602 yieldk a r = case cmp res a of
603 GT -> go (Just a) r
604 _ -> go (Just res) r
605 in foldStream defState yieldk single stop m1
607{-# INLINE maximum #-}
608maximum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
609maximum m = go Nothing m
610 where
611 go Nothing m1 =
612 let stop = return Nothing
613 single a = return (Just a)
614 yieldk a r = go (Just a) r
615 in foldStream defState yieldk single stop m1
617 go (Just res) m1 =
618 let stop = return (Just res)
619 single a =
620 if res <= a
621 then return (Just a)
622 else return (Just res)
623 yieldk a r =
624 if res <= a
625 then go (Just a) r
626 else go (Just res) r
627 in foldStream defState yieldk single stop m1
629{-# INLINE maximumBy #-}
630maximumBy :: (IsStream t, Monad m) => (a -> a -> Ordering) -> t m a -> m (Maybe a)
631maximumBy cmp m = go Nothing m
632 where
633 go Nothing m1 =
634 let stop = return Nothing
635 single a = return (Just a)
636 yieldk a r = go (Just a) r
637 in foldStream defState yieldk single stop m1
639 go (Just res) m1 =
640 let stop = return (Just res)
641 single a = case cmp res a of
642 GT -> return (Just res)
643 _ -> return (Just a)
644 yieldk a r = case cmp res a of
645 GT -> go (Just res) r
646 _ -> go (Just a) r
647 in foldStream defState yieldk single stop m1
649{-# INLINE (!!) #-}
650(!!) :: (IsStream t, Monad m) => t m a -> Int -> m (Maybe a)
651m !! i = go i m
652 where
653 go n m1 =
654 let single a | n == 0 = return $ Just a
655 | otherwise = return Nothing
656 yieldk a x | n < 0 = return Nothing
657 | n == 0 = return $ Just a
658 | otherwise = go (n - 1) x
659 in foldStream defState yieldk single (return Nothing) m1
661{-# INLINE lookup #-}
662lookup :: (IsStream t, Monad m, Eq a) => a -> t m (a, b) -> m (Maybe b)
663lookup e m = go m
664 where
665 go m1 =
666 let single (a, b) | a == e = return $ Just b
667 | otherwise = return Nothing
668 yieldk (a, b) x | a == e = return $ Just b
669 | otherwise = go x
670 in foldStream defState yieldk single (return Nothing) m1
672{-# INLINE findM #-}
673findM :: (IsStream t, Monad m) => (a -> m Bool) -> t m a -> m (Maybe a)
674findM p m = go m
675 where
676 go m1 =
677 let single a = do
678 b <- p a
679 if b then return $ Just a else return Nothing
680 yieldk a x = do
681 b <- p a
682 if b then return $ Just a else go x
683 in foldStream defState yieldk single (return Nothing) m1
685{-# INLINE find #-}
686find :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m (Maybe a)
687find p = findM (return . p)
689{-# INLINE findIndices #-}
690findIndices :: IsStream t => (a -> Bool) -> t m a -> t m Int
691findIndices p = go 0
692 where
693 go offset m1 = mkStream $ \st yld sng stp ->
694 let single a | p a = sng offset
695 | otherwise = stp
696 yieldk a x | p a = yld offset $ go (offset + 1) x
697 | otherwise = foldStream (adaptState st) yld sng stp $
698 go (offset + 1) x
699 in foldStream (adaptState st) yieldk single stp m1
701------------------------------------------------------------------------------
702-- Map and Fold
703------------------------------------------------------------------------------
705-- | Apply a monadic action to each element of the stream and discard the
706-- output of the action.
707{-# INLINE mapM_ #-}
708mapM_ :: (IsStream t, Monad m) => (a -> m b) -> t m a -> m ()
709mapM_ f m = go m
710 where
711 go m1 =
712 let stop = return ()
713 single a = void (f a)
714 yieldk a r = f a >> go r
715 in foldStream defState yieldk single stop m1
717------------------------------------------------------------------------------
718-- Converting folds
719------------------------------------------------------------------------------
721{-# INLINABLE toList #-}
722toList :: (IsStream t, Monad m) => t m a -> m [a]
723toList = foldr (:) []
725{-# INLINE toStreamK #-}
726toStreamK :: Stream m a -> Stream m a
727toStreamK = id
729-------------------------------------------------------------------------------
730-- Transformation by folding (Scans)
731-------------------------------------------------------------------------------
733{-# INLINE scanlx' #-}
734scanlx' :: IsStream t => (x -> a -> x) -> x -> (x -> b) -> t m a -> t m b
735scanlx' step begin done m =
736 cons (done begin) $ go m begin
737 where
738 go m1 !acc = mkStream $ \st yld sng stp ->
739 let single a = sng (done $ step acc a)
740 yieldk a r =
741 let s = step acc a
742 in yld (done s) (go r s)
743 in foldStream (adaptState st) yieldk single stp m1
745{-# INLINE scanl' #-}
746scanl' :: IsStream t => (b -> a -> b) -> b -> t m a -> t m b
747scanl' step begin = scanlx' step begin id
749-------------------------------------------------------------------------------
750-- Filtering
751-------------------------------------------------------------------------------
753{-# INLINE filter #-}
754filter :: IsStream t => (a -> Bool) -> t m a -> t m a
755filter p m = go m
756 where
757 go m1 = mkStream $ \st yld sng stp ->
758 let single a | p a = sng a
759 | otherwise = stp
760 yieldk a r | p a = yld a (go r)
761 | otherwise = foldStream st yieldk single stp r
762 in foldStream st yieldk single stp m1
764{-# INLINE take #-}
765take :: IsStream t => Int -> t m a -> t m a
766take n m = go n m
767 where
768 go n1 m1 = mkStream $ \st yld sng stp ->
769 let yieldk a r = yld a (go (n1 - 1) r)
770 in if n1 <= 0
771 then stp
772 else foldStream st yieldk sng stp m1
774{-# INLINE takeWhile #-}
775takeWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
776takeWhile p m = go m
777 where
778 go m1 = mkStream $ \st yld sng stp ->
779 let single a | p a = sng a
780 | otherwise = stp
781 yieldk a r | p a = yld a (go r)
782 | otherwise = stp
783 in foldStream st yieldk single stp m1
785{-# INLINE drop #-}
786drop :: IsStream t => Int -> t m a -> t m a
787drop n m = fromStream $ unShare (go n (toStream m))
788 where
789 go n1 m1 = mkStream $ \st yld sng stp ->
790 let single _ = stp
791 yieldk _ r = foldStreamShared st yld sng stp $ go (n1 - 1) r
792 -- Somehow "<=" check performs better than a ">"
793 in if n1 <= 0
794 then foldStreamShared st yld sng stp m1
795 else foldStreamShared st yieldk single stp m1
797{-# INLINE dropWhile #-}
798dropWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
799dropWhile p m = go m
800 where
801 go m1 = mkStream $ \st yld sng stp ->
802 let single a | p a = stp
803 | otherwise = sng a
804 yieldk a r | p a = foldStream st yieldk single stp r
805 | otherwise = yld a r
806 in foldStream st yieldk single stp m1
808-------------------------------------------------------------------------------
809-- Mapping
810-------------------------------------------------------------------------------
812-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
813-- other stream types to overload it.
814{-# INLINE sequence #-}
815sequence :: (IsStream t, MonadAsync m) => t m (m a) -> t m a
816sequence m = go m
817 where
818 go m1 = mkStream $ \st yld sng stp ->
819 let single ma = ma >>= sng
820 yieldk ma r = foldStreamShared st yld sng stp $ ma |: go r
821 in foldStream (adaptState st) yieldk single stp m1
823-------------------------------------------------------------------------------
824-- Inserting
825-------------------------------------------------------------------------------
827{-# INLINE intersperseM #-}
828intersperseM :: (IsStream t, MonadAsync m) => m a -> t m a -> t m a
829intersperseM a m = prependingStart m
830 where
831 prependingStart m1 = mkStream $ \st yld sng stp ->
832 let yieldk i x = foldStreamShared st yld sng stp $ return i |: go x
833 in foldStream st yieldk sng stp m1
834 go m2 = mkStream $ \st yld sng stp ->
835 let single i = foldStreamShared st yld sng stp $ a |: yield i
836 yieldk i x = foldStreamShared st yld sng stp $ a |: return i |: go x
837 in foldStream st yieldk single stp m2
839{-# INLINE intersperse #-}
840intersperse :: (IsStream t, MonadAsync m) => a -> t m a -> t m a
841intersperse a = intersperseM (return a)
843{-# INLINE insertBy #-}
844insertBy :: IsStream t => (a -> a -> Ordering) -> a -> t m a -> t m a
845insertBy cmp x m = go m
846 where
847 go m1 = mkStream $ \st yld _ _ ->
848 let single a = case cmp x a of
849 GT -> yld a (yield x)
850 _ -> yld x (yield a)
851 stop = yld x nil
852 yieldk a r = case cmp x a of
853 GT -> yld a (go r)
854 _ -> yld x (a `cons` r)
855 in foldStream st yieldk single stop m1
857------------------------------------------------------------------------------
858-- Deleting
859------------------------------------------------------------------------------
861{-# INLINE deleteBy #-}
862deleteBy :: IsStream t => (a -> a -> Bool) -> a -> t m a -> t m a
863deleteBy eq x m = go m
864 where
865 go m1 = mkStream $ \st yld sng stp ->
866 let single a = if eq x a then stp else sng a
867 yieldk a r = if eq x a
868 then foldStream st yld sng stp r
869 else yld a (go r)
870 in foldStream st yieldk single stp m1
872------------------------------------------------------------------------------
873-- Reordering
874------------------------------------------------------------------------------
876{-# INLINE reverse #-}
877reverse :: IsStream t => t m a -> t m a
878reverse = foldlS (flip cons) nil
880-------------------------------------------------------------------------------
881-- Map and Filter
882-------------------------------------------------------------------------------
884{-# INLINE mapMaybe #-}
885mapMaybe :: IsStream t => (a -> Maybe b) -> t m a -> t m b
886mapMaybe f m = go m
887 where
888 go m1 = mkStream $ \st yld sng stp ->
889 let single a = case f a of
890 Just b -> sng b
891 Nothing -> stp
892 yieldk a r = case f a of
893 Just b -> yld b $ go r
894 Nothing -> foldStream (adaptState st) yieldk single stp r
895 in foldStream (adaptState st) yieldk single stp m1
897------------------------------------------------------------------------------
898-- Serial Zipping
899------------------------------------------------------------------------------
901-- | Zip two streams serially using a pure zipping function.
902--
903-- @since 0.1.0
904{-# INLINABLE zipWith #-}
905zipWith :: IsStream t => (a -> b -> c) -> t m a -> t m b -> t m c
906zipWith f = go
907 where
908 go mx my = mkStream $ \st yld sng stp -> do
909 let merge a ra =
910 let single2 b = sng (f a b)
911 yield2 b rb = yld (f a b) (go ra rb)
912 in foldStream (adaptState st) yield2 single2 stp my
913 let single1 a = merge a nil
914 yield1 = merge
915 foldStream (adaptState st) yield1 single1 stp mx
917-- | Zip two streams serially using a monadic zipping function.
918--
919-- @since 0.1.0
920{-# INLINABLE zipWithM #-}
921zipWithM :: (IsStream t, Monad m) => (a -> b -> m c) -> t m a -> t m b -> t m c
922zipWithM f m1 m2 = go m1 m2
923 where
924 go mx my = mkStream $ \st yld sng stp -> do
925 let merge a ra =
926 let runIt x = foldStream st yld sng stp x
927 single2 b = f a b >>= sng
928 yield2 b rb = f a b >>= \x -> runIt (x `cons` go ra rb)
929 in foldStream (adaptState st) yield2 single2 stp my
930 let single1 a = merge a nil
931 yield1 = merge
932 foldStream (adaptState st) yield1 single1 stp mx
934------------------------------------------------------------------------------
935-- Merging
936------------------------------------------------------------------------------
938{-# INLINE mergeByM #-}
939mergeByM
940 :: (IsStream t, Monad m)
941 => (a -> a -> m Ordering) -> t m a -> t m a -> t m a
942mergeByM cmp = go
943 where
944 go mx my = mkStream $ \st yld sng stp -> do
945 let mergeWithY a ra =
946 let stop2 = foldStream st yld sng stp mx
947 single2 b = do
948 r <- cmp a b
949 case r of
950 GT -> yld b (go (a `cons` ra) nil)
951 _ -> yld a (go ra (b `cons` nil))
952 yield2 b rb = do
953 r <- cmp a b
954 case r of
955 GT -> yld b (go (a `cons` ra) rb)
956 _ -> yld a (go ra (b `cons` rb))
957 in foldStream st yield2 single2 stop2 my
958 let stopX = foldStream st yld sng stp my
959 singleX a = mergeWithY a nil
960 yieldX = mergeWithY
961 foldStream st yieldX singleX stopX mx
963{-# INLINABLE mergeBy #-}
964mergeBy
965 :: (IsStream t, Monad m)
966 => (a -> a -> Ordering) -> t m a -> t m a -> t m a
967mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
969------------------------------------------------------------------------------
970-- Transformation comprehensions
971------------------------------------------------------------------------------
973{-# INLINE the #-}
974the :: (Eq a, IsStream t, Monad m) => t m a -> m (Maybe a)
975the m = do
976 r <- uncons m
977 case r of
978 Nothing -> return Nothing
979 Just (h, t) -> go h t
980 where
981 go h m1 =
982 let single a | h == a = return $ Just h
983 | otherwise = return Nothing
984 yieldk a r | h == a = go h r
985 | otherwise = return Nothing
986 in foldStream defState yieldk single (return $ Just h) m1
988------------------------------------------------------------------------------
989-- Alternative & MonadPlus
990------------------------------------------------------------------------------
992_alt :: Stream m a -> Stream m a -> Stream m a
993_alt m1 m2 = mkStream $ \st yld sng stp ->
994 let stop = foldStream st yld sng stp m2
995 in foldStream st yld sng stop m1
997------------------------------------------------------------------------------
998-- MonadReader
999------------------------------------------------------------------------------
1001{-# INLINABLE withLocal #-}
1002withLocal :: MonadReader r m => (r -> r) -> Stream m a -> Stream m a
1003withLocal f m =
1004 mkStream $ \st yld sng stp ->
1005 let single = local f . sng
1006 yieldk a r = local f $ yld a (withLocal f r)
1007 in foldStream st yieldk single (local f stp) m
1009------------------------------------------------------------------------------
1010-- MonadError
1011------------------------------------------------------------------------------
1014-- XXX handle and test cross thread state transfer
1015withCatchError
1016 :: MonadError e m
1017 => Stream m a -> (e -> Stream m a) -> Stream m a
1018withCatchError m h =
1019 mkStream $ \_ stp sng yld ->
1020 let run x = unStream x Nothing stp sng yieldk
1021 handle r = r `catchError` \e -> run $ h e
1022 yieldk a r = yld a (withCatchError r h)
1023 in handle $ run m