git.iko.soy/streamly
on 9ccee33 help sign in
streamly/src/Streamly/Streams/StreamD.hs
1167 lines 37 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--
16-- License : BSD3
17-- Maintainer : harendra.kumar@gmail.com
18-- Stability : experimental
19-- Portability : GHC
20--
21-- Direct style re-implementation of CPS style stream in StreamK module. The
22-- symbol or suffix 'D' in this module denotes the "Direct" style. GHC is able
23-- to INLINE and fuse direct style better, providing better performance than
24-- CPS implementation.
25--
26-- @
27-- import qualified Streamly.Streams.StreamD as D
28-- @
30-- Some of functions in this file have been adapted from the vector
31-- library, https://hackage.haskell.org/package/vector.
33module Streamly.Streams.StreamD
34 (
35 -- * The stream type
36 Step (..)
37 , Stream (..)
39 -- * Construction
40 , nil
41 , cons
43 -- * Deconstruction
44 , uncons
46 -- * Generation
47 -- ** Unfolds
48 , unfoldr
49 , unfoldrM
51 -- ** Specialized Generation
52 -- | Generate a monadic stream from a seed.
53 , repeat
54 , enumFromStepN
56 -- ** Conversions
57 -- | Transform an input structure into a stream.
58 -- | Direct style stream does not support @fromFoldable@.
59 , yield
60 , yieldM
61 , fromList
62 , fromListM
63 , fromStreamK
64 , fromStreamD
66 -- * Elimination
67 -- ** General Folds
68 , foldr
69 , foldrM
70 , foldr1
71 , foldl'
72 , foldlM'
74 -- ** Specialized Folds
75 , runStream
76 , null
77 , head
78 , tail
79 , last
80 , elem
81 , notElem
82 , all
83 , any
84 , maximum
85 , maximumBy
86 , minimum
87 , minimumBy
88 , findIndices
89 , lookup
90 , find
91 , (!!)
93 -- ** Substreams
94 , isPrefixOf
95 , isSubsequenceOf
96 , stripPrefix
98 -- ** Map and Fold
99 , mapM_
101 -- ** Conversions
102 -- | Transform a stream into another type.
103 , toList
104 , toStreamK
105 , toStreamD
107 -- * Transformation
108 -- ** By folding (scans)
109 , scanlM'
110 , scanl'
111 , scanlM
112 , scanl
113 , scanl1M'
114 , scanl1'
115 , scanl1M
116 , scanl1
118 -- * Filtering
119 , filter
120 , filterM
121 , take
122 , takeWhile
123 , takeWhileM
124 , drop
125 , dropWhile
126 , dropWhileM
128 -- * Mapping
129 , map
130 , mapM
131 , sequence
133 -- * Inserting
134 , insertBy
136 -- * Deleting
137 , deleteBy
139 -- ** Map and Filter
140 , mapMaybe
141 , mapMaybeM
143 -- * Zipping
144 , zipWith
145 , zipWithM
147 -- * Comparisions
148 , eqBy
149 , cmpBy
151 -- * Merging
152 , mergeBy
154 -- * Transformation comprehensions
155 , the
156 )
157where
159import Data.Maybe (fromJust, isJust)
160import GHC.Types ( SPEC(..) )
161import Prelude
162 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
163 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
164 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
165 (!!), scanl, scanl1)
167import Streamly.SVar (MonadAsync, State(..), defState, rstState)
169import qualified Streamly.Streams.StreamK as K
171------------------------------------------------------------------------------
172-- The direct style stream type
173------------------------------------------------------------------------------
175-- | A stream is a succession of 'Step's. A 'Yield' produces a single value and
176-- the next state of the stream. 'Stop' indicates there are no more values in
177-- the stream.
178data Step s a = Yield a s | Skip s | Stop
180instance Functor (Step s) where
181 {-# INLINE fmap #-}
182 fmap f (Yield x s) = Yield (f x) s
183 fmap _ (Skip s) = Skip s
184 fmap _ Stop = Stop
186-- gst = global state
187-- | A stream consists of a step function that generates the next step given a
188-- current state, and the current state.
189data Stream m a = forall s. Stream (State K.Stream m a -> s -> m (Step s a)) s
191------------------------------------------------------------------------------
192-- Construction
193------------------------------------------------------------------------------
195-- | An empty 'Stream'.
196{-# INLINE_NORMAL nil #-}
197nil :: Monad m => Stream m a
198nil = Stream (\_ _ -> return Stop) ()
200-- | Can fuse but has O(n^2) complexity.
201{-# INLINE_NORMAL cons #-}
202cons :: Monad m => a -> Stream m a -> Stream m a
203cons x (Stream step state) = Stream step1 Nothing
204 where
205 {-# INLINE_LATE step1 #-}
206 step1 _ Nothing = return $ Yield x (Just state)
207 step1 gst (Just st) = do
208 r <- step (rstState gst) st
209 return $
210 case r of
211 Yield a s -> Yield a (Just s)
212 Skip s -> Skip (Just s)
213 Stop -> Stop
215-------------------------------------------------------------------------------
216-- Deconstruction
217-------------------------------------------------------------------------------
219-- Does not fuse, has the same performance as the StreamK version.
220{-# INLINE_NORMAL uncons #-}
221uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
222uncons (Stream step state) = go state
223 where
224 go st = do
225 r <- step defState st
226 case r of
227 Yield x s -> return $ Just (x, Stream step s)
228 Skip s -> go s
229 Stop -> return Nothing
231------------------------------------------------------------------------------
232-- Generation by unfold
233------------------------------------------------------------------------------
235{-# INLINE_NORMAL unfoldrM #-}
236unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
237unfoldrM next state = Stream step state
238 where
239 {-# INLINE_LATE step #-}
240 step _ st = do
241 r <- next st
242 return $ case r of
243 Just (x, s) -> Yield x s
244 Nothing -> Stop
246{-# INLINE_LATE unfoldr #-}
247unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
248unfoldr f = unfoldrM (return . f)
250------------------------------------------------------------------------------
251-- Specialized Generation
252------------------------------------------------------------------------------
254repeat :: Monad m => a -> Stream m a
255repeat x = Stream (\_ _ -> return $ Yield x ()) ()
257{-# INLINE_NORMAL enumFromStepN #-}
258enumFromStepN :: (Num a, Monad m) => a -> a -> Int -> Stream m a
259enumFromStepN from stride n =
260 from `seq` stride `seq` n `seq` Stream step (from, n)
261 where
262 {-# INLINE_LATE step #-}
263 step _ (x, i) | i > 0 = return $ Yield x (x + stride, i - 1)
264 | otherwise = return Stop
266-------------------------------------------------------------------------------
267-- Generation by Conversion
268-------------------------------------------------------------------------------
270-- | Create a singleton 'Stream' from a pure value.
271{-# INLINE_NORMAL yield #-}
272yield :: Monad m => a -> Stream m a
273yield x = Stream (\_ s -> return $ step undefined s) True
274 where
275 {-# INLINE_LATE step #-}
276 step _ True = Yield x False
277 step _ False = Stop
279-- | Create a singleton 'Stream' from a monadic action.
280{-# INLINE_NORMAL yieldM #-}
281yieldM :: Monad m => m a -> Stream m a
282yieldM m = Stream step True
283 where
284 {-# INLINE_LATE step #-}
285 step _ True = m >>= \x -> return $ Yield x False
286 step _ False = return Stop
288-- XXX we need the MonadAsync constraint because of a rewrite rule.
289-- | Convert a list of monadic actions to a 'Stream'
290{-# INLINE_LATE fromListM #-}
291fromListM :: MonadAsync m => [m a] -> Stream m a
292fromListM = Stream step
293 where
294 {-# INLINE_LATE step #-}
295 step _ (m:ms) = m >>= \x -> return $ Yield x ms
296 step _ [] = return Stop
298-- | Convert a list of pure values to a 'Stream'
299{-# INLINE_LATE fromList #-}
300fromList :: Monad m => [a] -> Stream m a
301fromList = Stream step
302 where
303 {-# INLINE_LATE step #-}
304 step _ (x:xs) = return $ Yield x xs
305 step _ [] = return Stop
307-- XXX pass the state to streamD
308{-# INLINE_LATE fromStreamK #-}
309fromStreamK :: Monad m => K.Stream m a -> Stream m a
310fromStreamK = Stream step
311 where
312 step gst m1 =
313 let stop = return Stop
314 single a = return $ Yield a K.nil
315 yieldk a r = return $ Yield a r
316 in K.unStream m1 gst stop single yieldk
318{-# INLINE toStreamD #-}
319toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
320toStreamD = fromStreamK . K.toStream
322------------------------------------------------------------------------------
323-- Elimination by Folds
324------------------------------------------------------------------------------
326{-# INLINE_NORMAL foldrM #-}
327foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b
328foldrM f z (Stream step state) = go SPEC state
329 where
330 go !_ st = do
331 r <- step defState st
332 case r of
333 Yield x s -> go SPEC s >>= f x
334 Skip s -> go SPEC s
335 Stop -> return z
337{-# INLINE_NORMAL foldr #-}
338foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
339foldr f = foldrM (\a b -> return (f a b))
341{-# INLINE_NORMAL foldr1 #-}
342foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
343foldr1 f m = do
344 r <- uncons m
345 case r of
346 Nothing -> return Nothing
347 Just (h, t) -> fmap Just (foldr f h t)
349{-# INLINE_NORMAL foldlM' #-}
350foldlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> m b
351foldlM' fstep begin (Stream step state) = go SPEC begin state
352 where
353 go !_ acc st = acc `seq` do
354 r <- step defState st
355 case r of
356 Yield x s -> do
357 acc' <- fstep acc x
358 go SPEC acc' s
359 Skip s -> go SPEC acc s
360 Stop -> return acc
362{-# INLINE foldl' #-}
363foldl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
364foldl' fstep = foldlM' (\b a -> return (fstep b a))
366------------------------------------------------------------------------------
367-- Specialized Folds
368------------------------------------------------------------------------------
370-- | Run a streaming composition, discard the results.
371{-# INLINE_LATE runStream #-}
372runStream :: Monad m => Stream m a -> m ()
373runStream (Stream step state) = go SPEC state
374 where
375 go !_ st = do
376 r <- step defState st
377 case r of
378 Yield _ s -> go SPEC s
379 Skip s -> go SPEC s
380 Stop -> return ()
382{-# INLINE_NORMAL null #-}
383null :: Monad m => Stream m a -> m Bool
384null (Stream step state) = go state
385 where
386 go st = do
387 r <- step defState st
388 case r of
389 Yield _ _ -> return False
390 Skip s -> go s
391 Stop -> return True
393-- XXX SPEC?
394{-# INLINE_NORMAL head #-}
395head :: Monad m => Stream m a -> m (Maybe a)
396head (Stream step state) = go state
397 where
398 go st = do
399 r <- step defState st
400 case r of
401 Yield x _ -> return (Just x)
402 Skip s -> go s
403 Stop -> return Nothing
405-- Does not fuse, has the same performance as the StreamK version.
406{-# INLINE_NORMAL tail #-}
407tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
408tail (Stream step state) = go state
409 where
410 go st = do
411 r <- step defState st
412 case r of
413 Yield _ s -> return (Just $ Stream step s)
414 Skip s -> go s
415 Stop -> return Nothing
417-- XXX will it fuse? need custom impl?
418{-# INLINE_NORMAL last #-}
419last :: Monad m => Stream m a -> m (Maybe a)
420last = foldl' (\_ y -> Just y) Nothing
422{-# INLINE_NORMAL elem #-}
423elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
424elem e (Stream step state) = go state
425 where
426 go st = do
427 r <- step defState st
428 case r of
429 Yield x s
430 | x == e -> return True
431 | otherwise -> go s
432 Skip s -> go s
433 Stop -> return False
435{-# INLINE_NORMAL notElem #-}
436notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
437notElem e s = fmap not (elem e s)
439{-# INLINE_NORMAL all #-}
440all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
441all p (Stream step state) = go state
442 where
443 go st = do
444 r <- step defState st
445 case r of
446 Yield x s
447 | p x -> go s
448 | otherwise -> return False
449 Skip s -> go s
450 Stop -> return True
452{-# INLINE_NORMAL any #-}
453any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
454any p (Stream step state) = go state
455 where
456 go st = do
457 r <- step defState st
458 case r of
459 Yield x s
460 | p x -> return True
461 | otherwise -> go s
462 Skip s -> go s
463 Stop -> return False
465{-# INLINE_NORMAL maximum #-}
466maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
467maximum (Stream step state) = go Nothing state
468 where
469 go Nothing st = do
470 r <- step defState st
471 case r of
472 Yield x s -> go (Just x) s
473 Skip s -> go Nothing s
474 Stop -> return Nothing
475 go (Just acc) st = do
476 r <- step defState st
477 case r of
478 Yield x s
479 | acc <= x -> go (Just x) s
480 | otherwise -> go (Just acc) s
481 Skip s -> go (Just acc) s
482 Stop -> return (Just acc)
484{-# INLINE_NORMAL maximumBy #-}
485maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
486maximumBy cmp (Stream step state) = go Nothing state
487 where
488 go Nothing st = do
489 r <- step defState st
490 case r of
491 Yield x s -> go (Just x) s
492 Skip s -> go Nothing s
493 Stop -> return Nothing
494 go (Just acc) st = do
495 r <- step defState st
496 case r of
497 Yield x s -> case cmp acc x of
498 GT -> go (Just acc) s
499 _ -> go (Just x) s
500 Skip s -> go (Just acc) s
501 Stop -> return (Just acc)
503{-# INLINE_NORMAL minimum #-}
504minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
505minimum (Stream step state) = go Nothing state
506 where
507 go Nothing st = do
508 r <- step defState st
509 case r of
510 Yield x s -> go (Just x) s
511 Skip s -> go Nothing s
512 Stop -> return Nothing
513 go (Just acc) st = do
514 r <- step defState st
515 case r of
516 Yield x s
517 | acc <= x -> go (Just acc) s
518 | otherwise -> go (Just x) s
519 Skip s -> go (Just acc) s
520 Stop -> return (Just acc)
522{-# INLINE_NORMAL minimumBy #-}
523minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
524minimumBy cmp (Stream step state) = go Nothing state
525 where
526 go Nothing st = do
527 r <- step defState st
528 case r of
529 Yield x s -> go (Just x) s
530 Skip s -> go Nothing s
531 Stop -> return Nothing
532 go (Just acc) st = do
533 r <- step defState st
534 case r of
535 Yield x s -> case cmp acc x of
536 GT -> go (Just x) s
537 _ -> go (Just acc) s
538 Skip s -> go (Just acc) s
539 Stop -> return (Just acc)
541{-# INLINE_NORMAL (!!) #-}
542(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
543(Stream step state) !! i = go i state
544 where
545 go n st = do
546 r <- step defState st
547 case r of
548 Yield x s | n < 0 -> return Nothing
549 | n == 0 -> return $ Just x
550 | otherwise -> go (n - 1) s
551 Skip s -> go n s
552 Stop -> return Nothing
554{-# INLINE_NORMAL lookup #-}
555lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
556lookup e (Stream step state) = go state
557 where
558 go st = do
559 r <- step defState st
560 case r of
561 Yield (a, b) s -> if e == a then return (Just b) else go s
562 Skip s -> go s
563 Stop -> return Nothing
565{-# INLINE_NORMAL find #-}
566find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
567find p (Stream step state) = go state
568 where
569 go st = do
570 r <- step defState st
571 case r of
572 Yield x s -> if p x then return (Just x) else go s
573 Skip s -> go s
574 Stop -> return Nothing
576{-# INLINE_NORMAL findIndices #-}
577findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
578findIndices p (Stream step state) = Stream step' (state, 0)
579 where
580 {-# INLINE_LATE step' #-}
581 step' gst (st, i) = do
582 r <- step (rstState gst) st
583 return $ case r of
584 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
585 Skip s -> Skip (s, i+1)
586 Stop -> Stop
588------------------------------------------------------------------------------
589-- Substreams
590------------------------------------------------------------------------------
592{-# INLINE_NORMAL isPrefixOf #-}
593isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
594isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
595 where
596 go (sa, sb, Nothing) = do
597 r <- stepa defState sa
598 case r of
599 Yield x sa' -> go (sa', sb, Just x)
600 Skip sa' -> go (sa', sb, Nothing)
601 Stop -> return True
603 go (sa, sb, Just x) = do
604 r <- stepb defState sb
605 case r of
606 Yield y sb' ->
607 if x == y
608 then go (sa, sb', Nothing)
609 else return False
610 Skip sb' -> go (sa, sb', Just x)
611 Stop -> return False
613{-# INLINE_NORMAL isSubsequenceOf #-}
614isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
615isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
616 where
617 go (sa, sb, Nothing) = do
618 r <- stepa defState sa
619 case r of
620 Yield x sa' -> go (sa', sb, Just x)
621 Skip sa' -> go (sa', sb, Nothing)
622 Stop -> return True
624 go (sa, sb, Just x) = do
625 r <- stepb defState sb
626 case r of
627 Yield y sb' ->
628 if x == y
629 then go (sa, sb', Nothing)
630 else go (sa, sb', Just x)
631 Skip sb' -> go (sa, sb', Just x)
632 Stop -> return False
634{-# INLINE_NORMAL stripPrefix #-}
635stripPrefix :: (Eq a, Monad m) => Stream m a -> Stream m a -> m (Maybe (Stream m a))
636stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
637 where
638 go (sa, sb, Nothing) = do
639 r <- stepa defState sa
640 case r of
641 Yield x sa' -> go (sa', sb, Just x)
642 Skip sa' -> go (sa', sb, Nothing)
643 Stop -> return $ Just (Stream stepb sb)
645 go (sa, sb, Just x) = do
646 r <- stepb defState sb
647 case r of
648 Yield y sb' ->
649 if x == y
650 then go (sa, sb', Nothing)
651 else return Nothing
652 Skip sb' -> go (sa, sb', Just x)
653 Stop -> return Nothing
655------------------------------------------------------------------------------
656-- Map and Fold
657------------------------------------------------------------------------------
659-- | Execute a monadic action for each element of the 'Stream'
660{-# INLINE_NORMAL mapM_ #-}
661mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
662mapM_ m = runStream . mapM m
664------------------------------------------------------------------------------
665-- Converting folds
666------------------------------------------------------------------------------
668{-# INLINE toList #-}
669toList :: Monad m => Stream m a -> m [a]
670toList = foldr (:) []
672-- Convert a direct stream to and from CPS encoded stream
673{-# INLINE_LATE toStreamK #-}
674toStreamK :: Monad m => Stream m a -> K.Stream m a
675toStreamK (Stream step state) = go state
676 where
677 go st = K.Stream $ \gst stp sng yld -> do
678 r <- step gst st
679 case r of
680 Yield x s -> yld x (go s)
681 Skip s -> K.unStream (go s) gst stp sng yld
682 Stop -> stp
684#ifndef DISABLE_FUSION
685{-# RULES "fromStreamK/toStreamK fusion"
686 forall s. toStreamK (fromStreamK s) = s #-}
687{-# RULES "toStreamK/fromStreamK fusion"
688 forall s. fromStreamK (toStreamK s) = s #-}
689#endif
691{-# INLINE fromStreamD #-}
692fromStreamD :: (K.IsStream t, Monad m) => Stream m a -> t m a
693fromStreamD = K.fromStream . toStreamK
695------------------------------------------------------------------------------
696-- Transformation by Folding (Scans)
697------------------------------------------------------------------------------
699{-# INLINE_NORMAL postscanlM' #-}
700postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
701postscanlM' fstep begin (Stream step state) =
702 begin `seq` Stream step' (state, begin)
703 where
704 {-# INLINE_LATE step' #-}
705 step' gst (st, acc) = acc `seq` do
706 r <- step (rstState gst) st
707 case r of
708 Yield x s -> do
709 y <- fstep acc x
710 y `seq` return (Yield y (s, y))
711 Skip s -> return $ Skip (s, acc)
712 Stop -> return Stop
714{-# INLINE_NORMAL postscanlM #-}
715postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
716postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
717 where
718 {-# INLINE_LATE step' #-}
719 step' gst (st, acc) = do
720 r <- step (rstState gst) st
721 case r of
722 Yield x s -> do
723 y <- fstep acc x
724 return (Yield y (s, y))
725 Skip s -> return $ Skip (s, acc)
726 Stop -> return Stop
728{-# INLINE_NORMAL scanlM' #-}
729scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
730scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
732{-# INLINE scanl' #-}
733scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
734scanl' f = scanlM' (\a b -> return (f a b))
736{-# INLINE_NORMAL scanlM #-}
737scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
738scanlM fstep begin s = begin `cons` postscanlM fstep begin s
740{-# INLINE scanl #-}
741scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
742scanl f = scanlM (\a b -> return (f a b))
744{-# INLINE_NORMAL scanl1M #-}
745scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
746scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
747 where
748 {-# INLINE_LATE step' #-}
749 step' gst (st, Nothing) = do
750 r <- step (rstState gst) st
751 case r of
752 Yield x s -> return $ Yield x (s, Just x)
753 Skip s -> return $ Skip (s, Nothing)
754 Stop -> return Stop
756 step' gst (st, Just acc) = do
757 r <- step (rstState gst) st
758 case r of
759 Yield y s -> do
760 z <- fstep acc y
761 return $ Yield z (s, Just z)
762 Skip s -> return $ Skip (s, Just acc)
763 Stop -> return Stop
765{-# INLINE scanl1 #-}
766scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
767scanl1 f = scanl1M (\x y -> return (f x y))
769{-# INLINE_NORMAL scanl1M' #-}
770scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
771scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
772 where
773 {-# INLINE_LATE step' #-}
774 step' gst (st, Nothing) = do
775 r <- step (rstState gst) st
776 case r of
777 Yield x s -> x `seq` return $ Yield x (s, Just x)
778 Skip s -> return $ Skip (s, Nothing)
779 Stop -> return Stop
781 step' gst (st, Just acc) = acc `seq` do
782 r <- step (rstState gst) st
783 case r of
784 Yield y s -> do
785 z <- fstep acc y
786 z `seq` return $ Yield z (s, Just z)
787 Skip s -> return $ Skip (s, Just acc)
788 Stop -> return Stop
790{-# INLINE scanl1' #-}
791scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
792scanl1' f = scanl1M' (\x y -> return (f x y))
794-------------------------------------------------------------------------------
795-- Filtering
796-------------------------------------------------------------------------------
798{-# INLINE_NORMAL take #-}
799take :: Monad m => Int -> Stream m a -> Stream m a
800take n (Stream step state) = n `seq` Stream step' (state, 0)
801 where
802 {-# INLINE_LATE step' #-}
803 step' gst (st, i) | i < n = do
804 r <- step (rstState gst) st
805 return $ case r of
806 Yield x s -> Yield x (s, i + 1)
807 Skip s -> Skip (s, i)
808 Stop -> Stop
809 step' _ (_, _) = return Stop
811{-# INLINE_NORMAL takeWhileM #-}
812takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
813takeWhileM f (Stream step state) = Stream step' state
814 where
815 {-# INLINE_LATE step' #-}
816 step' gst st = do
817 r <- step (rstState gst) st
818 case r of
819 Yield x s -> do
820 b <- f x
821 return $ if b then Yield x s else Stop
822 Skip s -> return $ Skip s
823 Stop -> return Stop
825{-# INLINE takeWhile #-}
826takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
827takeWhile f = takeWhileM (return . f)
829{-# INLINE_NORMAL drop #-}
830drop :: Monad m => Int -> Stream m a -> Stream m a
831drop n (Stream step state) = Stream step' (state, Just n)
832 where
833 {-# INLINE_LATE step' #-}
834 step' gst (st, Just i)
835 | i > 0 = do
836 r <- step (rstState gst) st
837 return $
838 case r of
839 Yield _ s -> Skip (s, Just (i - 1))
840 Skip s -> Skip (s, Just i)
841 Stop -> Stop
842 | otherwise = return $ Skip (st, Nothing)
844 step' gst (st, Nothing) = do
845 r <- step (rstState gst) st
846 return $
847 case r of
848 Yield x s -> Yield x (s, Nothing)
849 Skip s -> Skip (s, Nothing)
850 Stop -> Stop
852data DropWhileState s a
853 = DropWhileDrop s
854 | DropWhileYield a s
855 | DropWhileNext s
857{-# INLINE_NORMAL dropWhileM #-}
858dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
859dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
860 where
861 {-# INLINE_LATE step' #-}
862 step' gst (DropWhileDrop st) = do
863 r <- step (rstState gst) st
864 case r of
865 Yield x s -> do
866 b <- f x
867 if b
868 then return $ Skip (DropWhileDrop s)
869 else return $ Skip (DropWhileYield x s)
870 Skip s -> return $ Skip (DropWhileDrop s)
871 Stop -> return Stop
873 step' gst (DropWhileNext st) = do
874 r <- step (rstState gst) st
875 case r of
876 Yield x s -> return $ Skip (DropWhileYield x s)
877 Skip s -> return $ Skip (DropWhileNext s)
878 Stop -> return Stop
880 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
882{-# INLINE dropWhile #-}
883dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
884dropWhile f = dropWhileM (return . f)
886{-# INLINE_NORMAL filterM #-}
887filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
888filterM f (Stream step state) = Stream step' state
889 where
890 {-# INLINE_LATE step' #-}
891 step' gst st = do
892 r <- step (rstState gst) st
893 case r of
894 Yield x s -> do
895 b <- f x
896 return $ if b
897 then Yield x s
898 else Skip s
899 Skip s -> return $ Skip s
900 Stop -> return Stop
902{-# INLINE filter #-}
903filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
904filter f = filterM (return . f)
906------------------------------------------------------------------------------
907-- Transformation by Mapping
908------------------------------------------------------------------------------
910-- | Map a monadic function over a 'Stream'
911{-# INLINE_NORMAL mapM #-}
912mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
913mapM f (Stream step state) = Stream step' state
914 where
915 {-# INLINE_LATE step' #-}
916 step' gst st = do
917 r <- step (rstState gst) st
918 case r of
919 Yield x s -> f x >>= \a -> return $ Yield a s
920 Skip s -> return $ Skip s
921 Stop -> return Stop
923{-# INLINE map #-}
924map :: Monad m => (a -> b) -> Stream m a -> Stream m b
925map f = mapM (return . f)
927{-# INLINE_NORMAL sequence #-}
928sequence :: Monad m => Stream m (m a) -> Stream m a
929sequence (Stream step state) = Stream step' state
930 where
931 {-# INLINE_LATE step' #-}
932 step' gst st = do
933 r <- step (rstState gst) st
934 case r of
935 Yield x s -> x >>= \a -> return (Yield a s)
936 Skip s -> return $ Skip s
937 Stop -> return Stop
939------------------------------------------------------------------------------
940-- Inserting
941------------------------------------------------------------------------------
943{-# INLINE_NORMAL insertBy #-}
944insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
945insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
946 where
947 {-# INLINE_LATE step' #-}
948 step' gst (st, False, _) = do
949 r <- step (rstState gst) st
950 case r of
951 Yield x s -> case cmp a x of
952 GT -> return $ Yield x (s, False, Nothing)
953 _ -> return $ Yield a (s, True, Just x)
954 Skip s -> return $ Skip (s, False, Nothing)
955 Stop -> return $ Yield a (st, True, Nothing)
957 step' _ (_, True, Nothing) = return Stop
959 step' gst (st, True, Just prev) = do
960 r <- step (rstState gst) st
961 case r of
962 Yield x s -> return $ Yield prev (s, True, Just x)
963 Skip s -> return $ Skip (s, True, Just prev)
964 Stop -> return $ Yield prev (st, True, Nothing)
966------------------------------------------------------------------------------
967-- Deleting
968------------------------------------------------------------------------------
970{-# INLINE_NORMAL deleteBy #-}
971deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
972deleteBy eq x (Stream step state) = Stream step' (state, False)
973 where
974 {-# INLINE_LATE step' #-}
975 step' gst (st, False) = do
976 r <- step (rstState gst) st
977 case r of
978 Yield y s -> return $
979 if eq x y then Skip (s, True) else Yield y (s, False)
980 Skip s -> return $ Skip (s, False)
981 Stop -> return Stop
983 step' gst (st, True) = do
984 r <- step (rstState gst) st
985 case r of
986 Yield y s -> return $ Yield y (s, True)
987 Skip s -> return $ Skip (s, True)
988 Stop -> return Stop
990------------------------------------------------------------------------------
991-- Transformation by Map and Filter
992------------------------------------------------------------------------------
994-- XXX Will this always fuse properly?
995{-# INLINE_NORMAL mapMaybe #-}
996mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
997mapMaybe f = fmap fromJust . filter isJust . map f
999{-# INLINE_NORMAL mapMaybeM #-}
1000mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
1001mapMaybeM f = fmap fromJust . filter isJust . mapM f
1003------------------------------------------------------------------------------
1004-- Zipping
1005------------------------------------------------------------------------------
1007{-# INLINE_NORMAL zipWithM #-}
1008zipWithM :: Monad m
1009 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
1010zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
1011 where
1012 {-# INLINE_LATE step #-}
1013 step gst (sa, sb, Nothing) = do
1014 r <- stepa (rstState gst) sa
1015 return $
1016 case r of
1017 Yield x sa' -> Skip (sa', sb, Just x)
1018 Skip sa' -> Skip (sa', sb, Nothing)
1019 Stop -> Stop
1021 step gst (sa, sb, Just x) = do
1022 r <- stepb (rstState gst) sb
1023 case r of
1024 Yield y sb' -> do
1025 z <- f x y
1026 return $ Yield z (sa, sb', Nothing)
1027 Skip sb' -> return $ Skip (sa, sb', Just x)
1028 Stop -> return Stop
1030{-# RULES "zipWithM xs xs"
1031 forall f xs. zipWithM f xs xs = mapM (\x -> f x x) xs #-}
1033{-# INLINE zipWith #-}
1034zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
1035zipWith f = zipWithM (\a b -> return (f a b))
1037------------------------------------------------------------------------------
1038-- Comparisions
1039------------------------------------------------------------------------------
1041{-# INLINE_NORMAL eqBy #-}
1042eqBy :: Monad m => (a -> b -> Bool) -> Stream m a -> Stream m b -> m Bool
1043eqBy eq (Stream step1 t1) (Stream step2 t2) = eq_loop0 SPEC t1 t2
1044 where
1045 eq_loop0 !_ s1 s2 = do
1046 r <- step1 defState s1
1047 case r of
1048 Yield x s1' -> eq_loop1 SPEC x s1' s2
1049 Skip s1' -> eq_loop0 SPEC s1' s2
1050 Stop -> eq_null s2
1052 eq_loop1 !_ x s1 s2 = do
1053 r <- step2 defState s2
1054 case r of
1055 Yield y s2'
1056 | eq x y -> eq_loop0 SPEC s1 s2'
1057 | otherwise -> return False
1058 Skip s2' -> eq_loop1 SPEC x s1 s2'
1059 Stop -> return False
1061 eq_null s2 = do
1062 r <- step2 defState s2
1063 case r of
1064 Yield _ _ -> return False
1065 Skip s2' -> eq_null s2'
1066 Stop -> return True
1068-- | Compare two streams lexicographically
1069{-# INLINE_NORMAL cmpBy #-}
1070cmpBy
1071 :: Monad m
1072 => (a -> b -> Ordering) -> Stream m a -> Stream m b -> m Ordering
1073cmpBy cmp (Stream step1 t1) (Stream step2 t2) = cmp_loop0 SPEC t1 t2
1074 where
1075 cmp_loop0 !_ s1 s2 = do
1076 r <- step1 defState s1
1077 case r of
1078 Yield x s1' -> cmp_loop1 SPEC x s1' s2
1079 Skip s1' -> cmp_loop0 SPEC s1' s2
1080 Stop -> cmp_null s2
1082 cmp_loop1 !_ x s1 s2 = do
1083 r <- step2 defState s2
1084 case r of
1085 Yield y s2' -> case x `cmp` y of
1086 EQ -> cmp_loop0 SPEC s1 s2'
1087 c -> return c
1088 Skip s2' -> cmp_loop1 SPEC x s1 s2'
1089 Stop -> return GT
1091 cmp_null s2 = do
1092 r <- step2 defState s2
1093 case r of
1094 Yield _ _ -> return LT
1095 Skip s2' -> cmp_null s2'
1096 Stop -> return EQ
1098------------------------------------------------------------------------------
1099-- Merging
1100------------------------------------------------------------------------------
1102{-# INLINE_NORMAL mergeBy #-}
1103mergeBy :: (Monad m) => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
1104mergeBy f (Stream stepa ta) (Stream stepb tb) =
1105 Stream step (Just ta, Just tb, Nothing, Nothing)
1106 where
1107 {-# INLINE_LATE step #-}
1109 -- one of the values is missing, and the corresponding stream is running
1110 step gst (Just sa, sb, Nothing, b) = do
1111 r <- stepa (rstState gst) sa
1112 return $ case r of
1113 Yield a sa' -> Skip (Just sa', sb, Just a, b)
1114 Skip sa' -> Skip (Just sa', sb, Nothing, b)
1115 Stop -> Skip (Nothing, sb, Nothing, b)
1117 step gst (sa, Just sb, a, Nothing) = do
1118 r <- stepb (rstState gst) sb
1119 return $ case r of
1120 Yield b sb' -> Skip (sa, Just sb', a, Just b)
1121 Skip sb' -> Skip (sa, Just sb', a, Nothing)
1122 Stop -> Skip (sa, Nothing, a, Nothing)
1124 -- both the values are available
1125 step _ (sa, sb, Just a, Just b) =
1126 return $ case f a b of
1127 LT -> Yield a (sa, sb, Nothing, Just b)
1128 _ -> Yield b (sa, sb, Just a, Nothing)
1130 -- one of the values is missing, corresponding stream is done
1131 step _ (Nothing, sb, Nothing, Just b) = do
1132 return $ Yield b (Nothing, sb, Nothing, Nothing)
1134 step _ (sa, Nothing, Just a, Nothing) = do
1135 return $ Yield a (sa, Nothing, Nothing, Nothing)
1137 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
1139------------------------------------------------------------------------------
1140-- Transformation comprehensions
1141------------------------------------------------------------------------------
1143{-# INLINE_NORMAL the #-}
1144the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
1145the (Stream step state) = go state
1146 where
1147 go st = do
1148 r <- step defState st
1149 case r of
1150 Yield x s -> go' x s
1151 Skip s -> go s
1152 Stop -> return Nothing
1153 go' n st = do
1154 r <- step defState st
1155 case r of
1156 Yield x s | x == n -> go' n s
1157 | otherwise -> return Nothing
1158 Skip s -> go' n s
1159 Stop -> return (Just n)
1161------------------------------------------------------------------------------
1162-- Instances
1163------------------------------------------------------------------------------
1165instance Monad m => Functor (Stream m) where
1166 {-# INLINE fmap #-}
1167 fmap = map