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