git.iko.soy/streamly
on 3093e79 help sign in
streamly/src/Streamly/Streams/StreamK.hs
1016 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 , replicate
66 , replicateM
67 , fromIndices
68 , fromIndicesM
70 -- ** Conversions
71 , yield
72 , yieldM
73 , fromFoldable
74 , fromList
75 , fromStreamK
77 -- * foldr/build
78 , foldrS
79 , foldrSM
80 , buildS
81 , buildM
82 , augmentS
83 , augmentSM
85 -- * Elimination
86 -- ** General Folds
87 , foldr
88 , foldr1
89 , foldrM
90 , foldrT
92 , foldl'
93 , foldlM'
94 , foldlS
95 , foldlT
96 , foldlx'
97 , foldlMx'
99 -- ** Specialized Folds
100 , drain
101 , null
102 , head
103 , tail
104 , init
105 , elem
106 , notElem
107 , all
108 , any
109 , last
110 , minimum
111 , minimumBy
112 , maximum
113 , maximumBy
114 , findIndices
115 , lookup
116 , findM
117 , find
118 , (!!)
120 -- ** Map and Fold
121 , mapM_
123 -- ** Conversions
124 , toList
125 , toStreamK
127 -- * Transformation
128 -- ** By folding (scans)
129 , scanl'
130 , scanlx'
132 -- ** Filtering
133 , filter
134 , take
135 , takeWhile
136 , drop
137 , dropWhile
139 -- ** Mapping
140 , map
141 , mapM
142 , mapMSerial
143 , sequence
145 -- ** Inserting
146 , intersperseM
147 , intersperse
148 , insertBy
150 -- ** Deleting
151 , deleteBy
153 -- ** Reordering
154 , reverse
156 -- ** Map and Filter
157 , mapMaybe
159 -- ** Zipping
160 , zipWith
161 , zipWithM
163 -- ** Merging
164 , mergeBy
165 , mergeByM
167 -- ** Nesting
168 , concatMapBy
169 , concatMap
170 , bindWith
172 -- ** Transformation comprehensions
173 , the
175 -- * Semigroup Style Composition
176 , serial
178 -- * Utilities
179 , consMStream
180 , withLocal
182 -- * Deprecated
183 , Streaming -- deprecated
184 , once -- deprecated
185 )
186where
188import Control.Monad.Trans (MonadTrans(lift))
189import Control.Monad (void)
190import Control.Monad.Reader.Class (MonadReader(..))
191import Prelude
192 hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,
193 take, filter, all, any, takeWhile, drop, dropWhile, minimum,
194 maximum, elem, notElem, null, head, tail, init, zipWith, lookup,
195 foldr1, (!!), replicate, reverse, concatMap)
196import qualified Prelude
198import Streamly.Internal.Data.SVar
199import Streamly.Streams.StreamK.Type
201-------------------------------------------------------------------------------
202-- Deconstruction
203-------------------------------------------------------------------------------
205{-# INLINE uncons #-}
206uncons :: (IsStream t, Monad m) => t m a -> m (Maybe (a, t m a))
207uncons m =
208 let stop = return Nothing
209 single a = return (Just (a, nil))
210 yieldk a r = return (Just (a, r))
211 in foldStream defState yieldk single stop m
213-------------------------------------------------------------------------------
214-- Generation
215-------------------------------------------------------------------------------
217{-# INLINE unfoldr #-}
218unfoldr :: IsStream t => (b -> Maybe (a, b)) -> b -> t m a
219unfoldr next s0 = build $ \yld stp ->
220 let go s =
221 case next s of
222 Just (a, b) -> yld a (go b)
223 Nothing -> stp
224 in go s0
226{-# INLINE unfoldrM #-}
227unfoldrM :: (IsStream t, MonadAsync m) => (b -> m (Maybe (a, b))) -> b -> t m a
228unfoldrM step = go
229 where
230 go s = sharedM $ \yld _ stp -> do
231 r <- step s
232 case r of
233 Just (a, b) -> yld a (go b)
234 Nothing -> stp
236{-
237-- Generalization of concurrent streams/SVar via unfoldr.
238--
239-- Unfold a value into monadic actions and then run the resulting monadic
240-- actions to generate a stream. Since the step of generating the monadic
241-- action and running them are decoupled we can run the monadic actions
242-- cooncurrently. For example, the seed could be a list of monadic actions or a
243-- pure stream of monadic actions.
244--
245-- We can have different flavors of this depending on the stream type t. The
246-- concurrent version could be async or ahead etc. Depending on how we queue
247-- back the feedback portion b, it could be DFS or BFS style.
248--
249unfoldrA :: (IsStream t, MonadAsync m) => (b -> Maybe (m a, b)) -> b -> t m a
250unfoldrA = undefined
251-}
253-------------------------------------------------------------------------------
254-- Special generation
255-------------------------------------------------------------------------------
257-- | Same as yieldM
258--
259-- @since 0.2.0
260{-# DEPRECATED once "Please use yieldM instead." #-}
261{-# INLINE once #-}
262once :: (Monad m, IsStream t) => m a -> t m a
263once = yieldM
265-- |
266-- @
267-- repeatM = fix . cons
268-- repeatM = cycle1 . yield
269-- @
270--
271-- Generate an infinite stream by repeating a pure value.
272--
273-- @since 0.4.0
274{-# INLINE repeat #-}
275repeat :: IsStream t => a -> t m a
276repeat a = let x = cons a x in x
278{-# INLINE replicateM #-}
279replicateM :: (IsStream t, MonadAsync m) => Int -> m a -> t m a
280replicateM n m = go n
281 where
282 go cnt = if cnt <= 0 then nil else m |: go (cnt - 1)
284{-# INLINE replicate #-}
285replicate :: IsStream t => Int -> a -> t m a
286replicate n a = go n
287 where
288 go cnt = if cnt <= 0 then nil else a `cons` go (cnt - 1)
290{-# INLINE fromIndicesM #-}
291fromIndicesM :: (IsStream t, MonadAsync m) => (Int -> m a) -> t m a
292fromIndicesM gen = go 0
293 where
294 go i = mkStream $ \st stp sng yld -> do
295 foldStreamShared st stp sng yld (gen i |: go (i + 1))
297{-# INLINE fromIndices #-}
298fromIndices :: IsStream t => (Int -> a) -> t m a
299fromIndices gen = go 0
300 where
301 go n = (gen n) `cons` go (n + 1)
303-------------------------------------------------------------------------------
304-- Conversions
305-------------------------------------------------------------------------------
307-- |
308-- @
309-- fromFoldable = 'Prelude.foldr' 'cons' 'nil'
310-- @
311--
312-- Construct a stream from a 'Foldable' containing pure values:
313--
314-- @since 0.2.0
315{-# INLINE fromFoldable #-}
316fromFoldable :: (IsStream t, Foldable f) => f a -> t m a
317fromFoldable = Prelude.foldr cons nil
319{-# INLINE fromList #-}
320fromList :: IsStream t => [a] -> t m a
321fromList = fromFoldable
323{-# INLINE fromStreamK #-}
324fromStreamK :: IsStream t => Stream m a -> t m a
325fromStreamK = fromStream
327-------------------------------------------------------------------------------
328-- Elimination by Folding
329-------------------------------------------------------------------------------
331-- | Lazy right associative fold.
332{-# INLINE foldr #-}
333foldr :: (IsStream t, Monad m) => (a -> b -> b) -> b -> t m a -> m b
334foldr step acc = foldrM (\x xs -> xs >>= \b -> return (step x b)) (return acc)
336-- | Right associative fold to an arbitrary transformer monad.
337{-# INLINE foldrT #-}
338foldrT :: (IsStream t, Monad m, Monad (s m), MonadTrans s)
339 => (a -> s m b -> s m b) -> s m b -> t m a -> s m b
340foldrT step final m = go m
341 where
342 go m1 = do
343 res <- lift $ uncons m1
344 case res of
345 Just (h, t) -> step h (go t)
346 Nothing -> final
348{-# INLINE foldr1 #-}
349foldr1 :: (IsStream t, Monad m) => (a -> a -> a) -> t m a -> m (Maybe a)
350foldr1 step m = do
351 r <- uncons m
352 case r of
353 Nothing -> return Nothing
354 Just (h, t) -> fmap Just (go h t)
355 where
356 go p m1 =
357 let stp = return p
358 single a = return $ step a p
359 yieldk a r = fmap (step p) (go a r)
360 in foldStream defState yieldk single stp m1
362-- | Strict left fold with an extraction function. Like the standard strict
363-- left fold, but applies a user supplied extraction function (the third
364-- argument) to the folded value at the end. This is designed to work with the
365-- @foldl@ library. The suffix @x@ is a mnemonic for extraction.
366--
367-- Note that the accumulator is always evaluated including the initial value.
368{-# INLINE foldlx' #-}
369foldlx' :: forall t m a b x. (IsStream t, Monad m)
370 => (x -> a -> x) -> x -> (x -> b) -> t m a -> m b
371foldlx' step begin done m = get $ go m begin
372 where
373 {-# NOINLINE get #-}
374 get :: t m x -> m b
375 get m1 =
376 -- XXX we are not strictly evaluating the accumulator here. Is this
377 -- okay?
378 let single = return . done
379 -- XXX this is foldSingleton. why foldStreamShared?
380 in foldStreamShared undefined undefined single undefined m1
382 -- Note, this can be implemented by making a recursive call to "go",
383 -- however that is more expensive because of unnecessary recursion
384 -- that cannot be tail call optimized. Unfolding recursion explicitly via
385 -- continuations is much more efficient.
386 go :: t m a -> x -> t m x
387 go m1 !acc = mkStream $ \_ yld sng _ ->
388 let stop = sng acc
389 single a = sng $ step acc a
390 -- XXX this is foldNonEmptyStream
391 yieldk a r = foldStream defState yld sng undefined $
392 go r (step acc a)
393 in foldStream defState yieldk single stop m1
395-- | Strict left associative fold.
396{-# INLINE foldl' #-}
397foldl' :: (IsStream t, Monad m) => (b -> a -> b) -> b -> t m a -> m b
398foldl' step begin = foldlx' step begin id
400-- XXX replace the recursive "go" with explicit continuations.
401-- | Like 'foldx', but with a monadic step function.
402{-# INLINABLE foldlMx' #-}
403foldlMx' :: (IsStream t, Monad m)
404 => (x -> a -> m x) -> m x -> (x -> m b) -> t m a -> m b
405foldlMx' step begin done m = go begin m
406 where
407 go !acc m1 =
408 let stop = acc >>= done
409 single a = acc >>= \b -> step b a >>= done
410 yieldk a r = acc >>= \b -> step b a >>= \x -> go (return x) r
411 in foldStream defState yieldk single stop m1
413-- | Like 'foldl'' but with a monadic step function.
414{-# INLINE foldlM' #-}
415foldlM' :: (IsStream t, Monad m) => (b -> a -> m b) -> b -> t m a -> m b
416foldlM' step begin = foldlMx' step (return begin) return
418-- | Lazy left fold to a stream.
419{-# INLINE foldlS #-}
420foldlS :: IsStream t => (t m b -> a -> t m b) -> t m b -> t m a -> t m b
421foldlS step begin m = go begin m
422 where
423 go acc rest = mkStream $ \st yld sng stp ->
424 let run x = foldStream st yld sng stp x
425 stop = run acc
426 single a = run $ step acc a
427 yieldk a r = run $ go (step acc a) r
428 in foldStream (adaptState st) yieldk single stop rest
430-- | Lazy left fold to an arbitrary transformer monad.
431{-# INLINE foldlT #-}
432foldlT :: (IsStream t, Monad m, Monad (s m), MonadTrans s)
433 => (s m b -> a -> s m b) -> s m b -> t m a -> s m b
434foldlT step begin m = go begin m
435 where
436 go acc m1 = do
437 res <- lift $ uncons m1
438 case res of
439 Just (h, t) -> go (step acc h) t
440 Nothing -> acc
442------------------------------------------------------------------------------
443-- Specialized folds
444------------------------------------------------------------------------------
446-- XXX use foldrM to implement folds where possible
447-- XXX This (commented) definition of drain and mapM_ perform much better on
448-- some benchmarks but worse on others. Need to investigate why, may there is
449-- an optimization opportunity that we can exploit.
450-- drain = foldrM (\_ xs -> return () >> xs) (return ())
452-- |
453-- > drain = foldl' (\_ _ -> ()) ()
454-- > drain = mapM_ (\_ -> return ())
455{-# INLINE drain #-}
456drain :: (Monad m, IsStream t) => t m a -> m ()
457drain = foldrM (\_ xs -> xs) (return ())
458{-
459drain = go
460 where
461 go m1 =
462 let stop = return ()
463 single _ = return ()
464 yieldk _ r = go r
465 in foldStream defState yieldk single stop m1
466-}
468{-# INLINE null #-}
469null :: (IsStream t, Monad m) => t m a -> m Bool
470-- null = foldrM (\_ _ -> return True) (return False)
471null m =
472 let stop = return True
473 single _ = return False
474 yieldk _ _ = return False
475 in foldStream defState yieldk single stop m
477{-# INLINE head #-}
478head :: (IsStream t, Monad m) => t m a -> m (Maybe a)
479-- head = foldrM (\x _ -> return $ Just x) (return Nothing)
480head m =
481 let stop = return Nothing
482 single a = return (Just a)
483 yieldk a _ = return (Just a)
484 in foldStream defState yieldk single stop m
486{-# INLINE tail #-}
487tail :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
488tail m =
489 let stop = return Nothing
490 single _ = return $ Just nil
491 yieldk _ r = return $ Just r
492 in foldStream defState yieldk single stop m
494{-# INLINE init #-}
495init :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
496init m = go1 m
497 where
498 go1 m1 = do
499 r <- uncons m1
500 case r of
501 Nothing -> return Nothing
502 Just (h, t) -> return . Just $ go h t
503 go p m1 = mkStream $ \_ yld sng stp ->
504 let single _ = sng p
505 yieldk a x = yld p $ go a x
506 in foldStream defState yieldk single stp m1
508{-# INLINE elem #-}
509elem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
510elem e m = go m
511 where
512 go m1 =
513 let stop = return False
514 single a = return (a == e)
515 yieldk a r = if a == e then return True else go r
516 in foldStream defState yieldk single stop m1
518{-# INLINE notElem #-}
519notElem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
520notElem e m = go m
521 where
522 go m1 =
523 let stop = return True
524 single a = return (a /= e)
525 yieldk a r = if a == e then return False else go r
526 in foldStream defState yieldk single stop m1
528{-# INLINABLE all #-}
529all :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
530all p m = go m
531 where
532 go m1 =
533 let single a | p a = return True
534 | otherwise = return False
535 yieldk a r | p a = go r
536 | otherwise = return False
537 in foldStream defState yieldk single (return True) m1
539{-# INLINABLE any #-}
540any :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
541any p m = go m
542 where
543 go m1 =
544 let single a | p a = return True
545 | otherwise = return False
546 yieldk a r | p a = return True
547 | otherwise = go r
548 in foldStream defState yieldk single (return False) m1
550-- | Extract the last element of the stream, if any.
551{-# INLINE last #-}
552last :: (IsStream t, Monad m) => t m a -> m (Maybe a)
553last = foldlx' (\_ y -> Just y) Nothing id
555{-# INLINE minimum #-}
556minimum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
557minimum m = go Nothing m
558 where
559 go Nothing m1 =
560 let stop = return Nothing
561 single a = return (Just a)
562 yieldk a r = go (Just a) r
563 in foldStream defState yieldk single stop m1
565 go (Just res) m1 =
566 let stop = return (Just res)
567 single a =
568 if res <= a
569 then return (Just res)
570 else return (Just a)
571 yieldk a r =
572 if res <= a
573 then go (Just res) r
574 else go (Just a) r
575 in foldStream defState yieldk single stop m1
577{-# INLINE minimumBy #-}
578minimumBy
579 :: (IsStream t, Monad m)
580 => (a -> a -> Ordering) -> t m a -> m (Maybe a)
581minimumBy cmp m = go Nothing m
582 where
583 go Nothing m1 =
584 let stop = return Nothing
585 single a = return (Just a)
586 yieldk a r = go (Just a) r
587 in foldStream defState yieldk single stop m1
589 go (Just res) m1 =
590 let stop = return (Just res)
591 single a = case cmp res a of
592 GT -> return (Just a)
593 _ -> return (Just res)
594 yieldk a r = case cmp res a of
595 GT -> go (Just a) r
596 _ -> go (Just res) r
597 in foldStream defState yieldk single stop m1
599{-# INLINE maximum #-}
600maximum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
601maximum m = go Nothing m
602 where
603 go Nothing m1 =
604 let stop = return Nothing
605 single a = return (Just a)
606 yieldk a r = go (Just a) r
607 in foldStream defState yieldk single stop m1
609 go (Just res) m1 =
610 let stop = return (Just res)
611 single a =
612 if res <= a
613 then return (Just a)
614 else return (Just res)
615 yieldk a r =
616 if res <= a
617 then go (Just a) r
618 else go (Just res) r
619 in foldStream defState yieldk single stop m1
621{-# INLINE maximumBy #-}
622maximumBy :: (IsStream t, Monad m) => (a -> a -> Ordering) -> t m a -> m (Maybe a)
623maximumBy cmp m = go Nothing m
624 where
625 go Nothing m1 =
626 let stop = return Nothing
627 single a = return (Just a)
628 yieldk a r = go (Just a) r
629 in foldStream defState yieldk single stop m1
631 go (Just res) m1 =
632 let stop = return (Just res)
633 single a = case cmp res a of
634 GT -> return (Just res)
635 _ -> return (Just a)
636 yieldk a r = case cmp res a of
637 GT -> go (Just res) r
638 _ -> go (Just a) r
639 in foldStream defState yieldk single stop m1
641{-# INLINE (!!) #-}
642(!!) :: (IsStream t, Monad m) => t m a -> Int -> m (Maybe a)
643m !! i = go i m
644 where
645 go n m1 =
646 let single a | n == 0 = return $ Just a
647 | otherwise = return Nothing
648 yieldk a x | n < 0 = return Nothing
649 | n == 0 = return $ Just a
650 | otherwise = go (n - 1) x
651 in foldStream defState yieldk single (return Nothing) m1
653{-# INLINE lookup #-}
654lookup :: (IsStream t, Monad m, Eq a) => a -> t m (a, b) -> m (Maybe b)
655lookup e m = go m
656 where
657 go m1 =
658 let single (a, b) | a == e = return $ Just b
659 | otherwise = return Nothing
660 yieldk (a, b) x | a == e = return $ Just b
661 | otherwise = go x
662 in foldStream defState yieldk single (return Nothing) m1
664{-# INLINE findM #-}
665findM :: (IsStream t, Monad m) => (a -> m Bool) -> t m a -> m (Maybe a)
666findM p m = go m
667 where
668 go m1 =
669 let single a = do
670 b <- p a
671 if b then return $ Just a else return Nothing
672 yieldk a x = do
673 b <- p a
674 if b then return $ Just a else go x
675 in foldStream defState yieldk single (return Nothing) m1
677{-# INLINE find #-}
678find :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m (Maybe a)
679find p = findM (return . p)
681{-# INLINE findIndices #-}
682findIndices :: IsStream t => (a -> Bool) -> t m a -> t m Int
683findIndices p = go 0
684 where
685 go offset m1 = mkStream $ \st yld sng stp ->
686 let single a | p a = sng offset
687 | otherwise = stp
688 yieldk a x | p a = yld offset $ go (offset + 1) x
689 | otherwise = foldStream (adaptState st) yld sng stp $
690 go (offset + 1) x
691 in foldStream (adaptState st) yieldk single stp m1
693------------------------------------------------------------------------------
694-- Map and Fold
695------------------------------------------------------------------------------
697-- | Apply a monadic action to each element of the stream and discard the
698-- output of the action.
699{-# INLINE mapM_ #-}
700mapM_ :: (IsStream t, Monad m) => (a -> m b) -> t m a -> m ()
701mapM_ f m = go m
702 where
703 go m1 =
704 let stop = return ()
705 single a = void (f a)
706 yieldk a r = f a >> go r
707 in foldStream defState yieldk single stop m1
709------------------------------------------------------------------------------
710-- Converting folds
711------------------------------------------------------------------------------
713{-# INLINABLE toList #-}
714toList :: (IsStream t, Monad m) => t m a -> m [a]
715toList = foldr (:) []
717{-# INLINE toStreamK #-}
718toStreamK :: Stream m a -> Stream m a
719toStreamK = id
721-------------------------------------------------------------------------------
722-- Transformation by folding (Scans)
723-------------------------------------------------------------------------------
725{-# INLINE scanlx' #-}
726scanlx' :: IsStream t => (x -> a -> x) -> x -> (x -> b) -> t m a -> t m b
727scanlx' step begin done m =
728 cons (done begin) $ go m begin
729 where
730 go m1 !acc = mkStream $ \st yld sng stp ->
731 let single a = sng (done $ step acc a)
732 yieldk a r =
733 let s = step acc a
734 in yld (done s) (go r s)
735 in foldStream (adaptState st) yieldk single stp m1
737{-# INLINE scanl' #-}
738scanl' :: IsStream t => (b -> a -> b) -> b -> t m a -> t m b
739scanl' step begin = scanlx' step begin id
741-------------------------------------------------------------------------------
742-- Filtering
743-------------------------------------------------------------------------------
745{-# INLINE filter #-}
746filter :: IsStream t => (a -> Bool) -> t m a -> t m a
747filter p m = go m
748 where
749 go m1 = mkStream $ \st yld sng stp ->
750 let single a | p a = sng a
751 | otherwise = stp
752 yieldk a r | p a = yld a (go r)
753 | otherwise = foldStream st yieldk single stp r
754 in foldStream st yieldk single stp m1
756{-# INLINE take #-}
757take :: IsStream t => Int -> t m a -> t m a
758take n m = go n m
759 where
760 go n1 m1 = mkStream $ \st yld sng stp ->
761 let yieldk a r = yld a (go (n1 - 1) r)
762 in if n1 <= 0
763 then stp
764 else foldStream st yieldk sng stp m1
766{-# INLINE takeWhile #-}
767takeWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
768takeWhile p m = go m
769 where
770 go m1 = mkStream $ \st yld sng stp ->
771 let single a | p a = sng a
772 | otherwise = stp
773 yieldk a r | p a = yld a (go r)
774 | otherwise = stp
775 in foldStream st yieldk single stp m1
777{-# INLINE drop #-}
778drop :: IsStream t => Int -> t m a -> t m a
779drop n m = fromStream $ unShare (go n (toStream m))
780 where
781 go n1 m1 = mkStream $ \st yld sng stp ->
782 let single _ = stp
783 yieldk _ r = foldStreamShared st yld sng stp $ go (n1 - 1) r
784 -- Somehow "<=" check performs better than a ">"
785 in if n1 <= 0
786 then foldStreamShared st yld sng stp m1
787 else foldStreamShared st yieldk single stp m1
789{-# INLINE dropWhile #-}
790dropWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
791dropWhile p m = go m
792 where
793 go m1 = mkStream $ \st yld sng stp ->
794 let single a | p a = stp
795 | otherwise = sng a
796 yieldk a r | p a = foldStream st yieldk single stp r
797 | otherwise = yld a r
798 in foldStream st yieldk single stp m1
800-------------------------------------------------------------------------------
801-- Mapping
802-------------------------------------------------------------------------------
804-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
805-- other stream types to overload it.
806{-# INLINE sequence #-}
807sequence :: (IsStream t, MonadAsync m) => t m (m a) -> t m a
808sequence m = go m
809 where
810 go m1 = mkStream $ \st yld sng stp ->
811 let single ma = ma >>= sng
812 yieldk ma r = foldStreamShared st yld sng stp $ ma |: go r
813 in foldStream (adaptState st) yieldk single stp m1
815-------------------------------------------------------------------------------
816-- Inserting
817-------------------------------------------------------------------------------
819{-# INLINE intersperseM #-}
820intersperseM :: (IsStream t, MonadAsync m) => m a -> t m a -> t m a
821intersperseM a m = prependingStart m
822 where
823 prependingStart m1 = mkStream $ \st yld sng stp ->
824 let yieldk i x = foldStreamShared st yld sng stp $ return i |: go x
825 in foldStream st yieldk sng stp m1
826 go m2 = mkStream $ \st yld sng stp ->
827 let single i = foldStreamShared st yld sng stp $ a |: yield i
828 yieldk i x = foldStreamShared st yld sng stp $ a |: return i |: go x
829 in foldStream st yieldk single stp m2
831{-# INLINE intersperse #-}
832intersperse :: (IsStream t, MonadAsync m) => a -> t m a -> t m a
833intersperse a = intersperseM (return a)
835{-# INLINE insertBy #-}
836insertBy :: IsStream t => (a -> a -> Ordering) -> a -> t m a -> t m a
837insertBy cmp x m = go m
838 where
839 go m1 = mkStream $ \st yld _ _ ->
840 let single a = case cmp x a of
841 GT -> yld a (yield x)
842 _ -> yld x (yield a)
843 stop = yld x nil
844 yieldk a r = case cmp x a of
845 GT -> yld a (go r)
846 _ -> yld x (a `cons` r)
847 in foldStream st yieldk single stop m1
849------------------------------------------------------------------------------
850-- Deleting
851------------------------------------------------------------------------------
853{-# INLINE deleteBy #-}
854deleteBy :: IsStream t => (a -> a -> Bool) -> a -> t m a -> t m a
855deleteBy eq x m = go m
856 where
857 go m1 = mkStream $ \st yld sng stp ->
858 let single a = if eq x a then stp else sng a
859 yieldk a r = if eq x a
860 then foldStream st yld sng stp r
861 else yld a (go r)
862 in foldStream st yieldk single stp m1
864------------------------------------------------------------------------------
865-- Reordering
866------------------------------------------------------------------------------
868{-# INLINE reverse #-}
869reverse :: IsStream t => t m a -> t m a
870reverse = foldlS (flip cons) nil
872-------------------------------------------------------------------------------
873-- Map and Filter
874-------------------------------------------------------------------------------
876{-# INLINE mapMaybe #-}
877mapMaybe :: IsStream t => (a -> Maybe b) -> t m a -> t m b
878mapMaybe f m = go m
879 where
880 go m1 = mkStream $ \st yld sng stp ->
881 let single a = case f a of
882 Just b -> sng b
883 Nothing -> stp
884 yieldk a r = case f a of
885 Just b -> yld b $ go r
886 Nothing -> foldStream (adaptState st) yieldk single stp r
887 in foldStream (adaptState st) yieldk single stp m1
889------------------------------------------------------------------------------
890-- Serial Zipping
891------------------------------------------------------------------------------
893-- | Zip two streams serially using a pure zipping function.
894--
895-- @since 0.1.0
896{-# INLINABLE zipWith #-}
897zipWith :: IsStream t => (a -> b -> c) -> t m a -> t m b -> t m c
898zipWith f = go
899 where
900 go mx my = mkStream $ \st yld sng stp -> do
901 let merge a ra =
902 let single2 b = sng (f a b)
903 yield2 b rb = yld (f a b) (go ra rb)
904 in foldStream (adaptState st) yield2 single2 stp my
905 let single1 a = merge a nil
906 yield1 = merge
907 foldStream (adaptState st) yield1 single1 stp mx
909-- | Zip two streams serially using a monadic zipping function.
910--
911-- @since 0.1.0
912{-# INLINABLE zipWithM #-}
913zipWithM :: (IsStream t, Monad m) => (a -> b -> m c) -> t m a -> t m b -> t m c
914zipWithM f m1 m2 = go m1 m2
915 where
916 go mx my = mkStream $ \st yld sng stp -> do
917 let merge a ra =
918 let runIt x = foldStream st yld sng stp x
919 single2 b = f a b >>= sng
920 yield2 b rb = f a b >>= \x -> runIt (x `cons` go ra rb)
921 in foldStream (adaptState st) yield2 single2 stp my
922 let single1 a = merge a nil
923 yield1 = merge
924 foldStream (adaptState st) yield1 single1 stp mx
926------------------------------------------------------------------------------
927-- Merging
928------------------------------------------------------------------------------
930{-# INLINE mergeByM #-}
931mergeByM
932 :: (IsStream t, Monad m)
933 => (a -> a -> m Ordering) -> t m a -> t m a -> t m a
934mergeByM cmp = go
935 where
936 go mx my = mkStream $ \st yld sng stp -> do
937 let mergeWithY a ra =
938 let stop2 = foldStream st yld sng stp mx
939 single2 b = do
940 r <- cmp a b
941 case r of
942 GT -> yld b (go (a `cons` ra) nil)
943 _ -> yld a (go ra (b `cons` nil))
944 yield2 b rb = do
945 r <- cmp a b
946 case r of
947 GT -> yld b (go (a `cons` ra) rb)
948 _ -> yld a (go ra (b `cons` rb))
949 in foldStream st yield2 single2 stop2 my
950 let stopX = foldStream st yld sng stp my
951 singleX a = mergeWithY a nil
952 yieldX = mergeWithY
953 foldStream st yieldX singleX stopX mx
955{-# INLINABLE mergeBy #-}
956mergeBy
957 :: (IsStream t, Monad m)
958 => (a -> a -> Ordering) -> t m a -> t m a -> t m a
959mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
961------------------------------------------------------------------------------
962-- Transformation comprehensions
963------------------------------------------------------------------------------
965{-# INLINE the #-}
966the :: (Eq a, IsStream t, Monad m) => t m a -> m (Maybe a)
967the m = do
968 r <- uncons m
969 case r of
970 Nothing -> return Nothing
971 Just (h, t) -> go h t
972 where
973 go h m1 =
974 let single a | h == a = return $ Just h
975 | otherwise = return Nothing
976 yieldk a r | h == a = go h r
977 | otherwise = return Nothing
978 in foldStream defState yieldk single (return $ Just h) m1
980------------------------------------------------------------------------------
981-- Alternative & MonadPlus
982------------------------------------------------------------------------------
984_alt :: Stream m a -> Stream m a -> Stream m a
985_alt m1 m2 = mkStream $ \st yld sng stp ->
986 let stop = foldStream st yld sng stp m2
987 in foldStream st yld sng stop m1
989------------------------------------------------------------------------------
990-- MonadReader
991------------------------------------------------------------------------------
993{-# INLINABLE withLocal #-}
994withLocal :: MonadReader r m => (r -> r) -> Stream m a -> Stream m a
995withLocal f m =
996 mkStream $ \st yld sng stp ->
997 let single = local f . sng
998 yieldk a r = local f $ yld a (withLocal f r)
999 in foldStream st yieldk single (local f stp) m
1001------------------------------------------------------------------------------
1002-- MonadError
1003------------------------------------------------------------------------------
1006-- XXX handle and test cross thread state transfer
1007withCatchError
1008 :: MonadError e m
1009 => Stream m a -> (e -> Stream m a) -> Stream m a
1010withCatchError m h =
1011 mkStream $ \_ stp sng yld ->
1012 let run x = unStream x Nothing stp sng yieldk
1013 handle r = r `catchError` \e -> run $ h e
1014 yieldk a r = yld a (withCatchError r h)
1015 in handle $ run m