git.iko.soy/streamly
on a84e6cf help sign in
streamly/src/Streamly/Internal/Data/Stream/StreamK.hs
1092 lines 34 kB history raw
1{-# LANGUAGE UndecidableInstances #-}
3#include "inline.hs"
5-- |
6-- Module : Streamly.Internal.Data.Stream.StreamK
7-- Copyright : (c) 2017 Harendra Kumar
8--
9-- License : BSD3
10-- Maintainer : streamly@composewell.com
11-- Stability : experimental
12-- Portability : GHC
13--
14--
15-- Continuation passing style (CPS) stream implementation. The symbol 'K' below
16-- denotes a function as well as a Kontinuation.
17--
18-- @
19-- import qualified Streamly.Internal.Data.Stream.StreamK as K
20-- @
21--
22module Streamly.Internal.Data.Stream.StreamK
23 (
24 -- * A class for streams
25 IsStream (..)
26 , adapt
28 -- * The stream type
29 , Stream(..)
31 -- * Construction Primitives
32 , mkStream
33 , nil
34 , nilM
35 , cons
36 , (.:)
38 -- * Elimination Primitives
39 , foldStream
40 , foldStreamShared
42 -- * Transformation Primitives
43 , unShare
45 -- * Deconstruction
46 , uncons
48 -- * Generation
49 -- ** Unfolds
50 , unfoldr
51 , unfoldrM
53 -- ** Specialized Generation
54 , repeat
55 , repeatM
56 , replicate
57 , replicateM
58 , fromIndices
59 , fromIndicesM
60 , iterate
61 , iterateM
63 -- ** Conversions
64 , yield
65 , yieldM
66 , fromFoldable
67 , fromList
68 , fromStreamK
70 -- * foldr/build
71 , foldrS
72 , foldrSM
73 , buildS
74 , buildM
75 , augmentS
76 , augmentSM
78 -- * Elimination
79 -- ** General Folds
80 , foldr
81 , foldr1
82 , foldrM
83 , foldrT
85 , foldl'
86 , foldlM'
87 , foldlS
88 , foldlT
89 , foldlx'
90 , foldlMx'
92 -- ** Specialized Folds
93 , drain
94 , null
95 , head
96 , tail
97 , init
98 , elem
99 , notElem
100 , all
101 , any
102 , last
103 , minimum
104 , minimumBy
105 , maximum
106 , maximumBy
107 , findIndices
108 , lookup
109 , findM
110 , find
111 , (!!)
113 -- ** Map and Fold
114 , mapM_
116 -- ** Conversions
117 , toList
118 , toStreamK
119 , hoist
121 -- * Transformation
122 -- ** By folding (scans)
123 , scanl'
124 , scanlx'
126 -- ** Filtering
127 , filter
128 , take
129 , takeWhile
130 , drop
131 , dropWhile
133 -- ** Mapping
134 , map
135 , mapM
136 , mapMSerial
137 , sequence
139 -- ** Inserting
140 , intersperseM
141 , intersperse
142 , insertBy
144 -- ** Deleting
145 , deleteBy
147 -- ** Reordering
148 , reverse
150 -- ** Map and Filter
151 , mapMaybe
153 -- ** Zipping
154 , zipWith
155 , zipWithM
157 -- ** Merging
158 , mergeBy
159 , mergeByM
161 -- ** Nesting
162 , concatMapBy
163 , concatMap
164 , bindWith
165 , apWith
166 , apSerial
167 , apSerialDiscardFst
168 , apSerialDiscardSnd
170 -- ** Transformation comprehensions
171 , the
173 -- * Semigroup Style Composition
174 , serial
176 -- * Utilities
177 , consMStream
178 , withLocal
179 , mfix
181 -- * Deprecated
182 , Streaming -- deprecated
183 , once -- deprecated
184 )
185where
187import Control.Monad.Trans.Class (MonadTrans(lift))
188import Control.Monad (void, join)
189import Control.Monad.Reader.Class (MonadReader(..))
190import Data.Function (fix)
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, iterate)
196import qualified Prelude
198import Streamly.Internal.Data.SVar
199import Streamly.Internal.Data.Stream.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 monadic value.
272--
273-- /Internal/
274repeatM :: (IsStream t, MonadAsync m) => m a -> t m a
275repeatM = go
276 where go m = m |: go m
278-- Generate an infinite stream by repeating a pure value.
279--
280-- /Internal/
281{-# INLINE repeat #-}
282repeat :: IsStream t => a -> t m a
283repeat a = let x = cons a x in x
285{-# INLINE replicateM #-}
286replicateM :: (IsStream t, MonadAsync m) => Int -> m a -> t m a
287replicateM n m = go n
288 where
289 go cnt = if cnt <= 0 then nil else m |: go (cnt - 1)
291{-# INLINE replicate #-}
292replicate :: IsStream t => Int -> a -> t m a
293replicate n a = go n
294 where
295 go cnt = if cnt <= 0 then nil else a `cons` go (cnt - 1)
297{-# INLINE fromIndicesM #-}
298fromIndicesM :: (IsStream t, MonadAsync m) => (Int -> m a) -> t m a
299fromIndicesM gen = go 0
300 where
301 go i = mkStream $ \st stp sng yld -> do
302 foldStreamShared st stp sng yld (gen i |: go (i + 1))
304{-# INLINE fromIndices #-}
305fromIndices :: IsStream t => (Int -> a) -> t m a
306fromIndices gen = go 0
307 where
308 go n = (gen n) `cons` go (n + 1)
310{-# INLINE iterate #-}
311iterate :: IsStream t => (a -> a) -> a -> t m a
312iterate step = fromStream . go
313 where
314 go s = cons s (go (step s))
316{-# INLINE iterateM #-}
317iterateM :: (IsStream t, MonadAsync m) => (a -> m a) -> m a -> t m a
318iterateM step = go
319 where
320 go s = mkStream $ \st stp sng yld -> do
321 next <- s
322 foldStreamShared st stp sng yld (return next |: go (step next))
324-------------------------------------------------------------------------------
325-- Conversions
326-------------------------------------------------------------------------------
328-- |
329-- @
330-- fromFoldable = 'Prelude.foldr' 'cons' 'nil'
331-- @
332--
333-- Construct a stream from a 'Foldable' containing pure values:
334--
335-- @since 0.2.0
336{-# INLINE fromFoldable #-}
337fromFoldable :: (IsStream t, Foldable f) => f a -> t m a
338fromFoldable = Prelude.foldr cons nil
340{-# INLINE fromList #-}
341fromList :: IsStream t => [a] -> t m a
342fromList = fromFoldable
344{-# INLINE fromStreamK #-}
345fromStreamK :: IsStream t => Stream m a -> t m a
346fromStreamK = fromStream
348-------------------------------------------------------------------------------
349-- Elimination by Folding
350-------------------------------------------------------------------------------
352-- | Lazy right associative fold.
353{-# INLINE foldr #-}
354foldr :: (IsStream t, Monad m) => (a -> b -> b) -> b -> t m a -> m b
355foldr step acc = foldrM (\x xs -> xs >>= \b -> return (step x b)) (return acc)
357-- | Right associative fold to an arbitrary transformer monad.
358{-# INLINE foldrT #-}
359foldrT :: (IsStream t, Monad m, Monad (s m), MonadTrans s)
360 => (a -> s m b -> s m b) -> s m b -> t m a -> s m b
361foldrT step final m = go m
362 where
363 go m1 = do
364 res <- lift $ uncons m1
365 case res of
366 Just (h, t) -> step h (go t)
367 Nothing -> final
369{-# INLINE foldr1 #-}
370foldr1 :: (IsStream t, Monad m) => (a -> a -> a) -> t m a -> m (Maybe a)
371foldr1 step m = do
372 r <- uncons m
373 case r of
374 Nothing -> return Nothing
375 Just (h, t) -> fmap Just (go h t)
376 where
377 go p m1 =
378 let stp = return p
379 single a = return $ step a p
380 yieldk a r = fmap (step p) (go a r)
381 in foldStream defState yieldk single stp m1
383-- | Strict left fold with an extraction function. Like the standard strict
384-- left fold, but applies a user supplied extraction function (the third
385-- argument) to the folded value at the end. This is designed to work with the
386-- @foldl@ library. The suffix @x@ is a mnemonic for extraction.
387--
388-- Note that the accumulator is always evaluated including the initial value.
389{-# INLINE foldlx' #-}
390foldlx' :: forall t m a b x. (IsStream t, Monad m)
391 => (x -> a -> x) -> x -> (x -> b) -> t m a -> m b
392foldlx' step begin done m = get $ go m begin
393 where
394 {-# NOINLINE get #-}
395 get :: t m x -> m b
396 get m1 =
397 -- XXX we are not strictly evaluating the accumulator here. Is this
398 -- okay?
399 let single = return . done
400 -- XXX this is foldSingleton. why foldStreamShared?
401 in foldStreamShared undefined undefined single undefined m1
403 -- Note, this can be implemented by making a recursive call to "go",
404 -- however that is more expensive because of unnecessary recursion
405 -- that cannot be tail call optimized. Unfolding recursion explicitly via
406 -- continuations is much more efficient.
407 go :: t m a -> x -> t m x
408 go m1 !acc = mkStream $ \_ yld sng _ ->
409 let stop = sng acc
410 single a = sng $ step acc a
411 -- XXX this is foldNonEmptyStream
412 yieldk a r = foldStream defState yld sng undefined $
413 go r (step acc a)
414 in foldStream defState yieldk single stop m1
416-- | Strict left associative fold.
417{-# INLINE foldl' #-}
418foldl' :: (IsStream t, Monad m) => (b -> a -> b) -> b -> t m a -> m b
419foldl' step begin = foldlx' step begin id
421-- XXX replace the recursive "go" with explicit continuations.
422-- | Like 'foldx', but with a monadic step function.
423{-# INLINABLE foldlMx' #-}
424foldlMx' :: (IsStream t, Monad m)
425 => (x -> a -> m x) -> m x -> (x -> m b) -> t m a -> m b
426foldlMx' step begin done m = go begin m
427 where
428 go !acc m1 =
429 let stop = acc >>= done
430 single a = acc >>= \b -> step b a >>= done
431 yieldk a r = acc >>= \b -> step b a >>= \x -> go (return x) r
432 in foldStream defState yieldk single stop m1
434-- | Like 'foldl'' but with a monadic step function.
435{-# INLINE foldlM' #-}
436foldlM' :: (IsStream t, Monad m) => (b -> a -> m b) -> m b -> t m a -> m b
437foldlM' step begin = foldlMx' step begin return
439-- | Lazy left fold to a stream.
440{-# INLINE foldlS #-}
441foldlS :: IsStream t => (t m b -> a -> t m b) -> t m b -> t m a -> t m b
442foldlS step begin m = go begin m
443 where
444 go acc rest = mkStream $ \st yld sng stp ->
445 let run x = foldStream st yld sng stp x
446 stop = run acc
447 single a = run $ step acc a
448 yieldk a r = run $ go (step acc a) r
449 in foldStream (adaptState st) yieldk single stop rest
451-- | Lazy left fold to an arbitrary transformer monad.
452{-# INLINE foldlT #-}
453foldlT :: (IsStream t, Monad m, Monad (s m), MonadTrans s)
454 => (s m b -> a -> s m b) -> s m b -> t m a -> s m b
455foldlT step begin m = go begin m
456 where
457 go acc m1 = do
458 res <- lift $ uncons m1
459 case res of
460 Just (h, t) -> go (step acc h) t
461 Nothing -> acc
463------------------------------------------------------------------------------
464-- Specialized folds
465------------------------------------------------------------------------------
467-- XXX use foldrM to implement folds where possible
468-- XXX This (commented) definition of drain and mapM_ perform much better on
469-- some benchmarks but worse on others. Need to investigate why, may there is
470-- an optimization opportunity that we can exploit.
471-- drain = foldrM (\_ xs -> return () >> xs) (return ())
473-- |
474-- > drain = foldl' (\_ _ -> ()) ()
475-- > drain = mapM_ (\_ -> return ())
476{-# INLINE drain #-}
477drain :: (Monad m, IsStream t) => t m a -> m ()
478drain = foldrM (\_ xs -> xs) (return ())
479{-
480drain = go
481 where
482 go m1 =
483 let stop = return ()
484 single _ = return ()
485 yieldk _ r = go r
486 in foldStream defState yieldk single stop m1
487-}
489{-# INLINE null #-}
490null :: (IsStream t, Monad m) => t m a -> m Bool
491-- null = foldrM (\_ _ -> return True) (return False)
492null m =
493 let stop = return True
494 single _ = return False
495 yieldk _ _ = return False
496 in foldStream defState yieldk single stop m
498{-# INLINE head #-}
499head :: (IsStream t, Monad m) => t m a -> m (Maybe a)
500-- head = foldrM (\x _ -> return $ Just x) (return Nothing)
501head m =
502 let stop = return Nothing
503 single a = return (Just a)
504 yieldk a _ = return (Just a)
505 in foldStream defState yieldk single stop m
507{-# INLINE tail #-}
508tail :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
509tail m =
510 let stop = return Nothing
511 single _ = return $ Just nil
512 yieldk _ r = return $ Just r
513 in foldStream defState yieldk single stop m
515{-# INLINE headPartial #-}
516headPartial :: (IsStream t, Monad m) => t m a -> m a
517headPartial = foldrM (\x _ -> return x) (error "head of nil")
519{-# INLINE tailPartial #-}
520tailPartial :: IsStream t => t m a -> t m a
521tailPartial m = mkStream $ \st yld sng stp ->
522 let stop = error "tail of nil"
523 single _ = stp
524 yieldk _ r = foldStream st yld sng stp r
525 in foldStream st yieldk single stop m
527-- | Iterate a lazy function `f` of the shape `m a -> t m a` until it gets
528-- fully defined i.e. becomes independent of its argument action, then return
529-- the resulting value of the function (`t m a`).
530--
531-- It can be used to construct a stream that uses a cyclic definition. For
532-- example:
533--
534-- @
535-- import Streamly.Internal.Data.Stream.IsStream as S
536-- import System.IO.Unsafe (unsafeInterleaveIO)
537--
538-- main = do
539-- S.mapM_ print $ S.mfix $ \x -> do
540-- a <- S.fromList [1,2]
541-- b <- S.fromListM [return 3, unsafeInterleaveIO (fmap fst x)]
542-- return (a, b)
543-- @
544--
545-- Note that the function `f` must be lazy in its argument, that's why we use
546-- 'unsafeInterleaveIO' because IO monad is strict.
547--
548-- /Internal/
550mfix :: (IsStream t, Monad m) => (m a -> t m a) -> t m a
551mfix f = mkStream $ \st yld sng stp ->
552 let single a = foldStream st yld sng stp $ a `cons` ys
553 yieldk a _ = foldStream st yld sng stp $ a `cons` ys
554 in foldStream st yieldk single stp xs
555 where xs = fix (f . headPartial)
556 ys = mfix (tailPartial . f)
558{-# INLINE init #-}
559init :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
560init m = go1 m
561 where
562 go1 m1 = do
563 r <- uncons m1
564 case r of
565 Nothing -> return Nothing
566 Just (h, t) -> return . Just $ go h t
567 go p m1 = mkStream $ \_ yld sng stp ->
568 let single _ = sng p
569 yieldk a x = yld p $ go a x
570 in foldStream defState yieldk single stp m1
572{-# INLINE elem #-}
573elem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
574elem e m = go m
575 where
576 go m1 =
577 let stop = return False
578 single a = return (a == e)
579 yieldk a r = if a == e then return True else go r
580 in foldStream defState yieldk single stop m1
582{-# INLINE notElem #-}
583notElem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
584notElem e m = go m
585 where
586 go m1 =
587 let stop = return True
588 single a = return (a /= e)
589 yieldk a r = if a == e then return False else go r
590 in foldStream defState yieldk single stop m1
592{-# INLINABLE all #-}
593all :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
594all p m = go m
595 where
596 go m1 =
597 let single a | p a = return True
598 | otherwise = return False
599 yieldk a r | p a = go r
600 | otherwise = return False
601 in foldStream defState yieldk single (return True) m1
603{-# INLINABLE any #-}
604any :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
605any p m = go m
606 where
607 go m1 =
608 let single a | p a = return True
609 | otherwise = return False
610 yieldk a r | p a = return True
611 | otherwise = go r
612 in foldStream defState yieldk single (return False) m1
614-- | Extract the last element of the stream, if any.
615{-# INLINE last #-}
616last :: (IsStream t, Monad m) => t m a -> m (Maybe a)
617last = foldlx' (\_ y -> Just y) Nothing id
619{-# INLINE minimum #-}
620minimum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
621minimum m = go Nothing m
622 where
623 go Nothing m1 =
624 let stop = return Nothing
625 single a = return (Just a)
626 yieldk a r = go (Just a) r
627 in foldStream defState yieldk single stop m1
629 go (Just res) m1 =
630 let stop = return (Just res)
631 single a =
632 if res <= a
633 then return (Just res)
634 else return (Just a)
635 yieldk a r =
636 if res <= a
637 then go (Just res) r
638 else go (Just a) r
639 in foldStream defState yieldk single stop m1
641{-# INLINE minimumBy #-}
642minimumBy
643 :: (IsStream t, Monad m)
644 => (a -> a -> Ordering) -> t m a -> m (Maybe a)
645minimumBy cmp m = go Nothing m
646 where
647 go Nothing m1 =
648 let stop = return Nothing
649 single a = return (Just a)
650 yieldk a r = go (Just a) r
651 in foldStream defState yieldk single stop m1
653 go (Just res) m1 =
654 let stop = return (Just res)
655 single a = case cmp res a of
656 GT -> return (Just a)
657 _ -> return (Just res)
658 yieldk a r = case cmp res a of
659 GT -> go (Just a) r
660 _ -> go (Just res) r
661 in foldStream defState yieldk single stop m1
663{-# INLINE maximum #-}
664maximum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
665maximum m = go Nothing m
666 where
667 go Nothing m1 =
668 let stop = return Nothing
669 single a = return (Just a)
670 yieldk a r = go (Just a) r
671 in foldStream defState yieldk single stop m1
673 go (Just res) m1 =
674 let stop = return (Just res)
675 single a =
676 if res <= a
677 then return (Just a)
678 else return (Just res)
679 yieldk a r =
680 if res <= a
681 then go (Just a) r
682 else go (Just res) r
683 in foldStream defState yieldk single stop m1
685{-# INLINE maximumBy #-}
686maximumBy :: (IsStream t, Monad m) => (a -> a -> Ordering) -> t m a -> m (Maybe a)
687maximumBy cmp m = go Nothing m
688 where
689 go Nothing m1 =
690 let stop = return Nothing
691 single a = return (Just a)
692 yieldk a r = go (Just a) r
693 in foldStream defState yieldk single stop m1
695 go (Just res) m1 =
696 let stop = return (Just res)
697 single a = case cmp res a of
698 GT -> return (Just res)
699 _ -> return (Just a)
700 yieldk a r = case cmp res a of
701 GT -> go (Just res) r
702 _ -> go (Just a) r
703 in foldStream defState yieldk single stop m1
705{-# INLINE (!!) #-}
706(!!) :: (IsStream t, Monad m) => t m a -> Int -> m (Maybe a)
707m !! i = go i m
708 where
709 go n m1 =
710 let single a | n == 0 = return $ Just a
711 | otherwise = return Nothing
712 yieldk a x | n < 0 = return Nothing
713 | n == 0 = return $ Just a
714 | otherwise = go (n - 1) x
715 in foldStream defState yieldk single (return Nothing) m1
717{-# INLINE lookup #-}
718lookup :: (IsStream t, Monad m, Eq a) => a -> t m (a, b) -> m (Maybe b)
719lookup e m = go m
720 where
721 go m1 =
722 let single (a, b) | a == e = return $ Just b
723 | otherwise = return Nothing
724 yieldk (a, b) x | a == e = return $ Just b
725 | otherwise = go x
726 in foldStream defState yieldk single (return Nothing) m1
728{-# INLINE findM #-}
729findM :: (IsStream t, Monad m) => (a -> m Bool) -> t m a -> m (Maybe a)
730findM p m = go m
731 where
732 go m1 =
733 let single a = do
734 b <- p a
735 if b then return $ Just a else return Nothing
736 yieldk a x = do
737 b <- p a
738 if b then return $ Just a else go x
739 in foldStream defState yieldk single (return Nothing) m1
741{-# INLINE find #-}
742find :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m (Maybe a)
743find p = findM (return . p)
745{-# INLINE findIndices #-}
746findIndices :: IsStream t => (a -> Bool) -> t m a -> t m Int
747findIndices p = go 0
748 where
749 go offset m1 = mkStream $ \st yld sng stp ->
750 let single a | p a = sng offset
751 | otherwise = stp
752 yieldk a x | p a = yld offset $ go (offset + 1) x
753 | otherwise = foldStream (adaptState st) yld sng stp $
754 go (offset + 1) x
755 in foldStream (adaptState st) yieldk single stp m1
757------------------------------------------------------------------------------
758-- Map and Fold
759------------------------------------------------------------------------------
761-- | Apply a monadic action to each element of the stream and discard the
762-- output of the action.
763{-# INLINE mapM_ #-}
764mapM_ :: (IsStream t, Monad m) => (a -> m b) -> t m a -> m ()
765mapM_ f m = go m
766 where
767 go m1 =
768 let stop = return ()
769 single a = void (f a)
770 yieldk a r = f a >> go r
771 in foldStream defState yieldk single stop m1
773------------------------------------------------------------------------------
774-- Converting folds
775------------------------------------------------------------------------------
777{-# INLINABLE toList #-}
778toList :: (IsStream t, Monad m) => t m a -> m [a]
779toList = foldr (:) []
781{-# INLINE toStreamK #-}
782toStreamK :: Stream m a -> Stream m a
783toStreamK = id
785-- Based on suggestions by David Feuer and Pranay Sashank
786{-# INLINE hoist #-}
787hoist :: (IsStream t, Monad m, Monad n)
788 => (forall x. m x -> n x) -> t m a -> t n a
789hoist f str =
790 mkStream $ \st yld sng stp ->
791 let single = return . sng
792 yieldk a s = return $ yld a (hoist f s)
793 stop = return stp
794 state = adaptState st
795 in join . f $ foldStreamShared state yieldk single stop str
797-------------------------------------------------------------------------------
798-- Transformation by folding (Scans)
799-------------------------------------------------------------------------------
801{-# INLINE scanlx' #-}
802scanlx' :: IsStream t => (x -> a -> x) -> x -> (x -> b) -> t m a -> t m b
803scanlx' step begin done m =
804 cons (done begin) $ go m begin
805 where
806 go m1 !acc = mkStream $ \st yld sng stp ->
807 let single a = sng (done $ step acc a)
808 yieldk a r =
809 let s = step acc a
810 in yld (done s) (go r s)
811 in foldStream (adaptState st) yieldk single stp m1
813{-# INLINE scanl' #-}
814scanl' :: IsStream t => (b -> a -> b) -> b -> t m a -> t m b
815scanl' step begin = scanlx' step begin id
817-------------------------------------------------------------------------------
818-- Filtering
819-------------------------------------------------------------------------------
821{-# INLINE filter #-}
822filter :: IsStream t => (a -> Bool) -> t m a -> t m a
823filter p m = go m
824 where
825 go m1 = mkStream $ \st yld sng stp ->
826 let single a | p a = sng a
827 | otherwise = stp
828 yieldk a r | p a = yld a (go r)
829 | otherwise = foldStream st yieldk single stp r
830 in foldStream st yieldk single stp m1
832{-# INLINE take #-}
833take :: IsStream t => Int -> t m a -> t m a
834take n m = go n m
835 where
836 go n1 m1 = mkStream $ \st yld sng stp ->
837 let yieldk a r = yld a (go (n1 - 1) r)
838 in if n1 <= 0
839 then stp
840 else foldStream st yieldk sng stp m1
842{-# INLINE takeWhile #-}
843takeWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
844takeWhile p m = go m
845 where
846 go m1 = mkStream $ \st yld sng stp ->
847 let single a | p a = sng a
848 | otherwise = stp
849 yieldk a r | p a = yld a (go r)
850 | otherwise = stp
851 in foldStream st yieldk single stp m1
853{-# INLINE drop #-}
854drop :: IsStream t => Int -> t m a -> t m a
855drop n m = fromStream $ unShare (go n (toStream m))
856 where
857 go n1 m1 = mkStream $ \st yld sng stp ->
858 let single _ = stp
859 yieldk _ r = foldStreamShared st yld sng stp $ go (n1 - 1) r
860 -- Somehow "<=" check performs better than a ">"
861 in if n1 <= 0
862 then foldStreamShared st yld sng stp m1
863 else foldStreamShared st yieldk single stp m1
865{-# INLINE dropWhile #-}
866dropWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
867dropWhile p m = go m
868 where
869 go m1 = mkStream $ \st yld sng stp ->
870 let single a | p a = stp
871 | otherwise = sng a
872 yieldk a r | p a = foldStream st yieldk single stp r
873 | otherwise = yld a r
874 in foldStream st yieldk single stp m1
876-------------------------------------------------------------------------------
877-- Mapping
878-------------------------------------------------------------------------------
880-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
881-- other stream types to overload it.
882{-# INLINE sequence #-}
883sequence :: (IsStream t, MonadAsync m) => t m (m a) -> t m a
884sequence m = go m
885 where
886 go m1 = mkStream $ \st yld sng stp ->
887 let single ma = ma >>= sng
888 yieldk ma r = foldStreamShared st yld sng stp $ ma |: go r
889 in foldStream (adaptState st) yieldk single stp m1
891-------------------------------------------------------------------------------
892-- Inserting
893-------------------------------------------------------------------------------
895{-# INLINE intersperseM #-}
896intersperseM :: (IsStream t, MonadAsync m) => m a -> t m a -> t m a
897intersperseM a m = prependingStart m
898 where
899 prependingStart m1 = mkStream $ \st yld sng stp ->
900 let yieldk i x = foldStreamShared st yld sng stp $ return i |: go x
901 in foldStream st yieldk sng stp m1
902 go m2 = mkStream $ \st yld sng stp ->
903 let single i = foldStreamShared st yld sng stp $ a |: yield i
904 yieldk i x = foldStreamShared st yld sng stp $ a |: return i |: go x
905 in foldStream st yieldk single stp m2
907{-# INLINE intersperse #-}
908intersperse :: (IsStream t, MonadAsync m) => a -> t m a -> t m a
909intersperse a = intersperseM (return a)
911{-# INLINE insertBy #-}
912insertBy :: IsStream t => (a -> a -> Ordering) -> a -> t m a -> t m a
913insertBy cmp x m = go m
914 where
915 go m1 = mkStream $ \st yld _ _ ->
916 let single a = case cmp x a of
917 GT -> yld a (yield x)
918 _ -> yld x (yield a)
919 stop = yld x nil
920 yieldk a r = case cmp x a of
921 GT -> yld a (go r)
922 _ -> yld x (a `cons` r)
923 in foldStream st yieldk single stop m1
925------------------------------------------------------------------------------
926-- Deleting
927------------------------------------------------------------------------------
929{-# INLINE deleteBy #-}
930deleteBy :: IsStream t => (a -> a -> Bool) -> a -> t m a -> t m a
931deleteBy eq x m = go m
932 where
933 go m1 = mkStream $ \st yld sng stp ->
934 let single a = if eq x a then stp else sng a
935 yieldk a r = if eq x a
936 then foldStream st yld sng stp r
937 else yld a (go r)
938 in foldStream st yieldk single stp m1
940------------------------------------------------------------------------------
941-- Reordering
942------------------------------------------------------------------------------
944{-# INLINE reverse #-}
945reverse :: IsStream t => t m a -> t m a
946reverse = foldlS (flip cons) nil
948-------------------------------------------------------------------------------
949-- Map and Filter
950-------------------------------------------------------------------------------
952{-# INLINE mapMaybe #-}
953mapMaybe :: IsStream t => (a -> Maybe b) -> t m a -> t m b
954mapMaybe f m = go m
955 where
956 go m1 = mkStream $ \st yld sng stp ->
957 let single a = case f a of
958 Just b -> sng b
959 Nothing -> stp
960 yieldk a r = case f a of
961 Just b -> yld b $ go r
962 Nothing -> foldStream (adaptState st) yieldk single stp r
963 in foldStream (adaptState st) yieldk single stp m1
965------------------------------------------------------------------------------
966-- Serial Zipping
967------------------------------------------------------------------------------
969-- | Zip two streams serially using a pure zipping function.
970--
971-- @since 0.1.0
972{-# INLINABLE zipWith #-}
973zipWith :: IsStream t => (a -> b -> c) -> t m a -> t m b -> t m c
974zipWith f = go
975 where
976 go mx my = mkStream $ \st yld sng stp -> do
977 let merge a ra =
978 let single2 b = sng (f a b)
979 yield2 b rb = yld (f a b) (go ra rb)
980 in foldStream (adaptState st) yield2 single2 stp my
981 let single1 a = merge a nil
982 yield1 = merge
983 foldStream (adaptState st) yield1 single1 stp mx
985-- | Zip two streams serially using a monadic zipping function.
986--
987-- @since 0.1.0
988{-# INLINABLE zipWithM #-}
989zipWithM :: (IsStream t, Monad m) => (a -> b -> m c) -> t m a -> t m b -> t m c
990zipWithM f m1 m2 = go m1 m2
991 where
992 go mx my = mkStream $ \st yld sng stp -> do
993 let merge a ra =
994 let runIt x = foldStream st yld sng stp x
995 single2 b = f a b >>= sng
996 yield2 b rb = f a b >>= \x -> runIt (x `cons` go ra rb)
997 in foldStream (adaptState st) yield2 single2 stp my
998 let single1 a = merge a nil
999 yield1 = merge
1000 foldStream (adaptState st) yield1 single1 stp mx
1002------------------------------------------------------------------------------
1003-- Merging
1004------------------------------------------------------------------------------
1006{-# INLINE mergeByM #-}
1007mergeByM
1008 :: (IsStream t, Monad m)
1009 => (a -> a -> m Ordering) -> t m a -> t m a -> t m a
1010mergeByM cmp = go
1011 where
1012 go mx my = mkStream $ \st yld sng stp -> do
1013 let mergeWithY a ra =
1014 let stop2 = foldStream st yld sng stp mx
1015 single2 b = do
1016 r <- cmp a b
1017 case r of
1018 GT -> yld b (go (a `cons` ra) nil)
1019 _ -> yld a (go ra (b `cons` nil))
1020 yield2 b rb = do
1021 r <- cmp a b
1022 case r of
1023 GT -> yld b (go (a `cons` ra) rb)
1024 _ -> yld a (go ra (b `cons` rb))
1025 in foldStream st yield2 single2 stop2 my
1026 let stopX = foldStream st yld sng stp my
1027 singleX a = mergeWithY a nil
1028 yieldX = mergeWithY
1029 foldStream st yieldX singleX stopX mx
1031{-# INLINABLE mergeBy #-}
1032mergeBy
1033 :: (IsStream t, Monad m)
1034 => (a -> a -> Ordering) -> t m a -> t m a -> t m a
1035mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
1037------------------------------------------------------------------------------
1038-- Transformation comprehensions
1039------------------------------------------------------------------------------
1041{-# INLINE the #-}
1042the :: (Eq a, IsStream t, Monad m) => t m a -> m (Maybe a)
1043the m = do
1044 r <- uncons m
1045 case r of
1046 Nothing -> return Nothing
1047 Just (h, t) -> go h t
1048 where
1049 go h m1 =
1050 let single a | h == a = return $ Just h
1051 | otherwise = return Nothing
1052 yieldk a r | h == a = go h r
1053 | otherwise = return Nothing
1054 in foldStream defState yieldk single (return $ Just h) m1
1056------------------------------------------------------------------------------
1057-- Alternative & MonadPlus
1058------------------------------------------------------------------------------
1060_alt :: Stream m a -> Stream m a -> Stream m a
1061_alt m1 m2 = mkStream $ \st yld sng stp ->
1062 let stop = foldStream st yld sng stp m2
1063 in foldStream st yld sng stop m1
1065------------------------------------------------------------------------------
1066-- MonadReader
1067------------------------------------------------------------------------------
1069{-# INLINABLE withLocal #-}
1070withLocal :: MonadReader r m => (r -> r) -> Stream m a -> Stream m a
1071withLocal f m =
1072 mkStream $ \st yld sng stp ->
1073 let single = local f . sng
1074 yieldk a r = local f $ yld a (withLocal f r)
1075 in foldStream st yieldk single (local f stp) m
1077------------------------------------------------------------------------------
1078-- MonadError
1079------------------------------------------------------------------------------
1082-- XXX handle and test cross thread state transfer
1083withCatchError
1084 :: MonadError e m
1085 => Stream m a -> (e -> Stream m a) -> Stream m a
1086withCatchError m h =
1087 mkStream $ \_ stp sng yld ->
1088 let run x = unStream x Nothing stp sng yieldk
1089 handle r = r `catchError` \e -> run $ h e
1090 yieldk a r = yld a (withCatchError r h)
1091 in handle $ run m