git.iko.soy/streamly
on cff0369 help sign in
streamly/src/Streamly/Internal/Data/Stream/StreamK.hs
1130 lines 36 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-- | We can define cyclic structures using @let@:
528--
529-- >>> let (a, b) = ([1, b], head a) in (a, b)
530-- ([1,1],1)
531--
532-- The function @fix@ defined as:
533--
534-- > fix f = let x = f x in x
535--
536-- ensures that the argument of a function and its output refer to the same
537-- lazy value @x@ i.e. the same location in memory. Thus @x@ can be defined
538-- in terms of itself, creating structures with cyclic references.
539--
540-- >>> f ~(a, b) = ([1, b], head a)
541-- >>> fix f
542-- ([1,1],1)
543--
544-- 'Control.Monad.mfix' is essentially the same as @fix@ but for monadic
545-- values.
546--
547-- Using 'mfix' for streams we can construct a stream in which each element of
548-- the stream is defined in a cyclic fashion. The argument of the function
549-- being fixed represents the current element of the stream which is being
550-- returned by the stream monad. Thus, we can use the argument to construct
551-- itself.
552--
553-- In the following example, the argument @action@ of the function @f@
554-- represents the tuple @(x,y)@ returned by it in a given iteration. We define
555-- the first element of the tuple in terms of the second.
556--
557-- @
558-- import Streamly.Internal.Data.Stream.IsStream as Stream
559-- import System.IO.Unsafe (unsafeInterleaveIO)
560--
561-- main = do
562-- Stream.mapM_ print $ Stream.mfix f
563--
564-- where
565--
566-- f action = do
567-- let incr n act = fmap ((+n) . snd) $ unsafeInterleaveIO act
568-- x <- Stream.fromListM [incr 1 action, incr 2 action]
569-- y <- Stream.fromList [4,5]
570-- return (x, y)
571-- @
572--
573-- Note: you cannot achieve this by just changing the order of the monad
574-- statements because that would change the order in which the stream elements
575-- are generated.
576--
577-- Note that the function @f@ must be lazy in its argument, that's why we use
578-- 'unsafeInterleaveIO' on @action@ because IO monad is strict.
579--
580-- /Internal/
582mfix :: (IsStream t, Monad m) => (m a -> t m a) -> t m a
583mfix f = mkStream $ \st yld sng stp ->
584 let single a = foldStream st yld sng stp $ a `cons` ys
585 yieldk a _ = foldStream st yld sng stp $ a `cons` ys
586 in foldStream st yieldk single stp xs
588 where
590 -- fix the head element of the stream
591 xs = fix (f . headPartial)
593 -- now fix the tail recursively
594 ys = mfix (tailPartial . f)
596{-# INLINE init #-}
597init :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
598init m = go1 m
599 where
600 go1 m1 = do
601 r <- uncons m1
602 case r of
603 Nothing -> return Nothing
604 Just (h, t) -> return . Just $ go h t
605 go p m1 = mkStream $ \_ yld sng stp ->
606 let single _ = sng p
607 yieldk a x = yld p $ go a x
608 in foldStream defState yieldk single stp m1
610{-# INLINE elem #-}
611elem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
612elem e m = go m
613 where
614 go m1 =
615 let stop = return False
616 single a = return (a == e)
617 yieldk a r = if a == e then return True else go r
618 in foldStream defState yieldk single stop m1
620{-# INLINE notElem #-}
621notElem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
622notElem e m = go m
623 where
624 go m1 =
625 let stop = return True
626 single a = return (a /= e)
627 yieldk a r = if a == e then return False else go r
628 in foldStream defState yieldk single stop m1
630{-# INLINABLE all #-}
631all :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
632all p m = go m
633 where
634 go m1 =
635 let single a | p a = return True
636 | otherwise = return False
637 yieldk a r | p a = go r
638 | otherwise = return False
639 in foldStream defState yieldk single (return True) m1
641{-# INLINABLE any #-}
642any :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
643any p m = go m
644 where
645 go m1 =
646 let single a | p a = return True
647 | otherwise = return False
648 yieldk a r | p a = return True
649 | otherwise = go r
650 in foldStream defState yieldk single (return False) m1
652-- | Extract the last element of the stream, if any.
653{-# INLINE last #-}
654last :: (IsStream t, Monad m) => t m a -> m (Maybe a)
655last = foldlx' (\_ y -> Just y) Nothing id
657{-# INLINE minimum #-}
658minimum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
659minimum m = go Nothing m
660 where
661 go Nothing m1 =
662 let stop = return Nothing
663 single a = return (Just a)
664 yieldk a r = go (Just a) r
665 in foldStream defState yieldk single stop m1
667 go (Just res) m1 =
668 let stop = return (Just res)
669 single a =
670 if res <= a
671 then return (Just res)
672 else return (Just a)
673 yieldk a r =
674 if res <= a
675 then go (Just res) r
676 else go (Just a) r
677 in foldStream defState yieldk single stop m1
679{-# INLINE minimumBy #-}
680minimumBy
681 :: (IsStream t, Monad m)
682 => (a -> a -> Ordering) -> t m a -> m (Maybe a)
683minimumBy cmp m = go Nothing m
684 where
685 go Nothing m1 =
686 let stop = return Nothing
687 single a = return (Just a)
688 yieldk a r = go (Just a) r
689 in foldStream defState yieldk single stop m1
691 go (Just res) m1 =
692 let stop = return (Just res)
693 single a = case cmp res a of
694 GT -> return (Just a)
695 _ -> return (Just res)
696 yieldk a r = case cmp res a of
697 GT -> go (Just a) r
698 _ -> go (Just res) r
699 in foldStream defState yieldk single stop m1
701{-# INLINE maximum #-}
702maximum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
703maximum m = go Nothing m
704 where
705 go Nothing m1 =
706 let stop = return Nothing
707 single a = return (Just a)
708 yieldk a r = go (Just a) r
709 in foldStream defState yieldk single stop m1
711 go (Just res) m1 =
712 let stop = return (Just res)
713 single a =
714 if res <= a
715 then return (Just a)
716 else return (Just res)
717 yieldk a r =
718 if res <= a
719 then go (Just a) r
720 else go (Just res) r
721 in foldStream defState yieldk single stop m1
723{-# INLINE maximumBy #-}
724maximumBy :: (IsStream t, Monad m) => (a -> a -> Ordering) -> t m a -> m (Maybe a)
725maximumBy cmp m = go Nothing m
726 where
727 go Nothing m1 =
728 let stop = return Nothing
729 single a = return (Just a)
730 yieldk a r = go (Just a) r
731 in foldStream defState yieldk single stop m1
733 go (Just res) m1 =
734 let stop = return (Just res)
735 single a = case cmp res a of
736 GT -> return (Just res)
737 _ -> return (Just a)
738 yieldk a r = case cmp res a of
739 GT -> go (Just res) r
740 _ -> go (Just a) r
741 in foldStream defState yieldk single stop m1
743{-# INLINE (!!) #-}
744(!!) :: (IsStream t, Monad m) => t m a -> Int -> m (Maybe a)
745m !! i = go i m
746 where
747 go n m1 =
748 let single a | n == 0 = return $ Just a
749 | otherwise = return Nothing
750 yieldk a x | n < 0 = return Nothing
751 | n == 0 = return $ Just a
752 | otherwise = go (n - 1) x
753 in foldStream defState yieldk single (return Nothing) m1
755{-# INLINE lookup #-}
756lookup :: (IsStream t, Monad m, Eq a) => a -> t m (a, b) -> m (Maybe b)
757lookup e m = go m
758 where
759 go m1 =
760 let single (a, b) | a == e = return $ Just b
761 | otherwise = return Nothing
762 yieldk (a, b) x | a == e = return $ Just b
763 | otherwise = go x
764 in foldStream defState yieldk single (return Nothing) m1
766{-# INLINE findM #-}
767findM :: (IsStream t, Monad m) => (a -> m Bool) -> t m a -> m (Maybe a)
768findM p m = go m
769 where
770 go m1 =
771 let single a = do
772 b <- p a
773 if b then return $ Just a else return Nothing
774 yieldk a x = do
775 b <- p a
776 if b then return $ Just a else go x
777 in foldStream defState yieldk single (return Nothing) m1
779{-# INLINE find #-}
780find :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m (Maybe a)
781find p = findM (return . p)
783{-# INLINE findIndices #-}
784findIndices :: IsStream t => (a -> Bool) -> t m a -> t m Int
785findIndices p = go 0
786 where
787 go offset m1 = mkStream $ \st yld sng stp ->
788 let single a | p a = sng offset
789 | otherwise = stp
790 yieldk a x | p a = yld offset $ go (offset + 1) x
791 | otherwise = foldStream (adaptState st) yld sng stp $
792 go (offset + 1) x
793 in foldStream (adaptState st) yieldk single stp m1
795------------------------------------------------------------------------------
796-- Map and Fold
797------------------------------------------------------------------------------
799-- | Apply a monadic action to each element of the stream and discard the
800-- output of the action.
801{-# INLINE mapM_ #-}
802mapM_ :: (IsStream t, Monad m) => (a -> m b) -> t m a -> m ()
803mapM_ f m = go m
804 where
805 go m1 =
806 let stop = return ()
807 single a = void (f a)
808 yieldk a r = f a >> go r
809 in foldStream defState yieldk single stop m1
811------------------------------------------------------------------------------
812-- Converting folds
813------------------------------------------------------------------------------
815{-# INLINABLE toList #-}
816toList :: (IsStream t, Monad m) => t m a -> m [a]
817toList = foldr (:) []
819{-# INLINE toStreamK #-}
820toStreamK :: Stream m a -> Stream m a
821toStreamK = id
823-- Based on suggestions by David Feuer and Pranay Sashank
824{-# INLINE hoist #-}
825hoist :: (IsStream t, Monad m, Monad n)
826 => (forall x. m x -> n x) -> t m a -> t n a
827hoist f str =
828 mkStream $ \st yld sng stp ->
829 let single = return . sng
830 yieldk a s = return $ yld a (hoist f s)
831 stop = return stp
832 state = adaptState st
833 in join . f $ foldStreamShared state yieldk single stop str
835-------------------------------------------------------------------------------
836-- Transformation by folding (Scans)
837-------------------------------------------------------------------------------
839{-# INLINE scanlx' #-}
840scanlx' :: IsStream t => (x -> a -> x) -> x -> (x -> b) -> t m a -> t m b
841scanlx' step begin done m =
842 cons (done begin) $ go m begin
843 where
844 go m1 !acc = mkStream $ \st yld sng stp ->
845 let single a = sng (done $ step acc a)
846 yieldk a r =
847 let s = step acc a
848 in yld (done s) (go r s)
849 in foldStream (adaptState st) yieldk single stp m1
851{-# INLINE scanl' #-}
852scanl' :: IsStream t => (b -> a -> b) -> b -> t m a -> t m b
853scanl' step begin = scanlx' step begin id
855-------------------------------------------------------------------------------
856-- Filtering
857-------------------------------------------------------------------------------
859{-# INLINE filter #-}
860filter :: IsStream t => (a -> Bool) -> t m a -> t m a
861filter p m = go m
862 where
863 go m1 = mkStream $ \st yld sng stp ->
864 let single a | p a = sng a
865 | otherwise = stp
866 yieldk a r | p a = yld a (go r)
867 | otherwise = foldStream st yieldk single stp r
868 in foldStream st yieldk single stp m1
870{-# INLINE take #-}
871take :: IsStream t => Int -> t m a -> t m a
872take n m = go n m
873 where
874 go n1 m1 = mkStream $ \st yld sng stp ->
875 let yieldk a r = yld a (go (n1 - 1) r)
876 in if n1 <= 0
877 then stp
878 else foldStream st yieldk sng stp m1
880{-# INLINE takeWhile #-}
881takeWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
882takeWhile p m = go m
883 where
884 go m1 = mkStream $ \st yld sng stp ->
885 let single a | p a = sng a
886 | otherwise = stp
887 yieldk a r | p a = yld a (go r)
888 | otherwise = stp
889 in foldStream st yieldk single stp m1
891{-# INLINE drop #-}
892drop :: IsStream t => Int -> t m a -> t m a
893drop n m = fromStream $ unShare (go n (toStream m))
894 where
895 go n1 m1 = mkStream $ \st yld sng stp ->
896 let single _ = stp
897 yieldk _ r = foldStreamShared st yld sng stp $ go (n1 - 1) r
898 -- Somehow "<=" check performs better than a ">"
899 in if n1 <= 0
900 then foldStreamShared st yld sng stp m1
901 else foldStreamShared st yieldk single stp m1
903{-# INLINE dropWhile #-}
904dropWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
905dropWhile p m = go m
906 where
907 go m1 = mkStream $ \st yld sng stp ->
908 let single a | p a = stp
909 | otherwise = sng a
910 yieldk a r | p a = foldStream st yieldk single stp r
911 | otherwise = yld a r
912 in foldStream st yieldk single stp m1
914-------------------------------------------------------------------------------
915-- Mapping
916-------------------------------------------------------------------------------
918-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
919-- other stream types to overload it.
920{-# INLINE sequence #-}
921sequence :: (IsStream t, MonadAsync m) => t m (m a) -> t m a
922sequence m = go m
923 where
924 go m1 = mkStream $ \st yld sng stp ->
925 let single ma = ma >>= sng
926 yieldk ma r = foldStreamShared st yld sng stp $ ma |: go r
927 in foldStream (adaptState st) yieldk single stp m1
929-------------------------------------------------------------------------------
930-- Inserting
931-------------------------------------------------------------------------------
933{-# INLINE intersperseM #-}
934intersperseM :: (IsStream t, MonadAsync m) => m a -> t m a -> t m a
935intersperseM a m = prependingStart m
936 where
937 prependingStart m1 = mkStream $ \st yld sng stp ->
938 let yieldk i x = foldStreamShared st yld sng stp $ return i |: go x
939 in foldStream st yieldk sng stp m1
940 go m2 = mkStream $ \st yld sng stp ->
941 let single i = foldStreamShared st yld sng stp $ a |: yield i
942 yieldk i x = foldStreamShared st yld sng stp $ a |: return i |: go x
943 in foldStream st yieldk single stp m2
945{-# INLINE intersperse #-}
946intersperse :: (IsStream t, MonadAsync m) => a -> t m a -> t m a
947intersperse a = intersperseM (return a)
949{-# INLINE insertBy #-}
950insertBy :: IsStream t => (a -> a -> Ordering) -> a -> t m a -> t m a
951insertBy cmp x m = go m
952 where
953 go m1 = mkStream $ \st yld _ _ ->
954 let single a = case cmp x a of
955 GT -> yld a (yield x)
956 _ -> yld x (yield a)
957 stop = yld x nil
958 yieldk a r = case cmp x a of
959 GT -> yld a (go r)
960 _ -> yld x (a `cons` r)
961 in foldStream st yieldk single stop m1
963------------------------------------------------------------------------------
964-- Deleting
965------------------------------------------------------------------------------
967{-# INLINE deleteBy #-}
968deleteBy :: IsStream t => (a -> a -> Bool) -> a -> t m a -> t m a
969deleteBy eq x m = go m
970 where
971 go m1 = mkStream $ \st yld sng stp ->
972 let single a = if eq x a then stp else sng a
973 yieldk a r = if eq x a
974 then foldStream st yld sng stp r
975 else yld a (go r)
976 in foldStream st yieldk single stp m1
978------------------------------------------------------------------------------
979-- Reordering
980------------------------------------------------------------------------------
982{-# INLINE reverse #-}
983reverse :: IsStream t => t m a -> t m a
984reverse = foldlS (flip cons) nil
986-------------------------------------------------------------------------------
987-- Map and Filter
988-------------------------------------------------------------------------------
990{-# INLINE mapMaybe #-}
991mapMaybe :: IsStream t => (a -> Maybe b) -> t m a -> t m b
992mapMaybe f m = go m
993 where
994 go m1 = mkStream $ \st yld sng stp ->
995 let single a = case f a of
996 Just b -> sng b
997 Nothing -> stp
998 yieldk a r = case f a of
999 Just b -> yld b $ go r
1000 Nothing -> foldStream (adaptState st) yieldk single stp r
1001 in foldStream (adaptState st) yieldk single stp m1
1003------------------------------------------------------------------------------
1004-- Serial Zipping
1005------------------------------------------------------------------------------
1007-- | Zip two streams serially using a pure zipping function.
1009-- @since 0.1.0
1010{-# INLINABLE zipWith #-}
1011zipWith :: IsStream t => (a -> b -> c) -> t m a -> t m b -> t m c
1012zipWith f = go
1013 where
1014 go mx my = mkStream $ \st yld sng stp -> do
1015 let merge a ra =
1016 let single2 b = sng (f a b)
1017 yield2 b rb = yld (f a b) (go ra rb)
1018 in foldStream (adaptState st) yield2 single2 stp my
1019 let single1 a = merge a nil
1020 yield1 = merge
1021 foldStream (adaptState st) yield1 single1 stp mx
1023-- | Zip two streams serially using a monadic zipping function.
1025-- @since 0.1.0
1026{-# INLINABLE zipWithM #-}
1027zipWithM :: (IsStream t, Monad m) => (a -> b -> m c) -> t m a -> t m b -> t m c
1028zipWithM f m1 m2 = go m1 m2
1029 where
1030 go mx my = mkStream $ \st yld sng stp -> do
1031 let merge a ra =
1032 let runIt x = foldStream st yld sng stp x
1033 single2 b = f a b >>= sng
1034 yield2 b rb = f a b >>= \x -> runIt (x `cons` go ra rb)
1035 in foldStream (adaptState st) yield2 single2 stp my
1036 let single1 a = merge a nil
1037 yield1 = merge
1038 foldStream (adaptState st) yield1 single1 stp mx
1040------------------------------------------------------------------------------
1041-- Merging
1042------------------------------------------------------------------------------
1044{-# INLINE mergeByM #-}
1045mergeByM
1046 :: (IsStream t, Monad m)
1047 => (a -> a -> m Ordering) -> t m a -> t m a -> t m a
1048mergeByM cmp = go
1049 where
1050 go mx my = mkStream $ \st yld sng stp -> do
1051 let mergeWithY a ra =
1052 let stop2 = foldStream st yld sng stp mx
1053 single2 b = do
1054 r <- cmp a b
1055 case r of
1056 GT -> yld b (go (a `cons` ra) nil)
1057 _ -> yld a (go ra (b `cons` nil))
1058 yield2 b rb = do
1059 r <- cmp a b
1060 case r of
1061 GT -> yld b (go (a `cons` ra) rb)
1062 _ -> yld a (go ra (b `cons` rb))
1063 in foldStream st yield2 single2 stop2 my
1064 let stopX = foldStream st yld sng stp my
1065 singleX a = mergeWithY a nil
1066 yieldX = mergeWithY
1067 foldStream st yieldX singleX stopX mx
1069{-# INLINABLE mergeBy #-}
1070mergeBy
1071 :: (IsStream t, Monad m)
1072 => (a -> a -> Ordering) -> t m a -> t m a -> t m a
1073mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
1075------------------------------------------------------------------------------
1076-- Transformation comprehensions
1077------------------------------------------------------------------------------
1079{-# INLINE the #-}
1080the :: (Eq a, IsStream t, Monad m) => t m a -> m (Maybe a)
1081the m = do
1082 r <- uncons m
1083 case r of
1084 Nothing -> return Nothing
1085 Just (h, t) -> go h t
1086 where
1087 go h m1 =
1088 let single a | h == a = return $ Just h
1089 | otherwise = return Nothing
1090 yieldk a r | h == a = go h r
1091 | otherwise = return Nothing
1092 in foldStream defState yieldk single (return $ Just h) m1
1094------------------------------------------------------------------------------
1095-- Alternative & MonadPlus
1096------------------------------------------------------------------------------
1098_alt :: Stream m a -> Stream m a -> Stream m a
1099_alt m1 m2 = mkStream $ \st yld sng stp ->
1100 let stop = foldStream st yld sng stp m2
1101 in foldStream st yld sng stop m1
1103------------------------------------------------------------------------------
1104-- MonadReader
1105------------------------------------------------------------------------------
1107{-# INLINABLE withLocal #-}
1108withLocal :: MonadReader r m => (r -> r) -> Stream m a -> Stream m a
1109withLocal f m =
1110 mkStream $ \st yld sng stp ->
1111 let single = local f . sng
1112 yieldk a r = local f $ yld a (withLocal f r)
1113 in foldStream st yieldk single (local f stp) m
1115------------------------------------------------------------------------------
1116-- MonadError
1117------------------------------------------------------------------------------
1120-- XXX handle and test cross thread state transfer
1121withCatchError
1122 :: MonadError e m
1123 => Stream m a -> (e -> Stream m a) -> Stream m a
1124withCatchError m h =
1125 mkStream $ \_ stp sng yld ->
1126 let run x = unStream x Nothing stp sng yieldk
1127 handle r = r `catchError` \e -> run $ h e
1128 yieldk a r = yld a (withCatchError r h)
1129 in handle $ run m