on 6b192fb help sign in
953 lines 31 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE CPP #-}
3{-# LANGUAGE ConstraintKinds #-}
4{-# LANGUAGE FlexibleContexts #-}
5{-# LANGUAGE FlexibleInstances #-}
6{-# LANGUAGE InstanceSigs #-}
7{-# LANGUAGE LambdaCase #-}
8{-# LANGUAGE MagicHash #-}
9{-# LANGUAGE MultiParamTypeClasses #-}
10{-# LANGUAGE RankNTypes #-}
11{-# LANGUAGE UnboxedTuples #-}
12{-# LANGUAGE UndecidableInstances #-} -- XXX
14-- |
15-- Module : Streamly.Streams.StreamK
16-- Copyright : (c) 2017 Harendra Kumar
17--
18-- License : BSD3
19-- Maintainer : harendra.kumar@gmail.com
20-- Stability : experimental
21-- Portability : GHC
22--
23--
24-- Continuation passing style (CPS) stream implementation. The symbol 'K' below
25-- denotes a function as well as a Kontinuation.
26--
27-- @
28-- import qualified Streamly.Streams.StreamK as K
29-- @
30--
31module Streamly.Streams.StreamK
32 (
33 -- * A class for streams
34 IsStream (..)
35 , adapt
37 -- * The stream type
38 , Stream (..)
40 -- * Construction
41 , mkStream
42 , nil
43 , cons
44 , (.:)
46 -- * Asynchronous construction
47 , nilK
48 , yieldK
49 , consK
51 -- * Deconstruction
52 , uncons
54 -- * Generation
55 -- ** Unfolds
56 , unfoldr
57 , unfoldrM
59 -- ** Specialized Generation
60 , repeat
62 -- ** Conversions
63 , yield
64 , yieldM
65 , fromFoldable
66 , fromList
67 , fromStreamK
69 -- * Elimination
70 -- ** General Folds
71 , foldStream
72 , foldr
73 , foldrM
74 , foldl'
75 , foldlM'
76 , foldx
77 , foldxM
79 -- ** Specialized Folds
80 , runStream
81 , null
82 , head
83 , tail
84 , elem
85 , notElem
86 , all
87 , any
88 , last
89 , minimum
90 , maximum
92 -- ** Map and Fold
93 , mapM_
95 -- ** Conversions
96 , toList
97 , toStreamK
99 -- * Transformation
100 -- ** By folding (scans)
101 , scanl'
102 , scanx
104 -- ** Filtering
105 , filter
106 , take
107 , takeWhile
108 , drop
109 , dropWhile
111 -- ** Mapping
112 , map
113 , mapM
114 , sequence
116 -- ** Map and Filter
117 , mapMaybe
119 -- ** Zipping
120 , zipWith
121 , zipWithM
123 -- * Semigroup Style Composition
124 , serial
126 -- * Utilities
127 , consMSerial
128 , bindWith
129 , withLocal
131 -- * Deprecated
132 , Streaming -- deprecated
133 , once -- deprecated
134 )
135where
137import Control.Monad (void)
138import Control.Monad.Reader.Class (MonadReader(..))
139import Control.Monad.Trans.Class (MonadTrans(lift))
140import Data.Semigroup (Semigroup(..))
141import Prelude
142 hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,
143 take, filter, all, any, takeWhile, drop, dropWhile, minimum,
144 maximum, elem, notElem, null, head, tail, zipWith)
145import qualified Prelude
147import Streamly.SVar
149------------------------------------------------------------------------------
150-- The basic stream type
151------------------------------------------------------------------------------
153-- | The type @Stream m a@ represents a monadic stream of values of type 'a'
154-- constructed using actions in monad 'm'. It uses stop, singleton and yield
155-- continuations equivalent to the following direct style type:
156--
157-- @
158-- data Stream m a = Stop | Singleton a | Yield a (Stream m a)
159-- @
160--
161-- To facilitate parallel composition we maintain a local state in an 'SVar'
162-- that is shared across and is used for synchronization of the streams being
163-- composed.
164--
165-- The singleton case can be expressed in terms of stop and yield but we have
166-- it as a separate case to optimize composition operations for streams with
167-- single element. We build singleton streams in the implementation of 'pure'
168-- for Applicative and Monad, and in 'lift' for MonadTrans.
169--
170newtype Stream m a =
171 Stream {
172 unStream :: forall r.
173 State Stream m a -- state
174 -> m r -- stop
175 -> (a -> m r) -- singleton
176 -> (a -> Stream m a -> m r) -- yield
177 -> m r
178 }
180------------------------------------------------------------------------------
181-- Types that can behave as a Stream
182------------------------------------------------------------------------------
184infixr 5 `consM`
185infixr 5 |:
187-- | Class of types that can represent a stream of elements of some type 'a' in
188-- some monad 'm'.
189--
190-- @since 0.2.0
191class IsStream t where
192 toStream :: t m a -> Stream m a
193 fromStream :: Stream m a -> t m a
194 -- | Constructs a stream by adding a monadic action at the head of an
195 -- existing stream. For example:
196 --
197 -- @
198 -- > toList $ getLine \`consM` getLine \`consM` nil
199 -- hello
200 -- world
201 -- ["hello","world"]
202 -- @
203 --
204 -- /Concurrent (do not use 'parallely' to construct infinite streams)/
205 --
206 -- @since 0.2.0
207 consM :: MonadAsync m => m a -> t m a -> t m a
208 -- | Operator equivalent of 'consM'. We can read it as "@parallel colon@"
209 -- to remember that @|@ comes before ':'.
210 --
211 -- @
212 -- > toList $ getLine |: getLine |: nil
213 -- hello
214 -- world
215 -- ["hello","world"]
216 -- @
217 --
218 -- @
219 -- let delay = threadDelay 1000000 >> print 1
220 -- runStream $ serially $ delay |: delay |: delay |: nil
221 -- runStream $ parallely $ delay |: delay |: delay |: nil
222 -- @
223 --
224 -- /Concurrent (do not use 'parallely' to construct infinite streams)/
225 --
226 -- @since 0.2.0
227 (|:) :: MonadAsync m => m a -> t m a -> t m a
228 -- We can define (|:) just as 'consM' but it is defined explicitly for each
229 -- type because we want to use SPECIALIZE pragma on the definition.
231-- | Same as 'IsStream'.
232--
233-- @since 0.1.0
234{-# DEPRECATED Streaming "Please use IsStream instead." #-}
235type Streaming = IsStream
237-------------------------------------------------------------------------------
238-- Type adapting combinators
239-------------------------------------------------------------------------------
241-- | Adapt any specific stream type to any other specific stream type.
242--
243-- @since 0.1.0
244adapt :: (IsStream t1, IsStream t2) => t1 m a -> t2 m a
245adapt = fromStream . toStream
247------------------------------------------------------------------------------
248-- Building a stream
249------------------------------------------------------------------------------
251-- | Build a stream from an 'SVar', a stop continuation, a singleton stream
252-- continuation and a yield continuation.
253mkStream:: IsStream t
254 => (forall r. State Stream m a
255 -> m r
256 -> (a -> m r)
257 -> (a -> t m a -> m r)
258 -> m r)
259 -> t m a
260mkStream k = fromStream $ Stream $ \st stp sng yld ->
261 let yieldk a r = yld a (toStream r)
262 in k (rstState st) stp sng yieldk
264------------------------------------------------------------------------------
265-- Construction
266------------------------------------------------------------------------------
268-- | An empty stream.
269--
270-- @
271-- > toList nil
272-- []
273-- @
274--
275-- @since 0.1.0
276nil :: IsStream t => t m a
277nil = fromStream $ Stream $ \_ stp _ _ -> stp
279infixr 5 `cons`
281-- faster than consM because there is no bind.
282-- | Construct a stream by adding a pure value at the head of an existing
283-- stream. For serial streams this is the same as @(return a) \`consM` r@ but
284-- more efficient. For concurrent streams this is not concurrent whereas
285-- 'consM' is concurrent. For example:
286--
287-- @
288-- > toList $ 1 \`cons` 2 \`cons` 3 \`cons` nil
289-- [1,2,3]
290-- @
291--
292-- @since 0.1.0
293cons :: IsStream t => a -> t m a -> t m a
294cons a r = fromStream $ Stream $ \_ _ _ yld -> yld a (toStream r)
296infixr 5 .:
298-- | Operator equivalent of 'cons'.
299--
300-- @
301-- > toList $ 1 .: 2 .: 3 .: nil
302-- [1,2,3]
303-- @
304--
305-- @since 0.1.1
306(.:) :: IsStream t => a -> t m a -> t m a
307(.:) = cons
309{-# INLINE consMSerial #-}
310consMSerial :: (Monad m) => m a -> Stream m a -> Stream m a
311consMSerial m r = Stream $ \_ _ _ yld -> m >>= \a -> yld a r
313------------------------------------------------------------------------------
314-- Asynchronous construction
315------------------------------------------------------------------------------
317-- | Make an empty stream from a callback function.
318nilK :: IsStream t => (forall r. m r -> m r) -> t m a
319nilK k = fromStream $ Stream $ \_ stp _ _ -> k stp
321-- | Make a singleton stream from a one shot callback function.
322yieldK :: IsStream t => (forall r. (a -> m r) -> m r) -> t m a
323yieldK k = fromStream $ Stream $ \_ _ sng _ -> k sng
325-- | Construct a stream from a callback function.
326consK :: IsStream t => (forall r. (a -> m r) -> m r) -> t m a -> t m a
327consK k r = fromStream $ Stream $ \_ _ _ yld -> k (\x -> yld x (toStream r))
329-- XXX consK with concurrent callbacks
330-- XXX Build a stream from a repeating callback function.
332-------------------------------------------------------------------------------
333-- IsStream Stream
334-------------------------------------------------------------------------------
336instance IsStream Stream where
337 toStream = id
338 fromStream = id
340 {-# INLINE consM #-}
341 {-# SPECIALIZE consM :: IO a -> Stream IO a -> Stream IO a #-}
342 consM :: Monad m => m a -> Stream m a -> Stream m a
343 consM = consMSerial
345 {-# INLINE (|:) #-}
346 {-# SPECIALIZE (|:) :: IO a -> Stream IO a -> Stream IO a #-}
347 (|:) :: Monad m => m a -> Stream m a -> Stream m a
348 (|:) = consMSerial
350-------------------------------------------------------------------------------
351-- Deconstruction
352-------------------------------------------------------------------------------
354{-# INLINE uncons #-}
355uncons :: (IsStream t, Monad m) => t m a -> m (Maybe (a, t m a))
356uncons m =
357 let stop = return Nothing
358 single a = return (Just (a, nil))
359 yieldk a r = return (Just (a, fromStream r))
360 in (unStream (toStream m)) defState stop single yieldk
362-------------------------------------------------------------------------------
363-- Generation
364-------------------------------------------------------------------------------
366{-# INLINE unfoldr #-}
367unfoldr :: IsStream t => (b -> Maybe (a, b)) -> b -> t m a
368unfoldr step = fromStream . go
369 where
370 go s = Stream $ \_ stp _ yld ->
371 case step s of
372 Nothing -> stp
373 Just (a, b) -> yld a (go b)
375{-# INLINE unfoldrM #-}
376unfoldrM :: (IsStream t, MonadAsync m) => (b -> m (Maybe (a, b))) -> b -> t m a
377unfoldrM step = go
378 where
379 go s = fromStream $ Stream $ \svr stp sng yld -> do
380 mayb <- step s
381 case mayb of
382 Nothing -> stp
383 Just (a, b) ->
384 unStream (toStream (return a |: go b)) svr stp sng yld
386-------------------------------------------------------------------------------
387-- Special generation
388-------------------------------------------------------------------------------
390-- Faster than yieldM because there is no bind. Usually we can construct a
391-- stream from a pure value using "pure" in an applicative, however in case of
392-- Zip streams pure creates an infinite stream.
393--
394-- | Create a singleton stream from a pure value. In monadic streams, 'pure' or
395-- 'return' can be used in place of 'yield', however, in Zip applicative
396-- streams 'pure' is equivalent to 'repeat'.
397--
398-- @since 0.4.0
399yield :: IsStream t => a -> t m a
400yield a = fromStream $ Stream $ \_ _ single _ -> single a
402-- | Create a singleton stream from a monadic action. Same as @m \`consM` nil@
403-- but more efficient.
404--
405-- @
406-- > toList $ yieldM getLine
407-- hello
408-- ["hello"]
409-- @
410--
411-- @since 0.4.0
412{-# INLINE yieldM #-}
413yieldM :: (Monad m, IsStream t) => m a -> t m a
414yieldM m = fromStream $ Stream $ \_ _ single _ -> m >>= single
416-- | Same as yieldM
417--
418-- @since 0.2.0
419{-# DEPRECATED once "Please use yieldM instead." #-}
420{-# INLINE once #-}
421once :: (Monad m, IsStream t) => m a -> t m a
422once = yieldM
424-- | Generate an infinite stream by repeating a pure value.
425--
426-- @since 0.4.0
427repeat :: IsStream t => a -> t m a
428repeat a = let x = cons a x in x
430-------------------------------------------------------------------------------
431-- Conversions
432-------------------------------------------------------------------------------
434-- | Construct a stream from a 'Foldable' containing pure values.
435--
436-- @since 0.2.0
437{-# INLINE fromFoldable #-}
438fromFoldable :: (IsStream t, Foldable f) => f a -> t m a
439fromFoldable = Prelude.foldr cons nil
441{-# INLINE fromList #-}
442fromList :: IsStream t => [a] -> t m a
443fromList = fromFoldable
445{-# INLINE fromStreamK #-}
446fromStreamK :: Stream m a -> Stream m a
447fromStreamK = id
449-------------------------------------------------------------------------------
450-- Elimination by Folding
451-------------------------------------------------------------------------------
453-- | Fold a stream by providing an SVar, a stop continuation, a singleton
454-- continuation and a yield continuation.
455foldStream
456 :: IsStream t
457 => State Stream m a
458 -> m r
459 -> (a -> m r)
460 -> (a -> t m a -> m r)
461 -> t m a
462 -> m r
463foldStream st blank single step m =
464 let yieldk a x = step a (fromStream x)
465 in (unStream (toStream m)) st blank single yieldk
467-- | Lazy right associative fold.
468foldr :: (IsStream t, Monad m) => (a -> b -> b) -> b -> t m a -> m b
469foldr step acc m = go (toStream m)
470 where
471 go m1 =
472 let stop = return acc
473 single a = return (step a acc)
474 yieldk a r = go r >>= \b -> return (step a b)
475 in (unStream m1) defState stop single yieldk
477-- | Lazy right fold with a monadic step function.
478{-# INLINE foldrM #-}
479foldrM :: (IsStream t, Monad m) => (a -> b -> m b) -> b -> t m a -> m b
480foldrM step acc m = go (toStream m)
481 where
482 go m1 =
483 let stop = return acc
484 single a = step a acc
485 yieldk a r = go r >>= step a
486 in (unStream m1) defState stop single yieldk
488-- | Strict left fold with an extraction function. Like the standard strict
489-- left fold, but applies a user supplied extraction function (the third
490-- argument) to the folded value at the end. This is designed to work with the
491-- @foldl@ library. The suffix @x@ is a mnemonic for extraction.
492{-# INLINE foldx #-}
493foldx :: (IsStream t, Monad m)
494 => (x -> a -> x) -> x -> (x -> b) -> t m a -> m b
495foldx step begin done m = get $ go (toStream m) begin
496 where
497 {-# NOINLINE get #-}
498 get m1 =
499 let single = return . done
500 in (unStream m1) undefined undefined single undefined
502 -- Note, this can be implemented by making a recursive call to "go",
503 -- however that is more expensive because of unnecessary recursion
504 -- that cannot be tail call optimized. Unfolding recursion explicitly via
505 -- continuations is much more efficient.
506 go m1 !acc = Stream $ \_ _ sng yld ->
507 let stop = sng acc
508 single a = sng $ step acc a
509 yieldk a r =
510 let stream = go r (step acc a)
511 in (unStream stream) defState undefined sng yld
512 in (unStream m1) defState stop single yieldk
514-- | Strict left associative fold.
515{-# INLINE foldl' #-}
516foldl' :: (IsStream t, Monad m) => (b -> a -> b) -> b -> t m a -> m b
517foldl' step begin m = foldx step begin id m
519-- XXX replace the recursive "go" with explicit continuations.
520-- | Like 'foldx', but with a monadic step function.
521foldxM :: (IsStream t, Monad m)
522 => (x -> a -> m x) -> m x -> (x -> m b) -> t m a -> m b
523foldxM step begin done m = go begin (toStream m)
524 where
525 go !acc m1 =
526 let stop = acc >>= done
527 single a = acc >>= \b -> step b a >>= done
528 yieldk a r = acc >>= \b -> step b a >>= \x -> go (return x) r
529 in (unStream m1) defState stop single yieldk
531-- | Like 'foldl'' but with a monadic step function.
532foldlM' :: (IsStream t, Monad m) => (b -> a -> m b) -> b -> t m a -> m b
533foldlM' step begin m = foldxM step (return begin) return m
535------------------------------------------------------------------------------
536-- Specialized folds
537------------------------------------------------------------------------------
539{-# INLINE runStream #-}
540runStream :: (Monad m, IsStream t) => t m a -> m ()
541runStream m = go (toStream m)
542 where
543 go m1 =
544 let stop = return ()
545 single _ = return ()
546 yieldk _ r = go (toStream r)
547 in unStream m1 defState stop single yieldk
549{-# INLINE null #-}
550null :: (IsStream t, Monad m) => t m a -> m Bool
551null m =
552 let stop = return True
553 single _ = return False
554 yieldk _ _ = return False
555 in unStream (toStream m) defState stop single yieldk
557{-# INLINE head #-}
558head :: (IsStream t, Monad m) => t m a -> m (Maybe a)
559head m =
560 let stop = return Nothing
561 single a = return (Just a)
562 yieldk a _ = return (Just a)
563 in unStream (toStream m) defState stop single yieldk
565{-# INLINE tail #-}
566tail :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
567tail m =
568 let stop = return Nothing
569 single _ = return $ Just nil
570 yieldk _ r = return $ Just $ fromStream r
571 in unStream (toStream m) defState stop single yieldk
573{-# INLINE elem #-}
574elem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
575elem e m = go (toStream m)
576 where
577 go m1 =
578 let stop = return False
579 single a = return (a == e)
580 yieldk a r = if a == e then return True else go r
581 in (unStream m1) defState stop single yieldk
583{-# INLINE notElem #-}
584notElem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
585notElem e m = go (toStream m)
586 where
587 go m1 =
588 let stop = return True
589 single a = return (a /= e)
590 yieldk a r = if a == e then return False else go r
591 in (unStream m1) defState stop single yieldk
593all :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
594all p m = go (toStream 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 unStream m1 defState (return True) single yieldk
603any :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
604any p m = go (toStream m)
605 where
606 go m1 =
607 let single a | p a = return True
608 | otherwise = return False
609 yieldk a r | p a = return True
610 | otherwise = go r
611 in unStream m1 defState (return False) single yieldk
613-- | Extract the last element of the stream, if any.
614{-# INLINE last #-}
615last :: (IsStream t, Monad m) => t m a -> m (Maybe a)
616last = foldx (\_ y -> Just y) Nothing id
618{-# INLINE minimum #-}
619minimum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
620minimum m = go Nothing (toStream m)
621 where
622 go Nothing m1 =
623 let stop = return Nothing
624 single a = return (Just a)
625 yieldk a r = go (Just a) r
626 in unStream m1 defState stop single yieldk
628 go (Just res) m1 =
629 let stop = return (Just res)
630 single a =
631 if res <= a
632 then return (Just res)
633 else return (Just a)
634 yieldk a r =
635 if res <= a
636 then go (Just res) r
637 else go (Just a) r
638 in unStream m1 defState stop single yieldk
640{-# INLINE maximum #-}
641maximum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
642maximum m = go Nothing (toStream m)
643 where
644 go Nothing m1 =
645 let stop = return Nothing
646 single a = return (Just a)
647 yieldk a r = go (Just a) r
648 in unStream m1 defState stop single yieldk
650 go (Just res) m1 =
651 let stop = return (Just res)
652 single a =
653 if res <= a
654 then return (Just a)
655 else return (Just res)
656 yieldk a r =
657 if res <= a
658 then go (Just a) r
659 else go (Just res) r
660 in unStream m1 defState stop single yieldk
662------------------------------------------------------------------------------
663-- Map and Fold
664------------------------------------------------------------------------------
666-- | Apply a monadic action to each element of the stream and discard the
667-- output of the action.
668mapM_ :: (IsStream t, Monad m) => (a -> m b) -> t m a -> m ()
669mapM_ f m = go (toStream m)
670 where
671 go m1 =
672 let stop = return ()
673 single a = void (f a)
674 yieldk a r = f a >> go r
675 in (unStream m1) defState stop single yieldk
677------------------------------------------------------------------------------
678-- Converting folds
679------------------------------------------------------------------------------
681{-# INLINABLE toList #-}
682toList :: (IsStream t, Monad m) => t m a -> m [a]
683toList = foldr (:) []
685{-# INLINE toStreamK #-}
686toStreamK :: Stream m a -> Stream m a
687toStreamK = id
689-------------------------------------------------------------------------------
690-- Transformation by folding (Scans)
691-------------------------------------------------------------------------------
693{-# INLINE scanx #-}
694scanx :: IsStream t => (x -> a -> x) -> x -> (x -> b) -> t m a -> t m b
695scanx step begin done m =
696 cons (done begin) $ fromStream $ go (toStream m) begin
697 where
698 go m1 !acc = Stream $ \st stp sng yld ->
699 let single a = sng (done $ step acc a)
700 yieldk a r =
701 let s = step acc a
702 in yld (done s) (go r s)
703 in unStream m1 (rstState st) stp single yieldk
705{-# INLINE scanl' #-}
706scanl' :: IsStream t => (b -> a -> b) -> b -> t m a -> t m b
707scanl' step begin m = scanx step begin id m
709-------------------------------------------------------------------------------
710-- Filtering
711-------------------------------------------------------------------------------
713{-# INLINE filter #-}
714filter :: IsStream t => (a -> Bool) -> t m a -> t m a
715filter p m = fromStream $ go (toStream m)
716 where
717 go m1 = Stream $ \st stp sng yld ->
718 let single a | p a = sng a
719 | otherwise = stp
720 yieldk a r | p a = yld a (go r)
721 | otherwise = (unStream r) (rstState st) stp single yieldk
722 in unStream m1 (rstState st) stp single yieldk
724{-# INLINE take #-}
725take :: IsStream t => Int -> t m a -> t m a
726take n m = fromStream $ go n (toStream m)
727 where
728 go n1 m1 = Stream $ \st stp sng yld ->
729 let yieldk a r = yld a (go (n1 - 1) r)
730 in if n1 <= 0
731 then stp
732 else unStream m1 (rstState st) stp sng yieldk
734{-# INLINE takeWhile #-}
735takeWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
736takeWhile p m = fromStream $ go (toStream m)
737 where
738 go m1 = Stream $ \st stp sng yld ->
739 let single a | p a = sng a
740 | otherwise = stp
741 yieldk a r | p a = yld a (go r)
742 | otherwise = stp
743 in unStream m1 (rstState st) stp single yieldk
745drop :: IsStream t => Int -> t m a -> t m a
746drop n m = fromStream $ Stream $ \st stp sng yld ->
747 unStream (go n (toStream m)) (rstState st) stp sng yld
748 where
749 go n1 m1 = Stream $ \st stp sng yld ->
750 let single _ = stp
751 yieldk _ r = (unStream $ go (n1 - 1) r) st stp sng yld
752 -- Somehow "<=" check performs better than a ">"
753 in if n1 <= 0
754 then unStream m1 st stp sng yld
755 else unStream m1 st stp single yieldk
757{-# INLINE dropWhile #-}
758dropWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
759dropWhile p m = fromStream $ go (toStream m)
760 where
761 go m1 = Stream $ \st stp sng yld ->
762 let single a | p a = stp
763 | otherwise = sng a
764 yieldk a r | p a = (unStream r) (rstState st) stp single yieldk
765 | otherwise = yld a r
766 in unStream m1 (rstState st) stp single yieldk
768-------------------------------------------------------------------------------
769-- Mapping
770-------------------------------------------------------------------------------
772{-# INLINE map #-}
773map :: (IsStream t, Monad m) => (a -> b) -> t m a -> t m b
774map f m = fromStream $ Stream $ \st stp sng yld ->
775 let single = sng . f
776 yieldk a r = yld (f a) (fmap f r)
777 in unStream (toStream m) (rstState st) stp single yieldk
779-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
780-- other stream types to overload it.
781{-# INLINE mapM #-}
782mapM :: (IsStream t, MonadAsync m) => (a -> m b) -> t m a -> t m b
783mapM f m = go (toStream m)
784 where
785 go m1 = fromStream $ Stream $ \st stp sng yld ->
786 let single a = f a >>= sng
787 yieldk a r = unStream (toStream (f a |: (go r))) st stp sng yld
788 in (unStream m1) (rstState st) stp single yieldk
790-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
791-- other stream types to overload it.
792{-# INLINE sequence #-}
793sequence :: (IsStream t, MonadAsync m) => t m (m a) -> t m a
794sequence m = go (toStream m)
795 where
796 go m1 = fromStream $ Stream $ \st stp sng yld ->
797 let single ma = ma >>= sng
798 yieldk ma r = unStream (toStream $ ma |: go r) st stp sng yld
799 in (unStream m1) (rstState st) stp single yieldk
801-------------------------------------------------------------------------------
802-- Map and Filter
803-------------------------------------------------------------------------------
805{-# INLINE mapMaybe #-}
806mapMaybe :: IsStream t => (a -> Maybe b) -> t m a -> t m b
807mapMaybe f m = go (toStream m)
808 where
809 go m1 = fromStream $ Stream $ \st stp sng yld ->
810 let single a = case f a of
811 Just b -> sng b
812 Nothing -> stp
813 yieldk a r = case f a of
814 Just b -> yld b (toStream $ go r)
815 Nothing -> (unStream r) (rstState st) stp single yieldk
816 in unStream m1 (rstState st) stp single yieldk
818------------------------------------------------------------------------------
819-- Serial Zipping
820------------------------------------------------------------------------------
822{-# INLINE zipWithS #-}
823zipWithS :: (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
824zipWithS f m1 m2 = go m1 m2
825 where
826 go mx my = Stream $ \st stp sng yld -> do
827 let merge a ra =
828 let single2 b = sng (f a b)
829 yield2 b rb = yld (f a b) (go ra rb)
830 in unStream my (rstState st) stp single2 yield2
831 let single1 a = merge a nil
832 yield1 a ra = merge a ra
833 unStream mx (rstState st) stp single1 yield1
835-- | Zip two streams serially using a pure zipping function.
836--
837-- @since 0.1.0
838{-# INLINABLE zipWith #-}
839zipWith :: IsStream t => (a -> b -> c) -> t m a -> t m b -> t m c
840zipWith f m1 m2 = fromStream $ zipWithS f (toStream m1) (toStream m2)
842-- | Zip two streams serially using a monadic zipping function.
843--
844-- @since 0.1.0
845zipWithM :: (IsStream t, Monad m) => (a -> b -> m c) -> t m a -> t m b -> t m c
846zipWithM f m1 m2 = fromStream $ go (toStream m1) (toStream m2)
847 where
848 go mx my = Stream $ \st stp sng yld -> do
849 let merge a ra =
850 let runIt x = unStream x (rstState st) stp sng yld
851 single2 b = f a b >>= sng
852 yield2 b rb = f a b >>= \x -> runIt (x `cons` go ra rb)
853 in unStream my (rstState st) stp single2 yield2
854 let single1 a = merge a nil
855 yield1 a ra = merge a ra
856 unStream mx (rstState st) stp single1 yield1
858------------------------------------------------------------------------------
859-- Semigroup
860------------------------------------------------------------------------------
862-- | Concatenates two streams sequentially i.e. the first stream is
863-- exhausted completely before yielding any element from the second stream.
864{-# INLINE serial #-}
865serial :: Stream m a -> Stream m a -> Stream m a
866serial m1 m2 = go m1
867 where
868 go (Stream m) = Stream $ \st stp sng yld ->
869 let stop = (unStream m2) (rstState st) stp sng yld
870 single a = yld a m2
871 yieldk a r = yld a (go r)
872 in m (rstState st) stop single yieldk
874instance Semigroup (Stream m a) where
875 (<>) = serial
877------------------------------------------------------------------------------
878-- Monoid
879------------------------------------------------------------------------------
881instance Monoid (Stream m a) where
882 mempty = nil
883 mappend = (<>)
885-------------------------------------------------------------------------------
886-- Functor
887-------------------------------------------------------------------------------
889instance Monad m => Functor (Stream m) where
890 fmap = map
892-------------------------------------------------------------------------------
893-- Bind utility
894-------------------------------------------------------------------------------
896{-# INLINE bindWith #-}
897bindWith
898 :: (forall c. Stream m c -> Stream m c -> Stream m c)
899 -> Stream m a
900 -> (a -> Stream m b)
901 -> Stream m b
902bindWith par m f = go m
903 where
904 go (Stream g) =
905 Stream $ \st stp sng yld ->
906 let run x = (unStream x) st stp sng yld
907 single a = run $ f a
908 yieldk a r = run $ f a `par` go r
909 in g (rstState st) stp single yieldk
911------------------------------------------------------------------------------
912-- Alternative & MonadPlus
913------------------------------------------------------------------------------
915_alt :: Stream m a -> Stream m a -> Stream m a
916_alt m1 m2 = Stream $ \st stp sng yld ->
917 let stop = unStream m2 (rstState st) stp sng yld
918 in unStream m1 (rstState st) stop sng yld
920------------------------------------------------------------------------------
921-- MonadReader
922------------------------------------------------------------------------------
924withLocal :: MonadReader r m => (r -> r) -> Stream m a -> Stream m a
925withLocal f m =
926 Stream $ \st stp sng yld ->
927 let single = local f . sng
928 yieldk a r = local f $ yld a (withLocal f r)
929 in (unStream m) (rstState st) (local f stp) single yieldk
931------------------------------------------------------------------------------
932-- MonadError
933------------------------------------------------------------------------------
935{-
936-- XXX handle and test cross thread state transfer
937withCatchError
938 :: MonadError e m
939 => Stream m a -> (e -> Stream m a) -> Stream m a
940withCatchError m h =
941 Stream $ \_ stp sng yld ->
942 let run x = unStream x Nothing stp sng yieldk
943 handle r = r `catchError` \e -> run $ h e
944 yieldk a r = yld a (withCatchError r h)
945 in handle $ run m
946-}
948-------------------------------------------------------------------------------
949-- Transformers
950-------------------------------------------------------------------------------
952instance MonadTrans Stream where
953 lift = yieldM