git.iko.soy/streamly
on e08ff86 help sign in
streamly/src/Streamly/Streams/StreamD.hs
1446 lines 46 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE CPP #-}
3{-# LANGUAGE ConstraintKinds #-}
4{-# LANGUAGE ExistentialQuantification #-}
5{-# LANGUAGE FlexibleContexts #-}
6{-# LANGUAGE FlexibleInstances #-}
7{-# LANGUAGE MultiParamTypeClasses #-}
8{-# LANGUAGE PatternSynonyms #-}
9{-# LANGUAGE ViewPatterns #-}
10{-# LANGUAGE RankNTypes #-}
12#include "inline.hs"
14-- |
15-- Module : Streamly.Streams.StreamD
16-- Copyright : (c) 2018 Harendra Kumar
17-- Copyright : (c) Roman Leshchinskiy 2008-2010
18--
19-- License : BSD3
20-- Maintainer : harendra.kumar@gmail.com
21-- Stability : experimental
22-- Portability : GHC
23--
24-- Direct style re-implementation of CPS style stream in StreamK module. The
25-- symbol or suffix 'D' in this module denotes the "Direct" style. GHC is able
26-- to INLINE and fuse direct style better, providing better performance than
27-- CPS implementation.
28--
29-- @
30-- import qualified Streamly.Streams.StreamD as D
31-- @
33-- Some of the functions in this file have been adapted from the vector
34-- library, https://hackage.haskell.org/package/vector.
36module Streamly.Streams.StreamD
37 (
38 -- * The stream type
39 Step (..)
40 , Stream (..)
42 -- * Construction
43 , nil
44 , cons
46 -- * Deconstruction
47 , uncons
49 -- * Generation
50 -- ** Unfolds
51 , unfoldr
52 , unfoldrM
54 -- ** Specialized Generation
55 -- | Generate a monadic stream from a seed.
56 , repeat
57 , replicate
58 , replicateM
59 , fromIndices
60 , fromIndicesM
61 , generate
62 , generateM
64 -- ** Enumerations
65 , intFromStep
66 , intFrom
67 , intFromThen
68 , intFromTo
69 , intFromThenTo
71 , numFromStep
72 , numFrom
73 , numFromThen
74 , fracFromTo
75 , fracFromThenTo
77 -- ** Conversions
78 -- | Transform an input structure into a stream.
79 -- | Direct style stream does not support @fromFoldable@.
80 , yield
81 , yieldM
82 , fromList
83 , fromListM
84 , fromStreamK
85 , fromStreamD
87 -- * Elimination
88 -- ** General Folds
89 , foldr
90 , foldrM
91 , foldr1
92 , foldl'
93 , foldlM'
95 -- ** Specialized Folds
96 , runStream
97 , null
98 , head
99 , tail
100 , last
101 , elem
102 , notElem
103 , all
104 , any
105 , maximum
106 , maximumBy
107 , minimum
108 , minimumBy
109 , findIndices
110 , lookup
111 , findM
112 , find
113 , (!!)
114 , concatMapM
115 , concatMap
117 -- ** Substreams
118 , isPrefixOf
119 , isSubsequenceOf
120 , stripPrefix
122 -- ** Map and Fold
123 , mapM_
125 -- ** Conversions
126 -- | Transform a stream into another type.
127 , toList
128 , toStreamK
129 , toStreamD
131 -- * Transformation
132 -- ** By folding (scans)
133 , scanlM'
134 , scanl'
135 , scanlM
136 , scanl
137 , scanl1M'
138 , scanl1'
139 , scanl1M
140 , scanl1
142 , prescanl'
143 , prescanlM'
144 , postscanl
145 , postscanlM
146 , postscanl'
147 , postscanlM'
149 -- * Filtering
150 , filter
151 , filterM
152 , uniq
153 , take
154 , takeWhile
155 , takeWhileM
156 , drop
157 , dropWhile
158 , dropWhileM
160 -- * Mapping
161 , map
162 , mapM
163 , sequence
165 -- * Inserting
166 , insertBy
168 -- * Deleting
169 , deleteBy
171 -- ** Map and Filter
172 , mapMaybe
173 , mapMaybeM
175 -- * Zipping
176 , indexed
177 , indexedR
178 , zipWith
179 , zipWithM
181 -- * Comparisions
182 , eqBy
183 , cmpBy
185 -- * Merging
186 , mergeBy
187 , mergeByM
189 -- * Transformation comprehensions
190 , the
191 )
192where
194import Data.Maybe (fromJust, isJust)
195import GHC.Types ( SPEC(..) )
196import Prelude
197 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
198 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
199 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
200 (!!), scanl, scanl1, concatMap, replicate, enumFromTo)
202import Streamly.SVar (MonadAsync, defState, adaptState)
204import Streamly.Streams.StreamDType
205import qualified Streamly.Streams.StreamK as K
207------------------------------------------------------------------------------
208-- Construction
209------------------------------------------------------------------------------
211-- | An empty 'Stream'.
212{-# INLINE_NORMAL nil #-}
213nil :: Monad m => Stream m a
214nil = Stream (\_ _ -> return Stop) ()
216-- | Can fuse but has O(n^2) complexity.
217{-# INLINE_NORMAL cons #-}
218cons :: Monad m => a -> Stream m a -> Stream m a
219cons x (Stream step state) = Stream step1 Nothing
220 where
221 {-# INLINE_LATE step1 #-}
222 step1 _ Nothing = return $ Yield x (Just state)
223 step1 gst (Just st) = do
224 r <- step gst st
225 return $
226 case r of
227 Yield a s -> Yield a (Just s)
228 Skip s -> Skip (Just s)
229 Stop -> Stop
231-------------------------------------------------------------------------------
232-- Deconstruction
233-------------------------------------------------------------------------------
235-- Does not fuse, has the same performance as the StreamK version.
236{-# INLINE_NORMAL uncons #-}
237uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
238uncons (Stream step state) = go state
239 where
240 go st = do
241 r <- step defState st
242 case r of
243 Yield x s -> return $ Just (x, Stream step s)
244 Skip s -> go s
245 Stop -> return Nothing
247------------------------------------------------------------------------------
248-- Generation by unfold
249------------------------------------------------------------------------------
251{-# INLINE_NORMAL unfoldrM #-}
252unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
253unfoldrM next state = Stream step state
254 where
255 {-# INLINE_LATE step #-}
256 step _ st = do
257 r <- next st
258 return $ case r of
259 Just (x, s) -> Yield x s
260 Nothing -> Stop
262{-# INLINE_LATE unfoldr #-}
263unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
264unfoldr f = unfoldrM (return . f)
266------------------------------------------------------------------------------
267-- Specialized Generation
268------------------------------------------------------------------------------
270repeat :: Monad m => a -> Stream m a
271repeat x = Stream (\_ _ -> return $ Yield x ()) ()
273{-# INLINE_NORMAL replicateM #-}
274replicateM :: Monad m => Int -> m a -> Stream m a
275replicateM n p = Stream step n
276 where
277 {-# INLINE_LATE step #-}
278 step _ i | i <= 0 = return Stop
279 | otherwise = do
280 x <- p
281 return $ Yield x (i - 1)
283{-# INLINE_NORMAL replicate #-}
284replicate :: Monad m => Int -> a -> Stream m a
285replicate n x = replicateM n (return x)
287-- This would not work properly for floats, therefore we put an Integral
288-- constraint.
289-- | Can be used to enumerate unbounded integrals. This does not check for
290-- overflow or underflow for bounded integrals.
291{-# INLINE_NORMAL intFromStep #-}
292intFromStep :: (Integral a, Monad m) => a -> a -> Stream m a
293intFromStep from stride =
294 from `seq` stride `seq` Stream step from
295 where
296 {-# INLINE_LATE step #-}
297 step _ !x = return $ Yield x $! (x + stride)
299-- We are assuming that "to" is constrained by the type to be within
300-- max/min bounds.
301{-# INLINE intFromTo #-}
302intFromTo :: (Monad m, Integral a) => a -> a -> Stream m a
303intFromTo from to = takeWhile (<= to) $ intFromStep from 1
305{-# INLINE intFrom #-}
306intFrom :: (Monad m, Integral a, Bounded a) => a -> Stream m a
307intFrom from = intFromTo from maxBound
309data EnumState a = EnumInit | EnumYield a a a | EnumStop
311{-# INLINE_NORMAL intFromThenToUp #-}
312intFromThenToUp :: (Monad m, Integral a) => a -> a -> a -> Stream m a
313intFromThenToUp from next to = Stream step EnumInit
314 where
315 {-# INLINE_LATE step #-}
316 step _ EnumInit =
317 return $
318 if to < next
319 then if to < from
320 then Stop
321 else Yield from EnumStop
322 else -- from <= next <= to
323 let stride = next - from
324 in Skip $ EnumYield from stride (to - stride)
326 step _ (EnumYield x stride toMinus) =
327 return $
328 if x > toMinus
329 then Yield x EnumStop
330 else Yield x $ EnumYield (x + stride) stride toMinus
332 step _ EnumStop = return Stop
334{-# INLINE_NORMAL intFromThenToDn #-}
335intFromThenToDn :: (Monad m, Integral a) => a -> a -> a -> Stream m a
336intFromThenToDn from next to = Stream step EnumInit
337 where
338 {-# INLINE_LATE step #-}
339 step _ EnumInit =
340 return $ if to > next
341 then if to > from
342 then Stop
343 else Yield from EnumStop
344 else -- from >= next >= to
345 let stride = next - from
346 in Skip $ EnumYield from stride (to - stride)
348 step _ (EnumYield x stride toMinus) =
349 return $
350 if x < toMinus
351 then Yield x EnumStop
352 else Yield x $ EnumYield (x + stride) stride toMinus
354 step _ EnumStop = return Stop
356{-# INLINE_NORMAL intFromThenTo #-}
357intFromThenTo :: (Monad m, Integral a) => a -> a -> a -> Stream m a
358intFromThenTo from next to
359 | next >= from = intFromThenToUp from next to
360 | otherwise = intFromThenToDn from next to
362{-# INLINE_NORMAL intFromThen #-}
363intFromThen :: (Monad m, Integral a, Bounded a) => a -> a -> Stream m a
364intFromThen from next =
365 if next > from
366 then intFromThenToUp from next maxBound
367 else intFromThenToDn from next minBound
369-- For floating point numbers if the increment is less than the precision then
370-- it just gets lost. Therefore we cannot always increment it correctly by just
371-- repeated addition.
372-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
373-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
375-- Instead we accumulate the increment counter and compute the increment
376-- everytime before adding it to the starting number.
377--
378-- This works for Integrals as well as floating point numbers, but intFromStep
379-- is faster for integrals.
380{-# INLINE_NORMAL numFromStep #-}
381numFromStep :: (Monad m, Num a) => a -> a -> Stream m a
382numFromStep from stride = Stream step 0
383 where
384 {-# INLINE_LATE step #-}
385 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
387{-# INLINE_NORMAL numFrom #-}
388numFrom :: (Monad m, Num a) => a -> Stream m a
389numFrom from = numFromStep from 1
391{-# INLINE_NORMAL numFromThen #-}
392numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
393numFromThen from next = numFromStep from (next - from)
395-- We cannot write a general function for Num. The only way to write code
396-- portable between the two is to use a 'Real' constraint and convert between
397-- Fractional and Integral using fromRational which is horribly slow.
398{-# INLINE_NORMAL fracFromTo #-}
399fracFromTo :: (Monad m, Fractional a, Ord a) => a -> a -> Stream m a
400fracFromTo from to = takeWhile (<= to + 1/2) $ numFromStep from 1
402{-# INLINE_NORMAL fracFromThenTo #-}
403fracFromThenTo :: (Monad m, Fractional a, Ord a) => a -> a -> a -> Stream m a
404fracFromThenTo from next to = takeWhile predicate $ numFromThen from next
405 where
406 mid = (next - from) / 2
407 predicate | next >= from = (<= to + mid)
408 | otherwise = (>= to + mid)
410-------------------------------------------------------------------------------
411-- Generation by Conversion
412-------------------------------------------------------------------------------
414-- | Create a singleton 'Stream' from a pure value.
415{-# INLINE_NORMAL yield #-}
416yield :: Monad m => a -> Stream m a
417yield x = Stream (\_ s -> return $ step undefined s) True
418 where
419 {-# INLINE_LATE step #-}
420 step _ True = Yield x False
421 step _ False = Stop
423-- | Create a singleton 'Stream' from a monadic action.
424{-# INLINE_NORMAL yieldM #-}
425yieldM :: Monad m => m a -> Stream m a
426yieldM m = Stream step True
427 where
428 {-# INLINE_LATE step #-}
429 step _ True = m >>= \x -> return $ Yield x False
430 step _ False = return Stop
432{-# INLINE_NORMAL fromIndicesM #-}
433fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
434fromIndicesM gen = Stream step 0
435 where
436 {-# INLINE_LATE step #-}
437 step _ i = do
438 x <- gen i
439 return $ Yield x (i + 1)
441{-# INLINE fromIndices #-}
442fromIndices :: Monad m => (Int -> a) -> Stream m a
443fromIndices gen = fromIndicesM (return . gen)
445{-# INLINE_NORMAL generateM #-}
446generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
447generateM n gen = n `seq` Stream step 0
448 where
449 {-# INLINE_LATE step #-}
450 step _ i | i < n = do
451 x <- gen i
452 return $ Yield x (i + 1)
453 | otherwise = return Stop
455{-# INLINE generate #-}
456generate :: Monad m => Int -> (Int -> a) -> Stream m a
457generate n gen = generateM n (return . gen)
459-- XXX we need the MonadAsync constraint because of a rewrite rule.
460-- | Convert a list of monadic actions to a 'Stream'
461{-# INLINE_LATE fromListM #-}
462fromListM :: MonadAsync m => [m a] -> Stream m a
463fromListM = Stream step
464 where
465 {-# INLINE_LATE step #-}
466 step _ (m:ms) = m >>= \x -> return $ Yield x ms
467 step _ [] = return Stop
469-- | Convert a list of pure values to a 'Stream'
470{-# INLINE_LATE fromList #-}
471fromList :: Monad m => [a] -> Stream m a
472fromList = Stream step
473 where
474 {-# INLINE_LATE step #-}
475 step _ (x:xs) = return $ Yield x xs
476 step _ [] = return Stop
478{-# INLINE_LATE fromStreamK #-}
479fromStreamK :: Monad m => K.Stream m a -> Stream m a
480fromStreamK = Stream step
481 where
482 step gst m1 =
483 let stop = return Stop
484 single a = return $ Yield a K.nil
485 yieldk a r = return $ Yield a r
486 in K.unStreamShared m1 gst stop single yieldk
488{-# INLINE toStreamD #-}
489toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
490toStreamD = fromStreamK . K.toStream
492------------------------------------------------------------------------------
493-- Elimination by Folds
494------------------------------------------------------------------------------
496{-# INLINE_NORMAL foldrM #-}
497foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b
498foldrM f z (Stream step state) = go SPEC state
499 where
500 go !_ st = do
501 r <- step defState st
502 case r of
503 Yield x s -> go SPEC s >>= f x
504 Skip s -> go SPEC s
505 Stop -> return z
507{-# INLINE_NORMAL foldr #-}
508foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
509foldr f = foldrM (\a b -> return (f a b))
511{-# INLINE_NORMAL foldr1 #-}
512foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
513foldr1 f m = do
514 r <- uncons m
515 case r of
516 Nothing -> return Nothing
517 Just (h, t) -> fmap Just (foldr f h t)
519{-# INLINE_NORMAL foldlM' #-}
520foldlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> m b
521foldlM' fstep begin (Stream step state) = go SPEC begin state
522 where
523 go !_ acc st = acc `seq` do
524 r <- step defState st
525 case r of
526 Yield x s -> do
527 acc' <- fstep acc x
528 go SPEC acc' s
529 Skip s -> go SPEC acc s
530 Stop -> return acc
532{-# INLINE foldl' #-}
533foldl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
534foldl' fstep = foldlM' (\b a -> return (fstep b a))
536------------------------------------------------------------------------------
537-- Specialized Folds
538------------------------------------------------------------------------------
540-- | Run a streaming composition, discard the results.
541{-# INLINE_LATE runStream #-}
542runStream :: Monad m => Stream m a -> m ()
543runStream (Stream step state) = go SPEC state
544 where
545 go !_ st = do
546 r <- step defState st
547 case r of
548 Yield _ s -> go SPEC s
549 Skip s -> go SPEC s
550 Stop -> return ()
552{-# INLINE_NORMAL null #-}
553null :: Monad m => Stream m a -> m Bool
554null (Stream step state) = go state
555 where
556 go st = do
557 r <- step defState st
558 case r of
559 Yield _ _ -> return False
560 Skip s -> go s
561 Stop -> return True
563-- XXX SPEC?
564{-# INLINE_NORMAL head #-}
565head :: Monad m => Stream m a -> m (Maybe a)
566head (Stream step state) = go state
567 where
568 go st = do
569 r <- step defState st
570 case r of
571 Yield x _ -> return (Just x)
572 Skip s -> go s
573 Stop -> return Nothing
575-- Does not fuse, has the same performance as the StreamK version.
576{-# INLINE_NORMAL tail #-}
577tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
578tail (Stream step state) = go state
579 where
580 go st = do
581 r <- step defState st
582 case r of
583 Yield _ s -> return (Just $ Stream step s)
584 Skip s -> go s
585 Stop -> return Nothing
587-- XXX will it fuse? need custom impl?
588{-# INLINE_NORMAL last #-}
589last :: Monad m => Stream m a -> m (Maybe a)
590last = foldl' (\_ y -> Just y) Nothing
592{-# INLINE_NORMAL elem #-}
593elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
594elem e (Stream step state) = go state
595 where
596 go st = do
597 r <- step defState st
598 case r of
599 Yield x s
600 | x == e -> return True
601 | otherwise -> go s
602 Skip s -> go s
603 Stop -> return False
605{-# INLINE_NORMAL notElem #-}
606notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
607notElem e s = fmap not (elem e s)
609{-# INLINE_NORMAL all #-}
610all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
611all p (Stream step state) = go state
612 where
613 go st = do
614 r <- step defState st
615 case r of
616 Yield x s
617 | p x -> go s
618 | otherwise -> return False
619 Skip s -> go s
620 Stop -> return True
622{-# INLINE_NORMAL any #-}
623any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
624any p (Stream step state) = go state
625 where
626 go st = do
627 r <- step defState st
628 case r of
629 Yield x s
630 | p x -> return True
631 | otherwise -> go s
632 Skip s -> go s
633 Stop -> return False
635{-# INLINE_NORMAL maximum #-}
636maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
637maximum (Stream step state) = go Nothing state
638 where
639 go Nothing st = do
640 r <- step defState st
641 case r of
642 Yield x s -> go (Just x) s
643 Skip s -> go Nothing s
644 Stop -> return Nothing
645 go (Just acc) st = do
646 r <- step defState st
647 case r of
648 Yield x s
649 | acc <= x -> go (Just x) s
650 | otherwise -> go (Just acc) s
651 Skip s -> go (Just acc) s
652 Stop -> return (Just acc)
654{-# INLINE_NORMAL maximumBy #-}
655maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
656maximumBy cmp (Stream step state) = go Nothing state
657 where
658 go Nothing st = do
659 r <- step defState st
660 case r of
661 Yield x s -> go (Just x) s
662 Skip s -> go Nothing s
663 Stop -> return Nothing
664 go (Just acc) st = do
665 r <- step defState st
666 case r of
667 Yield x s -> case cmp acc x of
668 GT -> go (Just acc) s
669 _ -> go (Just x) s
670 Skip s -> go (Just acc) s
671 Stop -> return (Just acc)
673{-# INLINE_NORMAL minimum #-}
674minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
675minimum (Stream step state) = go Nothing state
676 where
677 go Nothing st = do
678 r <- step defState st
679 case r of
680 Yield x s -> go (Just x) s
681 Skip s -> go Nothing s
682 Stop -> return Nothing
683 go (Just acc) st = do
684 r <- step defState st
685 case r of
686 Yield x s
687 | acc <= x -> go (Just acc) s
688 | otherwise -> go (Just x) s
689 Skip s -> go (Just acc) s
690 Stop -> return (Just acc)
692{-# INLINE_NORMAL minimumBy #-}
693minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
694minimumBy cmp (Stream step state) = go Nothing state
695 where
696 go Nothing st = do
697 r <- step defState st
698 case r of
699 Yield x s -> go (Just x) s
700 Skip s -> go Nothing s
701 Stop -> return Nothing
702 go (Just acc) st = do
703 r <- step defState st
704 case r of
705 Yield x s -> case cmp acc x of
706 GT -> go (Just x) s
707 _ -> go (Just acc) s
708 Skip s -> go (Just acc) s
709 Stop -> return (Just acc)
711{-# INLINE_NORMAL (!!) #-}
712(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
713(Stream step state) !! i = go i state
714 where
715 go n st = do
716 r <- step defState st
717 case r of
718 Yield x s | n < 0 -> return Nothing
719 | n == 0 -> return $ Just x
720 | otherwise -> go (n - 1) s
721 Skip s -> go n s
722 Stop -> return Nothing
724{-# INLINE_NORMAL lookup #-}
725lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
726lookup e (Stream step state) = go state
727 where
728 go st = do
729 r <- step defState st
730 case r of
731 Yield (a, b) s -> if e == a then return (Just b) else go s
732 Skip s -> go s
733 Stop -> return Nothing
735{-# INLINE_NORMAL findM #-}
736findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
737findM p (Stream step state) = go SPEC state
738 where
739 go !_ st = do
740 r <- step defState st
741 case r of
742 Yield x s -> do
743 b <- p x
744 if b then return (Just x) else go SPEC s
745 Skip s -> go SPEC s
746 Stop -> return Nothing
748{-# INLINE find #-}
749find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
750find p = findM (return . p)
752{-# INLINE_NORMAL findIndices #-}
753findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
754findIndices p (Stream step state) = Stream step' (state, 0)
755 where
756 {-# INLINE_LATE step' #-}
757 step' gst (st, i) = do
758 r <- step (adaptState gst) st
759 return $ case r of
760 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
761 Skip s -> Skip (s, i+1)
762 Stop -> Stop
764{-# INLINE_NORMAL concatMapM #-}
765concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b
766concatMapM f (Stream step state) = Stream step' (Left state)
767 where
768 {-# INLINE_LATE step' #-}
769 step' gst (Left st) = do
770 r <- step (adaptState gst) st
771 case r of
772 Yield a s -> do
773 b_stream <- f a
774 return $ Skip (Right (b_stream, s))
775 Skip s -> return $ Skip (Left s)
776 Stop -> return Stop
778 step' _ (Right (Stream inner_step inner_st, st)) = do
779 r <- inner_step defState inner_st
780 case r of
781 Yield b inner_s ->
782 return $ Yield b (Right (Stream inner_step inner_s, st))
783 Skip inner_s ->
784 return $ Skip (Right (Stream inner_step inner_s, st))
785 Stop -> return $ Skip (Left st)
787{-# INLINE concatMap #-}
788concatMap :: Monad m => (a -> Stream m b) -> Stream m a -> Stream m b
789concatMap f = concatMapM (return . f)
791------------------------------------------------------------------------------
792-- Substreams
793------------------------------------------------------------------------------
795{-# INLINE_NORMAL isPrefixOf #-}
796isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
797isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
798 where
799 go (sa, sb, Nothing) = do
800 r <- stepa defState sa
801 case r of
802 Yield x sa' -> go (sa', sb, Just x)
803 Skip sa' -> go (sa', sb, Nothing)
804 Stop -> return True
806 go (sa, sb, Just x) = do
807 r <- stepb defState sb
808 case r of
809 Yield y sb' ->
810 if x == y
811 then go (sa, sb', Nothing)
812 else return False
813 Skip sb' -> go (sa, sb', Just x)
814 Stop -> return False
816{-# INLINE_NORMAL isSubsequenceOf #-}
817isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
818isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
819 where
820 go (sa, sb, Nothing) = do
821 r <- stepa defState sa
822 case r of
823 Yield x sa' -> go (sa', sb, Just x)
824 Skip sa' -> go (sa', sb, Nothing)
825 Stop -> return True
827 go (sa, sb, Just x) = do
828 r <- stepb defState sb
829 case r of
830 Yield y sb' ->
831 if x == y
832 then go (sa, sb', Nothing)
833 else go (sa, sb', Just x)
834 Skip sb' -> go (sa, sb', Just x)
835 Stop -> return False
837{-# INLINE_NORMAL stripPrefix #-}
838stripPrefix
839 :: (Eq a, Monad m)
840 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
841stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
842 where
843 go (sa, sb, Nothing) = do
844 r <- stepa defState sa
845 case r of
846 Yield x sa' -> go (sa', sb, Just x)
847 Skip sa' -> go (sa', sb, Nothing)
848 Stop -> return $ Just (Stream stepb sb)
850 go (sa, sb, Just x) = do
851 r <- stepb defState sb
852 case r of
853 Yield y sb' ->
854 if x == y
855 then go (sa, sb', Nothing)
856 else return Nothing
857 Skip sb' -> go (sa, sb', Just x)
858 Stop -> return Nothing
860------------------------------------------------------------------------------
861-- Map and Fold
862------------------------------------------------------------------------------
864-- | Execute a monadic action for each element of the 'Stream'
865{-# INLINE_NORMAL mapM_ #-}
866mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
867mapM_ m = runStream . mapM m
869------------------------------------------------------------------------------
870-- Converting folds
871------------------------------------------------------------------------------
873{-# INLINE toList #-}
874toList :: Monad m => Stream m a -> m [a]
875toList = foldr (:) []
877-- Convert a direct stream to and from CPS encoded stream
878{-# INLINE_LATE toStreamK #-}
879toStreamK :: Monad m => Stream m a -> K.Stream m a
880toStreamK (Stream step state) = go state
881 where
882 go st = K.mkStream $ \gst stp sng yld -> do
883 r <- step gst st
884 case r of
885 Yield x s -> yld x (go s)
886 Skip s -> K.unStreamShared (go s) gst stp sng yld
887 Stop -> stp
889#ifndef DISABLE_FUSION
890{-# RULES "fromStreamK/toStreamK fusion"
891 forall s. toStreamK (fromStreamK s) = s #-}
892{-# RULES "toStreamK/fromStreamK fusion"
893 forall s. fromStreamK (toStreamK s) = s #-}
894#endif
896{-# INLINE fromStreamD #-}
897fromStreamD :: (K.IsStream t, Monad m) => Stream m a -> t m a
898fromStreamD = K.fromStream . toStreamK
900------------------------------------------------------------------------------
901-- Transformation by Folding (Scans)
902------------------------------------------------------------------------------
904-- XXX Is a prescan useful, discarding the last step does not sound useful? I
905-- am not sure about the utility of this function, so this is implemented but
906-- not exposed. We can expose it if someone provides good reasons why this is
907-- useful.
908--
909-- XXX We have to execute the stream one step ahead to know that we are at the
910-- last step. The vector implementation of prescan executes the last fold step
911-- but does not yield the result. This means we have executed the effect but
912-- discarded value. This does not sound right. In this implementation we are
913-- not executing the last fold step.
914{-# INLINE_NORMAL prescanlM' #-}
915prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
916prescanlM' f mz (Stream step state) = Stream step' (state, mz)
917 where
918 {-# INLINE_LATE step' #-}
919 step' gst (st, prev) = do
920 r <- step (adaptState gst) st
921 case r of
922 Yield x s -> do
923 acc <- prev
924 return $ Yield acc (s, f acc x)
925 Skip s -> return $ Skip (s, prev)
926 Stop -> return Stop
928{-# INLINE prescanl' #-}
929prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
930prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
932-- XXX if we make the initial value of the accumulator monadic then should we
933-- execute it even if the stream is empty? In that case we would have generated
934-- the effect but discarded the value, but that is what a fold does when the
935-- stream is empty. Whatever we decide, need to reconcile this with prescan.
936-- If we execute the initial value here without even using it then it is ok to
937-- execute the last step there as well without using the value.
938-- Looking at the duality with right fold, in case of right fold we always
939-- perform the action when the construction terminates, so in case of left fold
940-- we should perform it only when the reduction starts.
941{-# INLINE_NORMAL postscanlM' #-}
942postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
943postscanlM' fstep begin (Stream step state) =
944 begin `seq` Stream step' (state, begin)
945 where
946 {-# INLINE_LATE step' #-}
947 step' gst (st, acc) = acc `seq` do
948 r <- step (adaptState gst) st
949 case r of
950 Yield x s -> do
951 y <- fstep acc x
952 y `seq` return (Yield y (s, y))
953 Skip s -> return $ Skip (s, acc)
954 Stop -> return Stop
956{-# INLINE_NORMAL postscanl' #-}
957postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
958postscanl' f = postscanlM' (\a b -> return (f a b))
960{-# INLINE_NORMAL postscanlM #-}
961postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
962postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
963 where
964 {-# INLINE_LATE step' #-}
965 step' gst (st, acc) = do
966 r <- step (adaptState gst) st
967 case r of
968 Yield x s -> do
969 y <- fstep acc x
970 return (Yield y (s, y))
971 Skip s -> return $ Skip (s, acc)
972 Stop -> return Stop
974{-# INLINE_NORMAL postscanl #-}
975postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
976postscanl f = postscanlM (\a b -> return (f a b))
978{-# INLINE_NORMAL scanlM' #-}
979scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
980scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
982{-# INLINE scanl' #-}
983scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
984scanl' f = scanlM' (\a b -> return (f a b))
986{-# INLINE_NORMAL scanlM #-}
987scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
988scanlM fstep begin s = begin `cons` postscanlM fstep begin s
990{-# INLINE scanl #-}
991scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
992scanl f = scanlM (\a b -> return (f a b))
994{-# INLINE_NORMAL scanl1M #-}
995scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
996scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
997 where
998 {-# INLINE_LATE step' #-}
999 step' gst (st, Nothing) = do
1000 r <- step gst st
1001 case r of
1002 Yield x s -> return $ Yield x (s, Just x)
1003 Skip s -> return $ Skip (s, Nothing)
1004 Stop -> return Stop
1006 step' gst (st, Just acc) = do
1007 r <- step gst st
1008 case r of
1009 Yield y s -> do
1010 z <- fstep acc y
1011 return $ Yield z (s, Just z)
1012 Skip s -> return $ Skip (s, Just acc)
1013 Stop -> return Stop
1015{-# INLINE scanl1 #-}
1016scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
1017scanl1 f = scanl1M (\x y -> return (f x y))
1019{-# INLINE_NORMAL scanl1M' #-}
1020scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
1021scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
1022 where
1023 {-# INLINE_LATE step' #-}
1024 step' gst (st, Nothing) = do
1025 r <- step gst st
1026 case r of
1027 Yield x s -> x `seq` return $ Yield x (s, Just x)
1028 Skip s -> return $ Skip (s, Nothing)
1029 Stop -> return Stop
1031 step' gst (st, Just acc) = acc `seq` do
1032 r <- step gst st
1033 case r of
1034 Yield y s -> do
1035 z <- fstep acc y
1036 z `seq` return $ Yield z (s, Just z)
1037 Skip s -> return $ Skip (s, Just acc)
1038 Stop -> return Stop
1040{-# INLINE scanl1' #-}
1041scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
1042scanl1' f = scanl1M' (\x y -> return (f x y))
1044-------------------------------------------------------------------------------
1045-- Filtering
1046-------------------------------------------------------------------------------
1048{-# INLINE_NORMAL take #-}
1049take :: Monad m => Int -> Stream m a -> Stream m a
1050take n (Stream step state) = n `seq` Stream step' (state, 0)
1051 where
1052 {-# INLINE_LATE step' #-}
1053 step' gst (st, i) | i < n = do
1054 r <- step gst st
1055 return $ case r of
1056 Yield x s -> Yield x (s, i + 1)
1057 Skip s -> Skip (s, i)
1058 Stop -> Stop
1059 step' _ (_, _) = return Stop
1061{-# INLINE_NORMAL takeWhileM #-}
1062takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
1063takeWhileM f (Stream step state) = Stream step' state
1064 where
1065 {-# INLINE_LATE step' #-}
1066 step' gst st = do
1067 r <- step gst st
1068 case r of
1069 Yield x s -> do
1070 b <- f x
1071 return $ if b then Yield x s else Stop
1072 Skip s -> return $ Skip s
1073 Stop -> return Stop
1075{-# INLINE takeWhile #-}
1076takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
1077takeWhile f = takeWhileM (return . f)
1079{-# INLINE_NORMAL drop #-}
1080drop :: Monad m => Int -> Stream m a -> Stream m a
1081drop n (Stream step state) = Stream step' (state, Just n)
1082 where
1083 {-# INLINE_LATE step' #-}
1084 step' gst (st, Just i)
1085 | i > 0 = do
1086 r <- step gst st
1087 return $
1088 case r of
1089 Yield _ s -> Skip (s, Just (i - 1))
1090 Skip s -> Skip (s, Just i)
1091 Stop -> Stop
1092 | otherwise = return $ Skip (st, Nothing)
1094 step' gst (st, Nothing) = do
1095 r <- step gst st
1096 return $
1097 case r of
1098 Yield x s -> Yield x (s, Nothing)
1099 Skip s -> Skip (s, Nothing)
1100 Stop -> Stop
1102data DropWhileState s a
1103 = DropWhileDrop s
1104 | DropWhileYield a s
1105 | DropWhileNext s
1107{-# INLINE_NORMAL dropWhileM #-}
1108dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
1109dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
1110 where
1111 {-# INLINE_LATE step' #-}
1112 step' gst (DropWhileDrop st) = do
1113 r <- step gst st
1114 case r of
1115 Yield x s -> do
1116 b <- f x
1117 if b
1118 then return $ Skip (DropWhileDrop s)
1119 else return $ Skip (DropWhileYield x s)
1120 Skip s -> return $ Skip (DropWhileDrop s)
1121 Stop -> return Stop
1123 step' gst (DropWhileNext st) = do
1124 r <- step gst st
1125 case r of
1126 Yield x s -> return $ Skip (DropWhileYield x s)
1127 Skip s -> return $ Skip (DropWhileNext s)
1128 Stop -> return Stop
1130 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
1132{-# INLINE dropWhile #-}
1133dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
1134dropWhile f = dropWhileM (return . f)
1136{-# INLINE_NORMAL filterM #-}
1137filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
1138filterM f (Stream step state) = Stream step' state
1139 where
1140 {-# INLINE_LATE step' #-}
1141 step' gst st = do
1142 r <- step gst st
1143 case r of
1144 Yield x s -> do
1145 b <- f x
1146 return $ if b
1147 then Yield x s
1148 else Skip s
1149 Skip s -> return $ Skip s
1150 Stop -> return Stop
1152{-# INLINE filter #-}
1153filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
1154filter f = filterM (return . f)
1156{-# INLINE_NORMAL uniq #-}
1157uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
1158uniq (Stream step state) = Stream step' (Nothing, state)
1159 where
1160 {-# INLINE_LATE step' #-}
1161 step' gst (Nothing, st) = do
1162 r <- step gst st
1163 case r of
1164 Yield x s -> return $ Yield x (Just x, s)
1165 Skip s -> return $ Skip (Nothing, s)
1166 Stop -> return Stop
1167 step' gst (Just x, st) = do
1168 r <- step gst st
1169 case r of
1170 Yield y s | x == y -> return $ Skip (Just x, s)
1171 | otherwise -> return $ Yield x (Just y, s)
1172 Skip s -> return $ Skip (Just x, s)
1173 Stop -> return Stop
1175------------------------------------------------------------------------------
1176-- Transformation by Mapping
1177------------------------------------------------------------------------------
1179{-# INLINE_NORMAL sequence #-}
1180sequence :: Monad m => Stream m (m a) -> Stream m a
1181sequence (Stream step state) = Stream step' state
1182 where
1183 {-# INLINE_LATE step' #-}
1184 step' gst st = do
1185 r <- step (adaptState gst) st
1186 case r of
1187 Yield x s -> x >>= \a -> return (Yield a s)
1188 Skip s -> return $ Skip s
1189 Stop -> return Stop
1191------------------------------------------------------------------------------
1192-- Inserting
1193------------------------------------------------------------------------------
1195{-# INLINE_NORMAL insertBy #-}
1196insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
1197insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
1198 where
1199 {-# INLINE_LATE step' #-}
1200 step' gst (st, False, _) = do
1201 r <- step gst st
1202 case r of
1203 Yield x s -> case cmp a x of
1204 GT -> return $ Yield x (s, False, Nothing)
1205 _ -> return $ Yield a (s, True, Just x)
1206 Skip s -> return $ Skip (s, False, Nothing)
1207 Stop -> return $ Yield a (st, True, Nothing)
1209 step' _ (_, True, Nothing) = return Stop
1211 step' gst (st, True, Just prev) = do
1212 r <- step gst st
1213 case r of
1214 Yield x s -> return $ Yield prev (s, True, Just x)
1215 Skip s -> return $ Skip (s, True, Just prev)
1216 Stop -> return $ Yield prev (st, True, Nothing)
1218------------------------------------------------------------------------------
1219-- Deleting
1220------------------------------------------------------------------------------
1222{-# INLINE_NORMAL deleteBy #-}
1223deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
1224deleteBy eq x (Stream step state) = Stream step' (state, False)
1225 where
1226 {-# INLINE_LATE step' #-}
1227 step' gst (st, False) = do
1228 r <- step gst st
1229 case r of
1230 Yield y s -> return $
1231 if eq x y then Skip (s, True) else Yield y (s, False)
1232 Skip s -> return $ Skip (s, False)
1233 Stop -> return Stop
1235 step' gst (st, True) = do
1236 r <- step gst st
1237 case r of
1238 Yield y s -> return $ Yield y (s, True)
1239 Skip s -> return $ Skip (s, True)
1240 Stop -> return Stop
1242------------------------------------------------------------------------------
1243-- Transformation by Map and Filter
1244------------------------------------------------------------------------------
1246-- XXX Will this always fuse properly?
1247{-# INLINE_NORMAL mapMaybe #-}
1248mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
1249mapMaybe f = fmap fromJust . filter isJust . map f
1251{-# INLINE_NORMAL mapMaybeM #-}
1252mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
1253mapMaybeM f = fmap fromJust . filter isJust . mapM f
1255------------------------------------------------------------------------------
1256-- Zipping
1257------------------------------------------------------------------------------
1259{-# INLINE_NORMAL indexed #-}
1260indexed :: Monad m => Stream m a -> Stream m (Int, a)
1261indexed (Stream step state) = Stream step' (state, 0)
1262 where
1263 {-# INLINE_LATE step' #-}
1264 step' gst (st, i) = i `seq` do
1265 r <- step (adaptState gst) st
1266 case r of
1267 Yield x s -> return $ Yield (i, x) (s, i+1)
1268 Skip s -> return $ Skip (s, i)
1269 Stop -> return Stop
1271{-# INLINE_NORMAL indexedR #-}
1272indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
1273indexedR m (Stream step state) = Stream step' (state, m)
1274 where
1275 {-# INLINE_LATE step' #-}
1276 step' gst (st, i) = i `seq` do
1277 r <- step (adaptState gst) st
1278 case r of
1279 Yield x s -> let i' = i - 1
1280 in
1281 return $ Yield (i', x) (s, i')
1282 Skip s -> return $ Skip (s, i)
1283 Stop -> return Stop
1285{-# INLINE_NORMAL zipWithM #-}
1286zipWithM :: Monad m
1287 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
1288zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
1289 where
1290 {-# INLINE_LATE step #-}
1291 step gst (sa, sb, Nothing) = do
1292 r <- stepa (adaptState gst) sa
1293 return $
1294 case r of
1295 Yield x sa' -> Skip (sa', sb, Just x)
1296 Skip sa' -> Skip (sa', sb, Nothing)
1297 Stop -> Stop
1299 step gst (sa, sb, Just x) = do
1300 r <- stepb (adaptState gst) sb
1301 case r of
1302 Yield y sb' -> do
1303 z <- f x y
1304 return $ Yield z (sa, sb', Nothing)
1305 Skip sb' -> return $ Skip (sa, sb', Just x)
1306 Stop -> return Stop
1308{-# RULES "zipWithM xs xs"
1309 forall f xs. zipWithM f xs xs = mapM (\x -> f x x) xs #-}
1311{-# INLINE zipWith #-}
1312zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
1313zipWith f = zipWithM (\a b -> return (f a b))
1315------------------------------------------------------------------------------
1316-- Comparisions
1317------------------------------------------------------------------------------
1319{-# INLINE_NORMAL eqBy #-}
1320eqBy :: Monad m => (a -> b -> Bool) -> Stream m a -> Stream m b -> m Bool
1321eqBy eq (Stream step1 t1) (Stream step2 t2) = eq_loop0 SPEC t1 t2
1322 where
1323 eq_loop0 !_ s1 s2 = do
1324 r <- step1 defState s1
1325 case r of
1326 Yield x s1' -> eq_loop1 SPEC x s1' s2
1327 Skip s1' -> eq_loop0 SPEC s1' s2
1328 Stop -> eq_null s2
1330 eq_loop1 !_ x s1 s2 = do
1331 r <- step2 defState s2
1332 case r of
1333 Yield y s2'
1334 | eq x y -> eq_loop0 SPEC s1 s2'
1335 | otherwise -> return False
1336 Skip s2' -> eq_loop1 SPEC x s1 s2'
1337 Stop -> return False
1339 eq_null s2 = do
1340 r <- step2 defState s2
1341 case r of
1342 Yield _ _ -> return False
1343 Skip s2' -> eq_null s2'
1344 Stop -> return True
1346-- | Compare two streams lexicographically
1347{-# INLINE_NORMAL cmpBy #-}
1348cmpBy
1349 :: Monad m
1350 => (a -> b -> Ordering) -> Stream m a -> Stream m b -> m Ordering
1351cmpBy cmp (Stream step1 t1) (Stream step2 t2) = cmp_loop0 SPEC t1 t2
1352 where
1353 cmp_loop0 !_ s1 s2 = do
1354 r <- step1 defState s1
1355 case r of
1356 Yield x s1' -> cmp_loop1 SPEC x s1' s2
1357 Skip s1' -> cmp_loop0 SPEC s1' s2
1358 Stop -> cmp_null s2
1360 cmp_loop1 !_ x s1 s2 = do
1361 r <- step2 defState s2
1362 case r of
1363 Yield y s2' -> case x `cmp` y of
1364 EQ -> cmp_loop0 SPEC s1 s2'
1365 c -> return c
1366 Skip s2' -> cmp_loop1 SPEC x s1 s2'
1367 Stop -> return GT
1369 cmp_null s2 = do
1370 r <- step2 defState s2
1371 case r of
1372 Yield _ _ -> return LT
1373 Skip s2' -> cmp_null s2'
1374 Stop -> return EQ
1376------------------------------------------------------------------------------
1377-- Merging
1378------------------------------------------------------------------------------
1380{-# INLINE_NORMAL mergeByM #-}
1381mergeByM
1382 :: (Monad m)
1383 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
1384mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
1385 Stream step (Just ta, Just tb, Nothing, Nothing)
1386 where
1387 {-# INLINE_LATE step #-}
1389 -- one of the values is missing, and the corresponding stream is running
1390 step gst (Just sa, sb, Nothing, b) = do
1391 r <- stepa gst sa
1392 return $ case r of
1393 Yield a sa' -> Skip (Just sa', sb, Just a, b)
1394 Skip sa' -> Skip (Just sa', sb, Nothing, b)
1395 Stop -> Skip (Nothing, sb, Nothing, b)
1397 step gst (sa, Just sb, a, Nothing) = do
1398 r <- stepb gst sb
1399 return $ case r of
1400 Yield b sb' -> Skip (sa, Just sb', a, Just b)
1401 Skip sb' -> Skip (sa, Just sb', a, Nothing)
1402 Stop -> Skip (sa, Nothing, a, Nothing)
1404 -- both the values are available
1405 step _ (sa, sb, Just a, Just b) = do
1406 res <- cmp a b
1407 return $ case res of
1408 GT -> Yield b (sa, sb, Just a, Nothing)
1409 _ -> Yield a (sa, sb, Nothing, Just b)
1411 -- one of the values is missing, corresponding stream is done
1412 step _ (Nothing, sb, Nothing, Just b) =
1413 return $ Yield b (Nothing, sb, Nothing, Nothing)
1415 step _ (sa, Nothing, Just a, Nothing) =
1416 return $ Yield a (sa, Nothing, Nothing, Nothing)
1418 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
1420{-# INLINE mergeBy #-}
1421mergeBy
1422 :: (Monad m)
1423 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
1424mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
1426------------------------------------------------------------------------------
1427-- Transformation comprehensions
1428------------------------------------------------------------------------------
1430{-# INLINE_NORMAL the #-}
1431the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
1432the (Stream step state) = go state
1433 where
1434 go st = do
1435 r <- step defState st
1436 case r of
1437 Yield x s -> go' x s
1438 Skip s -> go s
1439 Stop -> return Nothing
1440 go' n st = do
1441 r <- step defState st
1442 case r of
1443 Yield x s | x == n -> go' n s
1444 | otherwise -> return Nothing
1445 Skip s -> go' n s
1446 Stop -> return (Just n)