git.iko.soy/streamly
on 32785ad help sign in
streamly/core/src/Streamly/Internal/Data/Stream/StreamK.hs
1409 lines 47 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Stream.StreamK
3-- Copyright : (c) 2017 Composewell Technologies
4--
5-- License : BSD3
6-- Maintainer : streamly@composewell.com
7-- Stability : experimental
8-- Portability : GHC
9--
10-- To run examples in this module:
11--
12-- >>> import qualified Streamly.Data.Fold as Fold
13-- >>> import qualified Streamly.Data.Stream as Stream
14-- >>> import qualified Streamly.Data.Stream.StreamK as StreamK
15--
16-- We will add some more imports in the examples as needed.
17--
18-- For effectful streams we will use the following IO action:
19--
20-- >>> effect n = print n >> return n
21--
22-- = Overview
23--
24-- Continuation passing style (CPS) stream implementation. The 'K' in 'StreamK'
25-- stands for Kontinuation.
26--
27-- StreamK can be constructed like lists, except that they use 'nil' instead of
28-- '[]' and 'cons' instead of ':'.
29--
30-- `cons` adds a pure value at the head of the stream:
31--
32-- >>> import Streamly.Data.Stream.StreamK (StreamK, cons, consM, nil)
33-- >>> stream = 1 `cons` 2 `cons` nil :: StreamK IO Int
34--
35-- Convert 'StreamK' to 'Stream' to use functions from the
36-- "Streamly.Data.Stream" module:
37--
38-- >>> Stream.fold Fold.toList $ StreamK.toStream stream -- IO [Int]
39-- [1,2]
40--
41-- `consM` adds an effect at the head of the stream:
42--
43-- >>> stream = effect 1 `consM` effect 2 `consM` nil
44-- >>> Stream.fold Fold.toList $ StreamK.toStream stream
45-- 1
46-- 2
47-- [1,2]
48--
49module Streamly.Internal.Data.Stream.StreamK
50 (
51 -- * The stream type
52 Stream
53 , StreamK(..)
54 , fromStream
55 , toStream
57 , CrossStreamK
58 , unCross
59 , mkCross
61 -- * Construction Primitives
62 , mkStream
63 , nil
64 , nilM
65 , cons
66 , (.:)
68 -- * Elimination Primitives
69 , foldStream
70 , foldStreamShared
72 -- * Transformation Primitives
73 , unShare
75 -- * Deconstruction
76 , uncons
78 -- * Generation
79 -- ** Unfolds
80 , unfoldr
81 , unfoldrM
83 -- ** Specialized Generation
84 , repeat
85 , repeatM
86 , replicate
87 , replicateM
88 , fromIndices
89 , fromIndicesM
90 , iterate
91 , iterateM
93 -- ** Conversions
94 , fromPure
95 , fromEffect
96 , fromFoldable
97 , fromList
99 -- * foldr/build
100 , foldrS
101 , foldrSM
102 , buildS
103 , augmentS
105 -- * Elimination
106 -- ** General Folds
107 , foldr
108 , foldr1
109 , foldrM
111 , foldl'
112 , foldlM'
113 , foldlS
114 , foldlx'
115 , foldlMx'
116 , fold
117 , foldBreak
118 , foldEither
119 , foldConcat
120 , parseDBreak
121 , parseD
122 , parseBreakChunks
123 , parseChunks
125 -- ** Specialized Folds
126 , drain
127 , null
128 , head
129 , tail
130 , init
131 , elem
132 , notElem
133 , all
134 , any
135 , last
136 , minimum
137 , minimumBy
138 , maximum
139 , maximumBy
140 , findIndices
141 , lookup
142 , findM
143 , find
144 , (!!)
146 -- ** Map and Fold
147 , mapM_
149 -- ** Conversions
150 , toList
151 , hoist
153 -- * Transformation
154 -- ** By folding (scans)
155 , scanl'
156 , scanlx'
158 -- ** Filtering
159 , filter
160 , take
161 , takeWhile
162 , drop
163 , dropWhile
165 -- ** Mapping
166 , map
167 , mapM
168 , sequence
170 -- ** Inserting
171 , intersperseM
172 , intersperse
173 , insertBy
175 -- ** Deleting
176 , deleteBy
178 -- ** Reordering
179 , reverse
180 , sortBy
182 -- ** Map and Filter
183 , mapMaybe
185 -- ** Zipping
186 , zipWith
187 , zipWithM
189 -- ** Merging
190 , mergeBy
191 , mergeByM
193 -- ** Nesting
194 , crossApplyWith
195 , crossApply
196 , crossApplySnd
197 , crossApplyFst
198 , crossWith
200 , concatMapWith
201 , concatMap
202 , concatEffect
203 , bindWith
204 , concatIterateWith
205 , concatIterateLeftsWith
206 , concatIterateScanWith
208 , mergeMapWith
209 , mergeIterateWith
211 -- ** Transformation comprehensions
212 , the
214 -- * Semigroup Style Composition
215 , append
216 , interleave
218 -- * Utilities
219 , consM
220 , mfix
221 )
222where
224#include "ArrayMacros.h"
225#include "inline.hs"
226#include "assert.hs"
228import Control.Monad (void, join)
229import Data.Proxy (Proxy(..))
230import GHC.Types (SPEC(..))
231import Streamly.Internal.Data.Array.Type (Array(..))
232import Streamly.Internal.Data.Fold.Type (Fold(..))
233import Streamly.Internal.Data.Producer.Type (Producer(..))
234import Streamly.Internal.Data.SVar.Type (adaptState, defState)
235import Streamly.Internal.Data.Unboxed (sizeOf, Unbox)
236import Streamly.Internal.Data.Parser.ParserK.Type (ParserK)
238import qualified Streamly.Internal.Data.Array.Type as Array
239import qualified Streamly.Internal.Data.Fold.Type as FL
240import qualified Streamly.Internal.Data.Parser as Parser
241import qualified Streamly.Internal.Data.Parser.ParserD.Type as PR
242import qualified Streamly.Internal.Data.Parser.ParserK.Type as ParserK
243import qualified Streamly.Internal.Data.Stream.StreamD as Stream
244import qualified Prelude
246import Prelude
247 hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,
248 take, filter, all, any, takeWhile, drop, dropWhile, minimum,
249 maximum, elem, notElem, null, head, tail, init, zipWith, lookup,
250 foldr1, (!!), replicate, reverse, concatMap, iterate, splitAt)
252import Streamly.Internal.Data.Stream.StreamK.Type
253import Streamly.Internal.Data.Parser.ParserD (ParseError(..))
255-- $setup
256-- >>> :m
257-- >>> import qualified Streamly.Data.Fold as Fold
258-- >>> import qualified Streamly.Data.Parser as Parser
259-- >>> import qualified Streamly.Data.Stream as Stream
260-- >>> import qualified Streamly.Data.Stream.StreamK as StreamK
261-- >>> import qualified Streamly.Internal.Data.Stream.StreamK as StreamK
262--
264{-# INLINE fromStream #-}
265fromStream :: Monad m => Stream.Stream m a -> StreamK m a
266fromStream = Stream.toStreamK
268{-# INLINE toStream #-}
269toStream :: Applicative m => StreamK m a -> Stream.Stream m a
270toStream = Stream.fromStreamK
272-------------------------------------------------------------------------------
273-- Generation
274-------------------------------------------------------------------------------
276{-
277-- Generalization of concurrent streams/SVar via unfoldr.
278--
279-- Unfold a value into monadic actions and then run the resulting monadic
280-- actions to generate a stream. Since the step of generating the monadic
281-- action and running them are decoupled we can run the monadic actions
282-- cooncurrently. For example, the seed could be a list of monadic actions or a
283-- pure stream of monadic actions.
284--
285-- We can have different flavors of this depending on the stream type t. The
286-- concurrent version could be async or ahead etc. Depending on how we queue
287-- back the feedback portion b, it could be DFS or BFS style.
288--
289unfoldrA :: (b -> Maybe (m a, b)) -> b -> StreamK m a
290unfoldrA = undefined
291-}
293-------------------------------------------------------------------------------
294-- Special generation
295-------------------------------------------------------------------------------
297repeatM :: Monad m => m a -> StreamK m a
298repeatM = repeatMWith consM
300{-# INLINE replicateM #-}
301replicateM :: Monad m => Int -> m a -> StreamK m a
302replicateM = replicateMWith consM
303{-# INLINE replicate #-}
304replicate :: Int -> a -> StreamK m a
305replicate n a = go n
306 where
307 go cnt = if cnt <= 0 then nil else a `cons` go (cnt - 1)
309{-# INLINE fromIndicesM #-}
310fromIndicesM :: Monad m => (Int -> m a) -> StreamK m a
311fromIndicesM = fromIndicesMWith consM
312{-# INLINE fromIndices #-}
313fromIndices :: (Int -> a) -> StreamK m a
314fromIndices gen = go 0
315 where
316 go n = gen n `cons` go (n + 1)
318{-# INLINE iterate #-}
319iterate :: (a -> a) -> a -> StreamK m a
320iterate step = go
321 where
322 go !s = cons s (go (step s))
324{-# INLINE iterateM #-}
325iterateM :: Monad m => (a -> m a) -> m a -> StreamK m a
326iterateM = iterateMWith consM
328-------------------------------------------------------------------------------
329-- Conversions
330-------------------------------------------------------------------------------
332{-# INLINE fromList #-}
333fromList :: [a] -> StreamK m a
334fromList = fromFoldable
336-------------------------------------------------------------------------------
337-- Elimination by Folding
338-------------------------------------------------------------------------------
340{-# INLINE foldr1 #-}
341foldr1 :: Monad m => (a -> a -> a) -> StreamK m a -> m (Maybe a)
342foldr1 step m = do
343 r <- uncons m
344 case r of
345 Nothing -> return Nothing
346 Just (h, t) -> fmap Just (go h t)
347 where
348 go p m1 =
349 let stp = return p
350 single a = return $ step a p
351 yieldk a r = fmap (step p) (go a r)
352 in foldStream defState yieldk single stp m1
354-- XXX replace the recursive "go" with explicit continuations.
355-- | Like 'foldx', but with a monadic step function.
356{-# INLINABLE foldlMx' #-}
357foldlMx' :: Monad m
358 => (x -> a -> m x) -> m x -> (x -> m b) -> StreamK m a -> m b
359foldlMx' step begin done = go begin
360 where
361 go !acc m1 =
362 let stop = acc >>= done
363 single a = acc >>= \b -> step b a >>= done
364 yieldk a r = acc >>= \b -> step b a >>= \x -> go (return x) r
365 in foldStream defState yieldk single stop m1
367-- | Fold a stream using the supplied left 'Fold' and reducing the resulting
368-- expression strictly at each step. The behavior is similar to 'foldl''. A
369-- 'Fold' can terminate early without consuming the full stream. See the
370-- documentation of individual 'Fold's for termination behavior.
371--
372-- Definitions:
373--
374-- >>> fold f = fmap fst . StreamK.foldBreak f
375-- >>> fold f = StreamK.parseD (Parser.fromFold f)
376--
377-- Example:
378--
379-- >>> StreamK.fold Fold.sum $ StreamK.fromStream $ Stream.enumerateFromTo 1 100
380-- 5050
381--
382{-# INLINABLE fold #-}
383fold :: Monad m => FL.Fold m a b -> StreamK m a -> m b
384fold (FL.Fold step begin done) m = do
385 res <- begin
386 case res of
387 FL.Partial fs -> go fs m
388 FL.Done fb -> return fb
390 where
391 go !acc m1 =
392 let stop = done acc
393 single a = step acc a
394 >>= \case
395 FL.Partial s -> done s
396 FL.Done b1 -> return b1
397 yieldk a r = step acc a
398 >>= \case
399 FL.Partial s -> go s r
400 FL.Done b1 -> return b1
401 in foldStream defState yieldk single stop m1
403-- | Fold resulting in either breaking the stream or continuation of the fold.
404-- Instead of supplying the input stream in one go we can run the fold multiple
405-- times, each time supplying the next segment of the input stream. If the fold
406-- has not yet finished it returns a fold that can be run again otherwise it
407-- returns the fold result and the residual stream.
408--
409-- /Internal/
410{-# INLINE foldEither #-}
411foldEither :: Monad m =>
412 Fold m a b -> StreamK m a -> m (Either (Fold m a b) (b, StreamK m a))
413foldEither (FL.Fold step begin done) m = do
414 res <- begin
415 case res of
416 FL.Partial fs -> go fs m
417 FL.Done fb -> return $ Right (fb, m)
419 where
421 go !acc m1 =
422 let stop = return $ Left (Fold step (return $ FL.Partial acc) done)
423 single a =
424 step acc a
425 >>= \case
426 FL.Partial s ->
427 return $ Left (Fold step (return $ FL.Partial s) done)
428 FL.Done b1 -> return $ Right (b1, nil)
429 yieldk a r =
430 step acc a
431 >>= \case
432 FL.Partial s -> go s r
433 FL.Done b1 -> return $ Right (b1, r)
434 in foldStream defState yieldk single stop m1
436-- | Like 'fold' but also returns the remaining stream. The resulting stream
437-- would be 'StreamK.nil' if the stream finished before the fold.
438--
439{-# INLINE foldBreak #-}
440foldBreak :: Monad m => Fold m a b -> StreamK m a -> m (b, StreamK m a)
441foldBreak fld strm = do
442 r <- foldEither fld strm
443 case r of
444 Right res -> return res
445 Left (Fold _ initial extract) -> do
446 res <- initial
447 case res of
448 FL.Done _ -> error "foldBreak: unreachable state"
449 FL.Partial s -> do
450 b <- extract s
451 return (b, nil)
453-- XXX Array folds can be implemented using this.
454-- foldContainers? Specialized to foldArrays.
456-- | Generate streams from individual elements of a stream and fold the
457-- concatenation of those streams using the supplied fold. Return the result of
458-- the fold and residual stream.
459--
460-- For example, this can be used to efficiently fold an Array Word8 stream
461-- using Word8 folds.
462--
463-- /Internal/
464{-# INLINE foldConcat #-}
465foldConcat :: Monad m =>
466 Producer m a b -> Fold m b c -> StreamK m a -> m (c, StreamK m a)
467foldConcat
468 (Producer pstep pinject pextract)
469 (Fold fstep begin done)
470 stream = do
472 res <- begin
473 case res of
474 FL.Partial fs -> go fs stream
475 FL.Done fb -> return (fb, stream)
477 where
479 go !acc m1 = do
480 let stop = do
481 r <- done acc
482 return (r, nil)
483 single a = do
484 st <- pinject a
485 res <- go1 SPEC acc st
486 case res of
487 Left fs -> do
488 r <- done fs
489 return (r, nil)
490 Right (b, s) -> do
491 x <- pextract s
492 return (b, fromPure x)
493 yieldk a r = do
494 st <- pinject a
495 res <- go1 SPEC acc st
496 case res of
497 Left fs -> go fs r
498 Right (b, s) -> do
499 x <- pextract s
500 return (b, x `cons` r)
501 in foldStream defState yieldk single stop m1
503 {-# INLINE go1 #-}
504 go1 !_ !fs st = do
505 r <- pstep st
506 case r of
507 Stream.Yield x s -> do
508 res <- fstep fs x
509 case res of
510 FL.Done b -> return $ Right (b, s)
511 FL.Partial fs1 -> go1 SPEC fs1 s
512 Stream.Skip s -> go1 SPEC fs s
513 Stream.Stop -> return $ Left fs
515-- | Like 'foldl'' but with a monadic step function.
516{-# INLINE foldlM' #-}
517foldlM' :: Monad m => (b -> a -> m b) -> m b -> StreamK m a -> m b
518foldlM' step begin = foldlMx' step begin return
520------------------------------------------------------------------------------
521-- Specialized folds
522------------------------------------------------------------------------------
524{-# INLINE head #-}
525head :: Monad m => StreamK m a -> m (Maybe a)
526-- head = foldrM (\x _ -> return $ Just x) (return Nothing)
527head m =
528 let stop = return Nothing
529 single a = return (Just a)
530 yieldk a _ = return (Just a)
531 in foldStream defState yieldk single stop m
533{-# INLINE elem #-}
534elem :: (Monad m, Eq a) => a -> StreamK m a -> m Bool
535elem e = go
536 where
537 go m1 =
538 let stop = return False
539 single a = return (a == e)
540 yieldk a r = if a == e then return True else go r
541 in foldStream defState yieldk single stop m1
543{-# INLINE notElem #-}
544notElem :: (Monad m, Eq a) => a -> StreamK m a -> m Bool
545notElem e = go
546 where
547 go m1 =
548 let stop = return True
549 single a = return (a /= e)
550 yieldk a r = if a == e then return False else go r
551 in foldStream defState yieldk single stop m1
553{-# INLINABLE all #-}
554all :: Monad m => (a -> Bool) -> StreamK m a -> m Bool
555all p = go
556 where
557 go m1 =
558 let single a | p a = return True
559 | otherwise = return False
560 yieldk a r | p a = go r
561 | otherwise = return False
562 in foldStream defState yieldk single (return True) m1
564{-# INLINABLE any #-}
565any :: Monad m => (a -> Bool) -> StreamK m a -> m Bool
566any p = go
567 where
568 go m1 =
569 let single a | p a = return True
570 | otherwise = return False
571 yieldk a r | p a = return True
572 | otherwise = go r
573 in foldStream defState yieldk single (return False) m1
575-- | Extract the last element of the stream, if any.
576{-# INLINE last #-}
577last :: Monad m => StreamK m a -> m (Maybe a)
578last = foldlx' (\_ y -> Just y) Nothing id
580{-# INLINE minimum #-}
581minimum :: (Monad m, Ord a) => StreamK m a -> m (Maybe a)
582minimum = go Nothing
583 where
584 go Nothing m1 =
585 let stop = return Nothing
586 single a = return (Just a)
587 yieldk a r = go (Just a) r
588 in foldStream defState yieldk single stop m1
590 go (Just res) m1 =
591 let stop = return (Just res)
592 single a =
593 if res <= a
594 then return (Just res)
595 else return (Just a)
596 yieldk a r =
597 if res <= a
598 then go (Just res) r
599 else go (Just a) r
600 in foldStream defState yieldk single stop m1
602{-# INLINE minimumBy #-}
603minimumBy
604 :: (Monad m)
605 => (a -> a -> Ordering) -> StreamK m a -> m (Maybe a)
606minimumBy cmp = go Nothing
607 where
608 go Nothing m1 =
609 let stop = return Nothing
610 single a = return (Just a)
611 yieldk a r = go (Just a) r
612 in foldStream defState yieldk single stop m1
614 go (Just res) m1 =
615 let stop = return (Just res)
616 single a = case cmp res a of
617 GT -> return (Just a)
618 _ -> return (Just res)
619 yieldk a r = case cmp res a of
620 GT -> go (Just a) r
621 _ -> go (Just res) r
622 in foldStream defState yieldk single stop m1
624{-# INLINE maximum #-}
625maximum :: (Monad m, Ord a) => StreamK m a -> m (Maybe a)
626maximum = go Nothing
627 where
628 go Nothing m1 =
629 let stop = return Nothing
630 single a = return (Just a)
631 yieldk a r = go (Just a) r
632 in foldStream defState yieldk single stop m1
634 go (Just res) m1 =
635 let stop = return (Just res)
636 single a =
637 if res <= a
638 then return (Just a)
639 else return (Just res)
640 yieldk a r =
641 if res <= a
642 then go (Just a) r
643 else go (Just res) r
644 in foldStream defState yieldk single stop m1
646{-# INLINE maximumBy #-}
647maximumBy :: Monad m => (a -> a -> Ordering) -> StreamK m a -> m (Maybe a)
648maximumBy cmp = go Nothing
649 where
650 go Nothing m1 =
651 let stop = return Nothing
652 single a = return (Just a)
653 yieldk a r = go (Just a) r
654 in foldStream defState yieldk single stop m1
656 go (Just res) m1 =
657 let stop = return (Just res)
658 single a = case cmp res a of
659 GT -> return (Just res)
660 _ -> return (Just a)
661 yieldk a r = case cmp res a of
662 GT -> go (Just res) r
663 _ -> go (Just a) r
664 in foldStream defState yieldk single stop m1
666{-# INLINE (!!) #-}
667(!!) :: Monad m => StreamK m a -> Int -> m (Maybe a)
668m !! i = go i m
669 where
670 go n m1 =
671 let single a | n == 0 = return $ Just a
672 | otherwise = return Nothing
673 yieldk a x | n < 0 = return Nothing
674 | n == 0 = return $ Just a
675 | otherwise = go (n - 1) x
676 in foldStream defState yieldk single (return Nothing) m1
678{-# INLINE lookup #-}
679lookup :: (Monad m, Eq a) => a -> StreamK m (a, b) -> m (Maybe b)
680lookup e = go
681 where
682 go m1 =
683 let single (a, b) | a == e = return $ Just b
684 | otherwise = return Nothing
685 yieldk (a, b) x | a == e = return $ Just b
686 | otherwise = go x
687 in foldStream defState yieldk single (return Nothing) m1
689{-# INLINE findM #-}
690findM :: Monad m => (a -> m Bool) -> StreamK m a -> m (Maybe a)
691findM p = go
692 where
693 go m1 =
694 let single a = do
695 b <- p a
696 if b then return $ Just a else return Nothing
697 yieldk a x = do
698 b <- p a
699 if b then return $ Just a else go x
700 in foldStream defState yieldk single (return Nothing) m1
702{-# INLINE find #-}
703find :: Monad m => (a -> Bool) -> StreamK m a -> m (Maybe a)
704find p = findM (return . p)
706{-# INLINE findIndices #-}
707findIndices :: (a -> Bool) -> StreamK m a -> StreamK m Int
708findIndices p = go 0
709 where
710 go offset m1 = mkStream $ \st yld sng stp ->
711 let single a | p a = sng offset
712 | otherwise = stp
713 yieldk a x | p a = yld offset $ go (offset + 1) x
714 | otherwise = foldStream (adaptState st) yld sng stp $
715 go (offset + 1) x
716 in foldStream (adaptState st) yieldk single stp m1
718------------------------------------------------------------------------------
719-- Map and Fold
720------------------------------------------------------------------------------
722-- | Apply a monadic action to each element of the stream and discard the
723-- output of the action.
724{-# INLINE mapM_ #-}
725mapM_ :: Monad m => (a -> m b) -> StreamK m a -> m ()
726mapM_ f = go
727 where
728 go m1 =
729 let stop = return ()
730 single a = void (f a)
731 yieldk a r = f a >> go r
732 in foldStream defState yieldk single stop m1
734{-# INLINE mapM #-}
735mapM :: Monad m => (a -> m b) -> StreamK m a -> StreamK m b
736mapM = mapMWith consM
738------------------------------------------------------------------------------
739-- Converting folds
740------------------------------------------------------------------------------
742{-# INLINABLE toList #-}
743toList :: Monad m => StreamK m a -> m [a]
744toList = foldr (:) []
746-- Based on suggestions by David Feuer and Pranay Sashank
747{-# INLINE hoist #-}
748hoist :: (Monad m, Monad n)
749 => (forall x. m x -> n x) -> StreamK m a -> StreamK n a
750hoist f str =
751 mkStream $ \st yld sng stp ->
752 let single = return . sng
753 yieldk a s = return $ yld a (hoist f s)
754 stop = return stp
755 state = adaptState st
756 in join . f $ foldStreamShared state yieldk single stop str
758-------------------------------------------------------------------------------
759-- Transformation by folding (Scans)
760-------------------------------------------------------------------------------
762{-# INLINE scanlx' #-}
763scanlx' :: (x -> a -> x) -> x -> (x -> b) -> StreamK m a -> StreamK m b
764scanlx' step begin done m =
765 cons (done begin) $ go m begin
766 where
767 go m1 !acc = mkStream $ \st yld sng stp ->
768 let single a = sng (done $ step acc a)
769 yieldk a r =
770 let s = step acc a
771 in yld (done s) (go r s)
772 in foldStream (adaptState st) yieldk single stp m1
774{-# INLINE scanl' #-}
775scanl' :: (b -> a -> b) -> b -> StreamK m a -> StreamK m b
776scanl' step begin = scanlx' step begin id
778-------------------------------------------------------------------------------
779-- Filtering
780-------------------------------------------------------------------------------
782{-# INLINE filter #-}
783filter :: (a -> Bool) -> StreamK m a -> StreamK m a
784filter p = go
785 where
786 go m1 = mkStream $ \st yld sng stp ->
787 let single a | p a = sng a
788 | otherwise = stp
789 yieldk a r | p a = yld a (go r)
790 | otherwise = foldStream st yieldk single stp r
791 in foldStream st yieldk single stp m1
793{-# INLINE take #-}
794take :: Int -> StreamK m a -> StreamK m a
795take = go
796 where
797 go n1 m1 = mkStream $ \st yld sng stp ->
798 let yieldk a r = yld a (go (n1 - 1) r)
799 in if n1 <= 0
800 then stp
801 else foldStream st yieldk sng stp m1
803{-# INLINE takeWhile #-}
804takeWhile :: (a -> Bool) -> StreamK m a -> StreamK m a
805takeWhile p = go
806 where
807 go m1 = mkStream $ \st yld sng stp ->
808 let single a | p a = sng a
809 | otherwise = stp
810 yieldk a r | p a = yld a (go r)
811 | otherwise = stp
812 in foldStream st yieldk single stp m1
814{-# INLINE drop #-}
815drop :: Int -> StreamK m a -> StreamK m a
816drop n m = unShare (go n m)
817 where
818 go n1 m1 = mkStream $ \st yld sng stp ->
819 let single _ = stp
820 yieldk _ r = foldStreamShared st yld sng stp $ go (n1 - 1) r
821 -- Somehow "<=" check performs better than a ">"
822 in if n1 <= 0
823 then foldStreamShared st yld sng stp m1
824 else foldStreamShared st yieldk single stp m1
826{-# INLINE dropWhile #-}
827dropWhile :: (a -> Bool) -> StreamK m a -> StreamK m a
828dropWhile p = go
829 where
830 go m1 = mkStream $ \st yld sng stp ->
831 let single a | p a = stp
832 | otherwise = sng a
833 yieldk a r | p a = foldStream st yieldk single stp r
834 | otherwise = yld a r
835 in foldStream st yieldk single stp m1
837-------------------------------------------------------------------------------
838-- Mapping
839-------------------------------------------------------------------------------
841-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
842-- other stream types to overload it.
843{-# INLINE sequence #-}
844sequence :: Monad m => StreamK m (m a) -> StreamK m a
845sequence = go
846 where
847 go m1 = mkStream $ \st yld sng stp ->
848 let single ma = ma >>= sng
849 yieldk ma r = foldStreamShared st yld sng stp $ ma `consM` go r
850 in foldStream (adaptState st) yieldk single stp m1
852-------------------------------------------------------------------------------
853-- Inserting
854-------------------------------------------------------------------------------
856{-# INLINE intersperseM #-}
857intersperseM :: Monad m => m a -> StreamK m a -> StreamK m a
858intersperseM a = prependingStart
859 where
860 prependingStart m1 = mkStream $ \st yld sng stp ->
861 let yieldk i x =
862 foldStreamShared st yld sng stp $ return i `consM` go x
863 in foldStream st yieldk sng stp m1
864 go m2 = mkStream $ \st yld sng stp ->
865 let single i = foldStreamShared st yld sng stp $ a `consM` fromPure i
866 yieldk i x =
867 foldStreamShared
868 st yld sng stp $ a `consM` return i `consM` go x
869 in foldStream st yieldk single stp m2
871{-# INLINE intersperse #-}
872intersperse :: Monad m => a -> StreamK m a -> StreamK m a
873intersperse a = intersperseM (return a)
875{-# INLINE insertBy #-}
876insertBy :: (a -> a -> Ordering) -> a -> StreamK m a -> StreamK m a
877insertBy cmp x = go
878 where
879 go m1 = mkStream $ \st yld _ _ ->
880 let single a = case cmp x a of
881 GT -> yld a (fromPure x)
882 _ -> yld x (fromPure a)
883 stop = yld x nil
884 yieldk a r = case cmp x a of
885 GT -> yld a (go r)
886 _ -> yld x (a `cons` r)
887 in foldStream st yieldk single stop m1
889------------------------------------------------------------------------------
890-- Deleting
891------------------------------------------------------------------------------
893{-# INLINE deleteBy #-}
894deleteBy :: (a -> a -> Bool) -> a -> StreamK m a -> StreamK m a
895deleteBy eq x = go
896 where
897 go m1 = mkStream $ \st yld sng stp ->
898 let single a = if eq x a then stp else sng a
899 yieldk a r = if eq x a
900 then foldStream st yld sng stp r
901 else yld a (go r)
902 in foldStream st yieldk single stp m1
904-------------------------------------------------------------------------------
905-- Map and Filter
906-------------------------------------------------------------------------------
908{-# INLINE mapMaybe #-}
909mapMaybe :: (a -> Maybe b) -> StreamK m a -> StreamK m b
910mapMaybe f = go
911 where
912 go m1 = mkStream $ \st yld sng stp ->
913 let single a = maybe stp sng (f a)
914 yieldk a r = case f a of
915 Just b -> yld b $ go r
916 Nothing -> foldStream (adaptState st) yieldk single stp r
917 in foldStream (adaptState st) yieldk single stp m1
919------------------------------------------------------------------------------
920-- Serial Zipping
921------------------------------------------------------------------------------
923-- | Zip two streams serially using a pure zipping function.
924--
925-- @since 0.1.0
926{-# INLINE zipWith #-}
927zipWith :: Monad m => (a -> b -> c) -> StreamK m a -> StreamK m b -> StreamK m c
928zipWith f = zipWithM (\a b -> return (f a b))
930-- | Zip two streams serially using a monadic zipping function.
931--
932-- @since 0.1.0
933{-# INLINE zipWithM #-}
934zipWithM :: Monad m =>
935 (a -> b -> m c) -> StreamK m a -> StreamK m b -> StreamK m c
936zipWithM f = go
938 where
940 go mx my = mkStream $ \st yld sng stp -> do
941 let merge a ra =
942 let single2 b = f a b >>= sng
943 yield2 b rb = f a b >>= \x -> yld x (go ra rb)
944 in foldStream (adaptState st) yield2 single2 stp my
945 let single1 a = merge a nil
946 yield1 = merge
947 foldStream (adaptState st) yield1 single1 stp mx
949------------------------------------------------------------------------------
950-- Merging
951------------------------------------------------------------------------------
953{-# INLINE mergeByM #-}
954mergeByM :: Monad m =>
955 (a -> a -> m Ordering) -> StreamK m a -> StreamK m a -> StreamK m a
956mergeByM cmp = go
958 where
960 go mx my = mkStream $ \st yld sng stp -> do
961 let stop = foldStream st yld sng stp my
962 single x = foldStream st yld sng stp (goX0 x my)
963 yield x rx = foldStream st yld sng stp (goX x rx my)
964 foldStream st yield single stop mx
966 goX0 x my = mkStream $ \st yld sng _ -> do
967 let stop = sng x
968 single y = do
969 r <- cmp x y
970 case r of
971 GT -> yld y (fromPure x)
972 _ -> yld x (fromPure y)
973 yield y ry = do
974 r <- cmp x y
975 case r of
976 GT -> yld y (goX0 x ry)
977 _ -> yld x (y `cons` ry)
978 in foldStream st yield single stop my
980 goX x mx my = mkStream $ \st yld _ _ -> do
981 let stop = yld x mx
982 single y = do
983 r <- cmp x y
984 case r of
985 GT -> yld y (x `cons` mx)
986 _ -> yld x (goY0 mx y)
987 yield y ry = do
988 r <- cmp x y
989 case r of
990 GT -> yld y (goX x mx ry)
991 _ -> yld x (goY mx y ry)
992 in foldStream st yield single stop my
994 goY0 mx y = mkStream $ \st yld sng _ -> do
995 let stop = sng y
996 single x = do
997 r <- cmp x y
998 case r of
999 GT -> yld y (fromPure x)
1000 _ -> yld x (fromPure y)
1001 yield x rx = do
1002 r <- cmp x y
1003 case r of
1004 GT -> yld y (x `cons` rx)
1005 _ -> yld x (goY0 rx y)
1006 in foldStream st yield single stop mx
1008 goY mx y my = mkStream $ \st yld _ _ -> do
1009 let stop = yld y my
1010 single x = do
1011 r <- cmp x y
1012 case r of
1013 GT -> yld y (goX0 x my)
1014 _ -> yld x (y `cons` my)
1015 yield x rx = do
1016 r <- cmp x y
1017 case r of
1018 GT -> yld y (goX x rx my)
1019 _ -> yld x (goY rx y my)
1020 in foldStream st yield single stop mx
1022{-# INLINE mergeBy #-}
1023mergeBy :: (a -> a -> Ordering) -> StreamK m a -> StreamK m a -> StreamK m a
1024-- XXX GHC: This has slightly worse performance than replacing "r <- cmp x y"
1025-- with "let r = cmp x y" in the monadic version. The definition below is
1026-- exactly the same as mergeByM except this change.
1027-- mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
1028mergeBy cmp = go
1030 where
1032 go mx my = mkStream $ \st yld sng stp -> do
1033 let stop = foldStream st yld sng stp my
1034 single x = foldStream st yld sng stp (goX0 x my)
1035 yield x rx = foldStream st yld sng stp (goX x rx my)
1036 foldStream st yield single stop mx
1038 goX0 x my = mkStream $ \st yld sng _ -> do
1039 let stop = sng x
1040 single y = do
1041 case cmp x y of
1042 GT -> yld y (fromPure x)
1043 _ -> yld x (fromPure y)
1044 yield y ry = do
1045 case cmp x y of
1046 GT -> yld y (goX0 x ry)
1047 _ -> yld x (y `cons` ry)
1048 in foldStream st yield single stop my
1050 goX x mx my = mkStream $ \st yld _ _ -> do
1051 let stop = yld x mx
1052 single y = do
1053 case cmp x y of
1054 GT -> yld y (x `cons` mx)
1055 _ -> yld x (goY0 mx y)
1056 yield y ry = do
1057 case cmp x y of
1058 GT -> yld y (goX x mx ry)
1059 _ -> yld x (goY mx y ry)
1060 in foldStream st yield single stop my
1062 goY0 mx y = mkStream $ \st yld sng _ -> do
1063 let stop = sng y
1064 single x = do
1065 case cmp x y of
1066 GT -> yld y (fromPure x)
1067 _ -> yld x (fromPure y)
1068 yield x rx = do
1069 case cmp x y of
1070 GT -> yld y (x `cons` rx)
1071 _ -> yld x (goY0 rx y)
1072 in foldStream st yield single stop mx
1074 goY mx y my = mkStream $ \st yld _ _ -> do
1075 let stop = yld y my
1076 single x = do
1077 case cmp x y of
1078 GT -> yld y (goX0 x my)
1079 _ -> yld x (y `cons` my)
1080 yield x rx = do
1081 case cmp x y of
1082 GT -> yld y (goX x rx my)
1083 _ -> yld x (goY rx y my)
1084 in foldStream st yield single stop mx
1086------------------------------------------------------------------------------
1087-- Transformation comprehensions
1088------------------------------------------------------------------------------
1090{-# INLINE the #-}
1091the :: (Eq a, Monad m) => StreamK m a -> m (Maybe a)
1092the m = do
1093 r <- uncons m
1094 case r of
1095 Nothing -> return Nothing
1096 Just (h, t) -> go h t
1097 where
1098 go h m1 =
1099 let single a | h == a = return $ Just h
1100 | otherwise = return Nothing
1101 yieldk a r | h == a = go h r
1102 | otherwise = return Nothing
1103 in foldStream defState yieldk single (return $ Just h) m1
1105------------------------------------------------------------------------------
1106-- Alternative & MonadPlus
1107------------------------------------------------------------------------------
1109_alt :: StreamK m a -> StreamK m a -> StreamK m a
1110_alt m1 m2 = mkStream $ \st yld sng stp ->
1111 let stop = foldStream st yld sng stp m2
1112 in foldStream st yld sng stop m1
1114------------------------------------------------------------------------------
1115-- MonadError
1116------------------------------------------------------------------------------
1119-- XXX handle and test cross thread state transfer
1120withCatchError
1121 :: MonadError e m
1122 => StreamK m a -> (e -> StreamK m a) -> StreamK m a
1123withCatchError m h =
1124 mkStream $ \_ stp sng yld ->
1125 let run x = unStream x Nothing stp sng yieldk
1126 handle r = r `catchError` \e -> run $ h e
1127 yieldk a r = yld a (withCatchError r h)
1128 in handle $ run m
1131-------------------------------------------------------------------------------
1132-- Parsing
1133-------------------------------------------------------------------------------
1135-- Inlined definition.
1136{-# INLINE splitAt #-}
1137splitAt :: Int -> [a] -> ([a],[a])
1138splitAt n ls
1139 | n <= 0 = ([], ls)
1140 | otherwise = splitAt' n ls
1141 where
1142 splitAt' :: Int -> [a] -> ([a], [a])
1143 splitAt' _ [] = ([], [])
1144 splitAt' 1 (x:xs) = ([x], xs)
1145 splitAt' m (x:xs) = (x:xs', xs'')
1146 where
1147 (xs', xs'') = splitAt' (m - 1) xs
1149-- | Run a 'Parser' over a stream and return rest of the Stream.
1150{-# INLINE_NORMAL parseDBreak #-}
1151parseDBreak
1152 :: Monad m
1153 => PR.Parser a m b
1154 -> StreamK m a
1155 -> m (Either ParseError b, StreamK m a)
1156parseDBreak (PR.Parser pstep initial extract) stream = do
1157 res <- initial
1158 case res of
1159 PR.IPartial s -> goStream stream [] s
1160 PR.IDone b -> return (Right b, stream)
1161 PR.IError err -> return (Left (ParseError err), stream)
1163 where
1165 -- "buf" contains last few items in the stream that we may have to
1166 -- backtrack to.
1167 --
1168 -- XXX currently we are using a dumb list based approach for backtracking
1169 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
1170 -- That will allow us more efficient random back and forth movement.
1171 goStream st buf !pst =
1172 let stop = do
1173 r <- extract pst
1174 case r of
1175 PR.Error err -> return (Left (ParseError err), nil)
1176 PR.Done n b -> do
1177 assertM(n <= length buf)
1178 let src0 = Prelude.take n buf
1179 src = Prelude.reverse src0
1180 return (Right b, fromList src)
1181 PR.Partial _ _ -> error "Bug: parseBreak: Partial in extract"
1182 PR.Continue 0 s -> goStream nil buf s
1183 PR.Continue n s -> do
1184 assertM(n <= length buf)
1185 let (src0, buf1) = splitAt n buf
1186 src = Prelude.reverse src0
1187 goBuf nil buf1 src s
1188 single x = yieldk x nil
1189 yieldk x r = do
1190 res <- pstep pst x
1191 case res of
1192 PR.Partial 0 s -> goStream r [] s
1193 PR.Partial n s -> do
1194 assertM(n <= length (x:buf))
1195 let src0 = Prelude.take n (x:buf)
1196 src = Prelude.reverse src0
1197 goBuf r [] src s
1198 PR.Continue 0 s -> goStream r (x:buf) s
1199 PR.Continue n s -> do
1200 assertM(n <= length (x:buf))
1201 let (src0, buf1) = splitAt n (x:buf)
1202 src = Prelude.reverse src0
1203 goBuf r buf1 src s
1204 PR.Done 0 b -> return (Right b, r)
1205 PR.Done n b -> do
1206 assertM(n <= length (x:buf))
1207 let src0 = Prelude.take n (x:buf)
1208 src = Prelude.reverse src0
1209 return (Right b, append (fromList src) r)
1210 PR.Error err -> return (Left (ParseError err), r)
1211 in foldStream defState yieldk single stop st
1213 goBuf st buf [] !pst = goStream st buf pst
1214 goBuf st buf (x:xs) !pst = do
1215 pRes <- pstep pst x
1216 case pRes of
1217 PR.Partial 0 s -> goBuf st [] xs s
1218 PR.Partial n s -> do
1219 assert (n <= length (x:buf)) (return ())
1220 let src0 = Prelude.take n (x:buf)
1221 src = Prelude.reverse src0 ++ xs
1222 goBuf st [] src s
1223 PR.Continue 0 s -> goBuf st (x:buf) xs s
1224 PR.Continue n s -> do
1225 assert (n <= length (x:buf)) (return ())
1226 let (src0, buf1) = splitAt n (x:buf)
1227 src = Prelude.reverse src0 ++ xs
1228 goBuf st buf1 src s
1229 PR.Done n b -> do
1230 assert (n <= length (x:buf)) (return ())
1231 let src0 = Prelude.take n (x:buf)
1232 src = Prelude.reverse src0
1233 return (Right b, append (fromList src) st)
1234 PR.Error err -> return (Left (ParseError err), nil)
1236-- Using ParserD or ParserK on StreamK may not make much difference. We should
1237-- perhaps use only chunked parsing on StreamK. We can always convert a stream
1238-- to chunks before parsing. Or just have a ParserK element parser for StreamK
1239-- and convert ParserD to ParserK for element parsing using StreamK.
1240{-# INLINE parseD #-}
1241parseD :: Monad m =>
1242 Parser.Parser a m b -> StreamK m a -> m (Either ParseError b)
1243parseD f = fmap fst . parseDBreak f
1245-------------------------------------------------------------------------------
1246-- Chunked parsing using ParserK
1247-------------------------------------------------------------------------------
1249-- The backracking buffer consists of arrays in the most-recent-first order. We
1250-- want to take a total of n array elements from this buffer. Note: when we
1251-- have to take an array partially, we must take the last part of the array.
1252{-# INLINE backTrack #-}
1253backTrack :: forall m a. Unbox a =>
1254 Int
1255 -> [Array a]
1256 -> StreamK m (Array a)
1257 -> (StreamK m (Array a), [Array a])
1258backTrack = go
1260 where
1262 go _ [] stream = (stream, [])
1263 go n xs stream | n <= 0 = (stream, xs)
1264 go n (x:xs) stream =
1265 let len = Array.length x
1266 in if n > len
1267 then go (n - len) xs (cons x stream)
1268 else if n == len
1269 then (cons x stream, xs)
1270 else let !(Array contents start end) = x
1271 !start1 = end - (n * SIZE_OF(a))
1272 arr1 = Array contents start1 end
1273 arr2 = Array contents start start1
1274 in (cons arr1 stream, arr2:xs)
1276-- | A continuation to extract the result when a CPS parser is done.
1277{-# INLINE parserDone #-}
1278parserDone :: Applicative m =>
1279 ParserK.ParseResult b -> Int -> ParserK.Input a -> m (ParserK.Step a m b)
1280parserDone (ParserK.Success n b) _ _ = pure $ ParserK.Done n b
1281parserDone (ParserK.Failure n e) _ _ = pure $ ParserK.Error n e
1283-- | Run a 'ParserK' over a chunked 'StreamK' and return the rest of the Stream.
1284{-# INLINE_NORMAL parseBreakChunks #-}
1285parseBreakChunks
1286 :: (Monad m, Unbox a)
1287 => ParserK a m b
1288 -> StreamK m (Array a)
1289 -> m (Either ParseError b, StreamK m (Array a))
1290parseBreakChunks parser input = do
1291 let parserk = ParserK.runParser parser parserDone 0 0
1292 in go [] parserk input
1294 where
1296 {-# INLINE goStop #-}
1297 goStop backBuf parserk = do
1298 pRes <- parserk ParserK.None
1299 case pRes of
1300 -- If we stop in an alternative, it will try calling the next
1301 -- parser, the next parser may call initial returning Partial and
1302 -- then immediately we have to call extract on it.
1303 ParserK.Partial 0 cont1 ->
1304 go [] cont1 nil
1305 ParserK.Partial n cont1 -> do
1306 let n1 = negate n
1307 assertM(n1 >= 0 && n1 <= sum (Prelude.map Array.length backBuf))
1308 let (s1, backBuf1) = backTrack n1 backBuf nil
1309 in go backBuf1 cont1 s1
1310 ParserK.Continue 0 cont1 ->
1311 go backBuf cont1 nil
1312 ParserK.Continue n cont1 -> do
1313 let n1 = negate n
1314 assertM(n1 >= 0 && n1 <= sum (Prelude.map Array.length backBuf))
1315 let (s1, backBuf1) = backTrack n1 backBuf nil
1316 in go backBuf1 cont1 s1
1317 ParserK.Done 0 b ->
1318 return (Right b, nil)
1319 ParserK.Done n b -> do
1320 let n1 = negate n
1321 assertM(n1 >= 0 && n1 <= sum (Prelude.map Array.length backBuf))
1322 let (s1, _) = backTrack n1 backBuf nil
1323 in return (Right b, s1)
1324 ParserK.Error _ err -> return (Left (ParseError err), nil)
1326 seekErr n len =
1327 error $ "parseBreak: Partial: forward seek not implemented n = "
1328 ++ show n ++ " len = " ++ show len
1330 yieldk backBuf parserk arr stream = do
1331 pRes <- parserk (ParserK.Chunk arr)
1332 let len = Array.length arr
1333 case pRes of
1334 ParserK.Partial n cont1 ->
1335 case compare n len of
1336 EQ -> go [] cont1 stream
1337 LT -> do
1338 if n >= 0
1339 then yieldk [] cont1 arr stream
1340 else do
1341 let n1 = negate n
1342 bufLen = sum (Prelude.map Array.length backBuf)
1343 s = cons arr stream
1344 assertM(n1 >= 0 && n1 <= bufLen)
1345 let (s1, _) = backTrack n1 backBuf s
1346 go [] cont1 s1
1347 GT -> seekErr n len
1348 ParserK.Continue n cont1 ->
1349 case compare n len of
1350 EQ -> go (arr:backBuf) cont1 stream
1351 LT -> do
1352 if n >= 0
1353 then yieldk backBuf cont1 arr stream
1354 else do
1355 let n1 = negate n
1356 bufLen = sum (Prelude.map Array.length backBuf)
1357 s = cons arr stream
1358 assertM(n1 >= 0 && n1 <= bufLen)
1359 let (s1, backBuf1) = backTrack n1 backBuf s
1360 go backBuf1 cont1 s1
1361 GT -> seekErr n len
1362 ParserK.Done n b -> do
1363 let n1 = len - n
1364 assertM(n1 <= sum (Prelude.map Array.length (arr:backBuf)))
1365 let (s1, _) = backTrack n1 (arr:backBuf) stream
1366 in return (Right b, s1)
1367 ParserK.Error _ err -> return (Left (ParseError err), nil)
1369 go backBuf parserk stream = do
1370 let stop = goStop backBuf parserk
1371 single a = yieldk backBuf parserk a nil
1372 in foldStream
1373 defState (yieldk backBuf parserk) single stop stream
1375{-# INLINE parseChunks #-}
1376parseChunks :: (Monad m, Unbox a) =>
1377 ParserK a m b -> StreamK m (Array a) -> m (Either ParseError b)
1378parseChunks f = fmap fst . parseBreakChunks f
1380-------------------------------------------------------------------------------
1381-- Sorting
1382-------------------------------------------------------------------------------
1384-- | Sort the input stream using a supplied comparison function.
1386-- Sorting can be achieved by simply:
1388-- >>> sortBy cmp = StreamK.mergeMapWith (StreamK.mergeBy cmp) StreamK.fromPure
1390-- However, this combinator uses a parser to first split the input stream into
1391-- down and up sorted segments and then merges them to optimize sorting when
1392-- pre-sorted sequences exist in the input stream.
1394-- /O(n) space/
1396{-# INLINE sortBy #-}
1397sortBy :: Monad m => (a -> a -> Ordering) -> StreamK m a -> StreamK m a
1398-- sortBy f = Stream.concatPairsWith (Stream.mergeBy f) Stream.fromPure
1399sortBy cmp =
1400 let p =
1401 Parser.groupByRollingEither
1402 (\x -> (< GT) . cmp x)
1403 FL.toStreamKRev
1404 FL.toStreamK
1405 in mergeMapWith (mergeBy cmp) id
1406 . Stream.toStreamK
1407 . Stream.catRights -- its a non-failing backtracking parser
1408 . Stream.parseMany (fmap (either id id) p)
1409 . Stream.fromStreamK