git.iko.soy/streamly
on bcd787e help sign in
streamly/src/Streamly/Streams/StreamD.hs
1464 lines 47 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 , enumerateFromStepIntegral
66 , enumerateFromIntegral
67 , enumerateFromThenIntegral
68 , enumerateFromToIntegral
69 , enumerateFromThenToIntegral
71 , enumerateFromStepNum
72 , numFrom
73 , numFromThen
74 , enumerateFromToFractional
75 , enumerateFromThenToFractional
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 enumerateFromStepIntegral #-}
292enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
293enumerateFromStepIntegral 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 enumerateFromToIntegral #-}
302enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
303enumerateFromToIntegral from to =
304 takeWhile (<= to) $ enumerateFromStepIntegral from 1
306{-# INLINE enumerateFromIntegral #-}
307enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
308enumerateFromIntegral from = enumerateFromToIntegral from maxBound
310data EnumState a = EnumInit | EnumYield a a a | EnumStop
312{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
313enumerateFromThenToIntegralUp
314 :: (Monad m, Integral a)
315 => a -> a -> a -> Stream m a
316enumerateFromThenToIntegralUp from next to = Stream step EnumInit
317 where
318 {-# INLINE_LATE step #-}
319 step _ EnumInit =
320 return $
321 if to < next
322 then if to < from
323 then Stop
324 else Yield from EnumStop
325 else -- from <= next <= to
326 let stride = next - from
327 in Skip $ EnumYield from stride (to - stride)
329 step _ (EnumYield x stride toMinus) =
330 return $
331 if x > toMinus
332 then Yield x EnumStop
333 else Yield x $ EnumYield (x + stride) stride toMinus
335 step _ EnumStop = return Stop
337{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
338enumerateFromThenToIntegralDn
339 :: (Monad m, Integral a)
340 => a -> a -> a -> Stream m a
341enumerateFromThenToIntegralDn from next to = Stream step EnumInit
342 where
343 {-# INLINE_LATE step #-}
344 step _ EnumInit =
345 return $ if to > next
346 then if to > from
347 then Stop
348 else Yield from EnumStop
349 else -- from >= next >= to
350 let stride = next - from
351 in Skip $ EnumYield from stride (to - stride)
353 step _ (EnumYield x stride toMinus) =
354 return $
355 if x < toMinus
356 then Yield x EnumStop
357 else Yield x $ EnumYield (x + stride) stride toMinus
359 step _ EnumStop = return Stop
361{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
362enumerateFromThenToIntegral
363 :: (Monad m, Integral a)
364 => a -> a -> a -> Stream m a
365enumerateFromThenToIntegral from next to
366 | next >= from = enumerateFromThenToIntegralUp from next to
367 | otherwise = enumerateFromThenToIntegralDn from next to
369{-# INLINE_NORMAL enumerateFromThenIntegral #-}
370enumerateFromThenIntegral
371 :: (Monad m, Integral a, Bounded a)
372 => a -> a -> Stream m a
373enumerateFromThenIntegral from next =
374 if next > from
375 then enumerateFromThenToIntegralUp from next maxBound
376 else enumerateFromThenToIntegralDn from next minBound
378-- For floating point numbers if the increment is less than the precision then
379-- it just gets lost. Therefore we cannot always increment it correctly by just
380-- repeated addition.
381-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
382-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
384-- Instead we accumulate the increment counter and compute the increment
385-- everytime before adding it to the starting number.
386--
387-- This works for Integrals as well as floating point numbers, but
388-- enumerateFromStepIntegral is faster for integrals.
389{-# INLINE_NORMAL enumerateFromStepNum #-}
390enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
391enumerateFromStepNum from stride = Stream step 0
392 where
393 {-# INLINE_LATE step #-}
394 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
396{-# INLINE_NORMAL numFrom #-}
397numFrom :: (Monad m, Num a) => a -> Stream m a
398numFrom from = enumerateFromStepNum from 1
400{-# INLINE_NORMAL numFromThen #-}
401numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
402numFromThen from next = enumerateFromStepNum from (next - from)
404-- We cannot write a general function for Num. The only way to write code
405-- portable between the two is to use a 'Real' constraint and convert between
406-- Fractional and Integral using fromRational which is horribly slow.
407{-# INLINE_NORMAL enumerateFromToFractional #-}
408enumerateFromToFractional
409 :: (Monad m, Fractional a, Ord a)
410 => a -> a -> Stream m a
411enumerateFromToFractional from to =
412 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
414{-# INLINE_NORMAL enumerateFromThenToFractional #-}
415enumerateFromThenToFractional
416 :: (Monad m, Fractional a, Ord a)
417 => a -> a -> a -> Stream m a
418enumerateFromThenToFractional from next to =
419 takeWhile predicate $ numFromThen from next
420 where
421 mid = (next - from) / 2
422 predicate | next >= from = (<= to + mid)
423 | otherwise = (>= to + mid)
425-------------------------------------------------------------------------------
426-- Generation by Conversion
427-------------------------------------------------------------------------------
429-- | Create a singleton 'Stream' from a pure value.
430{-# INLINE_NORMAL yield #-}
431yield :: Monad m => a -> Stream m a
432yield x = Stream (\_ s -> return $ step undefined s) True
433 where
434 {-# INLINE_LATE step #-}
435 step _ True = Yield x False
436 step _ False = Stop
438-- | Create a singleton 'Stream' from a monadic action.
439{-# INLINE_NORMAL yieldM #-}
440yieldM :: Monad m => m a -> Stream m a
441yieldM m = Stream step True
442 where
443 {-# INLINE_LATE step #-}
444 step _ True = m >>= \x -> return $ Yield x False
445 step _ False = return Stop
447{-# INLINE_NORMAL fromIndicesM #-}
448fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
449fromIndicesM gen = Stream step 0
450 where
451 {-# INLINE_LATE step #-}
452 step _ i = do
453 x <- gen i
454 return $ Yield x (i + 1)
456{-# INLINE fromIndices #-}
457fromIndices :: Monad m => (Int -> a) -> Stream m a
458fromIndices gen = fromIndicesM (return . gen)
460{-# INLINE_NORMAL generateM #-}
461generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
462generateM n gen = n `seq` Stream step 0
463 where
464 {-# INLINE_LATE step #-}
465 step _ i | i < n = do
466 x <- gen i
467 return $ Yield x (i + 1)
468 | otherwise = return Stop
470{-# INLINE generate #-}
471generate :: Monad m => Int -> (Int -> a) -> Stream m a
472generate n gen = generateM n (return . gen)
474-- XXX we need the MonadAsync constraint because of a rewrite rule.
475-- | Convert a list of monadic actions to a 'Stream'
476{-# INLINE_LATE fromListM #-}
477fromListM :: MonadAsync m => [m a] -> Stream m a
478fromListM = Stream step
479 where
480 {-# INLINE_LATE step #-}
481 step _ (m:ms) = m >>= \x -> return $ Yield x ms
482 step _ [] = return Stop
484-- | Convert a list of pure values to a 'Stream'
485{-# INLINE_LATE fromList #-}
486fromList :: Monad m => [a] -> Stream m a
487fromList = Stream step
488 where
489 {-# INLINE_LATE step #-}
490 step _ (x:xs) = return $ Yield x xs
491 step _ [] = return Stop
493{-# INLINE_LATE fromStreamK #-}
494fromStreamK :: Monad m => K.Stream m a -> Stream m a
495fromStreamK = Stream step
496 where
497 step gst m1 =
498 let stop = return Stop
499 single a = return $ Yield a K.nil
500 yieldk a r = return $ Yield a r
501 in K.foldStreamShared gst yieldk single stop m1
503{-# INLINE toStreamD #-}
504toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
505toStreamD = fromStreamK . K.toStream
507------------------------------------------------------------------------------
508-- Elimination by Folds
509------------------------------------------------------------------------------
511{-# INLINE_NORMAL foldrM #-}
512foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b
513foldrM f z (Stream step state) = go SPEC state
514 where
515 go !_ st = do
516 r <- step defState st
517 case r of
518 Yield x s -> go SPEC s >>= f x
519 Skip s -> go SPEC s
520 Stop -> return z
522{-# INLINE_NORMAL foldr #-}
523foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
524foldr f = foldrM (\a b -> return (f a b))
526{-# INLINE_NORMAL foldr1 #-}
527foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
528foldr1 f m = do
529 r <- uncons m
530 case r of
531 Nothing -> return Nothing
532 Just (h, t) -> fmap Just (foldr f h t)
534{-# INLINE_NORMAL foldlM' #-}
535foldlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> m b
536foldlM' fstep begin (Stream step state) = go SPEC begin state
537 where
538 go !_ acc st = acc `seq` do
539 r <- step defState st
540 case r of
541 Yield x s -> do
542 acc' <- fstep acc x
543 go SPEC acc' s
544 Skip s -> go SPEC acc s
545 Stop -> return acc
547{-# INLINE foldl' #-}
548foldl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
549foldl' fstep = foldlM' (\b a -> return (fstep b a))
551------------------------------------------------------------------------------
552-- Specialized Folds
553------------------------------------------------------------------------------
555-- | Run a streaming composition, discard the results.
556{-# INLINE_LATE runStream #-}
557runStream :: Monad m => Stream m a -> m ()
558runStream (Stream step state) = go SPEC state
559 where
560 go !_ st = do
561 r <- step defState st
562 case r of
563 Yield _ s -> go SPEC s
564 Skip s -> go SPEC s
565 Stop -> return ()
567{-# INLINE_NORMAL null #-}
568null :: Monad m => Stream m a -> m Bool
569null (Stream step state) = go state
570 where
571 go st = do
572 r <- step defState st
573 case r of
574 Yield _ _ -> return False
575 Skip s -> go s
576 Stop -> return True
578-- XXX SPEC?
579{-# INLINE_NORMAL head #-}
580head :: Monad m => Stream m a -> m (Maybe a)
581head (Stream step state) = go state
582 where
583 go st = do
584 r <- step defState st
585 case r of
586 Yield x _ -> return (Just x)
587 Skip s -> go s
588 Stop -> return Nothing
590-- Does not fuse, has the same performance as the StreamK version.
591{-# INLINE_NORMAL tail #-}
592tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
593tail (UnStream step state) = go state
594 where
595 go st = do
596 r <- step defState st
597 case r of
598 Yield _ s -> return (Just $ Stream step s)
599 Skip s -> go s
600 Stop -> return Nothing
602-- XXX will it fuse? need custom impl?
603{-# INLINE_NORMAL last #-}
604last :: Monad m => Stream m a -> m (Maybe a)
605last = foldl' (\_ y -> Just y) Nothing
607{-# INLINE_NORMAL elem #-}
608elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
609elem e (Stream step state) = go state
610 where
611 go st = do
612 r <- step defState st
613 case r of
614 Yield x s
615 | x == e -> return True
616 | otherwise -> go s
617 Skip s -> go s
618 Stop -> return False
620{-# INLINE_NORMAL notElem #-}
621notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
622notElem e s = fmap not (elem e s)
624{-# INLINE_NORMAL all #-}
625all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
626all p (Stream step state) = go state
627 where
628 go st = do
629 r <- step defState st
630 case r of
631 Yield x s
632 | p x -> go s
633 | otherwise -> return False
634 Skip s -> go s
635 Stop -> return True
637{-# INLINE_NORMAL any #-}
638any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
639any p (Stream step state) = go state
640 where
641 go st = do
642 r <- step defState st
643 case r of
644 Yield x s
645 | p x -> return True
646 | otherwise -> go s
647 Skip s -> go s
648 Stop -> return False
650{-# INLINE_NORMAL maximum #-}
651maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
652maximum (Stream step state) = go Nothing state
653 where
654 go Nothing st = do
655 r <- step defState st
656 case r of
657 Yield x s -> go (Just x) s
658 Skip s -> go Nothing s
659 Stop -> return Nothing
660 go (Just acc) st = do
661 r <- step defState st
662 case r of
663 Yield x s
664 | acc <= x -> go (Just x) s
665 | otherwise -> go (Just acc) s
666 Skip s -> go (Just acc) s
667 Stop -> return (Just acc)
669{-# INLINE_NORMAL maximumBy #-}
670maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
671maximumBy cmp (Stream step state) = go Nothing state
672 where
673 go Nothing st = do
674 r <- step defState st
675 case r of
676 Yield x s -> go (Just x) s
677 Skip s -> go Nothing s
678 Stop -> return Nothing
679 go (Just acc) st = do
680 r <- step defState st
681 case r of
682 Yield x s -> case cmp acc x of
683 GT -> go (Just acc) s
684 _ -> go (Just x) s
685 Skip s -> go (Just acc) s
686 Stop -> return (Just acc)
688{-# INLINE_NORMAL minimum #-}
689minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
690minimum (Stream step state) = go Nothing state
691 where
692 go Nothing st = do
693 r <- step defState st
694 case r of
695 Yield x s -> go (Just x) s
696 Skip s -> go Nothing s
697 Stop -> return Nothing
698 go (Just acc) st = do
699 r <- step defState st
700 case r of
701 Yield x s
702 | acc <= x -> go (Just acc) s
703 | otherwise -> go (Just x) s
704 Skip s -> go (Just acc) s
705 Stop -> return (Just acc)
707{-# INLINE_NORMAL minimumBy #-}
708minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
709minimumBy cmp (Stream step state) = go Nothing state
710 where
711 go Nothing st = do
712 r <- step defState st
713 case r of
714 Yield x s -> go (Just x) s
715 Skip s -> go Nothing s
716 Stop -> return Nothing
717 go (Just acc) st = do
718 r <- step defState st
719 case r of
720 Yield x s -> case cmp acc x of
721 GT -> go (Just x) s
722 _ -> go (Just acc) s
723 Skip s -> go (Just acc) s
724 Stop -> return (Just acc)
726{-# INLINE_NORMAL (!!) #-}
727(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
728(Stream step state) !! i = go i state
729 where
730 go n st = do
731 r <- step defState st
732 case r of
733 Yield x s | n < 0 -> return Nothing
734 | n == 0 -> return $ Just x
735 | otherwise -> go (n - 1) s
736 Skip s -> go n s
737 Stop -> return Nothing
739{-# INLINE_NORMAL lookup #-}
740lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
741lookup e (Stream step state) = go state
742 where
743 go st = do
744 r <- step defState st
745 case r of
746 Yield (a, b) s -> if e == a then return (Just b) else go s
747 Skip s -> go s
748 Stop -> return Nothing
750{-# INLINE_NORMAL findM #-}
751findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
752findM p (Stream step state) = go SPEC state
753 where
754 go !_ st = do
755 r <- step defState st
756 case r of
757 Yield x s -> do
758 b <- p x
759 if b then return (Just x) else go SPEC s
760 Skip s -> go SPEC s
761 Stop -> return Nothing
763{-# INLINE find #-}
764find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
765find p = findM (return . p)
767{-# INLINE_NORMAL findIndices #-}
768findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
769findIndices p (Stream step state) = Stream step' (state, 0)
770 where
771 {-# INLINE_LATE step' #-}
772 step' gst (st, i) = do
773 r <- step (adaptState gst) st
774 return $ case r of
775 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
776 Skip s -> Skip (s, i+1)
777 Stop -> Stop
779{-# INLINE_NORMAL concatMapM #-}
780concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b
781concatMapM f (Stream step state) = Stream step' (Left state)
782 where
783 {-# INLINE_LATE step' #-}
784 step' gst (Left st) = do
785 r <- step (adaptState gst) st
786 case r of
787 Yield a s -> do
788 b_stream <- f a
789 return $ Skip (Right (b_stream, s))
790 Skip s -> return $ Skip (Left s)
791 Stop -> return Stop
793 -- XXX using the pattern synonym Stream causes a major performance issue
794 -- here even if the synonym does not include a adaptState call. Need to
795 -- find out why. Is that something to be fixed in GHC?
796 step' _ (Right (UnStream inner_step inner_st, st)) = do
797 r <- inner_step defState inner_st
798 case r of
799 Yield b inner_s ->
800 return $ Yield b (Right (Stream inner_step inner_s, st))
801 Skip inner_s ->
802 return $ Skip (Right (Stream inner_step inner_s, st))
803 Stop -> return $ Skip (Left st)
805{-# INLINE concatMap #-}
806concatMap :: Monad m => (a -> Stream m b) -> Stream m a -> Stream m b
807concatMap f = concatMapM (return . f)
809------------------------------------------------------------------------------
810-- Substreams
811------------------------------------------------------------------------------
813{-# INLINE_NORMAL isPrefixOf #-}
814isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
815isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
816 where
817 go (sa, sb, Nothing) = do
818 r <- stepa defState sa
819 case r of
820 Yield x sa' -> go (sa', sb, Just x)
821 Skip sa' -> go (sa', sb, Nothing)
822 Stop -> return True
824 go (sa, sb, Just x) = do
825 r <- stepb defState sb
826 case r of
827 Yield y sb' ->
828 if x == y
829 then go (sa, sb', Nothing)
830 else return False
831 Skip sb' -> go (sa, sb', Just x)
832 Stop -> return False
834{-# INLINE_NORMAL isSubsequenceOf #-}
835isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
836isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
837 where
838 go (sa, sb, Nothing) = do
839 r <- stepa defState sa
840 case r of
841 Yield x sa' -> go (sa', sb, Just x)
842 Skip sa' -> go (sa', sb, Nothing)
843 Stop -> return True
845 go (sa, sb, Just x) = do
846 r <- stepb defState sb
847 case r of
848 Yield y sb' ->
849 if x == y
850 then go (sa, sb', Nothing)
851 else go (sa, sb', Just x)
852 Skip sb' -> go (sa, sb', Just x)
853 Stop -> return False
855{-# INLINE_NORMAL stripPrefix #-}
856stripPrefix
857 :: (Eq a, Monad m)
858 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
859stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
860 where
861 go (sa, sb, Nothing) = do
862 r <- stepa defState sa
863 case r of
864 Yield x sa' -> go (sa', sb, Just x)
865 Skip sa' -> go (sa', sb, Nothing)
866 Stop -> return $ Just (Stream stepb sb)
868 go (sa, sb, Just x) = do
869 r <- stepb defState sb
870 case r of
871 Yield y sb' ->
872 if x == y
873 then go (sa, sb', Nothing)
874 else return Nothing
875 Skip sb' -> go (sa, sb', Just x)
876 Stop -> return Nothing
878------------------------------------------------------------------------------
879-- Map and Fold
880------------------------------------------------------------------------------
882-- | Execute a monadic action for each element of the 'Stream'
883{-# INLINE_NORMAL mapM_ #-}
884mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
885mapM_ m = runStream . mapM m
887------------------------------------------------------------------------------
888-- Converting folds
889------------------------------------------------------------------------------
891{-# INLINE toList #-}
892toList :: Monad m => Stream m a -> m [a]
893toList = foldr (:) []
895-- Convert a direct stream to and from CPS encoded stream
896{-# INLINE_LATE toStreamK #-}
897toStreamK :: Monad m => Stream m a -> K.Stream m a
898toStreamK (Stream step state) = go state
899 where
900 go st = K.mkStream $ \gst yld sng stp -> do
901 r <- step gst st
902 case r of
903 Yield x s -> yld x (go s)
904 Skip s -> K.foldStreamShared gst yld sng stp $ go s
905 Stop -> stp
907#ifndef DISABLE_FUSION
908{-# RULES "fromStreamK/toStreamK fusion"
909 forall s. toStreamK (fromStreamK s) = s #-}
910{-# RULES "toStreamK/fromStreamK fusion"
911 forall s. fromStreamK (toStreamK s) = s #-}
912#endif
914{-# INLINE fromStreamD #-}
915fromStreamD :: (K.IsStream t, Monad m) => Stream m a -> t m a
916fromStreamD = K.fromStream . toStreamK
918------------------------------------------------------------------------------
919-- Transformation by Folding (Scans)
920------------------------------------------------------------------------------
922-- XXX Is a prescan useful, discarding the last step does not sound useful? I
923-- am not sure about the utility of this function, so this is implemented but
924-- not exposed. We can expose it if someone provides good reasons why this is
925-- useful.
926--
927-- XXX We have to execute the stream one step ahead to know that we are at the
928-- last step. The vector implementation of prescan executes the last fold step
929-- but does not yield the result. This means we have executed the effect but
930-- discarded value. This does not sound right. In this implementation we are
931-- not executing the last fold step.
932{-# INLINE_NORMAL prescanlM' #-}
933prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
934prescanlM' f mz (Stream step state) = Stream step' (state, mz)
935 where
936 {-# INLINE_LATE step' #-}
937 step' gst (st, prev) = do
938 r <- step (adaptState gst) st
939 case r of
940 Yield x s -> do
941 acc <- prev
942 return $ Yield acc (s, f acc x)
943 Skip s -> return $ Skip (s, prev)
944 Stop -> return Stop
946{-# INLINE prescanl' #-}
947prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
948prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
950-- XXX if we make the initial value of the accumulator monadic then should we
951-- execute it even if the stream is empty? In that case we would have generated
952-- the effect but discarded the value, but that is what a fold does when the
953-- stream is empty. Whatever we decide, need to reconcile this with prescan.
954-- If we execute the initial value here without even using it then it is ok to
955-- execute the last step there as well without using the value.
956-- Looking at the duality with right fold, in case of right fold we always
957-- perform the action when the construction terminates, so in case of left fold
958-- we should perform it only when the reduction starts.
959{-# INLINE_NORMAL postscanlM' #-}
960postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
961postscanlM' fstep begin (Stream step state) =
962 begin `seq` Stream step' (state, begin)
963 where
964 {-# INLINE_LATE step' #-}
965 step' gst (st, acc) = acc `seq` do
966 r <- step (adaptState gst) st
967 case r of
968 Yield x s -> do
969 y <- fstep acc x
970 y `seq` 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 postscanlM #-}
979postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
980postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
981 where
982 {-# INLINE_LATE step' #-}
983 step' gst (st, acc) = do
984 r <- step (adaptState gst) st
985 case r of
986 Yield x s -> do
987 y <- fstep acc x
988 return (Yield y (s, y))
989 Skip s -> return $ Skip (s, acc)
990 Stop -> return Stop
992{-# INLINE_NORMAL postscanl #-}
993postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
994postscanl f = postscanlM (\a b -> return (f a b))
996{-# INLINE_NORMAL scanlM' #-}
997scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
998scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
1000{-# INLINE scanl' #-}
1001scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
1002scanl' f = scanlM' (\a b -> return (f a b))
1004{-# INLINE_NORMAL scanlM #-}
1005scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1006scanlM fstep begin s = begin `cons` postscanlM fstep begin s
1008{-# INLINE scanl #-}
1009scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
1010scanl f = scanlM (\a b -> return (f a b))
1012{-# INLINE_NORMAL scanl1M #-}
1013scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
1014scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
1015 where
1016 {-# INLINE_LATE step' #-}
1017 step' gst (st, Nothing) = do
1018 r <- step gst st
1019 case r of
1020 Yield x s -> return $ Yield x (s, Just x)
1021 Skip s -> return $ Skip (s, Nothing)
1022 Stop -> return Stop
1024 step' gst (st, Just acc) = do
1025 r <- step gst st
1026 case r of
1027 Yield y s -> do
1028 z <- fstep acc y
1029 return $ Yield z (s, Just z)
1030 Skip s -> return $ Skip (s, Just acc)
1031 Stop -> return Stop
1033{-# INLINE scanl1 #-}
1034scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
1035scanl1 f = scanl1M (\x y -> return (f x y))
1037{-# INLINE_NORMAL scanl1M' #-}
1038scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
1039scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
1040 where
1041 {-# INLINE_LATE step' #-}
1042 step' gst (st, Nothing) = do
1043 r <- step gst st
1044 case r of
1045 Yield x s -> x `seq` return $ Yield x (s, Just x)
1046 Skip s -> return $ Skip (s, Nothing)
1047 Stop -> return Stop
1049 step' gst (st, Just acc) = acc `seq` do
1050 r <- step gst st
1051 case r of
1052 Yield y s -> do
1053 z <- fstep acc y
1054 z `seq` return $ Yield z (s, Just z)
1055 Skip s -> return $ Skip (s, Just acc)
1056 Stop -> return Stop
1058{-# INLINE scanl1' #-}
1059scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
1060scanl1' f = scanl1M' (\x y -> return (f x y))
1062-------------------------------------------------------------------------------
1063-- Filtering
1064-------------------------------------------------------------------------------
1066{-# INLINE_NORMAL take #-}
1067take :: Monad m => Int -> Stream m a -> Stream m a
1068take n (Stream step state) = n `seq` Stream step' (state, 0)
1069 where
1070 {-# INLINE_LATE step' #-}
1071 step' gst (st, i) | i < n = do
1072 r <- step gst st
1073 return $ case r of
1074 Yield x s -> Yield x (s, i + 1)
1075 Skip s -> Skip (s, i)
1076 Stop -> Stop
1077 step' _ (_, _) = return Stop
1079{-# INLINE_NORMAL takeWhileM #-}
1080takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
1081takeWhileM f (Stream step state) = Stream step' state
1082 where
1083 {-# INLINE_LATE step' #-}
1084 step' gst st = do
1085 r <- step gst st
1086 case r of
1087 Yield x s -> do
1088 b <- f x
1089 return $ if b then Yield x s else Stop
1090 Skip s -> return $ Skip s
1091 Stop -> return Stop
1093{-# INLINE takeWhile #-}
1094takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
1095takeWhile f = takeWhileM (return . f)
1097{-# INLINE_NORMAL drop #-}
1098drop :: Monad m => Int -> Stream m a -> Stream m a
1099drop n (Stream step state) = Stream step' (state, Just n)
1100 where
1101 {-# INLINE_LATE step' #-}
1102 step' gst (st, Just i)
1103 | i > 0 = do
1104 r <- step gst st
1105 return $
1106 case r of
1107 Yield _ s -> Skip (s, Just (i - 1))
1108 Skip s -> Skip (s, Just i)
1109 Stop -> Stop
1110 | otherwise = return $ Skip (st, Nothing)
1112 step' gst (st, Nothing) = do
1113 r <- step gst st
1114 return $
1115 case r of
1116 Yield x s -> Yield x (s, Nothing)
1117 Skip s -> Skip (s, Nothing)
1118 Stop -> Stop
1120data DropWhileState s a
1121 = DropWhileDrop s
1122 | DropWhileYield a s
1123 | DropWhileNext s
1125{-# INLINE_NORMAL dropWhileM #-}
1126dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
1127dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
1128 where
1129 {-# INLINE_LATE step' #-}
1130 step' gst (DropWhileDrop st) = do
1131 r <- step gst st
1132 case r of
1133 Yield x s -> do
1134 b <- f x
1135 if b
1136 then return $ Skip (DropWhileDrop s)
1137 else return $ Skip (DropWhileYield x s)
1138 Skip s -> return $ Skip (DropWhileDrop s)
1139 Stop -> return Stop
1141 step' gst (DropWhileNext st) = do
1142 r <- step gst st
1143 case r of
1144 Yield x s -> return $ Skip (DropWhileYield x s)
1145 Skip s -> return $ Skip (DropWhileNext s)
1146 Stop -> return Stop
1148 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
1150{-# INLINE dropWhile #-}
1151dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
1152dropWhile f = dropWhileM (return . f)
1154{-# INLINE_NORMAL filterM #-}
1155filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
1156filterM f (Stream step state) = Stream step' state
1157 where
1158 {-# INLINE_LATE step' #-}
1159 step' gst st = do
1160 r <- step gst st
1161 case r of
1162 Yield x s -> do
1163 b <- f x
1164 return $ if b
1165 then Yield x s
1166 else Skip s
1167 Skip s -> return $ Skip s
1168 Stop -> return Stop
1170{-# INLINE filter #-}
1171filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
1172filter f = filterM (return . f)
1174{-# INLINE_NORMAL uniq #-}
1175uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
1176uniq (Stream step state) = Stream step' (Nothing, state)
1177 where
1178 {-# INLINE_LATE step' #-}
1179 step' gst (Nothing, st) = do
1180 r <- step gst st
1181 case r of
1182 Yield x s -> return $ Yield x (Just x, s)
1183 Skip s -> return $ Skip (Nothing, s)
1184 Stop -> return Stop
1185 step' gst (Just x, st) = do
1186 r <- step gst st
1187 case r of
1188 Yield y s | x == y -> return $ Skip (Just x, s)
1189 | otherwise -> return $ Yield x (Just y, s)
1190 Skip s -> return $ Skip (Just x, s)
1191 Stop -> return Stop
1193------------------------------------------------------------------------------
1194-- Transformation by Mapping
1195------------------------------------------------------------------------------
1197{-# INLINE_NORMAL sequence #-}
1198sequence :: Monad m => Stream m (m a) -> Stream m a
1199sequence (Stream step state) = Stream step' state
1200 where
1201 {-# INLINE_LATE step' #-}
1202 step' gst st = do
1203 r <- step (adaptState gst) st
1204 case r of
1205 Yield x s -> x >>= \a -> return (Yield a s)
1206 Skip s -> return $ Skip s
1207 Stop -> return Stop
1209------------------------------------------------------------------------------
1210-- Inserting
1211------------------------------------------------------------------------------
1213{-# INLINE_NORMAL insertBy #-}
1214insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
1215insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
1216 where
1217 {-# INLINE_LATE step' #-}
1218 step' gst (st, False, _) = do
1219 r <- step gst st
1220 case r of
1221 Yield x s -> case cmp a x of
1222 GT -> return $ Yield x (s, False, Nothing)
1223 _ -> return $ Yield a (s, True, Just x)
1224 Skip s -> return $ Skip (s, False, Nothing)
1225 Stop -> return $ Yield a (st, True, Nothing)
1227 step' _ (_, True, Nothing) = return Stop
1229 step' gst (st, True, Just prev) = do
1230 r <- step gst st
1231 case r of
1232 Yield x s -> return $ Yield prev (s, True, Just x)
1233 Skip s -> return $ Skip (s, True, Just prev)
1234 Stop -> return $ Yield prev (st, True, Nothing)
1236------------------------------------------------------------------------------
1237-- Deleting
1238------------------------------------------------------------------------------
1240{-# INLINE_NORMAL deleteBy #-}
1241deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
1242deleteBy eq x (Stream step state) = Stream step' (state, False)
1243 where
1244 {-# INLINE_LATE step' #-}
1245 step' gst (st, False) = do
1246 r <- step gst st
1247 case r of
1248 Yield y s -> return $
1249 if eq x y then Skip (s, True) else Yield y (s, False)
1250 Skip s -> return $ Skip (s, False)
1251 Stop -> return Stop
1253 step' gst (st, True) = do
1254 r <- step gst st
1255 case r of
1256 Yield y s -> return $ Yield y (s, True)
1257 Skip s -> return $ Skip (s, True)
1258 Stop -> return Stop
1260------------------------------------------------------------------------------
1261-- Transformation by Map and Filter
1262------------------------------------------------------------------------------
1264-- XXX Will this always fuse properly?
1265{-# INLINE_NORMAL mapMaybe #-}
1266mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
1267mapMaybe f = fmap fromJust . filter isJust . map f
1269{-# INLINE_NORMAL mapMaybeM #-}
1270mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
1271mapMaybeM f = fmap fromJust . filter isJust . mapM f
1273------------------------------------------------------------------------------
1274-- Zipping
1275------------------------------------------------------------------------------
1277{-# INLINE_NORMAL indexed #-}
1278indexed :: Monad m => Stream m a -> Stream m (Int, a)
1279indexed (Stream step state) = Stream step' (state, 0)
1280 where
1281 {-# INLINE_LATE step' #-}
1282 step' gst (st, i) = i `seq` do
1283 r <- step (adaptState gst) st
1284 case r of
1285 Yield x s -> return $ Yield (i, x) (s, i+1)
1286 Skip s -> return $ Skip (s, i)
1287 Stop -> return Stop
1289{-# INLINE_NORMAL indexedR #-}
1290indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
1291indexedR m (Stream step state) = Stream step' (state, m)
1292 where
1293 {-# INLINE_LATE step' #-}
1294 step' gst (st, i) = i `seq` do
1295 r <- step (adaptState gst) st
1296 case r of
1297 Yield x s -> let i' = i - 1
1298 in
1299 return $ Yield (i', x) (s, i')
1300 Skip s -> return $ Skip (s, i)
1301 Stop -> return Stop
1303{-# INLINE_NORMAL zipWithM #-}
1304zipWithM :: Monad m
1305 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
1306zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
1307 where
1308 {-# INLINE_LATE step #-}
1309 step gst (sa, sb, Nothing) = do
1310 r <- stepa (adaptState gst) sa
1311 return $
1312 case r of
1313 Yield x sa' -> Skip (sa', sb, Just x)
1314 Skip sa' -> Skip (sa', sb, Nothing)
1315 Stop -> Stop
1317 step gst (sa, sb, Just x) = do
1318 r <- stepb (adaptState gst) sb
1319 case r of
1320 Yield y sb' -> do
1321 z <- f x y
1322 return $ Yield z (sa, sb', Nothing)
1323 Skip sb' -> return $ Skip (sa, sb', Just x)
1324 Stop -> return Stop
1326{-# RULES "zipWithM xs xs"
1327 forall f xs. zipWithM f xs xs = mapM (\x -> f x x) xs #-}
1329{-# INLINE zipWith #-}
1330zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
1331zipWith f = zipWithM (\a b -> return (f a b))
1333------------------------------------------------------------------------------
1334-- Comparisions
1335------------------------------------------------------------------------------
1337{-# INLINE_NORMAL eqBy #-}
1338eqBy :: Monad m => (a -> b -> Bool) -> Stream m a -> Stream m b -> m Bool
1339eqBy eq (Stream step1 t1) (Stream step2 t2) = eq_loop0 SPEC t1 t2
1340 where
1341 eq_loop0 !_ s1 s2 = do
1342 r <- step1 defState s1
1343 case r of
1344 Yield x s1' -> eq_loop1 SPEC x s1' s2
1345 Skip s1' -> eq_loop0 SPEC s1' s2
1346 Stop -> eq_null s2
1348 eq_loop1 !_ x s1 s2 = do
1349 r <- step2 defState s2
1350 case r of
1351 Yield y s2'
1352 | eq x y -> eq_loop0 SPEC s1 s2'
1353 | otherwise -> return False
1354 Skip s2' -> eq_loop1 SPEC x s1 s2'
1355 Stop -> return False
1357 eq_null s2 = do
1358 r <- step2 defState s2
1359 case r of
1360 Yield _ _ -> return False
1361 Skip s2' -> eq_null s2'
1362 Stop -> return True
1364-- | Compare two streams lexicographically
1365{-# INLINE_NORMAL cmpBy #-}
1366cmpBy
1367 :: Monad m
1368 => (a -> b -> Ordering) -> Stream m a -> Stream m b -> m Ordering
1369cmpBy cmp (Stream step1 t1) (Stream step2 t2) = cmp_loop0 SPEC t1 t2
1370 where
1371 cmp_loop0 !_ s1 s2 = do
1372 r <- step1 defState s1
1373 case r of
1374 Yield x s1' -> cmp_loop1 SPEC x s1' s2
1375 Skip s1' -> cmp_loop0 SPEC s1' s2
1376 Stop -> cmp_null s2
1378 cmp_loop1 !_ x s1 s2 = do
1379 r <- step2 defState s2
1380 case r of
1381 Yield y s2' -> case x `cmp` y of
1382 EQ -> cmp_loop0 SPEC s1 s2'
1383 c -> return c
1384 Skip s2' -> cmp_loop1 SPEC x s1 s2'
1385 Stop -> return GT
1387 cmp_null s2 = do
1388 r <- step2 defState s2
1389 case r of
1390 Yield _ _ -> return LT
1391 Skip s2' -> cmp_null s2'
1392 Stop -> return EQ
1394------------------------------------------------------------------------------
1395-- Merging
1396------------------------------------------------------------------------------
1398{-# INLINE_NORMAL mergeByM #-}
1399mergeByM
1400 :: (Monad m)
1401 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
1402mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
1403 Stream step (Just ta, Just tb, Nothing, Nothing)
1404 where
1405 {-# INLINE_LATE step #-}
1407 -- one of the values is missing, and the corresponding stream is running
1408 step gst (Just sa, sb, Nothing, b) = do
1409 r <- stepa gst sa
1410 return $ case r of
1411 Yield a sa' -> Skip (Just sa', sb, Just a, b)
1412 Skip sa' -> Skip (Just sa', sb, Nothing, b)
1413 Stop -> Skip (Nothing, sb, Nothing, b)
1415 step gst (sa, Just sb, a, Nothing) = do
1416 r <- stepb gst sb
1417 return $ case r of
1418 Yield b sb' -> Skip (sa, Just sb', a, Just b)
1419 Skip sb' -> Skip (sa, Just sb', a, Nothing)
1420 Stop -> Skip (sa, Nothing, a, Nothing)
1422 -- both the values are available
1423 step _ (sa, sb, Just a, Just b) = do
1424 res <- cmp a b
1425 return $ case res of
1426 GT -> Yield b (sa, sb, Just a, Nothing)
1427 _ -> Yield a (sa, sb, Nothing, Just b)
1429 -- one of the values is missing, corresponding stream is done
1430 step _ (Nothing, sb, Nothing, Just b) =
1431 return $ Yield b (Nothing, sb, Nothing, Nothing)
1433 step _ (sa, Nothing, Just a, Nothing) =
1434 return $ Yield a (sa, Nothing, Nothing, Nothing)
1436 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
1438{-# INLINE mergeBy #-}
1439mergeBy
1440 :: (Monad m)
1441 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
1442mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
1444------------------------------------------------------------------------------
1445-- Transformation comprehensions
1446------------------------------------------------------------------------------
1448{-# INLINE_NORMAL the #-}
1449the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
1450the (Stream step state) = go state
1451 where
1452 go st = do
1453 r <- step defState st
1454 case r of
1455 Yield x s -> go' x s
1456 Skip s -> go s
1457 Stop -> return Nothing
1458 go' n st = do
1459 r <- step defState st
1460 case r of
1461 Yield x s | x == n -> go' n s
1462 | otherwise -> return Nothing
1463 Skip s -> go' n s
1464 Stop -> return (Just n)