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