git.iko.soy/streamly
on 315abcf help sign in
streamly/src/Streamly/Streams/StreamD.hs
1507 lines 48 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, State(..), defState, adaptState)
204import qualified Streamly.Streams.StreamK as K
206------------------------------------------------------------------------------
207-- The direct style stream type
208------------------------------------------------------------------------------
210-- | A stream is a succession of 'Step's. A 'Yield' produces a single value and
211-- the next state of the stream. 'Stop' indicates there are no more values in
212-- the stream.
213data Step s a = Yield a s | Skip s | Stop
215instance Functor (Step s) where
216 {-# INLINE fmap #-}
217 fmap f (Yield x s) = Yield (f x) s
218 fmap _ (Skip s) = Skip s
219 fmap _ Stop = Stop
221-- gst = global state
222-- | A stream consists of a step function that generates the next step given a
223-- current state, and the current state.
224data Stream m a =
225 forall s. UnStream (State K.Stream m a -> s -> m (Step s a)) s
227-- Do not use the UnStream constructor directly, use the pattern synonym
228-- instead. XXX we should perhaps put the type in a separate module and do not
229-- export the constructor.
231detachStream :: Stream m a -> Stream m a
232detachStream (UnStream step state) = UnStream step' state
233 where step' gst = step (adaptState gst)
235pattern Stream :: (State K.Stream m a -> s -> m (Step s a)) -> s -> Stream m a
236pattern Stream step state <- (detachStream -> UnStream step state)
237 where Stream = UnStream
239#if __GLASGOW_HASKELL__ >= 802
240{-# COMPLETE Stream #-}
241#endif
243------------------------------------------------------------------------------
244-- Construction
245------------------------------------------------------------------------------
247-- | An empty 'Stream'.
248{-# INLINE_NORMAL nil #-}
249nil :: Monad m => Stream m a
250nil = Stream (\_ _ -> return Stop) ()
252-- | Can fuse but has O(n^2) complexity.
253{-# INLINE_NORMAL cons #-}
254cons :: Monad m => a -> Stream m a -> Stream m a
255cons x (Stream step state) = Stream step1 Nothing
256 where
257 {-# INLINE_LATE step1 #-}
258 step1 _ Nothing = return $ Yield x (Just state)
259 step1 gst (Just st) = do
260 r <- step gst st
261 return $
262 case r of
263 Yield a s -> Yield a (Just s)
264 Skip s -> Skip (Just s)
265 Stop -> Stop
267-------------------------------------------------------------------------------
268-- Deconstruction
269-------------------------------------------------------------------------------
271-- Does not fuse, has the same performance as the StreamK version.
272{-# INLINE_NORMAL uncons #-}
273uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
274uncons (Stream step state) = go state
275 where
276 go st = do
277 r <- step defState st
278 case r of
279 Yield x s -> return $ Just (x, Stream step s)
280 Skip s -> go s
281 Stop -> return Nothing
283------------------------------------------------------------------------------
284-- Generation by unfold
285------------------------------------------------------------------------------
287{-# INLINE_NORMAL unfoldrM #-}
288unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
289unfoldrM next state = Stream step state
290 where
291 {-# INLINE_LATE step #-}
292 step _ st = do
293 r <- next st
294 return $ case r of
295 Just (x, s) -> Yield x s
296 Nothing -> Stop
298{-# INLINE_LATE unfoldr #-}
299unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
300unfoldr f = unfoldrM (return . f)
302------------------------------------------------------------------------------
303-- Specialized Generation
304------------------------------------------------------------------------------
306repeat :: Monad m => a -> Stream m a
307repeat x = Stream (\_ _ -> return $ Yield x ()) ()
309{-# INLINE_NORMAL replicateM #-}
310replicateM :: Monad m => Int -> m a -> Stream m a
311replicateM n p = Stream step n
312 where
313 {-# INLINE_LATE step #-}
314 step _ i | i <= 0 = return Stop
315 | otherwise = do
316 x <- p
317 return $ Yield x (i - 1)
319{-# INLINE_NORMAL replicate #-}
320replicate :: Monad m => Int -> a -> Stream m a
321replicate n x = replicateM n (return x)
323-- This would not work properly for floats, therefore we put an Integral
324-- constraint.
325-- | Can be used to enumerate unbounded integrals. This does not check for
326-- overflow or underflow for bounded integrals.
327{-# INLINE_NORMAL intFromStep #-}
328intFromStep :: (Integral a, Monad m) => a -> a -> Stream m a
329intFromStep from stride =
330 from `seq` stride `seq` Stream step from
331 where
332 {-# INLINE_LATE step #-}
333 step _ !x = return $ Yield x $! (x + stride)
335-- We are assuming that "to" is constrained by the type to be within
336-- max/min bounds.
337{-# INLINE intFromTo #-}
338intFromTo :: (Monad m, Integral a) => a -> a -> Stream m a
339intFromTo from to = takeWhile (<= to) $ intFromStep from 1
341{-# INLINE intFrom #-}
342intFrom :: (Monad m, Integral a, Bounded a) => a -> Stream m a
343intFrom from = intFromTo from maxBound
345data EnumState a = EnumInit | EnumYield a a a | EnumStop
347{-# INLINE_NORMAL intFromThenToUp #-}
348intFromThenToUp :: (Monad m, Integral a) => a -> a -> a -> Stream m a
349intFromThenToUp from next to = Stream step EnumInit
350 where
351 {-# INLINE_LATE step #-}
352 step _ EnumInit =
353 return $
354 if to < next
355 then if to < from
356 then Stop
357 else Yield from EnumStop
358 else -- from <= next <= to
359 let stride = next - from
360 in Skip $ EnumYield from stride (to - stride)
362 step _ (EnumYield x stride toMinus) =
363 return $
364 if x > toMinus
365 then Yield x EnumStop
366 else Yield x $ EnumYield (x + stride) stride toMinus
368 step _ EnumStop = return Stop
370{-# INLINE_NORMAL intFromThenToDn #-}
371intFromThenToDn :: (Monad m, Integral a) => a -> a -> a -> Stream m a
372intFromThenToDn from next to = Stream step EnumInit
373 where
374 {-# INLINE_LATE step #-}
375 step _ EnumInit =
376 return $ if to > next
377 then if to > from
378 then Stop
379 else Yield from EnumStop
380 else -- from >= next >= to
381 let stride = next - from
382 in Skip $ EnumYield from stride (to - stride)
384 step _ (EnumYield x stride toMinus) =
385 return $
386 if x < toMinus
387 then Yield x EnumStop
388 else Yield x $ EnumYield (x + stride) stride toMinus
390 step _ EnumStop = return Stop
392{-# INLINE_NORMAL intFromThenTo #-}
393intFromThenTo :: (Monad m, Integral a) => a -> a -> a -> Stream m a
394intFromThenTo from next to
395 | next >= from = intFromThenToUp from next to
396 | otherwise = intFromThenToDn from next to
398{-# INLINE_NORMAL intFromThen #-}
399intFromThen :: (Monad m, Integral a, Bounded a) => a -> a -> Stream m a
400intFromThen from next =
401 if next > from
402 then intFromThenToUp from next maxBound
403 else intFromThenToDn from next minBound
405-- For floating point numbers if the increment is less than the precision then
406-- it just gets lost. Therefore we cannot always increment it correctly by just
407-- repeated addition.
408-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
409-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
411-- Instead we accumulate the increment counter and compute the increment
412-- everytime before adding it to the starting number.
413--
414-- This works for Integrals as well as floating point numbers, but intFromStep
415-- is faster for integrals.
416{-# INLINE_NORMAL numFromStep #-}
417numFromStep :: (Monad m, Num a) => a -> a -> Stream m a
418numFromStep from stride = Stream step 0
419 where
420 {-# INLINE_LATE step #-}
421 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
423{-# INLINE_NORMAL numFrom #-}
424numFrom :: (Monad m, Num a) => a -> Stream m a
425numFrom from = numFromStep from 1
427{-# INLINE_NORMAL numFromThen #-}
428numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
429numFromThen from next = numFromStep from (next - from)
431-- We cannot write a general function for Num. The only way to write code
432-- portable between the two is to use a 'Real' constraint and convert between
433-- Fractional and Integral using fromRational which is horribly slow.
434{-# INLINE_NORMAL fracFromTo #-}
435fracFromTo :: (Monad m, Fractional a, Ord a) => a -> a -> Stream m a
436fracFromTo from to = takeWhile (<= to + 1/2) $ numFromStep from 1
438{-# INLINE_NORMAL fracFromThenTo #-}
439fracFromThenTo :: (Monad m, Fractional a, Ord a) => a -> a -> a -> Stream m a
440fracFromThenTo from next to = takeWhile predicate $ numFromThen from next
441 where
442 mid = (next - from) / 2
443 predicate | next >= from = (<= to + mid)
444 | otherwise = (>= to + mid)
446-------------------------------------------------------------------------------
447-- Generation by Conversion
448-------------------------------------------------------------------------------
450-- | Create a singleton 'Stream' from a pure value.
451{-# INLINE_NORMAL yield #-}
452yield :: Monad m => a -> Stream m a
453yield x = Stream (\_ s -> return $ step undefined s) True
454 where
455 {-# INLINE_LATE step #-}
456 step _ True = Yield x False
457 step _ False = Stop
459-- | Create a singleton 'Stream' from a monadic action.
460{-# INLINE_NORMAL yieldM #-}
461yieldM :: Monad m => m a -> Stream m a
462yieldM m = Stream step True
463 where
464 {-# INLINE_LATE step #-}
465 step _ True = m >>= \x -> return $ Yield x False
466 step _ False = return Stop
468{-# INLINE_NORMAL fromIndicesM #-}
469fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
470fromIndicesM gen = Stream step 0
471 where
472 {-# INLINE_LATE step #-}
473 step _ i = do
474 x <- gen i
475 return $ Yield x (i + 1)
477{-# INLINE fromIndices #-}
478fromIndices :: Monad m => (Int -> a) -> Stream m a
479fromIndices gen = fromIndicesM (return . gen)
481{-# INLINE_NORMAL generateM #-}
482generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
483generateM n gen = n `seq` Stream step 0
484 where
485 {-# INLINE_LATE step #-}
486 step _ i | i < n = do
487 x <- gen i
488 return $ Yield x (i + 1)
489 | otherwise = return Stop
491{-# INLINE generate #-}
492generate :: Monad m => Int -> (Int -> a) -> Stream m a
493generate n gen = generateM n (return . gen)
495-- XXX we need the MonadAsync constraint because of a rewrite rule.
496-- | Convert a list of monadic actions to a 'Stream'
497{-# INLINE_LATE fromListM #-}
498fromListM :: MonadAsync m => [m a] -> Stream m a
499fromListM = Stream step
500 where
501 {-# INLINE_LATE step #-}
502 step _ (m:ms) = m >>= \x -> return $ Yield x ms
503 step _ [] = return Stop
505-- | Convert a list of pure values to a 'Stream'
506{-# INLINE_LATE fromList #-}
507fromList :: Monad m => [a] -> Stream m a
508fromList = Stream step
509 where
510 {-# INLINE_LATE step #-}
511 step _ (x:xs) = return $ Yield x xs
512 step _ [] = return Stop
514{-# INLINE_LATE fromStreamK #-}
515fromStreamK :: Monad m => K.Stream m a -> Stream m a
516fromStreamK = Stream step
517 where
518 step gst m1 =
519 let stop = return Stop
520 single a = return $ Yield a K.nil
521 yieldk a r = return $ Yield a r
522 in K.unStreamShared m1 gst stop single yieldk
524{-# INLINE toStreamD #-}
525toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
526toStreamD = fromStreamK . K.toStream
528------------------------------------------------------------------------------
529-- Elimination by Folds
530------------------------------------------------------------------------------
532{-# INLINE_NORMAL foldrM #-}
533foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b
534foldrM f z (Stream step state) = go SPEC state
535 where
536 go !_ st = do
537 r <- step defState st
538 case r of
539 Yield x s -> go SPEC s >>= f x
540 Skip s -> go SPEC s
541 Stop -> return z
543{-# INLINE_NORMAL foldr #-}
544foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
545foldr f = foldrM (\a b -> return (f a b))
547{-# INLINE_NORMAL foldr1 #-}
548foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
549foldr1 f m = do
550 r <- uncons m
551 case r of
552 Nothing -> return Nothing
553 Just (h, t) -> fmap Just (foldr f h t)
555{-# INLINE_NORMAL foldlM' #-}
556foldlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> m b
557foldlM' fstep begin (Stream step state) = go SPEC begin state
558 where
559 go !_ acc st = acc `seq` do
560 r <- step defState st
561 case r of
562 Yield x s -> do
563 acc' <- fstep acc x
564 go SPEC acc' s
565 Skip s -> go SPEC acc s
566 Stop -> return acc
568{-# INLINE foldl' #-}
569foldl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
570foldl' fstep = foldlM' (\b a -> return (fstep b a))
572------------------------------------------------------------------------------
573-- Specialized Folds
574------------------------------------------------------------------------------
576-- | Run a streaming composition, discard the results.
577{-# INLINE_LATE runStream #-}
578runStream :: Monad m => Stream m a -> m ()
579runStream (Stream step state) = go SPEC state
580 where
581 go !_ st = do
582 r <- step defState st
583 case r of
584 Yield _ s -> go SPEC s
585 Skip s -> go SPEC s
586 Stop -> return ()
588{-# INLINE_NORMAL null #-}
589null :: Monad m => Stream m a -> m Bool
590null (Stream step state) = go state
591 where
592 go st = do
593 r <- step defState st
594 case r of
595 Yield _ _ -> return False
596 Skip s -> go s
597 Stop -> return True
599-- XXX SPEC?
600{-# INLINE_NORMAL head #-}
601head :: Monad m => Stream m a -> m (Maybe a)
602head (Stream step state) = go state
603 where
604 go st = do
605 r <- step defState st
606 case r of
607 Yield x _ -> return (Just x)
608 Skip s -> go s
609 Stop -> return Nothing
611-- Does not fuse, has the same performance as the StreamK version.
612{-# INLINE_NORMAL tail #-}
613tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
614tail (Stream step state) = go state
615 where
616 go st = do
617 r <- step defState st
618 case r of
619 Yield _ s -> return (Just $ Stream step s)
620 Skip s -> go s
621 Stop -> return Nothing
623-- XXX will it fuse? need custom impl?
624{-# INLINE_NORMAL last #-}
625last :: Monad m => Stream m a -> m (Maybe a)
626last = foldl' (\_ y -> Just y) Nothing
628{-# INLINE_NORMAL elem #-}
629elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
630elem e (Stream step state) = go state
631 where
632 go st = do
633 r <- step defState st
634 case r of
635 Yield x s
636 | x == e -> return True
637 | otherwise -> go s
638 Skip s -> go s
639 Stop -> return False
641{-# INLINE_NORMAL notElem #-}
642notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
643notElem e s = fmap not (elem e s)
645{-# INLINE_NORMAL all #-}
646all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
647all p (Stream step state) = go state
648 where
649 go st = do
650 r <- step defState st
651 case r of
652 Yield x s
653 | p x -> go s
654 | otherwise -> return False
655 Skip s -> go s
656 Stop -> return True
658{-# INLINE_NORMAL any #-}
659any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
660any p (Stream step state) = go state
661 where
662 go st = do
663 r <- step defState st
664 case r of
665 Yield x s
666 | p x -> return True
667 | otherwise -> go s
668 Skip s -> go s
669 Stop -> return False
671{-# INLINE_NORMAL maximum #-}
672maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
673maximum (Stream step state) = go Nothing state
674 where
675 go Nothing st = do
676 r <- step defState st
677 case r of
678 Yield x s -> go (Just x) s
679 Skip s -> go Nothing s
680 Stop -> return Nothing
681 go (Just acc) st = do
682 r <- step defState st
683 case r of
684 Yield x s
685 | acc <= x -> go (Just x) s
686 | otherwise -> go (Just acc) s
687 Skip s -> go (Just acc) s
688 Stop -> return (Just acc)
690{-# INLINE_NORMAL maximumBy #-}
691maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
692maximumBy cmp (Stream step state) = go Nothing state
693 where
694 go Nothing st = do
695 r <- step defState st
696 case r of
697 Yield x s -> go (Just x) s
698 Skip s -> go Nothing s
699 Stop -> return Nothing
700 go (Just acc) st = do
701 r <- step defState st
702 case r of
703 Yield x s -> case cmp acc x of
704 GT -> go (Just acc) s
705 _ -> go (Just x) s
706 Skip s -> go (Just acc) s
707 Stop -> return (Just acc)
709{-# INLINE_NORMAL minimum #-}
710minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
711minimum (Stream step state) = go Nothing state
712 where
713 go Nothing st = do
714 r <- step defState st
715 case r of
716 Yield x s -> go (Just x) s
717 Skip s -> go Nothing s
718 Stop -> return Nothing
719 go (Just acc) st = do
720 r <- step defState st
721 case r of
722 Yield x s
723 | acc <= x -> go (Just acc) s
724 | otherwise -> go (Just x) s
725 Skip s -> go (Just acc) s
726 Stop -> return (Just acc)
728{-# INLINE_NORMAL minimumBy #-}
729minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
730minimumBy cmp (Stream step state) = go Nothing state
731 where
732 go Nothing st = do
733 r <- step defState st
734 case r of
735 Yield x s -> go (Just x) s
736 Skip s -> go Nothing s
737 Stop -> return Nothing
738 go (Just acc) st = do
739 r <- step defState st
740 case r of
741 Yield x s -> case cmp acc x of
742 GT -> go (Just x) s
743 _ -> go (Just acc) s
744 Skip s -> go (Just acc) s
745 Stop -> return (Just acc)
747{-# INLINE_NORMAL (!!) #-}
748(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
749(Stream step state) !! i = go i state
750 where
751 go n st = do
752 r <- step defState st
753 case r of
754 Yield x s | n < 0 -> return Nothing
755 | n == 0 -> return $ Just x
756 | otherwise -> go (n - 1) s
757 Skip s -> go n s
758 Stop -> return Nothing
760{-# INLINE_NORMAL lookup #-}
761lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
762lookup e (Stream step state) = go state
763 where
764 go st = do
765 r <- step defState st
766 case r of
767 Yield (a, b) s -> if e == a then return (Just b) else go s
768 Skip s -> go s
769 Stop -> return Nothing
771{-# INLINE_NORMAL findM #-}
772findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
773findM p (Stream step state) = go SPEC state
774 where
775 go !_ st = do
776 r <- step defState st
777 case r of
778 Yield x s -> do
779 b <- p x
780 if b then return (Just x) else go SPEC s
781 Skip s -> go SPEC s
782 Stop -> return Nothing
784{-# INLINE find #-}
785find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
786find p = findM (return . p)
788{-# INLINE_NORMAL findIndices #-}
789findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
790findIndices p (Stream step state) = Stream step' (state, 0)
791 where
792 {-# INLINE_LATE step' #-}
793 step' gst (st, i) = do
794 r <- step (adaptState gst) st
795 return $ case r of
796 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
797 Skip s -> Skip (s, i+1)
798 Stop -> Stop
800{-# INLINE_NORMAL concatMapM #-}
801concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b
802concatMapM f (Stream step state) = Stream step' (Left state)
803 where
804 {-# INLINE_LATE step' #-}
805 step' gst (Left st) = do
806 r <- step (adaptState gst) st
807 case r of
808 Yield a s -> do
809 b_stream <- f a
810 return $ Skip (Right (b_stream, s))
811 Skip s -> return $ Skip (Left s)
812 Stop -> return Stop
814 step' _ (Right (Stream inner_step inner_st, st)) = do
815 r <- inner_step defState inner_st
816 case r of
817 Yield b inner_s ->
818 return $ Yield b (Right (Stream inner_step inner_s, st))
819 Skip inner_s ->
820 return $ Skip (Right (Stream inner_step inner_s, st))
821 Stop -> return $ Skip (Left st)
823{-# INLINE concatMap #-}
824concatMap :: Monad m => (a -> Stream m b) -> Stream m a -> Stream m b
825concatMap f = concatMapM (return . f)
827------------------------------------------------------------------------------
828-- Substreams
829------------------------------------------------------------------------------
831{-# INLINE_NORMAL isPrefixOf #-}
832isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
833isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
834 where
835 go (sa, sb, Nothing) = do
836 r <- stepa defState sa
837 case r of
838 Yield x sa' -> go (sa', sb, Just x)
839 Skip sa' -> go (sa', sb, Nothing)
840 Stop -> return True
842 go (sa, sb, Just x) = do
843 r <- stepb defState sb
844 case r of
845 Yield y sb' ->
846 if x == y
847 then go (sa, sb', Nothing)
848 else return False
849 Skip sb' -> go (sa, sb', Just x)
850 Stop -> return False
852{-# INLINE_NORMAL isSubsequenceOf #-}
853isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
854isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
855 where
856 go (sa, sb, Nothing) = do
857 r <- stepa defState sa
858 case r of
859 Yield x sa' -> go (sa', sb, Just x)
860 Skip sa' -> go (sa', sb, Nothing)
861 Stop -> return True
863 go (sa, sb, Just x) = do
864 r <- stepb defState sb
865 case r of
866 Yield y sb' ->
867 if x == y
868 then go (sa, sb', Nothing)
869 else go (sa, sb', Just x)
870 Skip sb' -> go (sa, sb', Just x)
871 Stop -> return False
873{-# INLINE_NORMAL stripPrefix #-}
874stripPrefix
875 :: (Eq a, Monad m)
876 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
877stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
878 where
879 go (sa, sb, Nothing) = do
880 r <- stepa defState sa
881 case r of
882 Yield x sa' -> go (sa', sb, Just x)
883 Skip sa' -> go (sa', sb, Nothing)
884 Stop -> return $ Just (Stream stepb sb)
886 go (sa, sb, Just x) = do
887 r <- stepb defState sb
888 case r of
889 Yield y sb' ->
890 if x == y
891 then go (sa, sb', Nothing)
892 else return Nothing
893 Skip sb' -> go (sa, sb', Just x)
894 Stop -> return Nothing
896------------------------------------------------------------------------------
897-- Map and Fold
898------------------------------------------------------------------------------
900-- | Execute a monadic action for each element of the 'Stream'
901{-# INLINE_NORMAL mapM_ #-}
902mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
903mapM_ m = runStream . mapM m
905------------------------------------------------------------------------------
906-- Converting folds
907------------------------------------------------------------------------------
909{-# INLINE toList #-}
910toList :: Monad m => Stream m a -> m [a]
911toList = foldr (:) []
913-- Convert a direct stream to and from CPS encoded stream
914{-# INLINE_LATE toStreamK #-}
915toStreamK :: Monad m => Stream m a -> K.Stream m a
916toStreamK (Stream step state) = go state
917 where
918 go st = K.mkStream $ \gst stp sng yld -> do
919 r <- step gst st
920 case r of
921 Yield x s -> yld x (go s)
922 Skip s -> K.unStreamShared (go s) gst stp sng yld
923 Stop -> stp
925#ifndef DISABLE_FUSION
926{-# RULES "fromStreamK/toStreamK fusion"
927 forall s. toStreamK (fromStreamK s) = s #-}
928{-# RULES "toStreamK/fromStreamK fusion"
929 forall s. fromStreamK (toStreamK s) = s #-}
930#endif
932{-# INLINE fromStreamD #-}
933fromStreamD :: (K.IsStream t, Monad m) => Stream m a -> t m a
934fromStreamD = K.fromStream . toStreamK
936------------------------------------------------------------------------------
937-- Transformation by Folding (Scans)
938------------------------------------------------------------------------------
940-- XXX Is a prescan useful, discarding the last step does not sound useful? I
941-- am not sure about the utility of this function, so this is implemented but
942-- not exposed. We can expose it if someone provides good reasons why this is
943-- useful.
944--
945-- XXX We have to execute the stream one step ahead to know that we are at the
946-- last step. The vector implementation of prescan executes the last fold step
947-- but does not yield the result. This means we have executed the effect but
948-- discarded value. This does not sound right. In this implementation we are
949-- not executing the last fold step.
950{-# INLINE_NORMAL prescanlM' #-}
951prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
952prescanlM' f mz (Stream step state) = Stream step' (state, mz)
953 where
954 {-# INLINE_LATE step' #-}
955 step' gst (st, prev) = do
956 r <- step (adaptState gst) st
957 case r of
958 Yield x s -> do
959 acc <- prev
960 return $ Yield acc (s, f acc x)
961 Skip s -> return $ Skip (s, prev)
962 Stop -> return Stop
964{-# INLINE prescanl' #-}
965prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
966prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
968-- XXX if we make the initial value of the accumulator monadic then should we
969-- execute it even if the stream is empty? In that case we would have generated
970-- the effect but discarded the value, but that is what a fold does when the
971-- stream is empty. Whatever we decide, need to reconcile this with prescan.
972-- If we execute the initial value here without even using it then it is ok to
973-- execute the last step there as well without using the value.
974-- Looking at the duality with right fold, in case of right fold we always
975-- perform the action when the construction terminates, so in case of left fold
976-- we should perform it only when the reduction starts.
977{-# INLINE_NORMAL postscanlM' #-}
978postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
979postscanlM' fstep begin (Stream step state) =
980 begin `seq` Stream step' (state, begin)
981 where
982 {-# INLINE_LATE step' #-}
983 step' gst (st, acc) = acc `seq` do
984 r <- step (adaptState gst) st
985 case r of
986 Yield x s -> do
987 y <- fstep acc x
988 y `seq` 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 postscanlM #-}
997postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
998postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
999 where
1000 {-# INLINE_LATE step' #-}
1001 step' gst (st, acc) = do
1002 r <- step (adaptState gst) st
1003 case r of
1004 Yield x s -> do
1005 y <- fstep acc x
1006 return (Yield y (s, y))
1007 Skip s -> return $ Skip (s, acc)
1008 Stop -> return Stop
1010{-# INLINE_NORMAL postscanl #-}
1011postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
1012postscanl f = postscanlM (\a b -> return (f a b))
1014{-# INLINE_NORMAL scanlM' #-}
1015scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1016scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
1018{-# INLINE scanl' #-}
1019scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
1020scanl' f = scanlM' (\a b -> return (f a b))
1022{-# INLINE_NORMAL scanlM #-}
1023scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1024scanlM fstep begin s = begin `cons` postscanlM fstep begin s
1026{-# INLINE scanl #-}
1027scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
1028scanl f = scanlM (\a b -> return (f a b))
1030{-# INLINE_NORMAL scanl1M #-}
1031scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
1032scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
1033 where
1034 {-# INLINE_LATE step' #-}
1035 step' gst (st, Nothing) = do
1036 r <- step gst st
1037 case r of
1038 Yield x s -> return $ Yield x (s, Just x)
1039 Skip s -> return $ Skip (s, Nothing)
1040 Stop -> return Stop
1042 step' gst (st, Just acc) = do
1043 r <- step gst st
1044 case r of
1045 Yield y s -> do
1046 z <- fstep acc y
1047 return $ Yield z (s, Just z)
1048 Skip s -> return $ Skip (s, Just acc)
1049 Stop -> return Stop
1051{-# INLINE scanl1 #-}
1052scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
1053scanl1 f = scanl1M (\x y -> return (f x y))
1055{-# INLINE_NORMAL scanl1M' #-}
1056scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
1057scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
1058 where
1059 {-# INLINE_LATE step' #-}
1060 step' gst (st, Nothing) = do
1061 r <- step gst st
1062 case r of
1063 Yield x s -> x `seq` return $ Yield x (s, Just x)
1064 Skip s -> return $ Skip (s, Nothing)
1065 Stop -> return Stop
1067 step' gst (st, Just acc) = acc `seq` do
1068 r <- step gst st
1069 case r of
1070 Yield y s -> do
1071 z <- fstep acc y
1072 z `seq` return $ Yield z (s, Just z)
1073 Skip s -> return $ Skip (s, Just acc)
1074 Stop -> return Stop
1076{-# INLINE scanl1' #-}
1077scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
1078scanl1' f = scanl1M' (\x y -> return (f x y))
1080-------------------------------------------------------------------------------
1081-- Filtering
1082-------------------------------------------------------------------------------
1084{-# INLINE_NORMAL take #-}
1085take :: Monad m => Int -> Stream m a -> Stream m a
1086take n (Stream step state) = n `seq` Stream step' (state, 0)
1087 where
1088 {-# INLINE_LATE step' #-}
1089 step' gst (st, i) | i < n = do
1090 r <- step gst st
1091 return $ case r of
1092 Yield x s -> Yield x (s, i + 1)
1093 Skip s -> Skip (s, i)
1094 Stop -> Stop
1095 step' _ (_, _) = return Stop
1097{-# INLINE_NORMAL takeWhileM #-}
1098takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
1099takeWhileM f (Stream step state) = Stream step' state
1100 where
1101 {-# INLINE_LATE step' #-}
1102 step' gst st = do
1103 r <- step gst st
1104 case r of
1105 Yield x s -> do
1106 b <- f x
1107 return $ if b then Yield x s else Stop
1108 Skip s -> return $ Skip s
1109 Stop -> return Stop
1111{-# INLINE takeWhile #-}
1112takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
1113takeWhile f = takeWhileM (return . f)
1115{-# INLINE_NORMAL drop #-}
1116drop :: Monad m => Int -> Stream m a -> Stream m a
1117drop n (Stream step state) = Stream step' (state, Just n)
1118 where
1119 {-# INLINE_LATE step' #-}
1120 step' gst (st, Just i)
1121 | i > 0 = do
1122 r <- step gst st
1123 return $
1124 case r of
1125 Yield _ s -> Skip (s, Just (i - 1))
1126 Skip s -> Skip (s, Just i)
1127 Stop -> Stop
1128 | otherwise = return $ Skip (st, Nothing)
1130 step' gst (st, Nothing) = do
1131 r <- step gst st
1132 return $
1133 case r of
1134 Yield x s -> Yield x (s, Nothing)
1135 Skip s -> Skip (s, Nothing)
1136 Stop -> Stop
1138data DropWhileState s a
1139 = DropWhileDrop s
1140 | DropWhileYield a s
1141 | DropWhileNext s
1143{-# INLINE_NORMAL dropWhileM #-}
1144dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
1145dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
1146 where
1147 {-# INLINE_LATE step' #-}
1148 step' gst (DropWhileDrop st) = do
1149 r <- step gst st
1150 case r of
1151 Yield x s -> do
1152 b <- f x
1153 if b
1154 then return $ Skip (DropWhileDrop s)
1155 else return $ Skip (DropWhileYield x s)
1156 Skip s -> return $ Skip (DropWhileDrop s)
1157 Stop -> return Stop
1159 step' gst (DropWhileNext st) = do
1160 r <- step gst st
1161 case r of
1162 Yield x s -> return $ Skip (DropWhileYield x s)
1163 Skip s -> return $ Skip (DropWhileNext s)
1164 Stop -> return Stop
1166 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
1168{-# INLINE dropWhile #-}
1169dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
1170dropWhile f = dropWhileM (return . f)
1172{-# INLINE_NORMAL filterM #-}
1173filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
1174filterM f (Stream step state) = Stream step' state
1175 where
1176 {-# INLINE_LATE step' #-}
1177 step' gst st = do
1178 r <- step gst st
1179 case r of
1180 Yield x s -> do
1181 b <- f x
1182 return $ if b
1183 then Yield x s
1184 else Skip s
1185 Skip s -> return $ Skip s
1186 Stop -> return Stop
1188{-# INLINE filter #-}
1189filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
1190filter f = filterM (return . f)
1192{-# INLINE_NORMAL uniq #-}
1193uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
1194uniq (Stream step state) = Stream step' (Nothing, state)
1195 where
1196 {-# INLINE_LATE step' #-}
1197 step' gst (Nothing, st) = do
1198 r <- step gst st
1199 case r of
1200 Yield x s -> return $ Yield x (Just x, s)
1201 Skip s -> return $ Skip (Nothing, s)
1202 Stop -> return Stop
1203 step' gst (Just x, st) = do
1204 r <- step gst st
1205 case r of
1206 Yield y s | x == y -> return $ Skip (Just x, s)
1207 | otherwise -> return $ Yield x (Just y, s)
1208 Skip s -> return $ Skip (Just x, s)
1209 Stop -> return Stop
1211------------------------------------------------------------------------------
1212-- Transformation by Mapping
1213------------------------------------------------------------------------------
1215-- | Map a monadic function over a 'Stream'
1216{-# INLINE_NORMAL mapM #-}
1217mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
1218mapM f (Stream step state) = Stream step' state
1219 where
1220 {-# INLINE_LATE step' #-}
1221 step' gst st = do
1222 r <- step (adaptState gst) st
1223 case r of
1224 Yield x s -> f x >>= \a -> return $ Yield a s
1225 Skip s -> return $ Skip s
1226 Stop -> return Stop
1228{-# INLINE map #-}
1229map :: Monad m => (a -> b) -> Stream m a -> Stream m b
1230map f = mapM (return . f)
1232{-# INLINE_NORMAL sequence #-}
1233sequence :: Monad m => Stream m (m a) -> Stream m a
1234sequence (Stream step state) = Stream step' state
1235 where
1236 {-# INLINE_LATE step' #-}
1237 step' gst st = do
1238 r <- step (adaptState gst) st
1239 case r of
1240 Yield x s -> x >>= \a -> return (Yield a s)
1241 Skip s -> return $ Skip s
1242 Stop -> return Stop
1244------------------------------------------------------------------------------
1245-- Inserting
1246------------------------------------------------------------------------------
1248{-# INLINE_NORMAL insertBy #-}
1249insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
1250insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
1251 where
1252 {-# INLINE_LATE step' #-}
1253 step' gst (st, False, _) = do
1254 r <- step gst st
1255 case r of
1256 Yield x s -> case cmp a x of
1257 GT -> return $ Yield x (s, False, Nothing)
1258 _ -> return $ Yield a (s, True, Just x)
1259 Skip s -> return $ Skip (s, False, Nothing)
1260 Stop -> return $ Yield a (st, True, Nothing)
1262 step' _ (_, True, Nothing) = return Stop
1264 step' gst (st, True, Just prev) = do
1265 r <- step gst st
1266 case r of
1267 Yield x s -> return $ Yield prev (s, True, Just x)
1268 Skip s -> return $ Skip (s, True, Just prev)
1269 Stop -> return $ Yield prev (st, True, Nothing)
1271------------------------------------------------------------------------------
1272-- Deleting
1273------------------------------------------------------------------------------
1275{-# INLINE_NORMAL deleteBy #-}
1276deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
1277deleteBy eq x (Stream step state) = Stream step' (state, False)
1278 where
1279 {-# INLINE_LATE step' #-}
1280 step' gst (st, False) = do
1281 r <- step gst st
1282 case r of
1283 Yield y s -> return $
1284 if eq x y then Skip (s, True) else Yield y (s, False)
1285 Skip s -> return $ Skip (s, False)
1286 Stop -> return Stop
1288 step' gst (st, True) = do
1289 r <- step gst st
1290 case r of
1291 Yield y s -> return $ Yield y (s, True)
1292 Skip s -> return $ Skip (s, True)
1293 Stop -> return Stop
1295------------------------------------------------------------------------------
1296-- Transformation by Map and Filter
1297------------------------------------------------------------------------------
1299-- XXX Will this always fuse properly?
1300{-# INLINE_NORMAL mapMaybe #-}
1301mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
1302mapMaybe f = fmap fromJust . filter isJust . map f
1304{-# INLINE_NORMAL mapMaybeM #-}
1305mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
1306mapMaybeM f = fmap fromJust . filter isJust . mapM f
1308------------------------------------------------------------------------------
1309-- Zipping
1310------------------------------------------------------------------------------
1312{-# INLINE_NORMAL indexed #-}
1313indexed :: Monad m => Stream m a -> Stream m (Int, a)
1314indexed (Stream step state) = Stream step' (state, 0)
1315 where
1316 {-# INLINE_LATE step' #-}
1317 step' gst (st, i) = i `seq` do
1318 r <- step (adaptState gst) st
1319 case r of
1320 Yield x s -> return $ Yield (i, x) (s, i+1)
1321 Skip s -> return $ Skip (s, i)
1322 Stop -> return Stop
1324{-# INLINE_NORMAL indexedR #-}
1325indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
1326indexedR m (Stream step state) = Stream step' (state, m)
1327 where
1328 {-# INLINE_LATE step' #-}
1329 step' gst (st, i) = i `seq` do
1330 r <- step (adaptState gst) st
1331 case r of
1332 Yield x s -> let i' = i - 1
1333 in
1334 return $ Yield (i', x) (s, i')
1335 Skip s -> return $ Skip (s, i)
1336 Stop -> return Stop
1338{-# INLINE_NORMAL zipWithM #-}
1339zipWithM :: Monad m
1340 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
1341zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
1342 where
1343 {-# INLINE_LATE step #-}
1344 step gst (sa, sb, Nothing) = do
1345 r <- stepa (adaptState gst) sa
1346 return $
1347 case r of
1348 Yield x sa' -> Skip (sa', sb, Just x)
1349 Skip sa' -> Skip (sa', sb, Nothing)
1350 Stop -> Stop
1352 step gst (sa, sb, Just x) = do
1353 r <- stepb (adaptState gst) sb
1354 case r of
1355 Yield y sb' -> do
1356 z <- f x y
1357 return $ Yield z (sa, sb', Nothing)
1358 Skip sb' -> return $ Skip (sa, sb', Just x)
1359 Stop -> return Stop
1361{-# RULES "zipWithM xs xs"
1362 forall f xs. zipWithM f xs xs = mapM (\x -> f x x) xs #-}
1364{-# INLINE zipWith #-}
1365zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
1366zipWith f = zipWithM (\a b -> return (f a b))
1368------------------------------------------------------------------------------
1369-- Comparisions
1370------------------------------------------------------------------------------
1372{-# INLINE_NORMAL eqBy #-}
1373eqBy :: Monad m => (a -> b -> Bool) -> Stream m a -> Stream m b -> m Bool
1374eqBy eq (Stream step1 t1) (Stream step2 t2) = eq_loop0 SPEC t1 t2
1375 where
1376 eq_loop0 !_ s1 s2 = do
1377 r <- step1 defState s1
1378 case r of
1379 Yield x s1' -> eq_loop1 SPEC x s1' s2
1380 Skip s1' -> eq_loop0 SPEC s1' s2
1381 Stop -> eq_null s2
1383 eq_loop1 !_ x s1 s2 = do
1384 r <- step2 defState s2
1385 case r of
1386 Yield y s2'
1387 | eq x y -> eq_loop0 SPEC s1 s2'
1388 | otherwise -> return False
1389 Skip s2' -> eq_loop1 SPEC x s1 s2'
1390 Stop -> return False
1392 eq_null s2 = do
1393 r <- step2 defState s2
1394 case r of
1395 Yield _ _ -> return False
1396 Skip s2' -> eq_null s2'
1397 Stop -> return True
1399-- | Compare two streams lexicographically
1400{-# INLINE_NORMAL cmpBy #-}
1401cmpBy
1402 :: Monad m
1403 => (a -> b -> Ordering) -> Stream m a -> Stream m b -> m Ordering
1404cmpBy cmp (Stream step1 t1) (Stream step2 t2) = cmp_loop0 SPEC t1 t2
1405 where
1406 cmp_loop0 !_ s1 s2 = do
1407 r <- step1 defState s1
1408 case r of
1409 Yield x s1' -> cmp_loop1 SPEC x s1' s2
1410 Skip s1' -> cmp_loop0 SPEC s1' s2
1411 Stop -> cmp_null s2
1413 cmp_loop1 !_ x s1 s2 = do
1414 r <- step2 defState s2
1415 case r of
1416 Yield y s2' -> case x `cmp` y of
1417 EQ -> cmp_loop0 SPEC s1 s2'
1418 c -> return c
1419 Skip s2' -> cmp_loop1 SPEC x s1 s2'
1420 Stop -> return GT
1422 cmp_null s2 = do
1423 r <- step2 defState s2
1424 case r of
1425 Yield _ _ -> return LT
1426 Skip s2' -> cmp_null s2'
1427 Stop -> return EQ
1429------------------------------------------------------------------------------
1430-- Merging
1431------------------------------------------------------------------------------
1433{-# INLINE_NORMAL mergeByM #-}
1434mergeByM
1435 :: (Monad m)
1436 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
1437mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
1438 Stream step (Just ta, Just tb, Nothing, Nothing)
1439 where
1440 {-# INLINE_LATE step #-}
1442 -- one of the values is missing, and the corresponding stream is running
1443 step gst (Just sa, sb, Nothing, b) = do
1444 r <- stepa gst sa
1445 return $ case r of
1446 Yield a sa' -> Skip (Just sa', sb, Just a, b)
1447 Skip sa' -> Skip (Just sa', sb, Nothing, b)
1448 Stop -> Skip (Nothing, sb, Nothing, b)
1450 step gst (sa, Just sb, a, Nothing) = do
1451 r <- stepb gst sb
1452 return $ case r of
1453 Yield b sb' -> Skip (sa, Just sb', a, Just b)
1454 Skip sb' -> Skip (sa, Just sb', a, Nothing)
1455 Stop -> Skip (sa, Nothing, a, Nothing)
1457 -- both the values are available
1458 step _ (sa, sb, Just a, Just b) = do
1459 res <- cmp a b
1460 return $ case res of
1461 GT -> Yield b (sa, sb, Just a, Nothing)
1462 _ -> Yield a (sa, sb, Nothing, Just b)
1464 -- one of the values is missing, corresponding stream is done
1465 step _ (Nothing, sb, Nothing, Just b) =
1466 return $ Yield b (Nothing, sb, Nothing, Nothing)
1468 step _ (sa, Nothing, Just a, Nothing) =
1469 return $ Yield a (sa, Nothing, Nothing, Nothing)
1471 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
1473{-# INLINE mergeBy #-}
1474mergeBy
1475 :: (Monad m)
1476 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
1477mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
1479------------------------------------------------------------------------------
1480-- Transformation comprehensions
1481------------------------------------------------------------------------------
1483{-# INLINE_NORMAL the #-}
1484the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
1485the (Stream step state) = go state
1486 where
1487 go st = do
1488 r <- step defState st
1489 case r of
1490 Yield x s -> go' x s
1491 Skip s -> go s
1492 Stop -> return Nothing
1493 go' n st = do
1494 r <- step defState st
1495 case r of
1496 Yield x s | x == n -> go' n s
1497 | otherwise -> return Nothing
1498 Skip s -> go' n s
1499 Stop -> return (Just n)
1501------------------------------------------------------------------------------
1502-- Instances
1503------------------------------------------------------------------------------
1505instance Monad m => Functor (Stream m) where
1506 {-# INLINE fmap #-}
1507 fmap = map