git.iko.soy/streamly
on 769ce84 help sign in
streamly/src/Streamly/Streams/StreamD.hs
1489 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 RankNTypes #-}
10#include "inline.hs"
12-- |
13-- Module : Streamly.Streams.StreamD
14-- Copyright : (c) 2018 Harendra Kumar
15-- Copyright : (c) Roman Leshchinskiy 2008-2010
16--
17-- License : BSD3
18-- Maintainer : harendra.kumar@gmail.com
19-- Stability : experimental
20-- Portability : GHC
21--
22-- Direct style re-implementation of CPS style stream in StreamK module. The
23-- symbol or suffix 'D' in this module denotes the "Direct" style. GHC is able
24-- to INLINE and fuse direct style better, providing better performance than
25-- CPS implementation.
26--
27-- @
28-- import qualified Streamly.Streams.StreamD as D
29-- @
31-- Some of the functions in this file have been adapted from the vector
32-- library, https://hackage.haskell.org/package/vector.
34module Streamly.Streams.StreamD
35 (
36 -- * The stream type
37 Step (..)
38 , Stream (..)
40 -- * Construction
41 , nil
42 , cons
44 -- * Deconstruction
45 , uncons
47 -- * Generation
48 -- ** Unfolds
49 , unfoldr
50 , unfoldrM
52 -- ** Specialized Generation
53 -- | Generate a monadic stream from a seed.
54 , repeat
55 , replicate
56 , replicateM
57 , fromIndices
58 , fromIndicesM
59 , generate
60 , generateM
62 -- ** Enumerations
63 , intFromStep
64 , intFrom
65 , intFromThen
66 , intFromTo
67 , intFromThenTo
69 , numFromStep
70 , numFrom
71 , numFromThen
72 , fracFromTo
73 , fracFromThenTo
75 -- ** Conversions
76 -- | Transform an input structure into a stream.
77 -- | Direct style stream does not support @fromFoldable@.
78 , yield
79 , yieldM
80 , fromList
81 , fromListM
82 , fromStreamK
83 , fromStreamD
85 -- * Elimination
86 -- ** General Folds
87 , foldr
88 , foldrM
89 , foldr1
90 , foldl'
91 , foldlM'
93 -- ** Specialized Folds
94 , runStream
95 , null
96 , head
97 , tail
98 , last
99 , elem
100 , notElem
101 , all
102 , any
103 , maximum
104 , maximumBy
105 , minimum
106 , minimumBy
107 , findIndices
108 , lookup
109 , findM
110 , find
111 , (!!)
112 , concatMapM
113 , concatMap
115 -- ** Substreams
116 , isPrefixOf
117 , isSubsequenceOf
118 , stripPrefix
120 -- ** Map and Fold
121 , mapM_
123 -- ** Conversions
124 -- | Transform a stream into another type.
125 , toList
126 , toStreamK
127 , toStreamD
129 -- * Transformation
130 -- ** By folding (scans)
131 , scanlM'
132 , scanl'
133 , scanlM
134 , scanl
135 , scanl1M'
136 , scanl1'
137 , scanl1M
138 , scanl1
140 , prescanl'
141 , prescanlM'
142 , postscanl
143 , postscanlM
144 , postscanl'
145 , postscanlM'
147 -- * Filtering
148 , filter
149 , filterM
150 , uniq
151 , take
152 , takeWhile
153 , takeWhileM
154 , drop
155 , dropWhile
156 , dropWhileM
158 -- * Mapping
159 , map
160 , mapM
161 , sequence
163 -- * Inserting
164 , insertBy
166 -- * Deleting
167 , deleteBy
169 -- ** Map and Filter
170 , mapMaybe
171 , mapMaybeM
173 -- * Zipping
174 , indexed
175 , indexedR
176 , zipWith
177 , zipWithM
179 -- * Comparisions
180 , eqBy
181 , cmpBy
183 -- * Merging
184 , mergeBy
185 , mergeByM
187 -- * Transformation comprehensions
188 , the
189 )
190where
192import Data.Maybe (fromJust, isJust)
193import GHC.Types ( SPEC(..) )
194import Prelude
195 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
196 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
197 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
198 (!!), scanl, scanl1, concatMap, replicate, enumFromTo)
200import Streamly.SVar (MonadAsync, State(..), defState, rstState)
202import qualified Streamly.Streams.StreamK as K
204------------------------------------------------------------------------------
205-- The direct style stream type
206------------------------------------------------------------------------------
208-- | A stream is a succession of 'Step's. A 'Yield' produces a single value and
209-- the next state of the stream. 'Stop' indicates there are no more values in
210-- the stream.
211data Step s a = Yield a s | Skip s | Stop
213instance Functor (Step s) where
214 {-# INLINE fmap #-}
215 fmap f (Yield x s) = Yield (f x) s
216 fmap _ (Skip s) = Skip s
217 fmap _ Stop = Stop
219-- gst = global state
220-- | A stream consists of a step function that generates the next step given a
221-- current state, and the current state.
222data Stream m a = forall s. Stream (State K.Stream m a -> s -> m (Step s a)) s
224------------------------------------------------------------------------------
225-- Construction
226------------------------------------------------------------------------------
228-- | An empty 'Stream'.
229{-# INLINE_NORMAL nil #-}
230nil :: Monad m => Stream m a
231nil = Stream (\_ _ -> return Stop) ()
233-- | Can fuse but has O(n^2) complexity.
234{-# INLINE_NORMAL cons #-}
235cons :: Monad m => a -> Stream m a -> Stream m a
236cons x (Stream step state) = Stream step1 Nothing
237 where
238 {-# INLINE_LATE step1 #-}
239 step1 _ Nothing = return $ Yield x (Just state)
240 step1 gst (Just st) = do
241 r <- step (rstState gst) st
242 return $
243 case r of
244 Yield a s -> Yield a (Just s)
245 Skip s -> Skip (Just s)
246 Stop -> Stop
248-------------------------------------------------------------------------------
249-- Deconstruction
250-------------------------------------------------------------------------------
252-- Does not fuse, has the same performance as the StreamK version.
253{-# INLINE_NORMAL uncons #-}
254uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
255uncons (Stream step state) = go state
256 where
257 go st = do
258 r <- step defState st
259 case r of
260 Yield x s -> return $ Just (x, Stream step s)
261 Skip s -> go s
262 Stop -> return Nothing
264------------------------------------------------------------------------------
265-- Generation by unfold
266------------------------------------------------------------------------------
268{-# INLINE_NORMAL unfoldrM #-}
269unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
270unfoldrM next state = Stream step state
271 where
272 {-# INLINE_LATE step #-}
273 step _ st = do
274 r <- next st
275 return $ case r of
276 Just (x, s) -> Yield x s
277 Nothing -> Stop
279{-# INLINE_LATE unfoldr #-}
280unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
281unfoldr f = unfoldrM (return . f)
283------------------------------------------------------------------------------
284-- Specialized Generation
285------------------------------------------------------------------------------
287repeat :: Monad m => a -> Stream m a
288repeat x = Stream (\_ _ -> return $ Yield x ()) ()
290{-# INLINE_NORMAL replicateM #-}
291replicateM :: Monad m => Int -> m a -> Stream m a
292replicateM n p = Stream step n
293 where
294 {-# INLINE_LATE step #-}
295 step _ i | i <= 0 = return Stop
296 | otherwise = do
297 x <- p
298 return $ Yield x (i - 1)
300{-# INLINE_NORMAL replicate #-}
301replicate :: Monad m => Int -> a -> Stream m a
302replicate n x = replicateM n (return x)
304-- This would not work properly for floats, therefore we put an Integral
305-- constraint.
306-- | Can be used to enumerate unbounded integrals. This does not check for
307-- overflow or underflow for bounded integrals.
308{-# INLINE_NORMAL intFromStep #-}
309intFromStep :: (Integral a, Monad m) => a -> a -> Stream m a
310intFromStep from stride =
311 from `seq` stride `seq` Stream step from
312 where
313 {-# INLINE_LATE step #-}
314 step _ !x = return $ Yield x $! (x + stride)
316-- We are assuming that "to" is constrained by the type to be within
317-- max/min bounds.
318{-# INLINE intFromTo #-}
319intFromTo :: (Monad m, Integral a) => a -> a -> Stream m a
320intFromTo from to = takeWhile (<= to) $ intFromStep from 1
322{-# INLINE intFrom #-}
323intFrom :: (Monad m, Integral a, Bounded a) => a -> Stream m a
324intFrom from = intFromTo from maxBound
326data EnumState a = EnumInit | EnumYield a a a | EnumStop
328{-# INLINE_NORMAL intFromThenToUp #-}
329intFromThenToUp :: (Monad m, Integral a) => a -> a -> a -> Stream m a
330intFromThenToUp from next to = Stream step EnumInit
331 where
332 {-# INLINE_LATE step #-}
333 step _ EnumInit =
334 return $
335 if to < next
336 then if to < from
337 then Stop
338 else Yield from EnumStop
339 else -- from <= next <= to
340 let stride = next - from
341 in Skip $ EnumYield from stride (to - stride)
343 step _ (EnumYield x stride toMinus) =
344 return $
345 if x > toMinus
346 then Yield x EnumStop
347 else Yield x $ EnumYield (x + stride) stride toMinus
349 step _ EnumStop = return Stop
351{-# INLINE_NORMAL intFromThenToDn #-}
352intFromThenToDn :: (Monad m, Integral a) => a -> a -> a -> Stream m a
353intFromThenToDn from next to = Stream step EnumInit
354 where
355 {-# INLINE_LATE step #-}
356 step _ EnumInit =
357 return $ if to > next
358 then if to > from
359 then Stop
360 else Yield from EnumStop
361 else -- from >= next >= to
362 let stride = next - from
363 in Skip $ EnumYield from stride (to - stride)
365 step _ (EnumYield x stride toMinus) =
366 return $
367 if x < toMinus
368 then Yield x EnumStop
369 else Yield x $ EnumYield (x + stride) stride toMinus
371 step _ EnumStop = return Stop
373{-# INLINE_NORMAL intFromThenTo #-}
374intFromThenTo :: (Monad m, Integral a) => a -> a -> a -> Stream m a
375intFromThenTo from next to
376 | next >= from = intFromThenToUp from next to
377 | otherwise = intFromThenToDn from next to
379{-# INLINE_NORMAL intFromThen #-}
380intFromThen :: (Monad m, Integral a, Bounded a) => a -> a -> Stream m a
381intFromThen from next =
382 if next > from
383 then intFromThenToUp from next maxBound
384 else intFromThenToDn from next minBound
386-- For floating point numbers if the increment is less than the precision then
387-- it just gets lost. Therefore we cannot always increment it correctly by just
388-- repeated addition.
389-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
390-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
392-- Instead we accumulate the increment counter and compute the increment
393-- everytime before adding it to the starting number.
394--
395-- This works for Integrals as well as floating point numbers, but intFromStep
396-- is faster for integrals.
397{-# INLINE_NORMAL numFromStep #-}
398numFromStep :: (Monad m, Num a) => a -> a -> Stream m a
399numFromStep from stride = Stream step 0
400 where
401 {-# INLINE_LATE step #-}
402 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
404{-# INLINE_NORMAL numFrom #-}
405numFrom :: (Monad m, Num a) => a -> Stream m a
406numFrom from = numFromStep from 1
408{-# INLINE_NORMAL numFromThen #-}
409numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
410numFromThen from next = numFromStep from (next - from)
412-- We cannot write a general function for Num. The only way to write code
413-- portable between the two is to use a 'Real' constraint and convert between
414-- Fractional and Integral using fromRational which is horribly slow.
415{-# INLINE_NORMAL fracFromTo #-}
416fracFromTo :: (Monad m, Fractional a, Ord a) => a -> a -> Stream m a
417fracFromTo from to = takeWhile (<= to + 1/2) $ numFromStep from 1
419{-# INLINE_NORMAL fracFromThenTo #-}
420fracFromThenTo :: (Monad m, Fractional a, Ord a) => a -> a -> a -> Stream m a
421fracFromThenTo from next to = takeWhile predicate $ numFromThen from next
422 where
423 mid = (next - from) / 2
424 predicate | next >= from = (<= to + mid)
425 | otherwise = (>= to + mid)
427-------------------------------------------------------------------------------
428-- Generation by Conversion
429-------------------------------------------------------------------------------
431-- | Create a singleton 'Stream' from a pure value.
432{-# INLINE_NORMAL yield #-}
433yield :: Monad m => a -> Stream m a
434yield x = Stream (\_ s -> return $ step undefined s) True
435 where
436 {-# INLINE_LATE step #-}
437 step _ True = Yield x False
438 step _ False = Stop
440-- | Create a singleton 'Stream' from a monadic action.
441{-# INLINE_NORMAL yieldM #-}
442yieldM :: Monad m => m a -> Stream m a
443yieldM m = Stream step True
444 where
445 {-# INLINE_LATE step #-}
446 step _ True = m >>= \x -> return $ Yield x False
447 step _ False = return Stop
449{-# INLINE_NORMAL fromIndicesM #-}
450fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
451fromIndicesM gen = Stream step 0
452 where
453 {-# INLINE_LATE step #-}
454 step _ i = do
455 x <- gen i
456 return $ Yield x (i + 1)
458{-# INLINE fromIndices #-}
459fromIndices :: Monad m => (Int -> a) -> Stream m a
460fromIndices gen = fromIndicesM (return . gen)
462{-# INLINE_NORMAL generateM #-}
463generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
464generateM n gen = n `seq` Stream step 0
465 where
466 {-# INLINE_LATE step #-}
467 step _ i | i < n = do
468 x <- gen i
469 return $ Yield x (i + 1)
470 | otherwise = return Stop
472{-# INLINE generate #-}
473generate :: Monad m => Int -> (Int -> a) -> Stream m a
474generate n gen = generateM n (return . gen)
476-- XXX we need the MonadAsync constraint because of a rewrite rule.
477-- | Convert a list of monadic actions to a 'Stream'
478{-# INLINE_LATE fromListM #-}
479fromListM :: MonadAsync m => [m a] -> Stream m a
480fromListM = Stream step
481 where
482 {-# INLINE_LATE step #-}
483 step _ (m:ms) = m >>= \x -> return $ Yield x ms
484 step _ [] = return Stop
486-- | Convert a list of pure values to a 'Stream'
487{-# INLINE_LATE fromList #-}
488fromList :: Monad m => [a] -> Stream m a
489fromList = Stream step
490 where
491 {-# INLINE_LATE step #-}
492 step _ (x:xs) = return $ Yield x xs
493 step _ [] = return Stop
495-- XXX pass the state to streamD
496{-# INLINE_LATE fromStreamK #-}
497fromStreamK :: Monad m => K.Stream m a -> Stream m a
498fromStreamK = Stream step
499 where
500 step gst m1 =
501 let stop = return Stop
502 single a = return $ Yield a K.nil
503 yieldk a r = return $ Yield a r
504 in K.unStreamShared m1 gst stop single yieldk
506{-# INLINE toStreamD #-}
507toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
508toStreamD = fromStreamK . K.toStream
510------------------------------------------------------------------------------
511-- Elimination by Folds
512------------------------------------------------------------------------------
514{-# INLINE_NORMAL foldrM #-}
515foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b
516foldrM f z (Stream step state) = go SPEC state
517 where
518 go !_ st = do
519 r <- step defState st
520 case r of
521 Yield x s -> go SPEC s >>= f x
522 Skip s -> go SPEC s
523 Stop -> return z
525{-# INLINE_NORMAL foldr #-}
526foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
527foldr f = foldrM (\a b -> return (f a b))
529{-# INLINE_NORMAL foldr1 #-}
530foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
531foldr1 f m = do
532 r <- uncons m
533 case r of
534 Nothing -> return Nothing
535 Just (h, t) -> fmap Just (foldr f h t)
537{-# INLINE_NORMAL foldlM' #-}
538foldlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> m b
539foldlM' fstep begin (Stream step state) = go SPEC begin state
540 where
541 go !_ acc st = acc `seq` do
542 r <- step defState st
543 case r of
544 Yield x s -> do
545 acc' <- fstep acc x
546 go SPEC acc' s
547 Skip s -> go SPEC acc s
548 Stop -> return acc
550{-# INLINE foldl' #-}
551foldl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
552foldl' fstep = foldlM' (\b a -> return (fstep b a))
554------------------------------------------------------------------------------
555-- Specialized Folds
556------------------------------------------------------------------------------
558-- | Run a streaming composition, discard the results.
559{-# INLINE_LATE runStream #-}
560runStream :: Monad m => Stream m a -> m ()
561runStream (Stream step state) = go SPEC state
562 where
563 go !_ st = do
564 r <- step defState st
565 case r of
566 Yield _ s -> go SPEC s
567 Skip s -> go SPEC s
568 Stop -> return ()
570{-# INLINE_NORMAL null #-}
571null :: Monad m => Stream m a -> m Bool
572null (Stream step state) = go state
573 where
574 go st = do
575 r <- step defState st
576 case r of
577 Yield _ _ -> return False
578 Skip s -> go s
579 Stop -> return True
581-- XXX SPEC?
582{-# INLINE_NORMAL head #-}
583head :: Monad m => Stream m a -> m (Maybe a)
584head (Stream step state) = go state
585 where
586 go st = do
587 r <- step defState st
588 case r of
589 Yield x _ -> return (Just x)
590 Skip s -> go s
591 Stop -> return Nothing
593-- Does not fuse, has the same performance as the StreamK version.
594{-# INLINE_NORMAL tail #-}
595tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
596tail (Stream step state) = go state
597 where
598 go st = do
599 r <- step defState st
600 case r of
601 Yield _ s -> return (Just $ Stream step s)
602 Skip s -> go s
603 Stop -> return Nothing
605-- XXX will it fuse? need custom impl?
606{-# INLINE_NORMAL last #-}
607last :: Monad m => Stream m a -> m (Maybe a)
608last = foldl' (\_ y -> Just y) Nothing
610{-# INLINE_NORMAL elem #-}
611elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
612elem e (Stream step state) = go state
613 where
614 go st = do
615 r <- step defState st
616 case r of
617 Yield x s
618 | x == e -> return True
619 | otherwise -> go s
620 Skip s -> go s
621 Stop -> return False
623{-# INLINE_NORMAL notElem #-}
624notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
625notElem e s = fmap not (elem e s)
627{-# INLINE_NORMAL all #-}
628all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
629all p (Stream step state) = go state
630 where
631 go st = do
632 r <- step defState st
633 case r of
634 Yield x s
635 | p x -> go s
636 | otherwise -> return False
637 Skip s -> go s
638 Stop -> return True
640{-# INLINE_NORMAL any #-}
641any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
642any p (Stream step state) = go state
643 where
644 go st = do
645 r <- step defState st
646 case r of
647 Yield x s
648 | p x -> return True
649 | otherwise -> go s
650 Skip s -> go s
651 Stop -> return False
653{-# INLINE_NORMAL maximum #-}
654maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
655maximum (Stream step state) = go Nothing state
656 where
657 go Nothing st = do
658 r <- step defState st
659 case r of
660 Yield x s -> go (Just x) s
661 Skip s -> go Nothing s
662 Stop -> return Nothing
663 go (Just acc) st = do
664 r <- step defState st
665 case r of
666 Yield x s
667 | acc <= x -> go (Just x) s
668 | otherwise -> go (Just acc) s
669 Skip s -> go (Just acc) s
670 Stop -> return (Just acc)
672{-# INLINE_NORMAL maximumBy #-}
673maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
674maximumBy cmp (Stream step state) = go Nothing state
675 where
676 go Nothing st = do
677 r <- step defState st
678 case r of
679 Yield x s -> go (Just x) s
680 Skip s -> go Nothing s
681 Stop -> return Nothing
682 go (Just acc) st = do
683 r <- step defState st
684 case r of
685 Yield x s -> case cmp acc x of
686 GT -> go (Just acc) s
687 _ -> go (Just x) s
688 Skip s -> go (Just acc) s
689 Stop -> return (Just acc)
691{-# INLINE_NORMAL minimum #-}
692minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
693minimum (Stream step state) = go Nothing state
694 where
695 go Nothing st = do
696 r <- step defState st
697 case r of
698 Yield x s -> go (Just x) s
699 Skip s -> go Nothing s
700 Stop -> return Nothing
701 go (Just acc) st = do
702 r <- step defState st
703 case r of
704 Yield x s
705 | acc <= x -> go (Just acc) s
706 | otherwise -> go (Just x) s
707 Skip s -> go (Just acc) s
708 Stop -> return (Just acc)
710{-# INLINE_NORMAL minimumBy #-}
711minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
712minimumBy cmp (Stream step state) = go Nothing state
713 where
714 go Nothing st = do
715 r <- step defState st
716 case r of
717 Yield x s -> go (Just x) s
718 Skip s -> go Nothing s
719 Stop -> return Nothing
720 go (Just acc) st = do
721 r <- step defState st
722 case r of
723 Yield x s -> case cmp acc x of
724 GT -> go (Just x) s
725 _ -> go (Just acc) s
726 Skip s -> go (Just acc) s
727 Stop -> return (Just acc)
729{-# INLINE_NORMAL (!!) #-}
730(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
731(Stream step state) !! i = go i state
732 where
733 go n st = do
734 r <- step defState st
735 case r of
736 Yield x s | n < 0 -> return Nothing
737 | n == 0 -> return $ Just x
738 | otherwise -> go (n - 1) s
739 Skip s -> go n s
740 Stop -> return Nothing
742{-# INLINE_NORMAL lookup #-}
743lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
744lookup e (Stream step state) = go state
745 where
746 go st = do
747 r <- step defState st
748 case r of
749 Yield (a, b) s -> if e == a then return (Just b) else go s
750 Skip s -> go s
751 Stop -> return Nothing
753{-# INLINE_NORMAL findM #-}
754findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
755findM p (Stream step state) = go SPEC state
756 where
757 go !_ st = do
758 r <- step defState st
759 case r of
760 Yield x s -> do
761 b <- p x
762 if b then return (Just x) else go SPEC s
763 Skip s -> go SPEC s
764 Stop -> return Nothing
766{-# INLINE find #-}
767find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
768find p = findM (return . p)
770{-# INLINE_NORMAL findIndices #-}
771findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
772findIndices p (Stream step state) = Stream step' (state, 0)
773 where
774 {-# INLINE_LATE step' #-}
775 step' gst (st, i) = do
776 r <- step (rstState gst) st
777 return $ case r of
778 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
779 Skip s -> Skip (s, i+1)
780 Stop -> Stop
782{-# INLINE_NORMAL concatMapM #-}
783concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b
784concatMapM f (Stream step state) = Stream step' (Left state)
785 where
786 {-# INLINE_LATE step' #-}
787 step' gst (Left st) = do
788 r <- step (rstState gst) st
789 case r of
790 Yield a s -> do
791 b_stream <- f a
792 return $ Skip (Right (b_stream, s))
793 Skip s -> return $ Skip (Left s)
794 Stop -> return Stop
796 step' _ (Right (Stream 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 stp sng yld -> do
901 r <- step gst st
902 case r of
903 Yield x s -> yld x (go s)
904 Skip s -> K.unStreamShared (go s) gst stp sng yld
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 (rstState 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 (rstState 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 (rstState 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 (rstState 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 (rstState 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 (rstState 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 (rstState 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 (rstState 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 (rstState 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 (rstState 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 (rstState 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 (rstState 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 (rstState 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 (rstState 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 (rstState 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 (rstState 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-- | Map a monadic function over a 'Stream'
1198{-# INLINE_NORMAL mapM #-}
1199mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
1200mapM f (Stream step state) = Stream step' state
1201 where
1202 {-# INLINE_LATE step' #-}
1203 step' gst st = do
1204 r <- step (rstState gst) st
1205 case r of
1206 Yield x s -> f x >>= \a -> return $ Yield a s
1207 Skip s -> return $ Skip s
1208 Stop -> return Stop
1210{-# INLINE map #-}
1211map :: Monad m => (a -> b) -> Stream m a -> Stream m b
1212map f = mapM (return . f)
1214{-# INLINE_NORMAL sequence #-}
1215sequence :: Monad m => Stream m (m a) -> Stream m a
1216sequence (Stream step state) = Stream step' state
1217 where
1218 {-# INLINE_LATE step' #-}
1219 step' gst st = do
1220 r <- step (rstState gst) st
1221 case r of
1222 Yield x s -> x >>= \a -> return (Yield a s)
1223 Skip s -> return $ Skip s
1224 Stop -> return Stop
1226------------------------------------------------------------------------------
1227-- Inserting
1228------------------------------------------------------------------------------
1230{-# INLINE_NORMAL insertBy #-}
1231insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
1232insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
1233 where
1234 {-# INLINE_LATE step' #-}
1235 step' gst (st, False, _) = do
1236 r <- step (rstState gst) st
1237 case r of
1238 Yield x s -> case cmp a x of
1239 GT -> return $ Yield x (s, False, Nothing)
1240 _ -> return $ Yield a (s, True, Just x)
1241 Skip s -> return $ Skip (s, False, Nothing)
1242 Stop -> return $ Yield a (st, True, Nothing)
1244 step' _ (_, True, Nothing) = return Stop
1246 step' gst (st, True, Just prev) = do
1247 r <- step (rstState gst) st
1248 case r of
1249 Yield x s -> return $ Yield prev (s, True, Just x)
1250 Skip s -> return $ Skip (s, True, Just prev)
1251 Stop -> return $ Yield prev (st, True, Nothing)
1253------------------------------------------------------------------------------
1254-- Deleting
1255------------------------------------------------------------------------------
1257{-# INLINE_NORMAL deleteBy #-}
1258deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
1259deleteBy eq x (Stream step state) = Stream step' (state, False)
1260 where
1261 {-# INLINE_LATE step' #-}
1262 step' gst (st, False) = do
1263 r <- step (rstState gst) st
1264 case r of
1265 Yield y s -> return $
1266 if eq x y then Skip (s, True) else Yield y (s, False)
1267 Skip s -> return $ Skip (s, False)
1268 Stop -> return Stop
1270 step' gst (st, True) = do
1271 r <- step (rstState gst) st
1272 case r of
1273 Yield y s -> return $ Yield y (s, True)
1274 Skip s -> return $ Skip (s, True)
1275 Stop -> return Stop
1277------------------------------------------------------------------------------
1278-- Transformation by Map and Filter
1279------------------------------------------------------------------------------
1281-- XXX Will this always fuse properly?
1282{-# INLINE_NORMAL mapMaybe #-}
1283mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
1284mapMaybe f = fmap fromJust . filter isJust . map f
1286{-# INLINE_NORMAL mapMaybeM #-}
1287mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
1288mapMaybeM f = fmap fromJust . filter isJust . mapM f
1290------------------------------------------------------------------------------
1291-- Zipping
1292------------------------------------------------------------------------------
1294{-# INLINE_NORMAL indexed #-}
1295indexed :: Monad m => Stream m a -> Stream m (Int, a)
1296indexed (Stream step state) = Stream step' (state, 0)
1297 where
1298 {-# INLINE_LATE step' #-}
1299 step' gst (st, i) = i `seq` do
1300 r <- step (rstState gst) st
1301 case r of
1302 Yield x s -> return $ Yield (i, x) (s, i+1)
1303 Skip s -> return $ Skip (s, i)
1304 Stop -> return Stop
1306{-# INLINE_NORMAL indexedR #-}
1307indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
1308indexedR m (Stream step state) = Stream step' (state, m)
1309 where
1310 {-# INLINE_LATE step' #-}
1311 step' gst (st, i) = i `seq` do
1312 r <- step (rstState gst) st
1313 case r of
1314 Yield x s -> let i' = i - 1
1315 in
1316 return $ Yield (i', x) (s, i')
1317 Skip s -> return $ Skip (s, i)
1318 Stop -> return Stop
1320{-# INLINE_NORMAL zipWithM #-}
1321zipWithM :: Monad m
1322 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
1323zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
1324 where
1325 {-# INLINE_LATE step #-}
1326 step gst (sa, sb, Nothing) = do
1327 r <- stepa (rstState gst) sa
1328 return $
1329 case r of
1330 Yield x sa' -> Skip (sa', sb, Just x)
1331 Skip sa' -> Skip (sa', sb, Nothing)
1332 Stop -> Stop
1334 step gst (sa, sb, Just x) = do
1335 r <- stepb (rstState gst) sb
1336 case r of
1337 Yield y sb' -> do
1338 z <- f x y
1339 return $ Yield z (sa, sb', Nothing)
1340 Skip sb' -> return $ Skip (sa, sb', Just x)
1341 Stop -> return Stop
1343{-# RULES "zipWithM xs xs"
1344 forall f xs. zipWithM f xs xs = mapM (\x -> f x x) xs #-}
1346{-# INLINE zipWith #-}
1347zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
1348zipWith f = zipWithM (\a b -> return (f a b))
1350------------------------------------------------------------------------------
1351-- Comparisions
1352------------------------------------------------------------------------------
1354{-# INLINE_NORMAL eqBy #-}
1355eqBy :: Monad m => (a -> b -> Bool) -> Stream m a -> Stream m b -> m Bool
1356eqBy eq (Stream step1 t1) (Stream step2 t2) = eq_loop0 SPEC t1 t2
1357 where
1358 eq_loop0 !_ s1 s2 = do
1359 r <- step1 defState s1
1360 case r of
1361 Yield x s1' -> eq_loop1 SPEC x s1' s2
1362 Skip s1' -> eq_loop0 SPEC s1' s2
1363 Stop -> eq_null s2
1365 eq_loop1 !_ x s1 s2 = do
1366 r <- step2 defState s2
1367 case r of
1368 Yield y s2'
1369 | eq x y -> eq_loop0 SPEC s1 s2'
1370 | otherwise -> return False
1371 Skip s2' -> eq_loop1 SPEC x s1 s2'
1372 Stop -> return False
1374 eq_null s2 = do
1375 r <- step2 defState s2
1376 case r of
1377 Yield _ _ -> return False
1378 Skip s2' -> eq_null s2'
1379 Stop -> return True
1381-- | Compare two streams lexicographically
1382{-# INLINE_NORMAL cmpBy #-}
1383cmpBy
1384 :: Monad m
1385 => (a -> b -> Ordering) -> Stream m a -> Stream m b -> m Ordering
1386cmpBy cmp (Stream step1 t1) (Stream step2 t2) = cmp_loop0 SPEC t1 t2
1387 where
1388 cmp_loop0 !_ s1 s2 = do
1389 r <- step1 defState s1
1390 case r of
1391 Yield x s1' -> cmp_loop1 SPEC x s1' s2
1392 Skip s1' -> cmp_loop0 SPEC s1' s2
1393 Stop -> cmp_null s2
1395 cmp_loop1 !_ x s1 s2 = do
1396 r <- step2 defState s2
1397 case r of
1398 Yield y s2' -> case x `cmp` y of
1399 EQ -> cmp_loop0 SPEC s1 s2'
1400 c -> return c
1401 Skip s2' -> cmp_loop1 SPEC x s1 s2'
1402 Stop -> return GT
1404 cmp_null s2 = do
1405 r <- step2 defState s2
1406 case r of
1407 Yield _ _ -> return LT
1408 Skip s2' -> cmp_null s2'
1409 Stop -> return EQ
1411------------------------------------------------------------------------------
1412-- Merging
1413------------------------------------------------------------------------------
1415{-# INLINE_NORMAL mergeByM #-}
1416mergeByM
1417 :: (Monad m)
1418 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
1419mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
1420 Stream step (Just ta, Just tb, Nothing, Nothing)
1421 where
1422 {-# INLINE_LATE step #-}
1424 -- one of the values is missing, and the corresponding stream is running
1425 step gst (Just sa, sb, Nothing, b) = do
1426 r <- stepa (rstState gst) sa
1427 return $ case r of
1428 Yield a sa' -> Skip (Just sa', sb, Just a, b)
1429 Skip sa' -> Skip (Just sa', sb, Nothing, b)
1430 Stop -> Skip (Nothing, sb, Nothing, b)
1432 step gst (sa, Just sb, a, Nothing) = do
1433 r <- stepb (rstState gst) sb
1434 return $ case r of
1435 Yield b sb' -> Skip (sa, Just sb', a, Just b)
1436 Skip sb' -> Skip (sa, Just sb', a, Nothing)
1437 Stop -> Skip (sa, Nothing, a, Nothing)
1439 -- both the values are available
1440 step _ (sa, sb, Just a, Just b) = do
1441 res <- cmp a b
1442 return $ case res of
1443 GT -> Yield b (sa, sb, Just a, Nothing)
1444 _ -> Yield a (sa, sb, Nothing, Just b)
1446 -- one of the values is missing, corresponding stream is done
1447 step _ (Nothing, sb, Nothing, Just b) =
1448 return $ Yield b (Nothing, sb, Nothing, Nothing)
1450 step _ (sa, Nothing, Just a, Nothing) =
1451 return $ Yield a (sa, Nothing, Nothing, Nothing)
1453 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
1455{-# INLINE mergeBy #-}
1456mergeBy
1457 :: (Monad m)
1458 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
1459mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
1461------------------------------------------------------------------------------
1462-- Transformation comprehensions
1463------------------------------------------------------------------------------
1465{-# INLINE_NORMAL the #-}
1466the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
1467the (Stream step state) = go state
1468 where
1469 go st = do
1470 r <- step defState st
1471 case r of
1472 Yield x s -> go' x s
1473 Skip s -> go s
1474 Stop -> return Nothing
1475 go' n st = do
1476 r <- step defState st
1477 case r of
1478 Yield x s | x == n -> go' n s
1479 | otherwise -> return Nothing
1480 Skip s -> go' n s
1481 Stop -> return (Just n)
1483------------------------------------------------------------------------------
1484-- Instances
1485------------------------------------------------------------------------------
1487instance Monad m => Functor (Stream m) where
1488 {-# INLINE fmap #-}
1489 fmap = map