git.iko.soy/streamly
on 9ccee33 help sign in
streamly/src/Streamly/Streams/StreamK.hs
1237 lines 41 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE CPP #-}
3{-# LANGUAGE ConstraintKinds #-}
4{-# LANGUAGE FlexibleContexts #-}
5{-# LANGUAGE FlexibleInstances #-}
6{-# LANGUAGE InstanceSigs #-}
7{-# LANGUAGE MultiParamTypeClasses #-}
8{-# LANGUAGE RankNTypes #-}
9{-# LANGUAGE UndecidableInstances #-} -- XXX
11-- |
12-- Module : Streamly.Streams.StreamK
13-- Copyright : (c) 2017 Harendra Kumar
14--
15-- License : BSD3
16-- Maintainer : harendra.kumar@gmail.com
17-- Stability : experimental
18-- Portability : GHC
19--
20--
21-- Continuation passing style (CPS) stream implementation. The symbol 'K' below
22-- denotes a function as well as a Kontinuation.
23--
24-- @
25-- import qualified Streamly.Streams.StreamK as K
26-- @
27--
28module Streamly.Streams.StreamK
29 (
30 -- * A class for streams
31 IsStream (..)
32 , adapt
34 -- * The stream type
35 , Stream (..)
36 , unStreamIsolated
37 , isolateStream
38 , unstreamShared
39 , runStreamSVar
41 -- * Construction
42 , mkStream
43 , nil
44 , cons
45 , (.:)
47 -- * Asynchronous construction
48 , nilK
49 , yieldK
50 , consK
52 -- * Deconstruction
53 , uncons
55 -- * Generation
56 -- ** Unfolds
57 , unfoldr
58 , unfoldrM
60 -- ** Specialized Generation
61 , repeat
63 -- ** Conversions
64 , yield
65 , yieldM
66 , fromFoldable
67 , fromList
68 , fromStreamK
70 -- * Elimination
71 -- ** General Folds
72 , foldStream
73 , foldr
74 , foldrM
75 , foldr1
76 , foldl'
77 , foldlM'
78 , foldx
79 , foldxM
81 -- ** Specialized Folds
82 , runStream
83 , null
84 , head
85 , tail
86 , init
87 , elem
88 , notElem
89 , all
90 , any
91 , last
92 , minimum
93 , minimumBy
94 , maximum
95 , maximumBy
96 , findIndices
97 , lookup
98 , find
99 , (!!)
101 -- ** Map and Fold
102 , mapM_
104 -- ** Conversions
105 , toList
106 , toStreamK
108 -- * Transformation
109 -- ** By folding (scans)
110 , scanl'
111 , scanx
113 -- ** Filtering
114 , filter
115 , take
116 , takeWhile
117 , drop
118 , dropWhile
120 -- ** Mapping
121 , map
122 , mapM
123 , sequence
125 -- ** Inserting
126 , intersperseM
127 , insertBy
129 -- ** Deleting
130 , deleteBy
132 -- ** Map and Filter
133 , mapMaybe
135 -- ** Zipping
136 , zipWith
137 , zipWithM
139 -- ** Merging
140 , mergeBy
142 -- ** Transformation comprehensions
143 , the
145 -- * Semigroup Style Composition
146 , serial
148 -- * Utilities
149 , consMSerial
150 , bindWith
151 , withLocal
153 -- * Deprecated
154 , Streaming -- deprecated
155 , once -- deprecated
156 )
157where
159import Control.Monad (void)
160import Control.Monad.IO.Class (MonadIO(liftIO))
161import Control.Monad.Reader.Class (MonadReader(..))
162import Control.Monad.Trans.Class (MonadTrans(lift))
163import Data.Semigroup (Semigroup(..))
164import Prelude
165 hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,
166 take, filter, all, any, takeWhile, drop, dropWhile, minimum,
167 maximum, elem, notElem, null, head, tail, init, zipWith, lookup,
168 foldr1, (!!))
169import qualified Prelude
171import Streamly.SVar
173------------------------------------------------------------------------------
174-- The basic stream type
175------------------------------------------------------------------------------
177-- | The type @Stream m a@ represents a monadic stream of values of type 'a'
178-- constructed using actions in monad 'm'. It uses stop, singleton and yield
179-- continuations equivalent to the following direct style type:
180--
181-- @
182-- data Stream m a = Stop | Singleton a | Yield a (Stream m a)
183-- @
184--
185-- To facilitate parallel composition we maintain a local state in an 'SVar'
186-- that is shared across and is used for synchronization of the streams being
187-- composed.
188--
189-- The singleton case can be expressed in terms of stop and yield but we have
190-- it as a separate case to optimize composition operations for streams with
191-- single element. We build singleton streams in the implementation of 'pure'
192-- for Applicative and Monad, and in 'lift' for MonadTrans.
193--
194newtype Stream m a =
195 Stream {
196 unStream :: forall r.
197 State Stream m a -- state
198 -> m r -- stop
199 -> (a -> m r) -- singleton
200 -> (a -> Stream m a -> m r) -- yield
201 -> m r
202 }
204-- XXX make this the default "unStream"
205-- | unwraps the Stream type producing the stream function that can be run with
206-- continuations.
207{-# INLINE unStreamIsolated #-}
208unStreamIsolated ::
209 Stream m a
210 -> State Stream m a -- state
211 -> m r -- stop
212 -> (a -> m r) -- singleton
213 -> (a -> Stream m a -> m r) -- yield
214 -> m r
215unStreamIsolated x st = unStream x (rstState st)
217{-# INLINE isolateStream #-}
218isolateStream :: Stream m a -> Stream m a
219isolateStream x = Stream $ \st stp sng yld ->
220 unStreamIsolated x st stp sng yld
222-- | Like unstream, but passes a shared SVar across continuations.
223{-# INLINE unstreamShared #-}
224unstreamShared ::
225 Stream m a
226 -> State Stream m a -- state
227 -> m r -- stop
228 -> (a -> m r) -- singleton
229 -> (a -> Stream m a -> m r) -- yield
230 -> m r
231unstreamShared = unStream
233-- Run the stream using a run function associated with the SVar that runs the
234-- streams with a captured snapshot of the monadic state.
235{-# INLINE runStreamSVar #-}
236runStreamSVar
237 :: MonadIO m
238 => SVar Stream m a
239 -> Stream m a
240 -> State Stream m a -- state
241 -> m r -- stop
242 -> (a -> m r) -- singleton
243 -> (a -> Stream m a -> m r) -- yield
244 -> m ()
245runStreamSVar sv m st stp sng yld =
246 let mrun = runInIO $ svarMrun sv
247 in void $ liftIO $ mrun $ unStream m st stp sng yld
249------------------------------------------------------------------------------
250-- Types that can behave as a Stream
251------------------------------------------------------------------------------
253infixr 5 `consM`
254infixr 5 |:
256-- | Class of types that can represent a stream of elements of some type 'a' in
257-- some monad 'm'.
258--
259-- @since 0.2.0
260class IsStream t where
261 toStream :: t m a -> Stream m a
262 fromStream :: Stream m a -> t m a
263 -- | Constructs a stream by adding a monadic action at the head of an
264 -- existing stream. For example:
265 --
266 -- @
267 -- > toList $ getLine \`consM` getLine \`consM` nil
268 -- hello
269 -- world
270 -- ["hello","world"]
271 -- @
272 --
273 -- /Concurrent (do not use 'parallely' to construct infinite streams)/
274 --
275 -- @since 0.2.0
276 consM :: MonadAsync m => m a -> t m a -> t m a
277 -- | Operator equivalent of 'consM'. We can read it as "@parallel colon@"
278 -- to remember that @|@ comes before ':'.
279 --
280 -- @
281 -- > toList $ getLine |: getLine |: nil
282 -- hello
283 -- world
284 -- ["hello","world"]
285 -- @
286 --
287 -- @
288 -- let delay = threadDelay 1000000 >> print 1
289 -- runStream $ serially $ delay |: delay |: delay |: nil
290 -- runStream $ parallely $ delay |: delay |: delay |: nil
291 -- @
292 --
293 -- /Concurrent (do not use 'parallely' to construct infinite streams)/
294 --
295 -- @since 0.2.0
296 (|:) :: MonadAsync m => m a -> t m a -> t m a
297 -- We can define (|:) just as 'consM' but it is defined explicitly for each
298 -- type because we want to use SPECIALIZE pragma on the definition.
300-- | Same as 'IsStream'.
301--
302-- @since 0.1.0
303{-# DEPRECATED Streaming "Please use IsStream instead." #-}
304type Streaming = IsStream
306-------------------------------------------------------------------------------
307-- Type adapting combinators
308-------------------------------------------------------------------------------
310-- | Adapt any specific stream type to any other specific stream type.
311--
312-- @since 0.1.0
313adapt :: (IsStream t1, IsStream t2) => t1 m a -> t2 m a
314adapt = fromStream . toStream
316------------------------------------------------------------------------------
317-- Building a stream
318------------------------------------------------------------------------------
320-- | Build a stream from an 'SVar', a stop continuation, a singleton stream
321-- continuation and a yield continuation.
322mkStream:: IsStream t
323 => (forall r. State Stream m a
324 -> m r
325 -> (a -> m r)
326 -> (a -> t m a -> m r)
327 -> m r)
328 -> t m a
329mkStream k = fromStream $ Stream $ \st stp sng yld ->
330 let yieldk a r = yld a (toStream r)
331 in k (rstState st) stp sng yieldk
333------------------------------------------------------------------------------
334-- Construction
335------------------------------------------------------------------------------
337-- | An empty stream.
338--
339-- @
340-- > toList nil
341-- []
342-- @
343--
344-- @since 0.1.0
345nil :: IsStream t => t m a
346nil = fromStream $ Stream $ \_ stp _ _ -> stp
348infixr 5 `cons`
350-- faster than consM because there is no bind.
351-- | Construct a stream by adding a pure value at the head of an existing
352-- stream. For serial streams this is the same as @(return a) \`consM` r@ but
353-- more efficient. For concurrent streams this is not concurrent whereas
354-- 'consM' is concurrent. For example:
355--
356-- @
357-- > toList $ 1 \`cons` 2 \`cons` 3 \`cons` nil
358-- [1,2,3]
359-- @
360--
361-- @since 0.1.0
362cons :: IsStream t => a -> t m a -> t m a
363cons a r = fromStream $ Stream $ \_ _ _ yld -> yld a (toStream r)
365infixr 5 .:
367-- | Operator equivalent of 'cons'.
368--
369-- @
370-- > toList $ 1 .: 2 .: 3 .: nil
371-- [1,2,3]
372-- @
373--
374-- @since 0.1.1
375(.:) :: IsStream t => a -> t m a -> t m a
376(.:) = cons
378{-# INLINE consMSerial #-}
379consMSerial :: (Monad m) => m a -> Stream m a -> Stream m a
380consMSerial m r = Stream $ \_ _ _ yld -> m >>= \a -> yld a r
382------------------------------------------------------------------------------
383-- Asynchronous construction
384------------------------------------------------------------------------------
386-- | Make an empty stream from a callback function.
387nilK :: IsStream t => (forall r. m r -> m r) -> t m a
388nilK k = fromStream $ Stream $ \_ stp _ _ -> k stp
390-- | Make a singleton stream from a one shot callback function.
391yieldK :: IsStream t => (forall r. (a -> m r) -> m r) -> t m a
392yieldK k = fromStream $ Stream $ \_ _ sng _ -> k sng
394-- | Construct a stream from a callback function.
395consK :: IsStream t => (forall r. (a -> m r) -> m r) -> t m a -> t m a
396consK k r = fromStream $ Stream $ \_ _ _ yld -> k (\x -> yld x (toStream r))
398-- XXX consK with concurrent callbacks
399-- XXX Build a stream from a repeating callback function.
401-------------------------------------------------------------------------------
402-- IsStream Stream
403-------------------------------------------------------------------------------
405instance IsStream Stream where
406 toStream = id
407 fromStream = id
409 {-# INLINE consM #-}
410 {-# SPECIALIZE consM :: IO a -> Stream IO a -> Stream IO a #-}
411 consM :: Monad m => m a -> Stream m a -> Stream m a
412 consM = consMSerial
414 {-# INLINE (|:) #-}
415 {-# SPECIALIZE (|:) :: IO a -> Stream IO a -> Stream IO a #-}
416 (|:) :: Monad m => m a -> Stream m a -> Stream m a
417 (|:) = consMSerial
419-------------------------------------------------------------------------------
420-- Deconstruction
421-------------------------------------------------------------------------------
423{-# INLINE uncons #-}
424uncons :: (IsStream t, Monad m) => t m a -> m (Maybe (a, t m a))
425uncons m =
426 let stop = return Nothing
427 single a = return (Just (a, nil))
428 yieldk a r = return (Just (a, fromStream r))
429 in unStream (toStream m) defState stop single yieldk
431-------------------------------------------------------------------------------
432-- Generation
433-------------------------------------------------------------------------------
435{-# INLINE unfoldr #-}
436unfoldr :: IsStream t => (b -> Maybe (a, b)) -> b -> t m a
437unfoldr step = fromStream . go
438 where
439 go s = Stream $ \_ stp _ yld ->
440 case step s of
441 Nothing -> stp
442 Just (a, b) -> yld a (go b)
444{-# INLINE unfoldrM #-}
445unfoldrM :: (IsStream t, MonadAsync m) => (b -> m (Maybe (a, b))) -> b -> t m a
446unfoldrM step = go
447 where
448 go s = fromStream $ Stream $ \svr stp sng yld -> do
449 mayb <- step s
450 case mayb of
451 Nothing -> stp
452 Just (a, b) ->
453 unStream (toStream (return a |: go b)) svr stp sng yld
455-------------------------------------------------------------------------------
456-- Special generation
457-------------------------------------------------------------------------------
459yield :: IsStream t => a -> t m a
460yield a = fromStream $ Stream $ \_ _ single _ -> single a
462{-# INLINE yieldM #-}
463yieldM :: (Monad m, IsStream t) => m a -> t m a
464yieldM m = fromStream $ Stream $ \_ _ single _ -> m >>= single
466-- | Same as yieldM
467--
468-- @since 0.2.0
469{-# DEPRECATED once "Please use yieldM instead." #-}
470{-# INLINE once #-}
471once :: (Monad m, IsStream t) => m a -> t m a
472once = yieldM
474-- | Generate an infinite stream by repeating a pure value.
475-- Can be expressed as @cycle1 . yield@.
476--
477-- @since 0.4.0
478repeat :: IsStream t => a -> t m a
479repeat a = let x = cons a x in x
481-------------------------------------------------------------------------------
482-- Conversions
483-------------------------------------------------------------------------------
485-- | Construct a stream from a 'Foldable' containing pure values. Same as
486-- @'Prelude.foldr' 'cons' 'nil'@.
487--
488-- @since 0.2.0
489{-# INLINE fromFoldable #-}
490fromFoldable :: (IsStream t, Foldable f) => f a -> t m a
491fromFoldable = Prelude.foldr cons nil
493{-# INLINE fromList #-}
494fromList :: IsStream t => [a] -> t m a
495fromList = fromFoldable
497{-# INLINE fromStreamK #-}
498fromStreamK :: Stream m a -> Stream m a
499fromStreamK = id
501-------------------------------------------------------------------------------
502-- Elimination by Folding
503-------------------------------------------------------------------------------
505-- | Fold a stream by providing an SVar, a stop continuation, a singleton
506-- continuation and a yield continuation.
507foldStream
508 :: IsStream t
509 => State Stream m a
510 -> m r
511 -> (a -> m r)
512 -> (a -> t m a -> m r)
513 -> t m a
514 -> m r
515foldStream st blank single step m =
516 let yieldk a x = step a (fromStream x)
517 in unStream (toStream m) st blank single yieldk
519-- | Lazy right associative fold.
520foldr :: (IsStream t, Monad m) => (a -> b -> b) -> b -> t m a -> m b
521foldr step acc m = go (toStream m)
522 where
523 go m1 =
524 let stop = return acc
525 single a = return (step a acc)
526 yieldk a r = go r >>= \b -> return (step a b)
527 in unStream m1 defState stop single yieldk
529-- | Lazy right fold with a monadic step function.
530{-# INLINE foldrM #-}
531foldrM :: (IsStream t, Monad m) => (a -> b -> m b) -> b -> t m a -> m b
532foldrM step acc m = go (toStream m)
533 where
534 go m1 =
535 let stop = return acc
536 single a = step a acc
537 yieldk a r = go r >>= step a
538 in unStream m1 defState stop single yieldk
540{-# INLINE foldr1 #-}
541foldr1 :: (IsStream t, Monad m) => (a -> a -> a) -> t m a -> m (Maybe a)
542foldr1 step m = do
543 r <- uncons m
544 case r of
545 Nothing -> return Nothing
546 Just (h, t) -> fmap Just (go h (toStream t))
547 where
548 go p m1 =
549 let stp = return p
550 single a = return $ step a p
551 yieldk a r = fmap (step p) (go a r)
552 in unStream m1 defState stp single yieldk
554-- | Strict left fold with an extraction function. Like the standard strict
555-- left fold, but applies a user supplied extraction function (the third
556-- argument) to the folded value at the end. This is designed to work with the
557-- @foldl@ library. The suffix @x@ is a mnemonic for extraction.
558{-# INLINE foldx #-}
559foldx :: (IsStream t, Monad m)
560 => (x -> a -> x) -> x -> (x -> b) -> t m a -> m b
561foldx step begin done m = get $ go (toStream m) begin
562 where
563 {-# NOINLINE get #-}
564 get m1 =
565 let single = return . done
566 in unStream m1 undefined undefined single undefined
568 -- Note, this can be implemented by making a recursive call to "go",
569 -- however that is more expensive because of unnecessary recursion
570 -- that cannot be tail call optimized. Unfolding recursion explicitly via
571 -- continuations is much more efficient.
572 go m1 !acc = Stream $ \_ _ sng yld ->
573 let stop = sng acc
574 single a = sng $ step acc a
575 yieldk a r =
576 let stream = go r (step acc a)
577 in unStream stream defState undefined sng yld
578 in unStream m1 defState stop single yieldk
580-- | Strict left associative fold.
581{-# INLINE foldl' #-}
582foldl' :: (IsStream t, Monad m) => (b -> a -> b) -> b -> t m a -> m b
583foldl' step begin = foldx step begin id
585-- XXX replace the recursive "go" with explicit continuations.
586-- | Like 'foldx', but with a monadic step function.
587foldxM :: (IsStream t, Monad m)
588 => (x -> a -> m x) -> m x -> (x -> m b) -> t m a -> m b
589foldxM step begin done m = go begin (toStream m)
590 where
591 go !acc m1 =
592 let stop = acc >>= done
593 single a = acc >>= \b -> step b a >>= done
594 yieldk a r = acc >>= \b -> step b a >>= \x -> go (return x) r
595 in unStream m1 defState stop single yieldk
597-- | Like 'foldl'' but with a monadic step function.
598foldlM' :: (IsStream t, Monad m) => (b -> a -> m b) -> b -> t m a -> m b
599foldlM' step begin = foldxM step (return begin) return
601------------------------------------------------------------------------------
602-- Specialized folds
603------------------------------------------------------------------------------
605{-# INLINE runStream #-}
606runStream :: (Monad m, IsStream t) => t m a -> m ()
607runStream m = go (toStream m)
608 where
609 go m1 =
610 let stop = return ()
611 single _ = return ()
612 yieldk _ r = go (toStream r)
613 in unStream m1 defState stop single yieldk
615{-# INLINE null #-}
616null :: (IsStream t, Monad m) => t m a -> m Bool
617null m =
618 let stop = return True
619 single _ = return False
620 yieldk _ _ = return False
621 in unStream (toStream m) defState stop single yieldk
623{-# INLINE head #-}
624head :: (IsStream t, Monad m) => t m a -> m (Maybe a)
625head m =
626 let stop = return Nothing
627 single a = return (Just a)
628 yieldk a _ = return (Just a)
629 in unStream (toStream m) defState stop single yieldk
631{-# INLINE tail #-}
632tail :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
633tail m =
634 let stop = return Nothing
635 single _ = return $ Just nil
636 yieldk _ r = return $ Just $ fromStream r
637 in unStream (toStream m) defState stop single yieldk
639{-# INLINE init #-}
640init :: (IsStream t, Monad m) => t m a -> m (Maybe (t m a))
641init m = go1 (toStream m)
642 where
643 go1 m1 = do
644 r <- uncons m1
645 case r of
646 Nothing -> return Nothing
647 Just (h, t) -> return . Just . fromStream $ go h t
648 go p m1 = Stream $ \_ stp sng yld ->
649 let single _ = sng p
650 yieldk a x = yld p $ go a x
651 in unStream m1 defState stp single yieldk
653{-# INLINE elem #-}
654elem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
655elem e m = go (toStream m)
656 where
657 go m1 =
658 let stop = return False
659 single a = return (a == e)
660 yieldk a r = if a == e then return True else go r
661 in unStream m1 defState stop single yieldk
663{-# INLINE notElem #-}
664notElem :: (IsStream t, Monad m, Eq a) => a -> t m a -> m Bool
665notElem e m = go (toStream m)
666 where
667 go m1 =
668 let stop = return True
669 single a = return (a /= e)
670 yieldk a r = if a == e then return False else go r
671 in unStream m1 defState stop single yieldk
673all :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
674all p m = go (toStream m)
675 where
676 go m1 =
677 let single a | p a = return True
678 | otherwise = return False
679 yieldk a r | p a = go r
680 | otherwise = return False
681 in unStream m1 defState (return True) single yieldk
683any :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m Bool
684any p m = go (toStream m)
685 where
686 go m1 =
687 let single a | p a = return True
688 | otherwise = return False
689 yieldk a r | p a = return True
690 | otherwise = go r
691 in unStream m1 defState (return False) single yieldk
693-- | Extract the last element of the stream, if any.
694{-# INLINE last #-}
695last :: (IsStream t, Monad m) => t m a -> m (Maybe a)
696last = foldx (\_ y -> Just y) Nothing id
698{-# INLINE minimum #-}
699minimum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
700minimum m = go Nothing (toStream m)
701 where
702 go Nothing m1 =
703 let stop = return Nothing
704 single a = return (Just a)
705 yieldk a r = go (Just a) r
706 in unStream m1 defState stop single yieldk
708 go (Just res) m1 =
709 let stop = return (Just res)
710 single a =
711 if res <= a
712 then return (Just res)
713 else return (Just a)
714 yieldk a r =
715 if res <= a
716 then go (Just res) r
717 else go (Just a) r
718 in unStream m1 defState stop single yieldk
720{-# INLINE minimumBy #-}
721minimumBy :: (IsStream t, Monad m) => (a -> a -> Ordering) -> t m a -> m (Maybe a)
722minimumBy cmp m = go Nothing (toStream m)
723 where
724 go Nothing m1 =
725 let stop = return Nothing
726 single a = return (Just a)
727 yieldk a r = go (Just a) r
728 in unStream m1 defState stop single yieldk
730 go (Just res) m1 =
731 let stop = return (Just res)
732 single a = case cmp res a of
733 GT -> return (Just a)
734 _ -> return (Just res)
735 yieldk a r = case cmp res a of
736 GT -> go (Just a) r
737 _ -> go (Just res) r
738 in unStream m1 defState stop single yieldk
740{-# INLINE maximum #-}
741maximum :: (IsStream t, Monad m, Ord a) => t m a -> m (Maybe a)
742maximum m = go Nothing (toStream m)
743 where
744 go Nothing m1 =
745 let stop = return Nothing
746 single a = return (Just a)
747 yieldk a r = go (Just a) r
748 in unStream m1 defState stop single yieldk
750 go (Just res) m1 =
751 let stop = return (Just res)
752 single a =
753 if res <= a
754 then return (Just a)
755 else return (Just res)
756 yieldk a r =
757 if res <= a
758 then go (Just a) r
759 else go (Just res) r
760 in unStream m1 defState stop single yieldk
762{-# INLINE maximumBy #-}
763maximumBy :: (IsStream t, Monad m) => (a -> a -> Ordering) -> t m a -> m (Maybe a)
764maximumBy cmp m = go Nothing (toStream m)
765 where
766 go Nothing m1 =
767 let stop = return Nothing
768 single a = return (Just a)
769 yieldk a r = go (Just a) r
770 in unStream m1 defState stop single yieldk
772 go (Just res) m1 =
773 let stop = return (Just res)
774 single a = case cmp res a of
775 GT -> return (Just res)
776 _ -> return (Just a)
777 yieldk a r = case cmp res a of
778 GT -> go (Just res) r
779 _ -> go (Just a) r
780 in unStream m1 defState stop single yieldk
782{-# INLINE (!!) #-}
783(!!) :: (IsStream t, Monad m) => t m a -> Int -> m (Maybe a)
784m !! i = go i (toStream m)
785 where
786 go n m1 =
787 let single a | n == 0 = return $ Just a
788 | otherwise = return Nothing
789 yieldk a x | n < 0 = return Nothing
790 | n == 0 = return $ Just a
791 | otherwise = go (n - 1) x
792 in unStream m1 defState (return Nothing) single yieldk
794{-# INLINE lookup #-}
795lookup :: (IsStream t, Monad m, Eq a) => a -> t m (a, b) -> m (Maybe b)
796lookup e m = go (toStream m)
797 where
798 go m1 =
799 let single (a, b) | a == e = return $ Just b
800 | otherwise = return Nothing
801 yieldk (a, b) x | a == e = return $ Just b
802 | otherwise = go x
803 in unStream m1 defState (return Nothing) single yieldk
805{-# INLINE find #-}
806find :: (IsStream t, Monad m) => (a -> Bool) -> t m a -> m (Maybe a)
807find p m = go (toStream m)
808 where
809 go m1 =
810 let single a | p a = return $ Just a
811 | otherwise = return Nothing
812 yieldk a x | p a = return $ Just a
813 | otherwise = go x
814 in unStream m1 defState (return Nothing) single yieldk
816{-# INLINE findIndices #-}
817findIndices :: IsStream t => (a -> Bool) -> t m a -> t m Int
818findIndices p = fromStream . go 0 . toStream
819 where
820 go offset m1 = Stream $ \st stp sng yld ->
821 let single a | p a = sng offset
822 | otherwise = stp
823 yieldk a x | p a = yld offset $ go (offset + 1) x
824 | otherwise = unStream (go (offset + 1) x) st stp sng yld
825 in unStream m1 (rstState st) stp single yieldk
827------------------------------------------------------------------------------
828-- Map and Fold
829------------------------------------------------------------------------------
831-- | Apply a monadic action to each element of the stream and discard the
832-- output of the action.
833mapM_ :: (IsStream t, Monad m) => (a -> m b) -> t m a -> m ()
834mapM_ f m = go (toStream m)
835 where
836 go m1 =
837 let stop = return ()
838 single a = void (f a)
839 yieldk a r = f a >> go r
840 in unStream m1 defState stop single yieldk
842------------------------------------------------------------------------------
843-- Converting folds
844------------------------------------------------------------------------------
846{-# INLINABLE toList #-}
847toList :: (IsStream t, Monad m) => t m a -> m [a]
848toList = foldr (:) []
850{-# INLINE toStreamK #-}
851toStreamK :: Stream m a -> Stream m a
852toStreamK = id
854-------------------------------------------------------------------------------
855-- Transformation by folding (Scans)
856-------------------------------------------------------------------------------
858{-# INLINE scanx #-}
859scanx :: IsStream t => (x -> a -> x) -> x -> (x -> b) -> t m a -> t m b
860scanx step begin done m =
861 cons (done begin) $ fromStream $ go (toStream m) begin
862 where
863 go m1 !acc = Stream $ \st stp sng yld ->
864 let single a = sng (done $ step acc a)
865 yieldk a r =
866 let s = step acc a
867 in yld (done s) (go r s)
868 in unStream m1 (rstState st) stp single yieldk
870{-# INLINE scanl' #-}
871scanl' :: IsStream t => (b -> a -> b) -> b -> t m a -> t m b
872scanl' step begin = scanx step begin id
874-------------------------------------------------------------------------------
875-- Filtering
876-------------------------------------------------------------------------------
878{-# INLINE filter #-}
879filter :: IsStream t => (a -> Bool) -> t m a -> t m a
880filter p m = fromStream $ go (toStream m)
881 where
882 go m1 = Stream $ \st stp sng yld ->
883 let single a | p a = sng a
884 | otherwise = stp
885 yieldk a r | p a = yld a (go r)
886 | otherwise = unStream r (rstState st) stp single yieldk
887 in unStream m1 (rstState st) stp single yieldk
889{-# INLINE take #-}
890take :: IsStream t => Int -> t m a -> t m a
891take n m = fromStream $ go n (toStream m)
892 where
893 go n1 m1 = Stream $ \st stp sng yld ->
894 let yieldk a r = yld a (go (n1 - 1) r)
895 in if n1 <= 0
896 then stp
897 else unStream m1 (rstState st) stp sng yieldk
899{-# INLINE takeWhile #-}
900takeWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
901takeWhile p m = fromStream $ go (toStream m)
902 where
903 go m1 = Stream $ \st stp sng yld ->
904 let single a | p a = sng a
905 | otherwise = stp
906 yieldk a r | p a = yld a (go r)
907 | otherwise = stp
908 in unStream m1 (rstState st) stp single yieldk
910drop :: IsStream t => Int -> t m a -> t m a
911drop n m = fromStream $ Stream $ \st stp sng yld ->
912 unStream (go n (toStream m)) (rstState st) stp sng yld
913 where
914 go n1 m1 = Stream $ \st stp sng yld ->
915 let single _ = stp
916 yieldk _ r = (unStream $ go (n1 - 1) r) st stp sng yld
917 -- Somehow "<=" check performs better than a ">"
918 in if n1 <= 0
919 then unStream m1 st stp sng yld
920 else unStream m1 st stp single yieldk
922{-# INLINE dropWhile #-}
923dropWhile :: IsStream t => (a -> Bool) -> t m a -> t m a
924dropWhile p m = fromStream $ go (toStream m)
925 where
926 go m1 = Stream $ \st stp sng yld ->
927 let single a | p a = stp
928 | otherwise = sng a
929 yieldk a r | p a = unStream r (rstState st) stp single yieldk
930 | otherwise = yld a r
931 in unStream m1 (rstState st) stp single yieldk
933-------------------------------------------------------------------------------
934-- Mapping
935-------------------------------------------------------------------------------
937{-# INLINE map #-}
938map :: (IsStream t, Monad m) => (a -> b) -> t m a -> t m b
939map f m = fromStream $ Stream $ \st stp sng yld ->
940 let single = sng . f
941 yieldk a r = yld (f a) (fmap f r)
942 in unStream (toStream m) (rstState st) stp single yieldk
944-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
945-- other stream types to overload it.
946{-# INLINE mapM #-}
947mapM :: (IsStream t, MonadAsync m) => (a -> m b) -> t m a -> t m b
948mapM f m = go (toStream m)
949 where
950 go m1 = fromStream $ Stream $ \st stp sng yld ->
951 let single a = f a >>= sng
952 yieldk a r = unStream (toStream (f a |: go r)) st stp sng yld
953 in unStream m1 (rstState st) stp single yieldk
955-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
956-- other stream types to overload it.
957{-# INLINE sequence #-}
958sequence :: (IsStream t, MonadAsync m) => t m (m a) -> t m a
959sequence m = go (toStream m)
960 where
961 go m1 = fromStream $ Stream $ \st stp sng yld ->
962 let single ma = ma >>= sng
963 yieldk ma r = unStream (toStream $ ma |: go r) st stp sng yld
964 in unStream m1 (rstState st) stp single yieldk
966-------------------------------------------------------------------------------
967-- Inserting
968-------------------------------------------------------------------------------
970{-# INLINE intersperseM #-}
971intersperseM :: (IsStream t, MonadAsync m) => m a -> t m a -> t m a
972intersperseM a m = fromStream $ prependingStart (toStream m)
973 where
974 prependingStart m1 = Stream $ \st stp sng yld ->
975 let yieldk i x = unStream (return i |: go x) st stp sng yld
976 in unStream m1 (rstState st) stp sng yieldk
977 go m2 = fromStream $ Stream $ \st stp sng yld ->
978 let single i = unStream (a |: yield i) st stp sng yld
979 yieldk i x = unStream (a |: return i |: go x) st stp sng yld
980 in unStream m2 (rstState st) stp single yieldk
982{-# INLINE insertBy #-}
983insertBy :: IsStream t => (a -> a -> Ordering) -> a -> t m a -> t m a
984insertBy cmp x m = fromStream $ go (toStream m)
985 where
986 go m1 = Stream $ \st _ _ yld ->
987 let single a = case cmp x a of
988 GT -> yld a (yield x)
989 _ -> yld x (yield a)
990 stop = yld x nil
991 yieldk a r = case cmp x a of
992 GT -> yld a (go r)
993 _ -> yld x (a `cons` r)
994 in unStream m1 (rstState st) stop single yieldk
996------------------------------------------------------------------------------
997-- Deleting
998------------------------------------------------------------------------------
1000{-# INLINE deleteBy #-}
1001deleteBy :: IsStream t => (a -> a -> Bool) -> a -> t m a -> t m a
1002deleteBy eq x m = fromStream $ go (toStream m)
1003 where
1004 go m1 = Stream $ \st stp sng yld ->
1005 let single a = if eq x a then stp else sng a
1006 yieldk a r = if eq x a
1007 then unStream r (rstState st) stp sng yld
1008 else yld a (go r)
1009 in unStream m1 (rstState st) stp single yieldk
1011-------------------------------------------------------------------------------
1012-- Map and Filter
1013-------------------------------------------------------------------------------
1015{-# INLINE mapMaybe #-}
1016mapMaybe :: IsStream t => (a -> Maybe b) -> t m a -> t m b
1017mapMaybe f m = go (toStream m)
1018 where
1019 go m1 = fromStream $ Stream $ \st stp sng yld ->
1020 let single a = case f a of
1021 Just b -> sng b
1022 Nothing -> stp
1023 yieldk a r = case f a of
1024 Just b -> yld b (toStream $ go r)
1025 Nothing -> unStream r (rstState st) stp single yieldk
1026 in unStream m1 (rstState st) stp single yieldk
1028------------------------------------------------------------------------------
1029-- Serial Zipping
1030------------------------------------------------------------------------------
1032{-# INLINE zipWithS #-}
1033zipWithS :: (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
1034zipWithS f = go
1035 where
1036 go mx my = Stream $ \st stp sng yld -> do
1037 let merge a ra =
1038 let single2 b = sng (f a b)
1039 yield2 b rb = yld (f a b) (go ra rb)
1040 in unStream my (rstState st) stp single2 yield2
1041 let single1 a = merge a nil
1042 yield1 = merge
1043 unStream mx (rstState st) stp single1 yield1
1045-- | Zip two streams serially using a pure zipping function.
1047-- @since 0.1.0
1048{-# INLINABLE zipWith #-}
1049zipWith :: IsStream t => (a -> b -> c) -> t m a -> t m b -> t m c
1050zipWith f m1 m2 = fromStream $ zipWithS f (toStream m1) (toStream m2)
1052-- | Zip two streams serially using a monadic zipping function.
1054-- @since 0.1.0
1055zipWithM :: (IsStream t, Monad m) => (a -> b -> m c) -> t m a -> t m b -> t m c
1056zipWithM f m1 m2 = fromStream $ go (toStream m1) (toStream m2)
1057 where
1058 go mx my = Stream $ \st stp sng yld -> do
1059 let merge a ra =
1060 let runIt x = unStream x (rstState st) stp sng yld
1061 single2 b = f a b >>= sng
1062 yield2 b rb = f a b >>= \x -> runIt (x `cons` go ra rb)
1063 in unStream my (rstState st) stp single2 yield2
1064 let single1 a = merge a nil
1065 yield1 = merge
1066 unStream mx (rstState st) stp single1 yield1
1068------------------------------------------------------------------------------
1069-- Merging
1070------------------------------------------------------------------------------
1072data MergeWithState l r o
1073 = MergeBoth l r
1074 | MergeWithLeft l r o
1075 | MergeWithRight l r o
1077{-# INLINE mergeByS #-}
1078mergeByS :: (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
1079mergeByS cmp m1 m2 = go (MergeBoth m1 m2)
1080 where
1081 go (MergeBoth mx my) =
1082 Stream $ \st stp sng yld -> do
1083 let yield1 a ra =
1084 let single2 b = sng (cond a b)
1085 stop2 = yld a ra
1086 yield2 b rb =
1087 case cmp a b of
1088 LT -> yld a (go (MergeWithRight ra rb b))
1089 _ -> yld b (go (MergeWithLeft ra rb a))
1090 in unStream my (rstState st) stop2 single2 yield2
1091 single1 a = sng a
1092 stop1 = unStream my (rstState st) stp sng yld
1093 unStream mx (rstState st) stop1 single1 yield1
1094 go (MergeWithLeft mx my x) =
1095 Stream $ \st _ sng yld -> do
1096 let single2 b = sng (cond x b)
1097 stop2 = yld x mx
1098 yield2 b rb =
1099 case cmp x b of
1100 LT -> yld x (go (MergeWithRight mx rb b))
1101 _ -> yld b (go (MergeWithLeft mx rb x))
1102 unStream my (rstState st) stop2 single2 yield2
1103 go (MergeWithRight mx my x) =
1104 Stream $ \st _ sng yld -> do
1105 let single1 a = sng (cond a x)
1106 stop1 = yld x my
1107 yield1 a ra =
1108 case cmp a x of
1109 LT -> yld a (go (MergeWithRight ra my x))
1110 _ -> yld x (go (MergeWithLeft ra my a))
1111 unStream mx (rstState st) stop1 single1 yield1
1112 cond x y =
1113 if cmp x y == LT
1114 then x
1115 else y
1117{-# INLINABLE mergeBy #-}
1118mergeBy :: IsStream t => (a -> a -> Ordering) -> t m a -> t m a -> t m a
1119mergeBy f m1 m2 = fromStream $ mergeByS f (toStream m1) (toStream m2)
1121------------------------------------------------------------------------------
1122-- Transformation comprehensions
1123------------------------------------------------------------------------------
1125{-# INLINE the #-}
1126the :: (Eq a, IsStream t, Monad m) => t m a -> m (Maybe a)
1127the m = do
1128 r <- uncons m
1129 case r of
1130 Nothing -> return Nothing
1131 Just (h, t) -> go h (toStream t)
1132 where
1133 go h m1 =
1134 let single a | h == a = return $ Just h
1135 | otherwise = return Nothing
1136 yieldk a r | h == a = go h r
1137 | otherwise = return Nothing
1138 in unStream m1 defState (return $ Just h) single yieldk
1140------------------------------------------------------------------------------
1141-- Semigroup
1142------------------------------------------------------------------------------
1144-- | Concatenates two streams sequentially i.e. the first stream is
1145-- exhausted completely before yielding any element from the second stream.
1146{-# INLINE serial #-}
1147serial :: Stream m a -> Stream m a -> Stream m a
1148serial m1 m2 = go m1
1149 where
1150 go (Stream m) = Stream $ \st stp sng yld ->
1151 let stop = unStream m2 (rstState st) stp sng yld
1152 single a = yld a m2
1153 yieldk a r = yld a (go r)
1154 in m (rstState st) stop single yieldk
1156instance Semigroup (Stream m a) where
1157 (<>) = serial
1159------------------------------------------------------------------------------
1160-- Monoid
1161------------------------------------------------------------------------------
1163instance Monoid (Stream m a) where
1164 mempty = nil
1165 mappend = (<>)
1167-------------------------------------------------------------------------------
1168-- Functor
1169-------------------------------------------------------------------------------
1171instance Monad m => Functor (Stream m) where
1172 fmap = map
1174-------------------------------------------------------------------------------
1175-- Bind utility
1176-------------------------------------------------------------------------------
1178{-# INLINE bindWith #-}
1179bindWith
1180 :: (forall c. Stream m c -> Stream m c -> Stream m c)
1181 -> Stream m a
1182 -> (a -> Stream m b)
1183 -> Stream m b
1184bindWith par m f = go m
1185 where
1186 go (Stream g) =
1187 Stream $ \st stp sng yld ->
1188 let runShared x = unstreamShared x st stp sng yld
1189 runIsolated x = unStreamIsolated x st stp sng yld
1191 single a = runIsolated $ f a
1192 yieldk a r = runShared $ isolateStream (f a) `par` go r
1193 in g (rstState st) stp single yieldk
1195------------------------------------------------------------------------------
1196-- Alternative & MonadPlus
1197------------------------------------------------------------------------------
1199_alt :: Stream m a -> Stream m a -> Stream m a
1200_alt m1 m2 = Stream $ \st stp sng yld ->
1201 let stop = unStream m2 (rstState st) stp sng yld
1202 in unStream m1 (rstState st) stop sng yld
1204------------------------------------------------------------------------------
1205-- MonadReader
1206------------------------------------------------------------------------------
1208withLocal :: MonadReader r m => (r -> r) -> Stream m a -> Stream m a
1209withLocal f m =
1210 Stream $ \st stp sng yld ->
1211 let single = local f . sng
1212 yieldk a r = local f $ yld a (withLocal f r)
1213 in unStream m (rstState st) (local f stp) single yieldk
1215------------------------------------------------------------------------------
1216-- MonadError
1217------------------------------------------------------------------------------
1220-- XXX handle and test cross thread state transfer
1221withCatchError
1222 :: MonadError e m
1223 => Stream m a -> (e -> Stream m a) -> Stream m a
1224withCatchError m h =
1225 Stream $ \_ stp sng yld ->
1226 let run x = unStream x Nothing stp sng yieldk
1227 handle r = r `catchError` \e -> run $ h e
1228 yieldk a r = yld a (withCatchError r h)
1229 in handle $ run m
1232-------------------------------------------------------------------------------
1233-- Transformers
1234-------------------------------------------------------------------------------
1236instance MonadTrans Stream where
1237 lift = yieldM