git.iko.soy/streamly
on ede29df help sign in
streamly/src/Streamly/Streams/StreamD.hs
1449 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.StreamD.Type
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 (UnStream 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.foldStreamShared gst yieldk single stop m1
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 (UnStream 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 -- XXX using the pattern synonym Stream causes a major performance issue
779 -- here even if the synonym does not include a adaptState call. Need to
780 -- find out why. Is that something to be fixed in GHC?
781 step' _ (Right (UnStream inner_step inner_st, st)) = do
782 r <- inner_step defState inner_st
783 case r of
784 Yield b inner_s ->
785 return $ Yield b (Right (Stream inner_step inner_s, st))
786 Skip inner_s ->
787 return $ Skip (Right (Stream inner_step inner_s, st))
788 Stop -> return $ Skip (Left st)
790{-# INLINE concatMap #-}
791concatMap :: Monad m => (a -> Stream m b) -> Stream m a -> Stream m b
792concatMap f = concatMapM (return . f)
794------------------------------------------------------------------------------
795-- Substreams
796------------------------------------------------------------------------------
798{-# INLINE_NORMAL isPrefixOf #-}
799isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
800isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
801 where
802 go (sa, sb, Nothing) = do
803 r <- stepa defState sa
804 case r of
805 Yield x sa' -> go (sa', sb, Just x)
806 Skip sa' -> go (sa', sb, Nothing)
807 Stop -> return True
809 go (sa, sb, Just x) = do
810 r <- stepb defState sb
811 case r of
812 Yield y sb' ->
813 if x == y
814 then go (sa, sb', Nothing)
815 else return False
816 Skip sb' -> go (sa, sb', Just x)
817 Stop -> return False
819{-# INLINE_NORMAL isSubsequenceOf #-}
820isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
821isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
822 where
823 go (sa, sb, Nothing) = do
824 r <- stepa defState sa
825 case r of
826 Yield x sa' -> go (sa', sb, Just x)
827 Skip sa' -> go (sa', sb, Nothing)
828 Stop -> return True
830 go (sa, sb, Just x) = do
831 r <- stepb defState sb
832 case r of
833 Yield y sb' ->
834 if x == y
835 then go (sa, sb', Nothing)
836 else go (sa, sb', Just x)
837 Skip sb' -> go (sa, sb', Just x)
838 Stop -> return False
840{-# INLINE_NORMAL stripPrefix #-}
841stripPrefix
842 :: (Eq a, Monad m)
843 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
844stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
845 where
846 go (sa, sb, Nothing) = do
847 r <- stepa defState sa
848 case r of
849 Yield x sa' -> go (sa', sb, Just x)
850 Skip sa' -> go (sa', sb, Nothing)
851 Stop -> return $ Just (Stream stepb sb)
853 go (sa, sb, Just x) = do
854 r <- stepb defState sb
855 case r of
856 Yield y sb' ->
857 if x == y
858 then go (sa, sb', Nothing)
859 else return Nothing
860 Skip sb' -> go (sa, sb', Just x)
861 Stop -> return Nothing
863------------------------------------------------------------------------------
864-- Map and Fold
865------------------------------------------------------------------------------
867-- | Execute a monadic action for each element of the 'Stream'
868{-# INLINE_NORMAL mapM_ #-}
869mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
870mapM_ m = runStream . mapM m
872------------------------------------------------------------------------------
873-- Converting folds
874------------------------------------------------------------------------------
876{-# INLINE toList #-}
877toList :: Monad m => Stream m a -> m [a]
878toList = foldr (:) []
880-- Convert a direct stream to and from CPS encoded stream
881{-# INLINE_LATE toStreamK #-}
882toStreamK :: Monad m => Stream m a -> K.Stream m a
883toStreamK (Stream step state) = go state
884 where
885 go st = K.mkStream $ \gst yld sng stp -> do
886 r <- step gst st
887 case r of
888 Yield x s -> yld x (go s)
889 Skip s -> K.foldStreamShared gst yld sng stp $ go s
890 Stop -> stp
892#ifndef DISABLE_FUSION
893{-# RULES "fromStreamK/toStreamK fusion"
894 forall s. toStreamK (fromStreamK s) = s #-}
895{-# RULES "toStreamK/fromStreamK fusion"
896 forall s. fromStreamK (toStreamK s) = s #-}
897#endif
899{-# INLINE fromStreamD #-}
900fromStreamD :: (K.IsStream t, Monad m) => Stream m a -> t m a
901fromStreamD = K.fromStream . toStreamK
903------------------------------------------------------------------------------
904-- Transformation by Folding (Scans)
905------------------------------------------------------------------------------
907-- XXX Is a prescan useful, discarding the last step does not sound useful? I
908-- am not sure about the utility of this function, so this is implemented but
909-- not exposed. We can expose it if someone provides good reasons why this is
910-- useful.
911--
912-- XXX We have to execute the stream one step ahead to know that we are at the
913-- last step. The vector implementation of prescan executes the last fold step
914-- but does not yield the result. This means we have executed the effect but
915-- discarded value. This does not sound right. In this implementation we are
916-- not executing the last fold step.
917{-# INLINE_NORMAL prescanlM' #-}
918prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
919prescanlM' f mz (Stream step state) = Stream step' (state, mz)
920 where
921 {-# INLINE_LATE step' #-}
922 step' gst (st, prev) = do
923 r <- step (adaptState gst) st
924 case r of
925 Yield x s -> do
926 acc <- prev
927 return $ Yield acc (s, f acc x)
928 Skip s -> return $ Skip (s, prev)
929 Stop -> return Stop
931{-# INLINE prescanl' #-}
932prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
933prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
935-- XXX if we make the initial value of the accumulator monadic then should we
936-- execute it even if the stream is empty? In that case we would have generated
937-- the effect but discarded the value, but that is what a fold does when the
938-- stream is empty. Whatever we decide, need to reconcile this with prescan.
939-- If we execute the initial value here without even using it then it is ok to
940-- execute the last step there as well without using the value.
941-- Looking at the duality with right fold, in case of right fold we always
942-- perform the action when the construction terminates, so in case of left fold
943-- we should perform it only when the reduction starts.
944{-# INLINE_NORMAL postscanlM' #-}
945postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
946postscanlM' fstep begin (Stream step state) =
947 begin `seq` Stream step' (state, begin)
948 where
949 {-# INLINE_LATE step' #-}
950 step' gst (st, acc) = acc `seq` do
951 r <- step (adaptState gst) st
952 case r of
953 Yield x s -> do
954 y <- fstep acc x
955 y `seq` return (Yield y (s, y))
956 Skip s -> return $ Skip (s, acc)
957 Stop -> return Stop
959{-# INLINE_NORMAL postscanl' #-}
960postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
961postscanl' f = postscanlM' (\a b -> return (f a b))
963{-# INLINE_NORMAL postscanlM #-}
964postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
965postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
966 where
967 {-# INLINE_LATE step' #-}
968 step' gst (st, acc) = do
969 r <- step (adaptState gst) st
970 case r of
971 Yield x s -> do
972 y <- fstep acc x
973 return (Yield y (s, y))
974 Skip s -> return $ Skip (s, acc)
975 Stop -> return Stop
977{-# INLINE_NORMAL postscanl #-}
978postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
979postscanl f = postscanlM (\a b -> return (f a b))
981{-# INLINE_NORMAL scanlM' #-}
982scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
983scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
985{-# INLINE scanl' #-}
986scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
987scanl' f = scanlM' (\a b -> return (f a b))
989{-# INLINE_NORMAL scanlM #-}
990scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
991scanlM fstep begin s = begin `cons` postscanlM fstep begin s
993{-# INLINE scanl #-}
994scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
995scanl f = scanlM (\a b -> return (f a b))
997{-# INLINE_NORMAL scanl1M #-}
998scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
999scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
1000 where
1001 {-# INLINE_LATE step' #-}
1002 step' gst (st, Nothing) = do
1003 r <- step gst st
1004 case r of
1005 Yield x s -> return $ Yield x (s, Just x)
1006 Skip s -> return $ Skip (s, Nothing)
1007 Stop -> return Stop
1009 step' gst (st, Just acc) = do
1010 r <- step gst st
1011 case r of
1012 Yield y s -> do
1013 z <- fstep acc y
1014 return $ Yield z (s, Just z)
1015 Skip s -> return $ Skip (s, Just acc)
1016 Stop -> return Stop
1018{-# INLINE scanl1 #-}
1019scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
1020scanl1 f = scanl1M (\x y -> return (f x y))
1022{-# INLINE_NORMAL scanl1M' #-}
1023scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
1024scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
1025 where
1026 {-# INLINE_LATE step' #-}
1027 step' gst (st, Nothing) = do
1028 r <- step gst st
1029 case r of
1030 Yield x s -> x `seq` return $ Yield x (s, Just x)
1031 Skip s -> return $ Skip (s, Nothing)
1032 Stop -> return Stop
1034 step' gst (st, Just acc) = acc `seq` do
1035 r <- step gst st
1036 case r of
1037 Yield y s -> do
1038 z <- fstep acc y
1039 z `seq` return $ Yield z (s, Just z)
1040 Skip s -> return $ Skip (s, Just acc)
1041 Stop -> return Stop
1043{-# INLINE scanl1' #-}
1044scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
1045scanl1' f = scanl1M' (\x y -> return (f x y))
1047-------------------------------------------------------------------------------
1048-- Filtering
1049-------------------------------------------------------------------------------
1051{-# INLINE_NORMAL take #-}
1052take :: Monad m => Int -> Stream m a -> Stream m a
1053take n (Stream step state) = n `seq` Stream step' (state, 0)
1054 where
1055 {-# INLINE_LATE step' #-}
1056 step' gst (st, i) | i < n = do
1057 r <- step gst st
1058 return $ case r of
1059 Yield x s -> Yield x (s, i + 1)
1060 Skip s -> Skip (s, i)
1061 Stop -> Stop
1062 step' _ (_, _) = return Stop
1064{-# INLINE_NORMAL takeWhileM #-}
1065takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
1066takeWhileM f (Stream step state) = Stream step' state
1067 where
1068 {-# INLINE_LATE step' #-}
1069 step' gst st = do
1070 r <- step gst st
1071 case r of
1072 Yield x s -> do
1073 b <- f x
1074 return $ if b then Yield x s else Stop
1075 Skip s -> return $ Skip s
1076 Stop -> return Stop
1078{-# INLINE takeWhile #-}
1079takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
1080takeWhile f = takeWhileM (return . f)
1082{-# INLINE_NORMAL drop #-}
1083drop :: Monad m => Int -> Stream m a -> Stream m a
1084drop n (Stream step state) = Stream step' (state, Just n)
1085 where
1086 {-# INLINE_LATE step' #-}
1087 step' gst (st, Just i)
1088 | i > 0 = do
1089 r <- step gst st
1090 return $
1091 case r of
1092 Yield _ s -> Skip (s, Just (i - 1))
1093 Skip s -> Skip (s, Just i)
1094 Stop -> Stop
1095 | otherwise = return $ Skip (st, Nothing)
1097 step' gst (st, Nothing) = do
1098 r <- step gst st
1099 return $
1100 case r of
1101 Yield x s -> Yield x (s, Nothing)
1102 Skip s -> Skip (s, Nothing)
1103 Stop -> Stop
1105data DropWhileState s a
1106 = DropWhileDrop s
1107 | DropWhileYield a s
1108 | DropWhileNext s
1110{-# INLINE_NORMAL dropWhileM #-}
1111dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
1112dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
1113 where
1114 {-# INLINE_LATE step' #-}
1115 step' gst (DropWhileDrop st) = do
1116 r <- step gst st
1117 case r of
1118 Yield x s -> do
1119 b <- f x
1120 if b
1121 then return $ Skip (DropWhileDrop s)
1122 else return $ Skip (DropWhileYield x s)
1123 Skip s -> return $ Skip (DropWhileDrop s)
1124 Stop -> return Stop
1126 step' gst (DropWhileNext st) = do
1127 r <- step gst st
1128 case r of
1129 Yield x s -> return $ Skip (DropWhileYield x s)
1130 Skip s -> return $ Skip (DropWhileNext s)
1131 Stop -> return Stop
1133 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
1135{-# INLINE dropWhile #-}
1136dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
1137dropWhile f = dropWhileM (return . f)
1139{-# INLINE_NORMAL filterM #-}
1140filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
1141filterM f (Stream step state) = Stream step' state
1142 where
1143 {-# INLINE_LATE step' #-}
1144 step' gst st = do
1145 r <- step gst st
1146 case r of
1147 Yield x s -> do
1148 b <- f x
1149 return $ if b
1150 then Yield x s
1151 else Skip s
1152 Skip s -> return $ Skip s
1153 Stop -> return Stop
1155{-# INLINE filter #-}
1156filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
1157filter f = filterM (return . f)
1159{-# INLINE_NORMAL uniq #-}
1160uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
1161uniq (Stream step state) = Stream step' (Nothing, state)
1162 where
1163 {-# INLINE_LATE step' #-}
1164 step' gst (Nothing, st) = do
1165 r <- step gst st
1166 case r of
1167 Yield x s -> return $ Yield x (Just x, s)
1168 Skip s -> return $ Skip (Nothing, s)
1169 Stop -> return Stop
1170 step' gst (Just x, st) = do
1171 r <- step gst st
1172 case r of
1173 Yield y s | x == y -> return $ Skip (Just x, s)
1174 | otherwise -> return $ Yield x (Just y, s)
1175 Skip s -> return $ Skip (Just x, s)
1176 Stop -> return Stop
1178------------------------------------------------------------------------------
1179-- Transformation by Mapping
1180------------------------------------------------------------------------------
1182{-# INLINE_NORMAL sequence #-}
1183sequence :: Monad m => Stream m (m a) -> Stream m a
1184sequence (Stream step state) = Stream step' state
1185 where
1186 {-# INLINE_LATE step' #-}
1187 step' gst st = do
1188 r <- step (adaptState gst) st
1189 case r of
1190 Yield x s -> x >>= \a -> return (Yield a s)
1191 Skip s -> return $ Skip s
1192 Stop -> return Stop
1194------------------------------------------------------------------------------
1195-- Inserting
1196------------------------------------------------------------------------------
1198{-# INLINE_NORMAL insertBy #-}
1199insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
1200insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
1201 where
1202 {-# INLINE_LATE step' #-}
1203 step' gst (st, False, _) = do
1204 r <- step gst st
1205 case r of
1206 Yield x s -> case cmp a x of
1207 GT -> return $ Yield x (s, False, Nothing)
1208 _ -> return $ Yield a (s, True, Just x)
1209 Skip s -> return $ Skip (s, False, Nothing)
1210 Stop -> return $ Yield a (st, True, Nothing)
1212 step' _ (_, True, Nothing) = return Stop
1214 step' gst (st, True, Just prev) = do
1215 r <- step gst st
1216 case r of
1217 Yield x s -> return $ Yield prev (s, True, Just x)
1218 Skip s -> return $ Skip (s, True, Just prev)
1219 Stop -> return $ Yield prev (st, True, Nothing)
1221------------------------------------------------------------------------------
1222-- Deleting
1223------------------------------------------------------------------------------
1225{-# INLINE_NORMAL deleteBy #-}
1226deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
1227deleteBy eq x (Stream step state) = Stream step' (state, False)
1228 where
1229 {-# INLINE_LATE step' #-}
1230 step' gst (st, False) = do
1231 r <- step gst st
1232 case r of
1233 Yield y s -> return $
1234 if eq x y then Skip (s, True) else Yield y (s, False)
1235 Skip s -> return $ Skip (s, False)
1236 Stop -> return Stop
1238 step' gst (st, True) = do
1239 r <- step gst st
1240 case r of
1241 Yield y s -> return $ Yield y (s, True)
1242 Skip s -> return $ Skip (s, True)
1243 Stop -> return Stop
1245------------------------------------------------------------------------------
1246-- Transformation by Map and Filter
1247------------------------------------------------------------------------------
1249-- XXX Will this always fuse properly?
1250{-# INLINE_NORMAL mapMaybe #-}
1251mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
1252mapMaybe f = fmap fromJust . filter isJust . map f
1254{-# INLINE_NORMAL mapMaybeM #-}
1255mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
1256mapMaybeM f = fmap fromJust . filter isJust . mapM f
1258------------------------------------------------------------------------------
1259-- Zipping
1260------------------------------------------------------------------------------
1262{-# INLINE_NORMAL indexed #-}
1263indexed :: Monad m => Stream m a -> Stream m (Int, a)
1264indexed (Stream step state) = Stream step' (state, 0)
1265 where
1266 {-# INLINE_LATE step' #-}
1267 step' gst (st, i) = i `seq` do
1268 r <- step (adaptState gst) st
1269 case r of
1270 Yield x s -> return $ Yield (i, x) (s, i+1)
1271 Skip s -> return $ Skip (s, i)
1272 Stop -> return Stop
1274{-# INLINE_NORMAL indexedR #-}
1275indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
1276indexedR m (Stream step state) = Stream step' (state, m)
1277 where
1278 {-# INLINE_LATE step' #-}
1279 step' gst (st, i) = i `seq` do
1280 r <- step (adaptState gst) st
1281 case r of
1282 Yield x s -> let i' = i - 1
1283 in
1284 return $ Yield (i', x) (s, i')
1285 Skip s -> return $ Skip (s, i)
1286 Stop -> return Stop
1288{-# INLINE_NORMAL zipWithM #-}
1289zipWithM :: Monad m
1290 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
1291zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
1292 where
1293 {-# INLINE_LATE step #-}
1294 step gst (sa, sb, Nothing) = do
1295 r <- stepa (adaptState gst) sa
1296 return $
1297 case r of
1298 Yield x sa' -> Skip (sa', sb, Just x)
1299 Skip sa' -> Skip (sa', sb, Nothing)
1300 Stop -> Stop
1302 step gst (sa, sb, Just x) = do
1303 r <- stepb (adaptState gst) sb
1304 case r of
1305 Yield y sb' -> do
1306 z <- f x y
1307 return $ Yield z (sa, sb', Nothing)
1308 Skip sb' -> return $ Skip (sa, sb', Just x)
1309 Stop -> return Stop
1311{-# RULES "zipWithM xs xs"
1312 forall f xs. zipWithM f xs xs = mapM (\x -> f x x) xs #-}
1314{-# INLINE zipWith #-}
1315zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
1316zipWith f = zipWithM (\a b -> return (f a b))
1318------------------------------------------------------------------------------
1319-- Comparisions
1320------------------------------------------------------------------------------
1322{-# INLINE_NORMAL eqBy #-}
1323eqBy :: Monad m => (a -> b -> Bool) -> Stream m a -> Stream m b -> m Bool
1324eqBy eq (Stream step1 t1) (Stream step2 t2) = eq_loop0 SPEC t1 t2
1325 where
1326 eq_loop0 !_ s1 s2 = do
1327 r <- step1 defState s1
1328 case r of
1329 Yield x s1' -> eq_loop1 SPEC x s1' s2
1330 Skip s1' -> eq_loop0 SPEC s1' s2
1331 Stop -> eq_null s2
1333 eq_loop1 !_ x s1 s2 = do
1334 r <- step2 defState s2
1335 case r of
1336 Yield y s2'
1337 | eq x y -> eq_loop0 SPEC s1 s2'
1338 | otherwise -> return False
1339 Skip s2' -> eq_loop1 SPEC x s1 s2'
1340 Stop -> return False
1342 eq_null s2 = do
1343 r <- step2 defState s2
1344 case r of
1345 Yield _ _ -> return False
1346 Skip s2' -> eq_null s2'
1347 Stop -> return True
1349-- | Compare two streams lexicographically
1350{-# INLINE_NORMAL cmpBy #-}
1351cmpBy
1352 :: Monad m
1353 => (a -> b -> Ordering) -> Stream m a -> Stream m b -> m Ordering
1354cmpBy cmp (Stream step1 t1) (Stream step2 t2) = cmp_loop0 SPEC t1 t2
1355 where
1356 cmp_loop0 !_ s1 s2 = do
1357 r <- step1 defState s1
1358 case r of
1359 Yield x s1' -> cmp_loop1 SPEC x s1' s2
1360 Skip s1' -> cmp_loop0 SPEC s1' s2
1361 Stop -> cmp_null s2
1363 cmp_loop1 !_ x s1 s2 = do
1364 r <- step2 defState s2
1365 case r of
1366 Yield y s2' -> case x `cmp` y of
1367 EQ -> cmp_loop0 SPEC s1 s2'
1368 c -> return c
1369 Skip s2' -> cmp_loop1 SPEC x s1 s2'
1370 Stop -> return GT
1372 cmp_null s2 = do
1373 r <- step2 defState s2
1374 case r of
1375 Yield _ _ -> return LT
1376 Skip s2' -> cmp_null s2'
1377 Stop -> return EQ
1379------------------------------------------------------------------------------
1380-- Merging
1381------------------------------------------------------------------------------
1383{-# INLINE_NORMAL mergeByM #-}
1384mergeByM
1385 :: (Monad m)
1386 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
1387mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
1388 Stream step (Just ta, Just tb, Nothing, Nothing)
1389 where
1390 {-# INLINE_LATE step #-}
1392 -- one of the values is missing, and the corresponding stream is running
1393 step gst (Just sa, sb, Nothing, b) = do
1394 r <- stepa gst sa
1395 return $ case r of
1396 Yield a sa' -> Skip (Just sa', sb, Just a, b)
1397 Skip sa' -> Skip (Just sa', sb, Nothing, b)
1398 Stop -> Skip (Nothing, sb, Nothing, b)
1400 step gst (sa, Just sb, a, Nothing) = do
1401 r <- stepb gst sb
1402 return $ case r of
1403 Yield b sb' -> Skip (sa, Just sb', a, Just b)
1404 Skip sb' -> Skip (sa, Just sb', a, Nothing)
1405 Stop -> Skip (sa, Nothing, a, Nothing)
1407 -- both the values are available
1408 step _ (sa, sb, Just a, Just b) = do
1409 res <- cmp a b
1410 return $ case res of
1411 GT -> Yield b (sa, sb, Just a, Nothing)
1412 _ -> Yield a (sa, sb, Nothing, Just b)
1414 -- one of the values is missing, corresponding stream is done
1415 step _ (Nothing, sb, Nothing, Just b) =
1416 return $ Yield b (Nothing, sb, Nothing, Nothing)
1418 step _ (sa, Nothing, Just a, Nothing) =
1419 return $ Yield a (sa, Nothing, Nothing, Nothing)
1421 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
1423{-# INLINE mergeBy #-}
1424mergeBy
1425 :: (Monad m)
1426 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
1427mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
1429------------------------------------------------------------------------------
1430-- Transformation comprehensions
1431------------------------------------------------------------------------------
1433{-# INLINE_NORMAL the #-}
1434the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
1435the (Stream step state) = go state
1436 where
1437 go st = do
1438 r <- step defState st
1439 case r of
1440 Yield x s -> go' x s
1441 Skip s -> go s
1442 Stop -> return Nothing
1443 go' n st = do
1444 r <- step defState st
1445 case r of
1446 Yield x s | x == n -> go' n s
1447 | otherwise -> return Nothing
1448 Skip s -> go' n s
1449 Stop -> return (Just n)