git.iko.soy/streamly
on de31277 help sign in
streamly/src/Streamly/Internal/Data/Stream/StreamK.hs
1116 lines 35 kB history raw
1{-# LANGUAGE UndecidableInstances #-}
3#include "inline.hs"
5-- |
6-- Module : Streamly.Internal.Data.Stream.StreamK
7-- Copyright : (c) 2017 Composewell Technologies
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'
91 , foldOnce
93 -- ** Specialized Folds
94 , drain
95 , null
96 , head
97 , tail
98 , init
99 , elem
100 , notElem
101 , all
102 , any
103 , last
104 , minimum
105 , minimumBy
106 , maximum
107 , maximumBy
108 , findIndices
109 , lookup
110 , findM
111 , find
112 , (!!)
114 -- ** Map and Fold
115 , mapM_
117 -- ** Conversions
118 , toList
119 , toStreamK
120 , hoist
122 -- * Transformation
123 -- ** By folding (scans)
124 , scanl'
125 , scanlx'
127 -- ** Filtering
128 , filter
129 , take
130 , takeWhile
131 , drop
132 , dropWhile
134 -- ** Mapping
135 , map
136 , mapM
137 , mapMSerial
138 , sequence
140 -- ** Inserting
141 , intersperseM
142 , intersperse
143 , insertBy
145 -- ** Deleting
146 , deleteBy
148 -- ** Reordering
149 , reverse
151 -- ** Map and Filter
152 , mapMaybe
154 -- ** Zipping
155 , zipWith
156 , zipWithM
158 -- ** Merging
159 , mergeBy
160 , mergeByM
162 -- ** Nesting
163 , concatMapBy
164 , concatMap
165 , bindWith
166 , apWith
167 , apSerial
168 , apSerialDiscardFst
169 , apSerialDiscardSnd
171 -- ** Transformation comprehensions
172 , the
174 -- * Semigroup Style Composition
175 , serial
177 -- * Utilities
178 , consMStream
179 , withLocal
180 , mfix
182 -- * Deprecated
183 , Streaming -- deprecated
184 , once -- deprecated
185 )
186where
188import Control.Monad.Trans.Class (MonadTrans(lift))
189import Control.Monad (void, join)
190import Control.Monad.Reader.Class (MonadReader(..))
191import Data.Function (fix)
192import Prelude
193 hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,
194 take, filter, all, any, takeWhile, drop, dropWhile, minimum,
195 maximum, elem, notElem, null, head, tail, init, zipWith, lookup,
196 foldr1, (!!), replicate, reverse, concatMap, iterate)
197import qualified Prelude
199import Streamly.Internal.Data.SVar
200import Streamly.Internal.Data.Stream.StreamK.Type
202import qualified Streamly.Internal.Data.Fold.Types as FL
204-------------------------------------------------------------------------------
205-- Deconstruction
206-------------------------------------------------------------------------------
208{-# INLINE uncons #-}
209uncons :: (IsStream t, Monad m) => t m a -> m (Maybe (a, t m a))
210uncons m =
211 let stop = return Nothing
212 single a = return (Just (a, nil))
213 yieldk a r = return (Just (a, r))
214 in foldStream defState yieldk single stop m
216-------------------------------------------------------------------------------
217-- Generation
218-------------------------------------------------------------------------------
220{-# INLINE unfoldr #-}
221unfoldr :: IsStream t => (b -> Maybe (a, b)) -> b -> t m a
222unfoldr next s0 = build $ \yld stp ->
223 let go s =
224 case next s of
225 Just (a, b) -> yld a (go b)
226 Nothing -> stp
227 in go s0
229{-# INLINE unfoldrM #-}
230unfoldrM :: (IsStream t, MonadAsync m) => (b -> m (Maybe (a, b))) -> b -> t m a
231unfoldrM step = go
232 where
233 go s = sharedM $ \yld _ stp -> do
234 r <- step s
235 case r of
236 Just (a, b) -> yld a (go b)
237 Nothing -> stp
239{-
240-- Generalization of concurrent streams/SVar via unfoldr.
241--
242-- Unfold a value into monadic actions and then run the resulting monadic
243-- actions to generate a stream. Since the step of generating the monadic
244-- action and running them are decoupled we can run the monadic actions
245-- cooncurrently. For example, the seed could be a list of monadic actions or a
246-- pure stream of monadic actions.
247--
248-- We can have different flavors of this depending on the stream type t. The
249-- concurrent version could be async or ahead etc. Depending on how we queue
250-- back the feedback portion b, it could be DFS or BFS style.
251--
252unfoldrA :: (IsStream t, MonadAsync m) => (b -> Maybe (m a, b)) -> b -> t m a
253unfoldrA = undefined
254-}
256-------------------------------------------------------------------------------
257-- Special generation
258-------------------------------------------------------------------------------
260-- | Same as yieldM
261--
262-- @since 0.2.0
263{-# DEPRECATED once "Please use yieldM instead." #-}
264{-# INLINE once #-}
265once :: (Monad m, IsStream t) => m a -> t m a
266once = yieldM
268-- |
269-- @
270-- repeatM = fix . cons
271-- repeatM = cycle1 . yield
272-- @
273--
274-- Generate an infinite stream by repeating a monadic value.
275--
276-- /Internal/
277repeatM :: (IsStream t, MonadAsync m) => m a -> t m a
278repeatM = go
279 where go m = m |: go m
281-- Generate an infinite stream by repeating a pure value.
282--
283-- /Internal/
284{-# INLINE repeat #-}
285repeat :: IsStream t => a -> t m a
286repeat a = let x = cons a x in x
288{-# INLINE replicateM #-}
289replicateM :: (IsStream t, MonadAsync m) => Int -> m a -> t m a
290replicateM n m = go n
291 where
292 go cnt = if cnt <= 0 then nil else m |: go (cnt - 1)
294{-# INLINE replicate #-}
295replicate :: IsStream t => Int -> a -> t m a
296replicate n a = go n
297 where
298 go cnt = if cnt <= 0 then nil else a `cons` go (cnt - 1)
300{-# INLINE fromIndicesM #-}
301fromIndicesM :: (IsStream t, MonadAsync m) => (Int -> m a) -> t m a
302fromIndicesM gen = go 0
303 where
304 go i = mkStream $ \st stp sng yld -> do
305 foldStreamShared st stp sng yld (gen i |: go (i + 1))
307{-# INLINE fromIndices #-}
308fromIndices :: IsStream t => (Int -> a) -> t m a
309fromIndices gen = go 0
310 where
311 go n = (gen n) `cons` go (n + 1)
313{-# INLINE iterate #-}
314iterate :: IsStream t => (a -> a) -> a -> t m a
315iterate step = fromStream . go
316 where
317 go s = cons s (go (step s))
319{-# INLINE iterateM #-}
320iterateM :: (IsStream t, MonadAsync m) => (a -> m a) -> m a -> t m a
321iterateM step = go
322 where
323 go s = mkStream $ \st stp sng yld -> do
324 next <- s
325 foldStreamShared st stp sng yld (return next |: go (step next))
327-------------------------------------------------------------------------------
328-- Conversions
329-------------------------------------------------------------------------------
331-- |
332-- @
333-- fromFoldable = 'Prelude.foldr' 'cons' 'nil'
334-- @
335--
336-- Construct a stream from a 'Foldable' containing pure values:
337--
338-- @since 0.2.0
339{-# INLINE fromFoldable #-}
340fromFoldable :: (IsStream t, Foldable f) => f a -> t m a
341fromFoldable = Prelude.foldr cons nil
343{-# INLINE fromList #-}
344fromList :: IsStream t => [a] -> t m a
345fromList = fromFoldable
347{-# INLINE fromStreamK #-}
348fromStreamK :: IsStream t => Stream m a -> t m a
349fromStreamK = fromStream
351-------------------------------------------------------------------------------
352-- Elimination by Folding
353-------------------------------------------------------------------------------
355-- | Lazy right associative fold.
356{-# INLINE foldr #-}
357foldr :: (IsStream t, Monad m) => (a -> b -> b) -> b -> t m a -> m b
358foldr step acc = foldrM (\x xs -> xs >>= \b -> return (step x b)) (return acc)
360-- | Right associative fold to an arbitrary transformer monad.
361{-# INLINE foldrT #-}
362foldrT :: (IsStream t, Monad m, Monad (s m), MonadTrans s)
363 => (a -> s m b -> s m b) -> s m b -> t m a -> s m b
364foldrT step final m = go m
365 where
366 go m1 = do
367 res <- lift $ uncons m1
368 case res of
369 Just (h, t) -> step h (go t)
370 Nothing -> final
372{-# INLINE foldr1 #-}
373foldr1 :: (IsStream t, Monad m) => (a -> a -> a) -> t m a -> m (Maybe a)
374foldr1 step m = do
375 r <- uncons m
376 case r of
377 Nothing -> return Nothing
378 Just (h, t) -> fmap Just (go h t)
379 where
380 go p m1 =
381 let stp = return p
382 single a = return $ step a p
383 yieldk a r = fmap (step p) (go a r)
384 in foldStream defState yieldk single stp m1
386-- XXX replace the recursive "go" with explicit continuations.
387-- | Like 'foldx', but with a monadic step function.
388{-# INLINABLE foldlMx' #-}
389foldlMx' :: (IsStream t, Monad m)
390 => (x -> a -> m x) -> m x -> (x -> m b) -> t m a -> m b
391foldlMx' step begin done m = go begin m
392 where
393 go !acc m1 =
394 let stop = acc >>= done
395 single a = acc >>= \b -> step b a >>= done
396 yieldk a r = acc >>= \b -> step b a >>= \x -> go (return x) r
397 in foldStream defState yieldk single stop m1
399{-# INLINABLE foldOnce #-}
400foldOnce :: (IsStream t, Monad m) => FL.Fold m a b -> t m a -> m b
401foldOnce (FL.Fold step begin done) m = do
402 res <- begin
403 case res of
404 FL.Partial fs -> go fs m
405 FL.Done fb -> return fb
407 where
408 go !acc m1 =
409 let stop = done acc
410 single a = step acc a
411 >>= \case
412 FL.Partial s -> done s
413 FL.Done b1 -> return b1
414 yieldk a r = step acc a
415 >>= \case
416 FL.Partial s -> go s r
417 FL.Done b1 -> return b1
418 in foldStream defState yieldk single stop m1
420-- | Like 'foldl'' but with a monadic step function.
421{-# INLINE foldlM' #-}
422foldlM' :: (IsStream t, Monad m) => (b -> a -> m b) -> m b -> t m a -> m b
423foldlM' step begin = foldlMx' step begin return
425-- | Lazy left fold to a stream.
426{-# INLINE foldlS #-}
427foldlS :: IsStream t => (t m b -> a -> t m b) -> t m b -> t m a -> t m b
428foldlS step begin m = go begin m
429 where
430 go acc rest = mkStream $ \st yld sng stp ->
431 let run x = foldStream st yld sng stp x
432 stop = run acc
433 single a = run $ step acc a
434 yieldk a r = run $ go (step acc a) r
435 in foldStream (adaptState st) yieldk single stop rest
437-- | Lazy left fold to an arbitrary transformer monad.
438{-# INLINE foldlT #-}
439foldlT :: (IsStream t, Monad m, Monad (s m), MonadTrans s)
440 => (s m b -> a -> s m b) -> s m b -> t m a -> s m b
441foldlT step begin m = go begin m
442 where
443 go acc m1 = do
444 res <- lift $ uncons m1
445 case res of
446 Just (h, t) -> go (step acc h) t
447 Nothing -> acc
449------------------------------------------------------------------------------
450-- Specialized folds
451------------------------------------------------------------------------------
453-- XXX use foldrM to implement folds where possible
454-- XXX This (commented) definition of drain and mapM_ perform much better on
455-- some benchmarks but worse on others. Need to investigate why, may there is
456-- an optimization opportunity that we can exploit.
457-- drain = foldrM (\_ xs -> return () >> xs) (return ())
459-- |
460-- > drain = foldl' (\_ _ -> ()) ()
461-- > drain = mapM_ (\_ -> return ())
462{-# INLINE drain #-}
463drain :: (Monad m, IsStream t) => t m a -> m ()
464drain = foldrM (\_ xs -> xs) (return ())
465{-
466drain = go
467 where
468 go m1 =
469 let stop = return ()
470 single _ = return ()
471 yieldk _ r = go r
472 in foldStream defState yieldk single stop m1
473-}
475{-# INLINE null #-}
476null :: (IsStream t, Monad m) => t m a -> m Bool
477-- null = foldrM (\_ _ -> return True) (return False)
478null m =
479 let stop = return True
480 single _ = return False
481 yieldk _ _ = return False
482 in foldStream defState yieldk single stop m
484{-# INLINE head #-}
485head :: (IsStream t, Monad m) => t m a -> m (Maybe a)
486-- head = foldrM (\x _ -> return $ Just x) (return Nothing)
487head m =
488 let stop = return Nothing
489 single a = return (Just a)
490 yieldk a _ = return (Just a)
491 in foldStream defState yieldk single stop m
493{-# INLINE tail #-}
494tail :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
495tail m =
496 let stop = return Nothing
497 single _ = return $ Just nil
498 yieldk _ r = return $ Just r
499 in foldStream defState yieldk single stop m
501{-# INLINE headPartial #-}
502headPartial :: (IsStream t, Monad m) => t m a -> m a
503headPartial = foldrM (\x _ -> return x) (error "head of nil")
505{-# INLINE tailPartial #-}
506tailPartial :: IsStream t => t m a -> t m a
507tailPartial m = mkStream $ \st yld sng stp ->
508 let stop = error "tail of nil"
509 single _ = stp
510 yieldk _ r = foldStream st yld sng stp r
511 in foldStream st yieldk single stop m
513-- | We can define cyclic structures using @let@:
514--
515-- >>> let (a, b) = ([1, b], head a) in (a, b)
516-- ([1,1],1)
517--
518-- The function @fix@ defined as:
519--
520-- > fix f = let x = f x in x
521--
522-- ensures that the argument of a function and its output refer to the same
523-- lazy value @x@ i.e. the same location in memory. Thus @x@ can be defined
524-- in terms of itself, creating structures with cyclic references.
525--
526-- >>> f ~(a, b) = ([1, b], head a)
527-- >>> fix f
528-- ([1,1],1)
529--
530-- 'Control.Monad.mfix' is essentially the same as @fix@ but for monadic
531-- values.
532--
533-- Using 'mfix' for streams we can construct a stream in which each element of
534-- the stream is defined in a cyclic fashion. The argument of the function
535-- being fixed represents the current element of the stream which is being
536-- returned by the stream monad. Thus, we can use the argument to construct
537-- itself.
538--
539-- In the following example, the argument @action@ of the function @f@
540-- represents the tuple @(x,y)@ returned by it in a given iteration. We define
541-- the first element of the tuple in terms of the second.
542--
543-- @
544-- import Streamly.Internal.Data.Stream.IsStream as Stream
545-- import System.IO.Unsafe (unsafeInterleaveIO)
546--
547-- main = do
548-- Stream.mapM_ print $ Stream.mfix f
549--
550-- where
551--
552-- f action = do
553-- let incr n act = fmap ((+n) . snd) $ unsafeInterleaveIO act
554-- x <- Stream.fromListM [incr 1 action, incr 2 action]
555-- y <- Stream.fromList [4,5]
556-- return (x, y)
557-- @
558--
559-- Note: you cannot achieve this by just changing the order of the monad
560-- statements because that would change the order in which the stream elements
561-- are generated.
562--
563-- Note that the function @f@ must be lazy in its argument, that's why we use
564-- 'unsafeInterleaveIO' on @action@ because IO monad is strict.
565--
566-- /Internal/
568mfix :: (IsStream t, Monad m) => (m a -> t m a) -> t m a
569mfix f = mkStream $ \st yld sng stp ->
570 let single a = foldStream st yld sng stp $ a `cons` ys
571 yieldk a _ = foldStream st yld sng stp $ a `cons` ys
572 in foldStream st yieldk single stp xs
574 where
576 -- fix the head element of the stream
577 xs = fix (f . headPartial)
579 -- now fix the tail recursively
580 ys = mfix (tailPartial . f)
582{-# INLINE init #-}
583init :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
584init m = go1 m
585 where
586 go1 m1 = do
587 r <- uncons m1
588 case r of
589 Nothing -> return Nothing
590 Just (h, t) -> return . Just $ go h t
591 go p m1 = mkStream $ \_ yld sng stp ->
592 let single _ = sng p
593 yieldk a x = yld p $ go a x
594 in foldStream defState yieldk single stp m1
596{-# INLINE elem #-}
597elem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
598elem e m = go m
599 where
600 go m1 =
601 let stop = return False
602 single a = return (a == e)
603 yieldk a r = if a == e then return True else go r
604 in foldStream defState yieldk single stop m1
606{-# INLINE notElem #-}
607notElem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
608notElem e m = go m
609 where
610 go m1 =
611 let stop = return True
612 single a = return (a /= e)
613 yieldk a r = if a == e then return False else go r
614 in foldStream defState yieldk single stop m1
616{-# INLINABLE all #-}
617all :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
618all p m = go m
619 where
620 go m1 =
621 let single a | p a = return True
622 | otherwise = return False
623 yieldk a r | p a = go r
624 | otherwise = return False
625 in foldStream defState yieldk single (return True) m1
627{-# INLINABLE any #-}
628any :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
629any p m = go m
630 where
631 go m1 =
632 let single a | p a = return True
633 | otherwise = return False
634 yieldk a r | p a = return True
635 | otherwise = go r
636 in foldStream defState yieldk single (return False) m1
638-- | Extract the last element of the stream, if any.
639{-# INLINE last #-}
640last :: (IsStream t, Monad m) => t m a -> m (Maybe a)
641last = foldlx' (\_ y -> Just y) Nothing id
643{-# INLINE minimum #-}
644minimum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
645minimum 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 =
656 if res <= a
657 then return (Just res)
658 else return (Just a)
659 yieldk a r =
660 if res <= a
661 then go (Just res) r
662 else go (Just a) r
663 in foldStream defState yieldk single stop m1
665{-# INLINE minimumBy #-}
666minimumBy
667 :: (IsStream t, Monad m)
668 => (a -> a -> Ordering) -> t m a -> m (Maybe a)
669minimumBy cmp m = go Nothing m
670 where
671 go Nothing m1 =
672 let stop = return Nothing
673 single a = return (Just a)
674 yieldk a r = go (Just a) r
675 in foldStream defState yieldk single stop m1
677 go (Just res) m1 =
678 let stop = return (Just res)
679 single a = case cmp res a of
680 GT -> return (Just a)
681 _ -> return (Just res)
682 yieldk a r = case cmp res a of
683 GT -> go (Just a) r
684 _ -> go (Just res) r
685 in foldStream defState yieldk single stop m1
687{-# INLINE maximum #-}
688maximum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
689maximum m = go Nothing m
690 where
691 go Nothing m1 =
692 let stop = return Nothing
693 single a = return (Just a)
694 yieldk a r = go (Just a) r
695 in foldStream defState yieldk single stop m1
697 go (Just res) m1 =
698 let stop = return (Just res)
699 single a =
700 if res <= a
701 then return (Just a)
702 else return (Just res)
703 yieldk a r =
704 if res <= a
705 then go (Just a) r
706 else go (Just res) r
707 in foldStream defState yieldk single stop m1
709{-# INLINE maximumBy #-}
710maximumBy :: (IsStream t, Monad m) => (a -> a -> Ordering) -> t m a -> m (Maybe a)
711maximumBy cmp m = go Nothing m
712 where
713 go Nothing m1 =
714 let stop = return Nothing
715 single a = return (Just a)
716 yieldk a r = go (Just a) r
717 in foldStream defState yieldk single stop m1
719 go (Just res) m1 =
720 let stop = return (Just res)
721 single a = case cmp res a of
722 GT -> return (Just res)
723 _ -> return (Just a)
724 yieldk a r = case cmp res a of
725 GT -> go (Just res) r
726 _ -> go (Just a) r
727 in foldStream defState yieldk single stop m1
729{-# INLINE (!!) #-}
730(!!) :: (IsStream t, Monad m) => t m a -> Int -> m (Maybe a)
731m !! i = go i m
732 where
733 go n m1 =
734 let single a | n == 0 = return $ Just a
735 | otherwise = return Nothing
736 yieldk a x | n < 0 = return Nothing
737 | n == 0 = return $ Just a
738 | otherwise = go (n - 1) x
739 in foldStream defState yieldk single (return Nothing) m1
741{-# INLINE lookup #-}
742lookup :: (IsStream t, Monad m, Eq a) => a -> t m (a, b) -> m (Maybe b)
743lookup e m = go m
744 where
745 go m1 =
746 let single (a, b) | a == e = return $ Just b
747 | otherwise = return Nothing
748 yieldk (a, b) x | a == e = return $ Just b
749 | otherwise = go x
750 in foldStream defState yieldk single (return Nothing) m1
752{-# INLINE findM #-}
753findM :: (IsStream t, Monad m) => (a -> m Bool) -> t m a -> m (Maybe a)
754findM p m = go m
755 where
756 go m1 =
757 let single a = do
758 b <- p a
759 if b then return $ Just a else return Nothing
760 yieldk a x = do
761 b <- p a
762 if b then return $ Just a else go x
763 in foldStream defState yieldk single (return Nothing) m1
765{-# INLINE find #-}
766find :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m (Maybe a)
767find p = findM (return . p)
769{-# INLINE findIndices #-}
770findIndices :: IsStream t => (a -> Bool) -> t m a -> t m Int
771findIndices p = go 0
772 where
773 go offset m1 = mkStream $ \st yld sng stp ->
774 let single a | p a = sng offset
775 | otherwise = stp
776 yieldk a x | p a = yld offset $ go (offset + 1) x
777 | otherwise = foldStream (adaptState st) yld sng stp $
778 go (offset + 1) x
779 in foldStream (adaptState st) yieldk single stp m1
781------------------------------------------------------------------------------
782-- Map and Fold
783------------------------------------------------------------------------------
785-- | Apply a monadic action to each element of the stream and discard the
786-- output of the action.
787{-# INLINE mapM_ #-}
788mapM_ :: (IsStream t, Monad m) => (a -> m b) -> t m a -> m ()
789mapM_ f m = go m
790 where
791 go m1 =
792 let stop = return ()
793 single a = void (f a)
794 yieldk a r = f a >> go r
795 in foldStream defState yieldk single stop m1
797------------------------------------------------------------------------------
798-- Converting folds
799------------------------------------------------------------------------------
801{-# INLINABLE toList #-}
802toList :: (IsStream t, Monad m) => t m a -> m [a]
803toList = foldr (:) []
805{-# INLINE toStreamK #-}
806toStreamK :: Stream m a -> Stream m a
807toStreamK = id
809-- Based on suggestions by David Feuer and Pranay Sashank
810{-# INLINE hoist #-}
811hoist :: (IsStream t, Monad m, Monad n)
812 => (forall x. m x -> n x) -> t m a -> t n a
813hoist f str =
814 mkStream $ \st yld sng stp ->
815 let single = return . sng
816 yieldk a s = return $ yld a (hoist f s)
817 stop = return stp
818 state = adaptState st
819 in join . f $ foldStreamShared state yieldk single stop str
821-------------------------------------------------------------------------------
822-- Transformation by folding (Scans)
823-------------------------------------------------------------------------------
825{-# INLINE scanlx' #-}
826scanlx' :: IsStream t => (x -> a -> x) -> x -> (x -> b) -> t m a -> t m b
827scanlx' step begin done m =
828 cons (done begin) $ go m begin
829 where
830 go m1 !acc = mkStream $ \st yld sng stp ->
831 let single a = sng (done $ step acc a)
832 yieldk a r =
833 let s = step acc a
834 in yld (done s) (go r s)
835 in foldStream (adaptState st) yieldk single stp m1
837{-# INLINE scanl' #-}
838scanl' :: IsStream t => (b -> a -> b) -> b -> t m a -> t m b
839scanl' step begin = scanlx' step begin id
841-------------------------------------------------------------------------------
842-- Filtering
843-------------------------------------------------------------------------------
845{-# INLINE filter #-}
846filter :: IsStream t => (a -> Bool) -> t m a -> t m a
847filter p m = go m
848 where
849 go m1 = mkStream $ \st yld sng stp ->
850 let single a | p a = sng a
851 | otherwise = stp
852 yieldk a r | p a = yld a (go r)
853 | otherwise = foldStream st yieldk single stp r
854 in foldStream st yieldk single stp m1
856{-# INLINE take #-}
857take :: IsStream t => Int -> t m a -> t m a
858take n m = go n m
859 where
860 go n1 m1 = mkStream $ \st yld sng stp ->
861 let yieldk a r = yld a (go (n1 - 1) r)
862 in if n1 <= 0
863 then stp
864 else foldStream st yieldk sng stp m1
866{-# INLINE takeWhile #-}
867takeWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
868takeWhile p m = go m
869 where
870 go m1 = mkStream $ \st yld sng stp ->
871 let single a | p a = sng a
872 | otherwise = stp
873 yieldk a r | p a = yld a (go r)
874 | otherwise = stp
875 in foldStream st yieldk single stp m1
877{-# INLINE drop #-}
878drop :: IsStream t => Int -> t m a -> t m a
879drop n m = fromStream $ unShare (go n (toStream m))
880 where
881 go n1 m1 = mkStream $ \st yld sng stp ->
882 let single _ = stp
883 yieldk _ r = foldStreamShared st yld sng stp $ go (n1 - 1) r
884 -- Somehow "<=" check performs better than a ">"
885 in if n1 <= 0
886 then foldStreamShared st yld sng stp m1
887 else foldStreamShared st yieldk single stp m1
889{-# INLINE dropWhile #-}
890dropWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
891dropWhile p m = go m
892 where
893 go m1 = mkStream $ \st yld sng stp ->
894 let single a | p a = stp
895 | otherwise = sng a
896 yieldk a r | p a = foldStream st yieldk single stp r
897 | otherwise = yld a r
898 in foldStream st yieldk single stp m1
900-------------------------------------------------------------------------------
901-- Mapping
902-------------------------------------------------------------------------------
904-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
905-- other stream types to overload it.
906{-# INLINE sequence #-}
907sequence :: (IsStream t, MonadAsync m) => t m (m a) -> t m a
908sequence m = go m
909 where
910 go m1 = mkStream $ \st yld sng stp ->
911 let single ma = ma >>= sng
912 yieldk ma r = foldStreamShared st yld sng stp $ ma |: go r
913 in foldStream (adaptState st) yieldk single stp m1
915-------------------------------------------------------------------------------
916-- Inserting
917-------------------------------------------------------------------------------
919{-# INLINE intersperseM #-}
920intersperseM :: (IsStream t, MonadAsync m) => m a -> t m a -> t m a
921intersperseM a m = prependingStart m
922 where
923 prependingStart m1 = mkStream $ \st yld sng stp ->
924 let yieldk i x = foldStreamShared st yld sng stp $ return i |: go x
925 in foldStream st yieldk sng stp m1
926 go m2 = mkStream $ \st yld sng stp ->
927 let single i = foldStreamShared st yld sng stp $ a |: yield i
928 yieldk i x = foldStreamShared st yld sng stp $ a |: return i |: go x
929 in foldStream st yieldk single stp m2
931{-# INLINE intersperse #-}
932intersperse :: (IsStream t, MonadAsync m) => a -> t m a -> t m a
933intersperse a = intersperseM (return a)
935{-# INLINE insertBy #-}
936insertBy :: IsStream t => (a -> a -> Ordering) -> a -> t m a -> t m a
937insertBy cmp x m = go m
938 where
939 go m1 = mkStream $ \st yld _ _ ->
940 let single a = case cmp x a of
941 GT -> yld a (yield x)
942 _ -> yld x (yield a)
943 stop = yld x nil
944 yieldk a r = case cmp x a of
945 GT -> yld a (go r)
946 _ -> yld x (a `cons` r)
947 in foldStream st yieldk single stop m1
949------------------------------------------------------------------------------
950-- Deleting
951------------------------------------------------------------------------------
953{-# INLINE deleteBy #-}
954deleteBy :: IsStream t => (a -> a -> Bool) -> a -> t m a -> t m a
955deleteBy eq x m = go m
956 where
957 go m1 = mkStream $ \st yld sng stp ->
958 let single a = if eq x a then stp else sng a
959 yieldk a r = if eq x a
960 then foldStream st yld sng stp r
961 else yld a (go r)
962 in foldStream st yieldk single stp m1
964------------------------------------------------------------------------------
965-- Reordering
966------------------------------------------------------------------------------
968{-# INLINE reverse #-}
969reverse :: IsStream t => t m a -> t m a
970reverse = foldlS (flip cons) nil
972-------------------------------------------------------------------------------
973-- Map and Filter
974-------------------------------------------------------------------------------
976{-# INLINE mapMaybe #-}
977mapMaybe :: IsStream t => (a -> Maybe b) -> t m a -> t m b
978mapMaybe f m = go m
979 where
980 go m1 = mkStream $ \st yld sng stp ->
981 let single a = case f a of
982 Just b -> sng b
983 Nothing -> stp
984 yieldk a r = case f a of
985 Just b -> yld b $ go r
986 Nothing -> foldStream (adaptState st) yieldk single stp r
987 in foldStream (adaptState st) yieldk single stp m1
989------------------------------------------------------------------------------
990-- Serial Zipping
991------------------------------------------------------------------------------
993-- | Zip two streams serially using a pure zipping function.
994--
995-- @since 0.1.0
996{-# INLINABLE zipWith #-}
997zipWith :: IsStream t => (a -> b -> c) -> t m a -> t m b -> t m c
998zipWith f = go
999 where
1000 go mx my = mkStream $ \st yld sng stp -> do
1001 let merge a ra =
1002 let single2 b = sng (f a b)
1003 yield2 b rb = yld (f a b) (go ra rb)
1004 in foldStream (adaptState st) yield2 single2 stp my
1005 let single1 a = merge a nil
1006 yield1 = merge
1007 foldStream (adaptState st) yield1 single1 stp mx
1009-- | Zip two streams serially using a monadic zipping function.
1011-- @since 0.1.0
1012{-# INLINABLE zipWithM #-}
1013zipWithM :: (IsStream t, Monad m) => (a -> b -> m c) -> t m a -> t m b -> t m c
1014zipWithM f m1 m2 = go m1 m2
1015 where
1016 go mx my = mkStream $ \st yld sng stp -> do
1017 let merge a ra =
1018 let runIt x = foldStream st yld sng stp x
1019 single2 b = f a b >>= sng
1020 yield2 b rb = f a b >>= \x -> runIt (x `cons` go ra rb)
1021 in foldStream (adaptState st) yield2 single2 stp my
1022 let single1 a = merge a nil
1023 yield1 = merge
1024 foldStream (adaptState st) yield1 single1 stp mx
1026------------------------------------------------------------------------------
1027-- Merging
1028------------------------------------------------------------------------------
1030{-# INLINE mergeByM #-}
1031mergeByM
1032 :: (IsStream t, Monad m)
1033 => (a -> a -> m Ordering) -> t m a -> t m a -> t m a
1034mergeByM cmp = go
1035 where
1036 go mx my = mkStream $ \st yld sng stp -> do
1037 let mergeWithY a ra =
1038 let stop2 = foldStream st yld sng stp mx
1039 single2 b = do
1040 r <- cmp a b
1041 case r of
1042 GT -> yld b (go (a `cons` ra) nil)
1043 _ -> yld a (go ra (b `cons` nil))
1044 yield2 b rb = do
1045 r <- cmp a b
1046 case r of
1047 GT -> yld b (go (a `cons` ra) rb)
1048 _ -> yld a (go ra (b `cons` rb))
1049 in foldStream st yield2 single2 stop2 my
1050 let stopX = foldStream st yld sng stp my
1051 singleX a = mergeWithY a nil
1052 yieldX = mergeWithY
1053 foldStream st yieldX singleX stopX mx
1055{-# INLINABLE mergeBy #-}
1056mergeBy
1057 :: (IsStream t, Monad m)
1058 => (a -> a -> Ordering) -> t m a -> t m a -> t m a
1059mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
1061------------------------------------------------------------------------------
1062-- Transformation comprehensions
1063------------------------------------------------------------------------------
1065{-# INLINE the #-}
1066the :: (Eq a, IsStream t, Monad m) => t m a -> m (Maybe a)
1067the m = do
1068 r <- uncons m
1069 case r of
1070 Nothing -> return Nothing
1071 Just (h, t) -> go h t
1072 where
1073 go h m1 =
1074 let single a | h == a = return $ Just h
1075 | otherwise = return Nothing
1076 yieldk a r | h == a = go h r
1077 | otherwise = return Nothing
1078 in foldStream defState yieldk single (return $ Just h) m1
1080------------------------------------------------------------------------------
1081-- Alternative & MonadPlus
1082------------------------------------------------------------------------------
1084_alt :: Stream m a -> Stream m a -> Stream m a
1085_alt m1 m2 = mkStream $ \st yld sng stp ->
1086 let stop = foldStream st yld sng stp m2
1087 in foldStream st yld sng stop m1
1089------------------------------------------------------------------------------
1090-- MonadReader
1091------------------------------------------------------------------------------
1093{-# INLINABLE withLocal #-}
1094withLocal :: MonadReader r m => (r -> r) -> Stream m a -> Stream m a
1095withLocal f m =
1096 mkStream $ \st yld sng stp ->
1097 let single = local f . sng
1098 yieldk a r = local f $ yld a (withLocal f r)
1099 in foldStream st yieldk single (local f stp) m
1101------------------------------------------------------------------------------
1102-- MonadError
1103------------------------------------------------------------------------------
1106-- XXX handle and test cross thread state transfer
1107withCatchError
1108 :: MonadError e m
1109 => Stream m a -> (e -> Stream m a) -> Stream m a
1110withCatchError m h =
1111 mkStream $ \_ stp sng yld ->
1112 let run x = unStream x Nothing stp sng yieldk
1113 handle r = r `catchError` \e -> run $ h e
1114 yieldk a r = yld a (withCatchError r h)
1115 in handle $ run m