git.iko.soy/streamly
on 879d159 help sign in
streamly/src/Streamly/Streams/StreamK.hs
1037 lines 33 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 : streamly@composewell.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
127 , hoist
129 -- * Transformation
130 -- ** By folding (scans)
131 , scanl'
132 , scanlx'
134 -- ** Filtering
135 , filter
136 , take
137 , takeWhile
138 , drop
139 , dropWhile
141 -- ** Mapping
142 , map
143 , mapM
144 , mapMSerial
145 , sequence
147 -- ** Inserting
148 , intersperseM
149 , intersperse
150 , insertBy
152 -- ** Deleting
153 , deleteBy
155 -- ** Reordering
156 , reverse
158 -- ** Map and Filter
159 , mapMaybe
161 -- ** Zipping
162 , zipWith
163 , zipWithM
165 -- ** Merging
166 , mergeBy
167 , mergeByM
169 -- ** Nesting
170 , concatMapBy
171 , concatMap
172 , bindWith
174 -- ** Transformation comprehensions
175 , the
177 -- * Semigroup Style Composition
178 , serial
180 -- * Utilities
181 , consMStream
182 , withLocal
184 -- * Deprecated
185 , Streaming -- deprecated
186 , once -- deprecated
187 )
188where
190import Control.Monad.Trans (MonadTrans(lift))
191import Control.Monad (void, join)
192import Control.Monad.Reader.Class (MonadReader(..))
193import Prelude
194 hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,
195 take, filter, all, any, takeWhile, drop, dropWhile, minimum,
196 maximum, elem, notElem, null, head, tail, init, zipWith, lookup,
197 foldr1, (!!), replicate, reverse, concatMap)
198import qualified Prelude
200import Streamly.Internal.Data.SVar
201import Streamly.Streams.StreamK.Type
203-------------------------------------------------------------------------------
204-- Deconstruction
205-------------------------------------------------------------------------------
207{-# INLINE uncons #-}
208uncons :: (IsStream t, Monad m) => t m a -> m (Maybe (a, t m a))
209uncons m =
210 let stop = return Nothing
211 single a = return (Just (a, nil))
212 yieldk a r = return (Just (a, r))
213 in foldStream defState yieldk single stop m
215-------------------------------------------------------------------------------
216-- Generation
217-------------------------------------------------------------------------------
219{-# INLINE unfoldr #-}
220unfoldr :: IsStream t => (b -> Maybe (a, b)) -> b -> t m a
221unfoldr next s0 = build $ \yld stp ->
222 let go s =
223 case next s of
224 Just (a, b) -> yld a (go b)
225 Nothing -> stp
226 in go s0
228{-# INLINE unfoldrM #-}
229unfoldrM :: (IsStream t, MonadAsync m) => (b -> m (Maybe (a, b))) -> b -> t m a
230unfoldrM step = go
231 where
232 go s = sharedM $ \yld _ stp -> do
233 r <- step s
234 case r of
235 Just (a, b) -> yld a (go b)
236 Nothing -> stp
238{-
239-- Generalization of concurrent streams/SVar via unfoldr.
240--
241-- Unfold a value into monadic actions and then run the resulting monadic
242-- actions to generate a stream. Since the step of generating the monadic
243-- action and running them are decoupled we can run the monadic actions
244-- cooncurrently. For example, the seed could be a list of monadic actions or a
245-- pure stream of monadic actions.
246--
247-- We can have different flavors of this depending on the stream type t. The
248-- concurrent version could be async or ahead etc. Depending on how we queue
249-- back the feedback portion b, it could be DFS or BFS style.
250--
251unfoldrA :: (IsStream t, MonadAsync m) => (b -> Maybe (m a, b)) -> b -> t m a
252unfoldrA = undefined
253-}
255-------------------------------------------------------------------------------
256-- Special generation
257-------------------------------------------------------------------------------
259-- | Same as yieldM
260--
261-- @since 0.2.0
262{-# DEPRECATED once "Please use yieldM instead." #-}
263{-# INLINE once #-}
264once :: (Monad m, IsStream t) => m a -> t m a
265once = yieldM
267-- |
268-- @
269-- repeatM = fix . cons
270-- repeatM = cycle1 . yield
271-- @
272--
273-- Generate an infinite stream by repeating a monadic value.
274--
275-- /Internal/
276repeatM :: (IsStream t, MonadAsync m) => m a -> t m a
277repeatM = go
278 where go m = m |: go m
280-- Generate an infinite stream by repeating a pure value.
281--
282-- /Internal/
283{-# INLINE repeat #-}
284repeat :: IsStream t => a -> t m a
285repeat a = let x = cons a x in x
287{-# INLINE replicateM #-}
288replicateM :: (IsStream t, MonadAsync m) => Int -> m a -> t m a
289replicateM n m = go n
290 where
291 go cnt = if cnt <= 0 then nil else m |: go (cnt - 1)
293{-# INLINE replicate #-}
294replicate :: IsStream t => Int -> a -> t m a
295replicate n a = go n
296 where
297 go cnt = if cnt <= 0 then nil else a `cons` go (cnt - 1)
299{-# INLINE fromIndicesM #-}
300fromIndicesM :: (IsStream t, MonadAsync m) => (Int -> m a) -> t m a
301fromIndicesM gen = go 0
302 where
303 go i = mkStream $ \st stp sng yld -> do
304 foldStreamShared st stp sng yld (gen i |: go (i + 1))
306{-# INLINE fromIndices #-}
307fromIndices :: IsStream t => (Int -> a) -> t m a
308fromIndices gen = go 0
309 where
310 go n = (gen n) `cons` go (n + 1)
312-------------------------------------------------------------------------------
313-- Conversions
314-------------------------------------------------------------------------------
316-- |
317-- @
318-- fromFoldable = 'Prelude.foldr' 'cons' 'nil'
319-- @
320--
321-- Construct a stream from a 'Foldable' containing pure values:
322--
323-- @since 0.2.0
324{-# INLINE fromFoldable #-}
325fromFoldable :: (IsStream t, Foldable f) => f a -> t m a
326fromFoldable = Prelude.foldr cons nil
328{-# INLINE fromList #-}
329fromList :: IsStream t => [a] -> t m a
330fromList = fromFoldable
332{-# INLINE fromStreamK #-}
333fromStreamK :: IsStream t => Stream m a -> t m a
334fromStreamK = fromStream
336-------------------------------------------------------------------------------
337-- Elimination by Folding
338-------------------------------------------------------------------------------
340-- | Lazy right associative fold.
341{-# INLINE foldr #-}
342foldr :: (IsStream t, Monad m) => (a -> b -> b) -> b -> t m a -> m b
343foldr step acc = foldrM (\x xs -> xs >>= \b -> return (step x b)) (return acc)
345-- | Right associative fold to an arbitrary transformer monad.
346{-# INLINE foldrT #-}
347foldrT :: (IsStream t, Monad m, Monad (s m), MonadTrans s)
348 => (a -> s m b -> s m b) -> s m b -> t m a -> s m b
349foldrT step final m = go m
350 where
351 go m1 = do
352 res <- lift $ uncons m1
353 case res of
354 Just (h, t) -> step h (go t)
355 Nothing -> final
357{-# INLINE foldr1 #-}
358foldr1 :: (IsStream t, Monad m) => (a -> a -> a) -> t m a -> m (Maybe a)
359foldr1 step m = do
360 r <- uncons m
361 case r of
362 Nothing -> return Nothing
363 Just (h, t) -> fmap Just (go h t)
364 where
365 go p m1 =
366 let stp = return p
367 single a = return $ step a p
368 yieldk a r = fmap (step p) (go a r)
369 in foldStream defState yieldk single stp m1
371-- | Strict left fold with an extraction function. Like the standard strict
372-- left fold, but applies a user supplied extraction function (the third
373-- argument) to the folded value at the end. This is designed to work with the
374-- @foldl@ library. The suffix @x@ is a mnemonic for extraction.
375--
376-- Note that the accumulator is always evaluated including the initial value.
377{-# INLINE foldlx' #-}
378foldlx' :: forall t m a b x. (IsStream t, Monad m)
379 => (x -> a -> x) -> x -> (x -> b) -> t m a -> m b
380foldlx' step begin done m = get $ go m begin
381 where
382 {-# NOINLINE get #-}
383 get :: t m x -> m b
384 get m1 =
385 -- XXX we are not strictly evaluating the accumulator here. Is this
386 -- okay?
387 let single = return . done
388 -- XXX this is foldSingleton. why foldStreamShared?
389 in foldStreamShared undefined undefined single undefined m1
391 -- Note, this can be implemented by making a recursive call to "go",
392 -- however that is more expensive because of unnecessary recursion
393 -- that cannot be tail call optimized. Unfolding recursion explicitly via
394 -- continuations is much more efficient.
395 go :: t m a -> x -> t m x
396 go m1 !acc = mkStream $ \_ yld sng _ ->
397 let stop = sng acc
398 single a = sng $ step acc a
399 -- XXX this is foldNonEmptyStream
400 yieldk a r = foldStream defState yld sng undefined $
401 go r (step acc a)
402 in foldStream defState yieldk single stop m1
404-- | Strict left associative fold.
405{-# INLINE foldl' #-}
406foldl' :: (IsStream t, Monad m) => (b -> a -> b) -> b -> t m a -> m b
407foldl' step begin = foldlx' step begin id
409-- XXX replace the recursive "go" with explicit continuations.
410-- | Like 'foldx', but with a monadic step function.
411{-# INLINABLE foldlMx' #-}
412foldlMx' :: (IsStream t, Monad m)
413 => (x -> a -> m x) -> m x -> (x -> m b) -> t m a -> m b
414foldlMx' step begin done m = go begin m
415 where
416 go !acc m1 =
417 let stop = acc >>= done
418 single a = acc >>= \b -> step b a >>= done
419 yieldk a r = acc >>= \b -> step b a >>= \x -> go (return x) r
420 in foldStream defState yieldk single stop m1
422-- | Like 'foldl'' but with a monadic step function.
423{-# INLINE foldlM' #-}
424foldlM' :: (IsStream t, Monad m) => (b -> a -> m b) -> b -> t m a -> m b
425foldlM' step begin = foldlMx' step (return begin) return
427-- | Lazy left fold to a stream.
428{-# INLINE foldlS #-}
429foldlS :: IsStream t => (t m b -> a -> t m b) -> t m b -> t m a -> t m b
430foldlS step begin m = go begin m
431 where
432 go acc rest = mkStream $ \st yld sng stp ->
433 let run x = foldStream st yld sng stp x
434 stop = run acc
435 single a = run $ step acc a
436 yieldk a r = run $ go (step acc a) r
437 in foldStream (adaptState st) yieldk single stop rest
439-- | Lazy left fold to an arbitrary transformer monad.
440{-# INLINE foldlT #-}
441foldlT :: (IsStream t, Monad m, Monad (s m), MonadTrans s)
442 => (s m b -> a -> s m b) -> s m b -> t m a -> s m b
443foldlT step begin m = go begin m
444 where
445 go acc m1 = do
446 res <- lift $ uncons m1
447 case res of
448 Just (h, t) -> go (step acc h) t
449 Nothing -> acc
451------------------------------------------------------------------------------
452-- Specialized folds
453------------------------------------------------------------------------------
455-- XXX use foldrM to implement folds where possible
456-- XXX This (commented) definition of drain and mapM_ perform much better on
457-- some benchmarks but worse on others. Need to investigate why, may there is
458-- an optimization opportunity that we can exploit.
459-- drain = foldrM (\_ xs -> return () >> xs) (return ())
461-- |
462-- > drain = foldl' (\_ _ -> ()) ()
463-- > drain = mapM_ (\_ -> return ())
464{-# INLINE drain #-}
465drain :: (Monad m, IsStream t) => t m a -> m ()
466drain = foldrM (\_ xs -> xs) (return ())
467{-
468drain = go
469 where
470 go m1 =
471 let stop = return ()
472 single _ = return ()
473 yieldk _ r = go r
474 in foldStream defState yieldk single stop m1
475-}
477{-# INLINE null #-}
478null :: (IsStream t, Monad m) => t m a -> m Bool
479-- null = foldrM (\_ _ -> return True) (return False)
480null m =
481 let stop = return True
482 single _ = return False
483 yieldk _ _ = return False
484 in foldStream defState yieldk single stop m
486{-# INLINE head #-}
487head :: (IsStream t, Monad m) => t m a -> m (Maybe a)
488-- head = foldrM (\x _ -> return $ Just x) (return Nothing)
489head m =
490 let stop = return Nothing
491 single a = return (Just a)
492 yieldk a _ = return (Just a)
493 in foldStream defState yieldk single stop m
495{-# INLINE tail #-}
496tail :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
497tail m =
498 let stop = return Nothing
499 single _ = return $ Just nil
500 yieldk _ r = return $ Just r
501 in foldStream defState yieldk single stop m
503{-# INLINE init #-}
504init :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
505init m = go1 m
506 where
507 go1 m1 = do
508 r <- uncons m1
509 case r of
510 Nothing -> return Nothing
511 Just (h, t) -> return . Just $ go h t
512 go p m1 = mkStream $ \_ yld sng stp ->
513 let single _ = sng p
514 yieldk a x = yld p $ go a x
515 in foldStream defState yieldk single stp m1
517{-# INLINE elem #-}
518elem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
519elem e m = go m
520 where
521 go m1 =
522 let stop = return False
523 single a = return (a == e)
524 yieldk a r = if a == e then return True else go r
525 in foldStream defState yieldk single stop m1
527{-# INLINE notElem #-}
528notElem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
529notElem e m = go m
530 where
531 go m1 =
532 let stop = return True
533 single a = return (a /= e)
534 yieldk a r = if a == e then return False else go r
535 in foldStream defState yieldk single stop m1
537{-# INLINABLE all #-}
538all :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
539all p m = go m
540 where
541 go m1 =
542 let single a | p a = return True
543 | otherwise = return False
544 yieldk a r | p a = go r
545 | otherwise = return False
546 in foldStream defState yieldk single (return True) m1
548{-# INLINABLE any #-}
549any :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
550any p m = go m
551 where
552 go m1 =
553 let single a | p a = return True
554 | otherwise = return False
555 yieldk a r | p a = return True
556 | otherwise = go r
557 in foldStream defState yieldk single (return False) m1
559-- | Extract the last element of the stream, if any.
560{-# INLINE last #-}
561last :: (IsStream t, Monad m) => t m a -> m (Maybe a)
562last = foldlx' (\_ y -> Just y) Nothing id
564{-# INLINE minimum #-}
565minimum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
566minimum m = go Nothing m
567 where
568 go Nothing m1 =
569 let stop = return Nothing
570 single a = return (Just a)
571 yieldk a r = go (Just a) r
572 in foldStream defState yieldk single stop m1
574 go (Just res) m1 =
575 let stop = return (Just res)
576 single a =
577 if res <= a
578 then return (Just res)
579 else return (Just a)
580 yieldk a r =
581 if res <= a
582 then go (Just res) r
583 else go (Just a) r
584 in foldStream defState yieldk single stop m1
586{-# INLINE minimumBy #-}
587minimumBy
588 :: (IsStream t, Monad m)
589 => (a -> a -> Ordering) -> t m a -> m (Maybe a)
590minimumBy cmp m = go Nothing m
591 where
592 go Nothing m1 =
593 let stop = return Nothing
594 single a = return (Just a)
595 yieldk a r = go (Just a) r
596 in foldStream defState yieldk single stop m1
598 go (Just res) m1 =
599 let stop = return (Just res)
600 single a = case cmp res a of
601 GT -> return (Just a)
602 _ -> return (Just res)
603 yieldk a r = case cmp res a of
604 GT -> go (Just a) r
605 _ -> go (Just res) r
606 in foldStream defState yieldk single stop m1
608{-# INLINE maximum #-}
609maximum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
610maximum m = go Nothing m
611 where
612 go Nothing m1 =
613 let stop = return Nothing
614 single a = return (Just a)
615 yieldk a r = go (Just a) r
616 in foldStream defState yieldk single stop m1
618 go (Just res) m1 =
619 let stop = return (Just res)
620 single a =
621 if res <= a
622 then return (Just a)
623 else return (Just res)
624 yieldk a r =
625 if res <= a
626 then go (Just a) r
627 else go (Just res) r
628 in foldStream defState yieldk single stop m1
630{-# INLINE maximumBy #-}
631maximumBy :: (IsStream t, Monad m) => (a -> a -> Ordering) -> t m a -> m (Maybe a)
632maximumBy cmp m = go Nothing m
633 where
634 go Nothing m1 =
635 let stop = return Nothing
636 single a = return (Just a)
637 yieldk a r = go (Just a) r
638 in foldStream defState yieldk single stop m1
640 go (Just res) m1 =
641 let stop = return (Just res)
642 single a = case cmp res a of
643 GT -> return (Just res)
644 _ -> return (Just a)
645 yieldk a r = case cmp res a of
646 GT -> go (Just res) r
647 _ -> go (Just a) r
648 in foldStream defState yieldk single stop m1
650{-# INLINE (!!) #-}
651(!!) :: (IsStream t, Monad m) => t m a -> Int -> m (Maybe a)
652m !! i = go i m
653 where
654 go n m1 =
655 let single a | n == 0 = return $ Just a
656 | otherwise = return Nothing
657 yieldk a x | n < 0 = return Nothing
658 | n == 0 = return $ Just a
659 | otherwise = go (n - 1) x
660 in foldStream defState yieldk single (return Nothing) m1
662{-# INLINE lookup #-}
663lookup :: (IsStream t, Monad m, Eq a) => a -> t m (a, b) -> m (Maybe b)
664lookup e m = go m
665 where
666 go m1 =
667 let single (a, b) | a == e = return $ Just b
668 | otherwise = return Nothing
669 yieldk (a, b) x | a == e = return $ Just b
670 | otherwise = go x
671 in foldStream defState yieldk single (return Nothing) m1
673{-# INLINE findM #-}
674findM :: (IsStream t, Monad m) => (a -> m Bool) -> t m a -> m (Maybe a)
675findM p m = go m
676 where
677 go m1 =
678 let single a = do
679 b <- p a
680 if b then return $ Just a else return Nothing
681 yieldk a x = do
682 b <- p a
683 if b then return $ Just a else go x
684 in foldStream defState yieldk single (return Nothing) m1
686{-# INLINE find #-}
687find :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m (Maybe a)
688find p = findM (return . p)
690{-# INLINE findIndices #-}
691findIndices :: IsStream t => (a -> Bool) -> t m a -> t m Int
692findIndices p = go 0
693 where
694 go offset m1 = mkStream $ \st yld sng stp ->
695 let single a | p a = sng offset
696 | otherwise = stp
697 yieldk a x | p a = yld offset $ go (offset + 1) x
698 | otherwise = foldStream (adaptState st) yld sng stp $
699 go (offset + 1) x
700 in foldStream (adaptState st) yieldk single stp m1
702------------------------------------------------------------------------------
703-- Map and Fold
704------------------------------------------------------------------------------
706-- | Apply a monadic action to each element of the stream and discard the
707-- output of the action.
708{-# INLINE mapM_ #-}
709mapM_ :: (IsStream t, Monad m) => (a -> m b) -> t m a -> m ()
710mapM_ f m = go m
711 where
712 go m1 =
713 let stop = return ()
714 single a = void (f a)
715 yieldk a r = f a >> go r
716 in foldStream defState yieldk single stop m1
718------------------------------------------------------------------------------
719-- Converting folds
720------------------------------------------------------------------------------
722{-# INLINABLE toList #-}
723toList :: (IsStream t, Monad m) => t m a -> m [a]
724toList = foldr (:) []
726{-# INLINE toStreamK #-}
727toStreamK :: Stream m a -> Stream m a
728toStreamK = id
730-- Based on suggestions by David Feuer and Pranay Sashank
731{-# INLINE hoist #-}
732hoist :: (IsStream t, Monad m, Monad n)
733 => (forall x. m x -> n x) -> t m a -> t n a
734hoist f str =
735 mkStream $ \st yld sng stp ->
736 let single = return . sng
737 yieldk a s = return $ yld a (hoist f s)
738 stop = return stp
739 state = adaptState st
740 in join . f $ foldStreamShared state yieldk single stop str
742-------------------------------------------------------------------------------
743-- Transformation by folding (Scans)
744-------------------------------------------------------------------------------
746{-# INLINE scanlx' #-}
747scanlx' :: IsStream t => (x -> a -> x) -> x -> (x -> b) -> t m a -> t m b
748scanlx' step begin done m =
749 cons (done begin) $ go m begin
750 where
751 go m1 !acc = mkStream $ \st yld sng stp ->
752 let single a = sng (done $ step acc a)
753 yieldk a r =
754 let s = step acc a
755 in yld (done s) (go r s)
756 in foldStream (adaptState st) yieldk single stp m1
758{-# INLINE scanl' #-}
759scanl' :: IsStream t => (b -> a -> b) -> b -> t m a -> t m b
760scanl' step begin = scanlx' step begin id
762-------------------------------------------------------------------------------
763-- Filtering
764-------------------------------------------------------------------------------
766{-# INLINE filter #-}
767filter :: IsStream t => (a -> Bool) -> t m a -> t m a
768filter 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 = foldStream st yieldk single stp r
775 in foldStream st yieldk single stp m1
777{-# INLINE take #-}
778take :: IsStream t => Int -> t m a -> t m a
779take n m = go n m
780 where
781 go n1 m1 = mkStream $ \st yld sng stp ->
782 let yieldk a r = yld a (go (n1 - 1) r)
783 in if n1 <= 0
784 then stp
785 else foldStream st yieldk sng stp m1
787{-# INLINE takeWhile #-}
788takeWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
789takeWhile p m = go m
790 where
791 go m1 = mkStream $ \st yld sng stp ->
792 let single a | p a = sng a
793 | otherwise = stp
794 yieldk a r | p a = yld a (go r)
795 | otherwise = stp
796 in foldStream st yieldk single stp m1
798{-# INLINE drop #-}
799drop :: IsStream t => Int -> t m a -> t m a
800drop n m = fromStream $ unShare (go n (toStream m))
801 where
802 go n1 m1 = mkStream $ \st yld sng stp ->
803 let single _ = stp
804 yieldk _ r = foldStreamShared st yld sng stp $ go (n1 - 1) r
805 -- Somehow "<=" check performs better than a ">"
806 in if n1 <= 0
807 then foldStreamShared st yld sng stp m1
808 else foldStreamShared st yieldk single stp m1
810{-# INLINE dropWhile #-}
811dropWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
812dropWhile p m = go m
813 where
814 go m1 = mkStream $ \st yld sng stp ->
815 let single a | p a = stp
816 | otherwise = sng a
817 yieldk a r | p a = foldStream st yieldk single stp r
818 | otherwise = yld a r
819 in foldStream st yieldk single stp m1
821-------------------------------------------------------------------------------
822-- Mapping
823-------------------------------------------------------------------------------
825-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
826-- other stream types to overload it.
827{-# INLINE sequence #-}
828sequence :: (IsStream t, MonadAsync m) => t m (m a) -> t m a
829sequence m = go m
830 where
831 go m1 = mkStream $ \st yld sng stp ->
832 let single ma = ma >>= sng
833 yieldk ma r = foldStreamShared st yld sng stp $ ma |: go r
834 in foldStream (adaptState st) yieldk single stp m1
836-------------------------------------------------------------------------------
837-- Inserting
838-------------------------------------------------------------------------------
840{-# INLINE intersperseM #-}
841intersperseM :: (IsStream t, MonadAsync m) => m a -> t m a -> t m a
842intersperseM a m = prependingStart m
843 where
844 prependingStart m1 = mkStream $ \st yld sng stp ->
845 let yieldk i x = foldStreamShared st yld sng stp $ return i |: go x
846 in foldStream st yieldk sng stp m1
847 go m2 = mkStream $ \st yld sng stp ->
848 let single i = foldStreamShared st yld sng stp $ a |: yield i
849 yieldk i x = foldStreamShared st yld sng stp $ a |: return i |: go x
850 in foldStream st yieldk single stp m2
852{-# INLINE intersperse #-}
853intersperse :: (IsStream t, MonadAsync m) => a -> t m a -> t m a
854intersperse a = intersperseM (return a)
856{-# INLINE insertBy #-}
857insertBy :: IsStream t => (a -> a -> Ordering) -> a -> t m a -> t m a
858insertBy cmp x m = go m
859 where
860 go m1 = mkStream $ \st yld _ _ ->
861 let single a = case cmp x a of
862 GT -> yld a (yield x)
863 _ -> yld x (yield a)
864 stop = yld x nil
865 yieldk a r = case cmp x a of
866 GT -> yld a (go r)
867 _ -> yld x (a `cons` r)
868 in foldStream st yieldk single stop m1
870------------------------------------------------------------------------------
871-- Deleting
872------------------------------------------------------------------------------
874{-# INLINE deleteBy #-}
875deleteBy :: IsStream t => (a -> a -> Bool) -> a -> t m a -> t m a
876deleteBy eq x m = go m
877 where
878 go m1 = mkStream $ \st yld sng stp ->
879 let single a = if eq x a then stp else sng a
880 yieldk a r = if eq x a
881 then foldStream st yld sng stp r
882 else yld a (go r)
883 in foldStream st yieldk single stp m1
885------------------------------------------------------------------------------
886-- Reordering
887------------------------------------------------------------------------------
889{-# INLINE reverse #-}
890reverse :: IsStream t => t m a -> t m a
891reverse = foldlS (flip cons) nil
893-------------------------------------------------------------------------------
894-- Map and Filter
895-------------------------------------------------------------------------------
897{-# INLINE mapMaybe #-}
898mapMaybe :: IsStream t => (a -> Maybe b) -> t m a -> t m b
899mapMaybe f m = go m
900 where
901 go m1 = mkStream $ \st yld sng stp ->
902 let single a = case f a of
903 Just b -> sng b
904 Nothing -> stp
905 yieldk a r = case f a of
906 Just b -> yld b $ go r
907 Nothing -> foldStream (adaptState st) yieldk single stp r
908 in foldStream (adaptState st) yieldk single stp m1
910------------------------------------------------------------------------------
911-- Serial Zipping
912------------------------------------------------------------------------------
914-- | Zip two streams serially using a pure zipping function.
915--
916-- @since 0.1.0
917{-# INLINABLE zipWith #-}
918zipWith :: IsStream t => (a -> b -> c) -> t m a -> t m b -> t m c
919zipWith f = go
920 where
921 go mx my = mkStream $ \st yld sng stp -> do
922 let merge a ra =
923 let single2 b = sng (f a b)
924 yield2 b rb = yld (f a b) (go ra rb)
925 in foldStream (adaptState st) yield2 single2 stp my
926 let single1 a = merge a nil
927 yield1 = merge
928 foldStream (adaptState st) yield1 single1 stp mx
930-- | Zip two streams serially using a monadic zipping function.
931--
932-- @since 0.1.0
933{-# INLINABLE zipWithM #-}
934zipWithM :: (IsStream t, Monad m) => (a -> b -> m c) -> t m a -> t m b -> t m c
935zipWithM f m1 m2 = go m1 m2
936 where
937 go mx my = mkStream $ \st yld sng stp -> do
938 let merge a ra =
939 let runIt x = foldStream st yld sng stp x
940 single2 b = f a b >>= sng
941 yield2 b rb = f a b >>= \x -> runIt (x `cons` go ra rb)
942 in foldStream (adaptState st) yield2 single2 stp my
943 let single1 a = merge a nil
944 yield1 = merge
945 foldStream (adaptState st) yield1 single1 stp mx
947------------------------------------------------------------------------------
948-- Merging
949------------------------------------------------------------------------------
951{-# INLINE mergeByM #-}
952mergeByM
953 :: (IsStream t, Monad m)
954 => (a -> a -> m Ordering) -> t m a -> t m a -> t m a
955mergeByM cmp = go
956 where
957 go mx my = mkStream $ \st yld sng stp -> do
958 let mergeWithY a ra =
959 let stop2 = foldStream st yld sng stp mx
960 single2 b = do
961 r <- cmp a b
962 case r of
963 GT -> yld b (go (a `cons` ra) nil)
964 _ -> yld a (go ra (b `cons` nil))
965 yield2 b rb = do
966 r <- cmp a b
967 case r of
968 GT -> yld b (go (a `cons` ra) rb)
969 _ -> yld a (go ra (b `cons` rb))
970 in foldStream st yield2 single2 stop2 my
971 let stopX = foldStream st yld sng stp my
972 singleX a = mergeWithY a nil
973 yieldX = mergeWithY
974 foldStream st yieldX singleX stopX mx
976{-# INLINABLE mergeBy #-}
977mergeBy
978 :: (IsStream t, Monad m)
979 => (a -> a -> Ordering) -> t m a -> t m a -> t m a
980mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
982------------------------------------------------------------------------------
983-- Transformation comprehensions
984------------------------------------------------------------------------------
986{-# INLINE the #-}
987the :: (Eq a, IsStream t, Monad m) => t m a -> m (Maybe a)
988the m = do
989 r <- uncons m
990 case r of
991 Nothing -> return Nothing
992 Just (h, t) -> go h t
993 where
994 go h m1 =
995 let single a | h == a = return $ Just h
996 | otherwise = return Nothing
997 yieldk a r | h == a = go h r
998 | otherwise = return Nothing
999 in foldStream defState yieldk single (return $ Just h) m1
1001------------------------------------------------------------------------------
1002-- Alternative & MonadPlus
1003------------------------------------------------------------------------------
1005_alt :: Stream m a -> Stream m a -> Stream m a
1006_alt m1 m2 = mkStream $ \st yld sng stp ->
1007 let stop = foldStream st yld sng stp m2
1008 in foldStream st yld sng stop m1
1010------------------------------------------------------------------------------
1011-- MonadReader
1012------------------------------------------------------------------------------
1014{-# INLINABLE withLocal #-}
1015withLocal :: MonadReader r m => (r -> r) -> Stream m a -> Stream m a
1016withLocal f m =
1017 mkStream $ \st yld sng stp ->
1018 let single = local f . sng
1019 yieldk a r = local f $ yld a (withLocal f r)
1020 in foldStream st yieldk single (local f stp) m
1022------------------------------------------------------------------------------
1023-- MonadError
1024------------------------------------------------------------------------------
1027-- XXX handle and test cross thread state transfer
1028withCatchError
1029 :: MonadError e m
1030 => Stream m a -> (e -> Stream m a) -> Stream m a
1031withCatchError m h =
1032 mkStream $ \_ stp sng yld ->
1033 let run x = unStream x Nothing stp sng yieldk
1034 handle r = r `catchError` \e -> run $ h e
1035 yieldk a r = yld a (withCatchError r h)
1036 in handle $ run m