git.iko.soy/streamly
on 162fc17 help sign in
streamly/core/src/Streamly/Internal/Data/Stream/StreamK.hs
1092 lines 35 kB history raw
1{-# LANGUAGE UndecidableInstances #-}
3#include "inline.hs"
5-- |
6-- Module : Streamly.Internal.Data.Stream.StreamK
7-- Copyright : (c) 2017 Composewell Technologies
8--
9-- License : BSD3
10-- Maintainer : streamly@composewell.com
11-- Stability : experimental
12-- Portability : GHC
13--
14--
15-- Continuation passing style (CPS) stream implementation. The symbol 'K' below
16-- denotes a function as well as a Kontinuation.
17--
18-- @
19-- import qualified Streamly.Internal.Data.Stream.StreamK as K
20-- @
21--
22module Streamly.Internal.Data.Stream.StreamK
23 (
24 -- * The stream type
25 Stream(..)
27 -- * Construction Primitives
28 , mkStream
29 , nil
30 , nilM
31 , cons
32 , (.:)
34 -- * Elimination Primitives
35 , foldStream
36 , foldStreamShared
38 -- * Transformation Primitives
39 , unShare
41 -- * Deconstruction
42 , uncons
44 -- * Generation
45 -- ** Unfolds
46 , unfoldr
47 , unfoldrM
49 -- ** Specialized Generation
50 , repeat
51 , repeatM
52 , replicate
53 , replicateM
54 , fromIndices
55 , fromIndicesM
56 , iterate
57 , iterateM
59 -- ** Conversions
60 , fromPure
61 , fromEffect
62 , fromFoldable
63 , fromList
65 -- * foldr/build
66 , foldrS
67 , foldrSM
68 , buildS
69 , augmentS
71 -- * Elimination
72 -- ** General Folds
73 , foldr
74 , foldr1
75 , foldrM
76 , foldrT
78 , foldl'
79 , foldlM'
80 , foldlS
81 , foldlT
82 , foldlx'
83 , foldlMx'
84 , fold
85 , foldBreak
86 , parseBreak
88 -- ** Specialized Folds
89 , drain
90 , null
91 , head
92 , tail
93 , init
94 , elem
95 , notElem
96 , all
97 , any
98 , last
99 , minimum
100 , minimumBy
101 , maximum
102 , maximumBy
103 , findIndices
104 , lookup
105 , findM
106 , find
107 , (!!)
109 -- ** Map and Fold
110 , mapM_
112 -- ** Conversions
113 , toList
114 , hoist
116 -- * Transformation
117 -- ** By folding (scans)
118 , scanl'
119 , scanlx'
121 -- ** Filtering
122 , filter
123 , take
124 , takeWhile
125 , drop
126 , dropWhile
128 -- ** Mapping
129 , map
130 , mapM
131 , sequence
133 -- ** Inserting
134 , intersperseM
135 , intersperse
136 , insertBy
138 -- ** Deleting
139 , deleteBy
141 -- ** Reordering
142 , reverse
144 -- ** Map and Filter
145 , mapMaybe
147 -- ** Zipping
148 , zipWith
149 , zipWithM
151 -- ** Merging
152 , mergeBy
153 , mergeByM
155 -- ** Nesting
156 , concatMapWith
157 , concatMap
158 , bindWith
159 , concatPairsWith
160 , apWith
161 , apSerial
162 , apSerialDiscardFst
163 , apSerialDiscardSnd
165 -- ** Transformation comprehensions
166 , the
168 -- * Semigroup Style Composition
169 , serial
171 -- * Utilities
172 , consM
173 , withLocal
174 , mfix
175 )
176where
178#include "assert.hs"
180import Control.Monad.Catch (MonadThrow, throwM)
181import Control.Monad.Trans.Class (MonadTrans(lift))
182import Control.Monad (void, join)
183import Streamly.Internal.Data.SVar.Type (adaptState, defState)
185import qualified Streamly.Internal.Data.Fold.Type as FL
186import qualified Streamly.Internal.Data.Parser.ParserD.Type as PR
187import qualified Prelude
189import Prelude
190 hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,
191 take, filter, all, any, takeWhile, drop, dropWhile, minimum,
192 maximum, elem, notElem, null, head, tail, init, zipWith, lookup,
193 foldr1, (!!), replicate, reverse, concatMap, iterate, splitAt)
195import Streamly.Internal.Data.Stream.StreamK.Type
197-- $setup
198-- >>> :m
200-------------------------------------------------------------------------------
201-- Generation
202-------------------------------------------------------------------------------
204{-# INLINE unfoldrM #-}
205unfoldrM :: Monad m => (b -> m (Maybe (a, b))) -> b -> Stream m a
206unfoldrM = unfoldrMWith consM
207{-
208-- Generalization of concurrent streams/SVar via unfoldr.
209--
210-- Unfold a value into monadic actions and then run the resulting monadic
211-- actions to generate a stream. Since the step of generating the monadic
212-- action and running them are decoupled we can run the monadic actions
213-- cooncurrently. For example, the seed could be a list of monadic actions or a
214-- pure stream of monadic actions.
215--
216-- We can have different flavors of this depending on the stream type t. The
217-- concurrent version could be async or ahead etc. Depending on how we queue
218-- back the feedback portion b, it could be DFS or BFS style.
219--
220unfoldrA :: (IsStream t, MonadAsync m) =>
221 (b -> Maybe (m a, b)) -> b -> Stream m a
222unfoldrA = undefined
223-}
225-------------------------------------------------------------------------------
226-- Special generation
227-------------------------------------------------------------------------------
229repeatM :: Monad m => m a -> Stream m a
230repeatM = repeatMWith consM
232{-# INLINE replicateM #-}
233replicateM :: Monad m => Int -> m a -> Stream m a
234replicateM = replicateMWith consM
235{-# INLINE replicate #-}
236replicate :: Int -> a -> Stream m a
237replicate n a = go n
238 where
239 go cnt = if cnt <= 0 then nil else a `cons` go (cnt - 1)
241{-# INLINE fromIndicesM #-}
242fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
243fromIndicesM = fromIndicesMWith consM
244{-# INLINE fromIndices #-}
245fromIndices :: (Int -> a) -> Stream m a
246fromIndices gen = go 0
247 where
248 go n = gen n `cons` go (n + 1)
250{-# INLINE iterate #-}
251iterate :: (a -> a) -> a -> Stream m a
252iterate step = go
253 where
254 go !s = cons s (go (step s))
256{-# INLINE iterateM #-}
257iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
258iterateM = iterateMWith consM
260-------------------------------------------------------------------------------
261-- Conversions
262-------------------------------------------------------------------------------
264{-# INLINE fromList #-}
265fromList :: [a] -> Stream m a
266fromList = fromFoldable
268-------------------------------------------------------------------------------
269-- Elimination by Folding
270-------------------------------------------------------------------------------
272-- | Lazy right associative fold.
273{-# INLINE foldr #-}
274foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
275foldr step acc = foldrM (\x xs -> xs >>= \b -> return (step x b)) (return acc)
277-- | Right associative fold to an arbitrary transformer monad.
278{-# INLINE foldrT #-}
279foldrT :: (Monad m, Monad (s m), MonadTrans s)
280 => (a -> s m b -> s m b) -> s m b -> Stream m a -> s m b
281foldrT step final = go
282 where
283 go m1 = do
284 res <- lift $ uncons m1
285 case res of
286 Just (h, t) -> step h (go t)
287 Nothing -> final
289{-# INLINE foldr1 #-}
290foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
291foldr1 step m = do
292 r <- uncons m
293 case r of
294 Nothing -> return Nothing
295 Just (h, t) -> fmap Just (go h t)
296 where
297 go p m1 =
298 let stp = return p
299 single a = return $ step a p
300 yieldk a r = fmap (step p) (go a r)
301 in foldStream defState yieldk single stp m1
303-- XXX replace the recursive "go" with explicit continuations.
304-- | Like 'foldx', but with a monadic step function.
305{-# INLINABLE foldlMx' #-}
306foldlMx' :: Monad m
307 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> m b
308foldlMx' step begin done = go begin
309 where
310 go !acc m1 =
311 let stop = acc >>= done
312 single a = acc >>= \b -> step b a >>= done
313 yieldk a r = acc >>= \b -> step b a >>= \x -> go (return x) r
314 in foldStream defState yieldk single stop m1
316{-# INLINABLE fold #-}
317fold :: Monad m => FL.Fold m a b -> Stream m a -> m b
318fold (FL.Fold step begin done) m = do
319 res <- begin
320 case res of
321 FL.Partial fs -> go fs m
322 FL.Done fb -> return fb
324 where
325 go !acc m1 =
326 let stop = done acc
327 single a = step acc a
328 >>= \case
329 FL.Partial s -> done s
330 FL.Done b1 -> return b1
331 yieldk a r = step acc a
332 >>= \case
333 FL.Partial s -> go s r
334 FL.Done b1 -> return b1
335 in foldStream defState yieldk single stop m1
337{-# INLINE foldBreak #-}
338foldBreak :: Monad m => FL.Fold m a b -> Stream m a -> m (b, Stream m a)
339foldBreak (FL.Fold step begin done) m = do
340 res <- begin
341 case res of
342 FL.Partial fs -> go fs m
343 FL.Done fb -> return (fb, m)
345 where
347 go !acc m1 =
348 let stop = (, nil) <$> done acc
349 single a =
350 step acc a
351 >>= \case
352 FL.Partial s -> (, nil) <$> done s
353 FL.Done b1 -> return (b1, nil)
354 yieldk a r =
355 step acc a
356 >>= \case
357 FL.Partial s -> go s r
358 FL.Done b1 -> return (b1, r)
359 in foldStream defState yieldk single stop m1
361-- | Like 'foldl'' but with a monadic step function.
362{-# INLINE foldlM' #-}
363foldlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> m b
364foldlM' step begin = foldlMx' step begin return
366-- | Lazy left fold to an arbitrary transformer monad.
367{-# INLINE foldlT #-}
368foldlT :: (Monad m, Monad (s m), MonadTrans s)
369 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
370foldlT step = go
371 where
372 go acc m1 = do
373 res <- lift $ uncons m1
374 case res of
375 Just (h, t) -> go (step acc h) t
376 Nothing -> acc
378------------------------------------------------------------------------------
379-- Specialized folds
380------------------------------------------------------------------------------
382{-# INLINE head #-}
383head :: Monad m => Stream m a -> m (Maybe a)
384-- head = foldrM (\x _ -> return $ Just x) (return Nothing)
385head m =
386 let stop = return Nothing
387 single a = return (Just a)
388 yieldk a _ = return (Just a)
389 in foldStream defState yieldk single stop m
391{-# INLINE elem #-}
392elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
393elem e = go
394 where
395 go m1 =
396 let stop = return False
397 single a = return (a == e)
398 yieldk a r = if a == e then return True else go r
399 in foldStream defState yieldk single stop m1
401{-# INLINE notElem #-}
402notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
403notElem e = go
404 where
405 go m1 =
406 let stop = return True
407 single a = return (a /= e)
408 yieldk a r = if a == e then return False else go r
409 in foldStream defState yieldk single stop m1
411{-# INLINABLE all #-}
412all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
413all p = go
414 where
415 go m1 =
416 let single a | p a = return True
417 | otherwise = return False
418 yieldk a r | p a = go r
419 | otherwise = return False
420 in foldStream defState yieldk single (return True) m1
422{-# INLINABLE any #-}
423any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
424any p = go
425 where
426 go m1 =
427 let single a | p a = return True
428 | otherwise = return False
429 yieldk a r | p a = return True
430 | otherwise = go r
431 in foldStream defState yieldk single (return False) m1
433-- | Extract the last element of the stream, if any.
434{-# INLINE last #-}
435last :: Monad m => Stream m a -> m (Maybe a)
436last = foldlx' (\_ y -> Just y) Nothing id
438{-# INLINE minimum #-}
439minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
440minimum = go Nothing
441 where
442 go Nothing m1 =
443 let stop = return Nothing
444 single a = return (Just a)
445 yieldk a r = go (Just a) r
446 in foldStream defState yieldk single stop m1
448 go (Just res) m1 =
449 let stop = return (Just res)
450 single a =
451 if res <= a
452 then return (Just res)
453 else return (Just a)
454 yieldk a r =
455 if res <= a
456 then go (Just res) r
457 else go (Just a) r
458 in foldStream defState yieldk single stop m1
460{-# INLINE minimumBy #-}
461minimumBy
462 :: (Monad m)
463 => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
464minimumBy cmp = go Nothing
465 where
466 go Nothing m1 =
467 let stop = return Nothing
468 single a = return (Just a)
469 yieldk a r = go (Just a) r
470 in foldStream defState yieldk single stop m1
472 go (Just res) m1 =
473 let stop = return (Just res)
474 single a = case cmp res a of
475 GT -> return (Just a)
476 _ -> return (Just res)
477 yieldk a r = case cmp res a of
478 GT -> go (Just a) r
479 _ -> go (Just res) r
480 in foldStream defState yieldk single stop m1
482{-# INLINE maximum #-}
483maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
484maximum = go Nothing
485 where
486 go Nothing m1 =
487 let stop = return Nothing
488 single a = return (Just a)
489 yieldk a r = go (Just a) r
490 in foldStream defState yieldk single stop m1
492 go (Just res) m1 =
493 let stop = return (Just res)
494 single a =
495 if res <= a
496 then return (Just a)
497 else return (Just res)
498 yieldk a r =
499 if res <= a
500 then go (Just a) r
501 else go (Just res) r
502 in foldStream defState yieldk single stop m1
504{-# INLINE maximumBy #-}
505maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
506maximumBy cmp = go Nothing
507 where
508 go Nothing m1 =
509 let stop = return Nothing
510 single a = return (Just a)
511 yieldk a r = go (Just a) r
512 in foldStream defState yieldk single stop m1
514 go (Just res) m1 =
515 let stop = return (Just res)
516 single a = case cmp res a of
517 GT -> return (Just res)
518 _ -> return (Just a)
519 yieldk a r = case cmp res a of
520 GT -> go (Just res) r
521 _ -> go (Just a) r
522 in foldStream defState yieldk single stop m1
524{-# INLINE (!!) #-}
525(!!) :: Monad m => Stream m a -> Int -> m (Maybe a)
526m !! i = go i m
527 where
528 go n m1 =
529 let single a | n == 0 = return $ Just a
530 | otherwise = return Nothing
531 yieldk a x | n < 0 = return Nothing
532 | n == 0 = return $ Just a
533 | otherwise = go (n - 1) x
534 in foldStream defState yieldk single (return Nothing) m1
536{-# INLINE lookup #-}
537lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
538lookup e = go
539 where
540 go m1 =
541 let single (a, b) | a == e = return $ Just b
542 | otherwise = return Nothing
543 yieldk (a, b) x | a == e = return $ Just b
544 | otherwise = go x
545 in foldStream defState yieldk single (return Nothing) m1
547{-# INLINE findM #-}
548findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
549findM p = go
550 where
551 go m1 =
552 let single a = do
553 b <- p a
554 if b then return $ Just a else return Nothing
555 yieldk a x = do
556 b <- p a
557 if b then return $ Just a else go x
558 in foldStream defState yieldk single (return Nothing) m1
560{-# INLINE find #-}
561find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
562find p = findM (return . p)
564{-# INLINE findIndices #-}
565findIndices :: (a -> Bool) -> Stream m a -> Stream m Int
566findIndices p = go 0
567 where
568 go offset m1 = mkStream $ \st yld sng stp ->
569 let single a | p a = sng offset
570 | otherwise = stp
571 yieldk a x | p a = yld offset $ go (offset + 1) x
572 | otherwise = foldStream (adaptState st) yld sng stp $
573 go (offset + 1) x
574 in foldStream (adaptState st) yieldk single stp m1
576------------------------------------------------------------------------------
577-- Map and Fold
578------------------------------------------------------------------------------
580-- | Apply a monadic action to each element of the stream and discard the
581-- output of the action.
582{-# INLINE mapM_ #-}
583mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
584mapM_ f = go
585 where
586 go m1 =
587 let stop = return ()
588 single a = void (f a)
589 yieldk a r = f a >> go r
590 in foldStream defState yieldk single stop m1
592{-# INLINE mapM #-}
593mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
594mapM = mapMWith consM
596------------------------------------------------------------------------------
597-- Converting folds
598------------------------------------------------------------------------------
600{-# INLINABLE toList #-}
601toList :: Monad m => Stream m a -> m [a]
602toList = foldr (:) []
604-- Based on suggestions by David Feuer and Pranay Sashank
605{-# INLINE hoist #-}
606hoist :: (Monad m, Monad n)
607 => (forall x. m x -> n x) -> Stream m a -> Stream n a
608hoist f str =
609 mkStream $ \st yld sng stp ->
610 let single = return . sng
611 yieldk a s = return $ yld a (hoist f s)
612 stop = return stp
613 state = adaptState st
614 in join . f $ foldStreamShared state yieldk single stop str
616-------------------------------------------------------------------------------
617-- Transformation by folding (Scans)
618-------------------------------------------------------------------------------
620{-# INLINE scanlx' #-}
621scanlx' :: (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
622scanlx' step begin done m =
623 cons (done begin) $ go m begin
624 where
625 go m1 !acc = mkStream $ \st yld sng stp ->
626 let single a = sng (done $ step acc a)
627 yieldk a r =
628 let s = step acc a
629 in yld (done s) (go r s)
630 in foldStream (adaptState st) yieldk single stp m1
632{-# INLINE scanl' #-}
633scanl' :: (b -> a -> b) -> b -> Stream m a -> Stream m b
634scanl' step begin = scanlx' step begin id
636-------------------------------------------------------------------------------
637-- Filtering
638-------------------------------------------------------------------------------
640{-# INLINE filter #-}
641filter :: (a -> Bool) -> Stream m a -> Stream m a
642filter p = go
643 where
644 go m1 = mkStream $ \st yld sng stp ->
645 let single a | p a = sng a
646 | otherwise = stp
647 yieldk a r | p a = yld a (go r)
648 | otherwise = foldStream st yieldk single stp r
649 in foldStream st yieldk single stp m1
651{-# INLINE take #-}
652take :: Int -> Stream m a -> Stream m a
653take = go
654 where
655 go n1 m1 = mkStream $ \st yld sng stp ->
656 let yieldk a r = yld a (go (n1 - 1) r)
657 in if n1 <= 0
658 then stp
659 else foldStream st yieldk sng stp m1
661{-# INLINE takeWhile #-}
662takeWhile :: (a -> Bool) -> Stream m a -> Stream m a
663takeWhile p = go
664 where
665 go m1 = mkStream $ \st yld sng stp ->
666 let single a | p a = sng a
667 | otherwise = stp
668 yieldk a r | p a = yld a (go r)
669 | otherwise = stp
670 in foldStream st yieldk single stp m1
672{-# INLINE drop #-}
673drop :: Int -> Stream m a -> Stream m a
674drop n m = unShare (go n m)
675 where
676 go n1 m1 = mkStream $ \st yld sng stp ->
677 let single _ = stp
678 yieldk _ r = foldStreamShared st yld sng stp $ go (n1 - 1) r
679 -- Somehow "<=" check performs better than a ">"
680 in if n1 <= 0
681 then foldStreamShared st yld sng stp m1
682 else foldStreamShared st yieldk single stp m1
684{-# INLINE dropWhile #-}
685dropWhile :: (a -> Bool) -> Stream m a -> Stream m a
686dropWhile p = go
687 where
688 go m1 = mkStream $ \st yld sng stp ->
689 let single a | p a = stp
690 | otherwise = sng a
691 yieldk a r | p a = foldStream st yieldk single stp r
692 | otherwise = yld a r
693 in foldStream st yieldk single stp m1
695-------------------------------------------------------------------------------
696-- Mapping
697-------------------------------------------------------------------------------
699-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
700-- other stream types to overload it.
701{-# INLINE sequence #-}
702sequence :: Monad m => Stream m (m a) -> Stream m a
703sequence = go
704 where
705 go m1 = mkStream $ \st yld sng stp ->
706 let single ma = ma >>= sng
707 yieldk ma r = foldStreamShared st yld sng stp $ ma `consM` go r
708 in foldStream (adaptState st) yieldk single stp m1
710-------------------------------------------------------------------------------
711-- Inserting
712-------------------------------------------------------------------------------
714{-# INLINE intersperseM #-}
715intersperseM :: Monad m => m a -> Stream m a -> Stream m a
716intersperseM a = prependingStart
717 where
718 prependingStart m1 = mkStream $ \st yld sng stp ->
719 let yieldk i x =
720 foldStreamShared st yld sng stp $ return i `consM` go x
721 in foldStream st yieldk sng stp m1
722 go m2 = mkStream $ \st yld sng stp ->
723 let single i = foldStreamShared st yld sng stp $ a `consM` fromPure i
724 yieldk i x =
725 foldStreamShared
726 st yld sng stp $ a `consM` return i `consM` go x
727 in foldStream st yieldk single stp m2
729{-# INLINE intersperse #-}
730intersperse :: Monad m => a -> Stream m a -> Stream m a
731intersperse a = intersperseM (return a)
733{-# INLINE insertBy #-}
734insertBy :: (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
735insertBy cmp x = go
736 where
737 go m1 = mkStream $ \st yld _ _ ->
738 let single a = case cmp x a of
739 GT -> yld a (fromPure x)
740 _ -> yld x (fromPure a)
741 stop = yld x nil
742 yieldk a r = case cmp x a of
743 GT -> yld a (go r)
744 _ -> yld x (a `cons` r)
745 in foldStream st yieldk single stop m1
747------------------------------------------------------------------------------
748-- Deleting
749------------------------------------------------------------------------------
751{-# INLINE deleteBy #-}
752deleteBy :: (a -> a -> Bool) -> a -> Stream m a -> Stream m a
753deleteBy eq x = go
754 where
755 go m1 = mkStream $ \st yld sng stp ->
756 let single a = if eq x a then stp else sng a
757 yieldk a r = if eq x a
758 then foldStream st yld sng stp r
759 else yld a (go r)
760 in foldStream st yieldk single stp m1
762-------------------------------------------------------------------------------
763-- Map and Filter
764-------------------------------------------------------------------------------
766{-# INLINE mapMaybe #-}
767mapMaybe :: (a -> Maybe b) -> Stream m a -> Stream m b
768mapMaybe f = go
769 where
770 go m1 = mkStream $ \st yld sng stp ->
771 let single a = maybe stp sng (f a)
772 yieldk a r = case f a of
773 Just b -> yld b $ go r
774 Nothing -> foldStream (adaptState st) yieldk single stp r
775 in foldStream (adaptState st) yieldk single stp m1
777------------------------------------------------------------------------------
778-- Serial Zipping
779------------------------------------------------------------------------------
781-- | Zip two streams serially using a pure zipping function.
782--
783-- @since 0.1.0
784{-# INLINE zipWith #-}
785zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
786zipWith f = zipWithM (\a b -> return (f a b))
788-- | Zip two streams serially using a monadic zipping function.
789--
790-- @since 0.1.0
791{-# INLINE zipWithM #-}
792zipWithM :: Monad m =>
793 (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
794zipWithM f = go
796 where
798 go mx my = mkStream $ \st yld sng stp -> do
799 let merge a ra =
800 let single2 b = f a b >>= sng
801 yield2 b rb = f a b >>= \x -> yld x (go ra rb)
802 in foldStream (adaptState st) yield2 single2 stp my
803 let single1 a = merge a nil
804 yield1 = merge
805 foldStream (adaptState st) yield1 single1 stp mx
807------------------------------------------------------------------------------
808-- Merging
809------------------------------------------------------------------------------
811{-# INLINE mergeByM #-}
812mergeByM :: Monad m =>
813 (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
814mergeByM cmp = go
816 where
818 go mx my = mkStream $ \st yld sng stp -> do
819 let stop = foldStream st yld sng stp my
820 single x = foldStream st yld sng stp (goX0 x my)
821 yield x rx = foldStream st yld sng stp (goX x rx my)
822 foldStream st yield single stop mx
824 goX0 x my = mkStream $ \st yld sng _ -> do
825 let stop = sng x
826 single y = do
827 r <- cmp x y
828 case r of
829 GT -> yld y (fromPure x)
830 _ -> yld x (fromPure y)
831 yield y ry = do
832 r <- cmp x y
833 case r of
834 GT -> yld y (goX0 x ry)
835 _ -> yld x (y `cons` ry)
836 in foldStream st yield single stop my
838 goX x mx my = mkStream $ \st yld _ _ -> do
839 let stop = yld x mx
840 single y = do
841 r <- cmp x y
842 case r of
843 GT -> yld y (x `cons` mx)
844 _ -> yld x (goY0 mx y)
845 yield y ry = do
846 r <- cmp x y
847 case r of
848 GT -> yld y (goX x mx ry)
849 _ -> yld x (goY mx y ry)
850 in foldStream st yield single stop my
852 goY0 mx y = mkStream $ \st yld sng _ -> do
853 let stop = sng y
854 single x = do
855 r <- cmp x y
856 case r of
857 GT -> yld y (fromPure x)
858 _ -> yld x (fromPure y)
859 yield x rx = do
860 r <- cmp x y
861 case r of
862 GT -> yld y (x `cons` rx)
863 _ -> yld x (goY0 rx y)
864 in foldStream st yield single stop mx
866 goY mx y my = mkStream $ \st yld _ _ -> do
867 let stop = yld y my
868 single x = do
869 r <- cmp x y
870 case r of
871 GT -> yld y (goX0 x my)
872 _ -> yld x (y `cons` my)
873 yield x rx = do
874 r <- cmp x y
875 case r of
876 GT -> yld y (goX x rx my)
877 _ -> yld x (goY rx y my)
878 in foldStream st yield single stop mx
880{-# INLINE mergeBy #-}
881mergeBy :: (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
882-- XXX GHC: This has slightly worse performance than replacing "r <- cmp x y"
883-- with "let r = cmp x y" in the monadic version. The definition below is
884-- exactly the same as mergeByM except this change.
885-- mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
886mergeBy cmp = go
888 where
890 go mx my = mkStream $ \st yld sng stp -> do
891 let stop = foldStream st yld sng stp my
892 single x = foldStream st yld sng stp (goX0 x my)
893 yield x rx = foldStream st yld sng stp (goX x rx my)
894 foldStream st yield single stop mx
896 goX0 x my = mkStream $ \st yld sng _ -> do
897 let stop = sng x
898 single y = do
899 case cmp x y of
900 GT -> yld y (fromPure x)
901 _ -> yld x (fromPure y)
902 yield y ry = do
903 case cmp x y of
904 GT -> yld y (goX0 x ry)
905 _ -> yld x (y `cons` ry)
906 in foldStream st yield single stop my
908 goX x mx my = mkStream $ \st yld _ _ -> do
909 let stop = yld x mx
910 single y = do
911 case cmp x y of
912 GT -> yld y (x `cons` mx)
913 _ -> yld x (goY0 mx y)
914 yield y ry = do
915 case cmp x y of
916 GT -> yld y (goX x mx ry)
917 _ -> yld x (goY mx y ry)
918 in foldStream st yield single stop my
920 goY0 mx y = mkStream $ \st yld sng _ -> do
921 let stop = sng y
922 single x = do
923 case cmp x y of
924 GT -> yld y (fromPure x)
925 _ -> yld x (fromPure y)
926 yield x rx = do
927 case cmp x y of
928 GT -> yld y (x `cons` rx)
929 _ -> yld x (goY0 rx y)
930 in foldStream st yield single stop mx
932 goY mx y my = mkStream $ \st yld _ _ -> do
933 let stop = yld y my
934 single x = do
935 case cmp x y of
936 GT -> yld y (goX0 x my)
937 _ -> yld x (y `cons` my)
938 yield x rx = do
939 case cmp x y of
940 GT -> yld y (goX x rx my)
941 _ -> yld x (goY rx y my)
942 in foldStream st yield single stop mx
944------------------------------------------------------------------------------
945-- Transformation comprehensions
946------------------------------------------------------------------------------
948{-# INLINE the #-}
949the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
950the m = do
951 r <- uncons m
952 case r of
953 Nothing -> return Nothing
954 Just (h, t) -> go h t
955 where
956 go h m1 =
957 let single a | h == a = return $ Just h
958 | otherwise = return Nothing
959 yieldk a r | h == a = go h r
960 | otherwise = return Nothing
961 in foldStream defState yieldk single (return $ Just h) m1
963------------------------------------------------------------------------------
964-- Alternative & MonadPlus
965------------------------------------------------------------------------------
967_alt :: Stream m a -> Stream m a -> Stream m a
968_alt m1 m2 = mkStream $ \st yld sng stp ->
969 let stop = foldStream st yld sng stp m2
970 in foldStream st yld sng stop m1
972------------------------------------------------------------------------------
973-- MonadError
974------------------------------------------------------------------------------
976{-
977-- XXX handle and test cross thread state transfer
978withCatchError
979 :: MonadError e m
980 => Stream m a -> (e -> Stream m a) -> Stream m a
981withCatchError m h =
982 mkStream $ \_ stp sng yld ->
983 let run x = unStream x Nothing stp sng yieldk
984 handle r = r `catchError` \e -> run $ h e
985 yieldk a r = yld a (withCatchError r h)
986 in handle $ run m
987-}
989-------------------------------------------------------------------------------
990-- Parsing
991-------------------------------------------------------------------------------
993-- Inlined definition.
994{-# INLINE splitAt #-}
995splitAt :: Int -> [a] -> ([a],[a])
996splitAt n ls
997 | n <= 0 = ([], ls)
998 | otherwise = splitAt' n ls
999 where
1000 splitAt' :: Int -> [a] -> ([a], [a])
1001 splitAt' _ [] = ([], [])
1002 splitAt' 1 (x:xs) = ([x], xs)
1003 splitAt' m (x:xs) = (x:xs', xs'')
1004 where
1005 (xs', xs'') = splitAt' (m - 1) xs
1007-- | Run a 'Parser' over a stream and return rest of the Stream.
1008{-# INLINE_NORMAL parseBreak #-}
1009parseBreak
1010 :: MonadThrow m
1011 => PR.Parser m a b
1012 -> Stream m a
1013 -> m (b, Stream m a)
1014parseBreak (PR.Parser pstep initial extract) stream = do
1015 res <- initial
1016 case res of
1017 PR.IPartial s -> goStream stream [] s
1018 PR.IDone b -> return (b, stream)
1019 PR.IError err -> throwM $ PR.ParseError err
1021 where
1023 -- "buf" contains last few items in the stream that we may have to
1024 -- backtrack to.
1025 --
1026 -- XXX currently we are using a dumb list based approach for backtracking
1027 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
1028 -- That will allow us more efficient random back and forth movement.
1029 goStream st buf !pst =
1030 let stop = do
1031 r <- extract pst
1032 case r of
1033 PR.Error err -> throwM $ PR.ParseError err
1034 PR.Done n b -> do
1035 assertM(n <= length buf)
1036 let src0 = Prelude.take n buf
1037 src = Prelude.reverse src0
1038 return (b, fromList src)
1039 PR.Partial _ _ -> error "Bug: parseBreak: Partial in extract"
1040 PR.Continue 0 _ -> error "parseBreak: extract, Continue 0 creates infinite loop"
1041 PR.Continue n s -> do
1042 assertM(n <= length buf)
1043 let (src0, buf1) = splitAt n buf
1044 src = Prelude.reverse src0
1045 goBuf nil buf1 src s
1046 single x = yieldk x nil
1047 yieldk x r = do
1048 res <- pstep pst x
1049 case res of
1050 PR.Partial 0 s -> goStream r [] s
1051 PR.Partial n s -> do
1052 assertM(n <= length (x:buf))
1053 let src0 = Prelude.take n (x:buf)
1054 src = Prelude.reverse src0
1055 goBuf r [] src s
1056 PR.Continue 0 s -> goStream r (x:buf) s
1057 PR.Continue n s -> do
1058 assertM(n <= length (x:buf))
1059 let (src0, buf1) = splitAt n (x:buf)
1060 src = Prelude.reverse src0
1061 goBuf r buf1 src s
1062 PR.Done 0 b -> return (b, r)
1063 PR.Done n b -> do
1064 assertM(n <= length (x:buf))
1065 let src0 = Prelude.take n (x:buf)
1066 src = Prelude.reverse src0
1067 return (b, serial (fromList src) r)
1068 PR.Error err -> throwM $ PR.ParseError err
1069 in foldStream defState yieldk single stop st
1071 goBuf st buf [] !pst = goStream st buf pst
1072 goBuf st buf (x:xs) !pst = do
1073 pRes <- pstep pst x
1074 case pRes of
1075 PR.Partial 0 s -> goBuf st [] xs s
1076 PR.Partial n s -> do
1077 assert (n <= length (x:buf)) (return ())
1078 let src0 = Prelude.take n (x:buf)
1079 src = Prelude.reverse src0 ++ xs
1080 goBuf st [] src s
1081 PR.Continue 0 s -> goBuf st (x:buf) xs s
1082 PR.Continue n s -> do
1083 assert (n <= length (x:buf)) (return ())
1084 let (src0, buf1) = splitAt n (x:buf)
1085 src = Prelude.reverse src0 ++ xs
1086 goBuf st buf1 src s
1087 PR.Done n b -> do
1088 assert (n <= length (x:buf)) (return ())
1089 let src0 = Prelude.take n (x:buf)
1090 src = Prelude.reverse src0
1091 return (b, serial (fromList src) st)
1092 PR.Error err -> throwM $ PR.ParseError err