git.iko.soy/streamly
on cedc81c help sign in
streamly/core/src/Streamly/Internal/Data/StreamK.hs
1735 lines 62 kB history raw
1{-# LANGUAGE CPP #-}
2-- |
3-- Module : Streamly.Internal.Data.StreamK
4-- Copyright : (c) 2017 Composewell Technologies
5--
6-- License : BSD3
7-- Maintainer : streamly@composewell.com
8-- Stability : experimental
9-- Portability : GHC
10--
11module Streamly.Internal.Data.StreamK
12 (
13 -- * Setup
14 -- | To execute the code examples provided in this module in ghci, please
15 -- run the following commands first.
16 --
17 -- $setup
19 -- * The stream type
20 module Streamly.Internal.Data.StreamK.Type
21 , module Streamly.Internal.Data.StreamK.Transformer
23 , StreamK(..)
24 , fromList
25 , fromStream
26 , toStream
28 -- ** Specialized Generation
29 , repeatM
30 , replicate
31 , replicateM
32 , fromIndices
33 , fromIndicesM
34 , iterate
35 , iterateM
37 -- * Elimination
38 -- ** General Folds
39 , foldr1
40 , foldlM'
41 , foldlMx'
42 , fold
43 , foldBreak
44 , foldEither
45 , foldConcat
46 , parseDBreak
47 , parseD
48 , parseBreakChunks
49 , parseChunks
50 , parseBreak
51 , parse
52 , parseBreakChunksGeneric
53 , parseChunksGeneric
55 -- ** Specialized Folds
56 , head
57 , elem
58 , notElem
59 , all
60 , any
61 , last
62 , minimum
63 , minimumBy
64 , maximum
65 , maximumBy
66 , findIndices
67 , lookup
68 , findM
69 , find
70 , (!!)
72 -- ** Map and Fold
73 , mapM_
75 -- ** Conversions
76 , toList
77 , hoist
79 -- * Transformation
80 -- ** By folding (scans)
81 , scanl'
82 , scanlx'
84 -- ** Filtering
85 , filter
86 , take
87 , takeWhile
88 , drop
89 , dropWhile
91 -- ** Mapping
92 , mapM
93 , sequence
95 -- ** Inserting
96 , intersperseM
97 , intersperse
98 , insertBy
100 -- ** Deleting
101 , deleteBy
103 -- ** Reordering
104 , sortBy
106 -- ** Map and Filter
107 , mapMaybe
109 -- ** Zipping
110 , zipWith
111 , zipWithM
113 -- ** Merging
114 , mergeBy
115 , mergeByM
117 -- ** Transformation comprehensions
118 , the
120 -- * Exceptions
121 , handle
123 -- * Resource Management
124 , bracketIO
125 )
126where
128#include "ArrayMacros.h"
129#include "inline.hs"
130#include "assert.hs"
132import Control.Exception (mask_, Exception)
133import Control.Monad (void, join)
134import Control.Monad.Catch (MonadCatch)
135import Control.Monad.IO.Class (MonadIO(..))
136import Data.Proxy (Proxy(..))
137import GHC.Types (SPEC(..))
138import Streamly.Internal.Data.Array.Type (Array(..))
139import Streamly.Internal.Data.Fold.Type (Fold(..))
140import Streamly.Internal.Data.IOFinalizer (newIOFinalizer, runIOFinalizer)
141import Streamly.Internal.Data.ParserK.Type (ParserK)
142import Streamly.Internal.Data.Producer.Type (Producer(..))
143import Streamly.Internal.Data.SVar.Type (adaptState, defState)
144import Streamly.Internal.Data.Unbox (sizeOf, Unbox)
146import qualified Control.Monad.Catch as MC
147import qualified Streamly.Internal.Data.Array.Type as Array
148import qualified Streamly.Internal.Data.Array.Generic as GenArr
149import qualified Streamly.Internal.Data.Fold.Type as FL
150import qualified Streamly.Internal.Data.Parser as Parser
151import qualified Streamly.Internal.Data.Parser.Type as PR
152import qualified Streamly.Internal.Data.ParserK.Type as ParserK
153import qualified Streamly.Internal.Data.Stream as Stream
154import qualified Prelude
156import Prelude
157 hiding (Foldable(..), last, map, mapM, mapM_, repeat, sequence,
158 take, filter, all, any, takeWhile, drop, dropWhile,
159 notElem, head, tail, init, zipWith, lookup,
160 (!!), replicate, reverse, concatMap, iterate, splitAt)
161import Data.Foldable (sum, length)
162import Streamly.Internal.Data.StreamK.Type
163import Streamly.Internal.Data.StreamK.Transformer
164import Streamly.Internal.Data.Parser (ParseError(..))
166#include "DocTestDataStreamK.hs"
168-- | Convert a fused 'Stream' to 'StreamK'.
169--
170-- For example:
171--
172-- >>> s1 = StreamK.fromStream $ Stream.fromList [1,2]
173-- >>> s2 = StreamK.fromStream $ Stream.fromList [3,4]
174-- >>> Stream.fold Fold.toList $ StreamK.toStream $ s1 `StreamK.append` s2
175-- [1,2,3,4]
176--
177{-# INLINE fromStream #-}
178fromStream :: Monad m => Stream.Stream m a -> StreamK m a
179fromStream = Stream.toStreamK
181-- | Convert a 'StreamK' to a fused 'Stream'.
182--
183{-# INLINE toStream #-}
184toStream :: Applicative m => StreamK m a -> Stream.Stream m a
185toStream = Stream.fromStreamK
187-------------------------------------------------------------------------------
188-- Generation
189-------------------------------------------------------------------------------
191{-
192-- Generalization of concurrent streams/SVar via unfoldr.
193--
194-- Unfold a value into monadic actions and then run the resulting monadic
195-- actions to generate a stream. Since the step of generating the monadic
196-- action and running them are decoupled we can run the monadic actions
197-- cooncurrently. For example, the seed could be a list of monadic actions or a
198-- pure stream of monadic actions.
199--
200-- We can have different flavors of this depending on the stream type t. The
201-- concurrent version could be async or ahead etc. Depending on how we queue
202-- back the feedback portion b, it could be DFS or BFS style.
203--
204unfoldrA :: (b -> Maybe (m a, b)) -> b -> StreamK m a
205unfoldrA = undefined
206-}
208-------------------------------------------------------------------------------
209-- Special generation
210-------------------------------------------------------------------------------
212-- |
213-- >>> repeatM = StreamK.sequence . StreamK.repeat
214-- >>> repeatM = fix . StreamK.consM
215-- >>> repeatM = cycle1 . StreamK.fromEffect
216--
217-- Generate a stream by repeatedly executing a monadic action forever.
218--
219-- >>> :{
220-- repeatAction =
221-- StreamK.repeatM (threadDelay 1000000 >> print 1)
222-- & StreamK.take 10
223-- & StreamK.fold Fold.drain
224-- :}
225--
226repeatM :: Monad m => m a -> StreamK m a
227repeatM = repeatMWith consM
229{-# INLINE replicateM #-}
230replicateM :: Monad m => Int -> m a -> StreamK m a
231replicateM = replicateMWith consM
232{-# INLINE replicate #-}
233replicate :: Int -> a -> StreamK m a
234replicate n a = go n
235 where
236 go cnt = if cnt <= 0 then nil else a `cons` go (cnt - 1)
238{-# INLINE fromIndicesM #-}
239fromIndicesM :: Monad m => (Int -> m a) -> StreamK m a
240fromIndicesM = fromIndicesMWith consM
241{-# INLINE fromIndices #-}
242fromIndices :: (Int -> a) -> StreamK m a
243fromIndices gen = go 0
244 where
245 go n = gen n `cons` go (n + 1)
247-- |
248-- >>> iterate f x = x `StreamK.cons` iterate f x
249--
250-- Generate an infinite stream with @x@ as the first element and each
251-- successive element derived by applying the function @f@ on the previous
252-- element.
253--
254-- >>> StreamK.toList $ StreamK.take 5 $ StreamK.iterate (+1) 1
255-- [1,2,3,4,5]
256--
257{-# INLINE iterate #-}
258iterate :: (a -> a) -> a -> StreamK m a
259iterate step = go
260 where
261 go !s = cons s (go (step s))
263-- |
264-- >>> iterateM f m = m >>= \a -> return a `StreamK.consM` iterateM f (f a)
265--
266-- Generate an infinite stream with the first element generated by the action
267-- @m@ and each successive element derived by applying the monadic function
268-- @f@ on the previous element.
269--
270-- >>> :{
271-- StreamK.iterateM (\x -> print x >> return (x + 1)) (return 0)
272-- & StreamK.take 3
273-- & StreamK.toList
274-- :}
275-- 0
276-- 1
277-- [0,1,2]
278--
279{-# INLINE iterateM #-}
280iterateM :: Monad m => (a -> m a) -> m a -> StreamK m a
281iterateM = iterateMWith consM
283-------------------------------------------------------------------------------
284-- Conversions
285-------------------------------------------------------------------------------
287{-# INLINE fromList #-}
288fromList :: [a] -> StreamK m a
289fromList = fromFoldable
291-------------------------------------------------------------------------------
292-- Elimination by Folding
293-------------------------------------------------------------------------------
295{-# INLINE foldr1 #-}
296foldr1 :: Monad m => (a -> a -> a) -> StreamK m a -> m (Maybe a)
297foldr1 step m = do
298 r <- uncons m
299 case r of
300 Nothing -> return Nothing
301 Just (h, t) -> fmap Just (go h t)
302 where
303 go p m1 =
304 let stp = return p
305 single a = return $ step a p
306 yieldk a r = fmap (step p) (go a r)
307 in foldStream defState yieldk single stp m1
309-- XXX replace the recursive "go" with explicit continuations.
310-- | Like 'foldx', but with a monadic step function.
311{-# INLINABLE foldlMx' #-}
312foldlMx' :: Monad m
313 => (x -> a -> m x) -> m x -> (x -> m b) -> StreamK m a -> m b
314foldlMx' step begin done = go begin
315 where
316 go !acc m1 =
317 let stop = acc >>= done
318 single a = acc >>= \b -> step b a >>= done
319 yieldk a r = acc >>= \b -> step b a >>= \x -> go (return x) r
320 in foldStream defState yieldk single stop m1
322-- | Fold a stream using the supplied left 'Fold' and reducing the resulting
323-- expression strictly at each step. The behavior is similar to 'foldl''. A
324-- 'Fold' can terminate early without consuming the full stream. See the
325-- documentation of individual 'Fold's for termination behavior.
326--
327-- Definitions:
328--
329-- >>> fold f = fmap fst . StreamK.foldBreak f
330-- >>> fold f = StreamK.parseD (Parser.fromFold f)
331--
332-- Example:
333--
334-- >>> StreamK.fold Fold.sum $ StreamK.fromStream $ Stream.enumerateFromTo 1 100
335-- 5050
336--
337{-# INLINABLE fold #-}
338fold :: Monad m => FL.Fold m a b -> StreamK m a -> m b
339fold (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
345 where
346 go !acc m1 =
347 let stop = done acc
348 single a = step acc a
349 >>= \case
350 FL.Partial s -> done s
351 FL.Done b1 -> return b1
352 yieldk a r = step acc a
353 >>= \case
354 FL.Partial s -> go s r
355 FL.Done b1 -> return b1
356 in foldStream defState yieldk single stop m1
358-- | Fold resulting in either breaking the stream or continuation of the fold.
359-- Instead of supplying the input stream in one go we can run the fold multiple
360-- times, each time supplying the next segment of the input stream. If the fold
361-- has not yet finished it returns a fold that can be run again otherwise it
362-- returns the fold result and the residual stream.
363--
364-- /Internal/
365{-# INLINE foldEither #-}
366foldEither :: Monad m =>
367 Fold m a b -> StreamK m a -> m (Either (Fold m a b) (b, StreamK m a))
368foldEither (FL.Fold step begin done) m = do
369 res <- begin
370 case res of
371 FL.Partial fs -> go fs m
372 FL.Done fb -> return $ Right (fb, m)
374 where
376 go !acc m1 =
377 let stop = return $ Left (Fold step (return $ FL.Partial acc) done)
378 single a =
379 step acc a
380 >>= \case
381 FL.Partial s ->
382 return $ Left (Fold step (return $ FL.Partial s) done)
383 FL.Done b1 -> return $ Right (b1, nil)
384 yieldk a r =
385 step acc a
386 >>= \case
387 FL.Partial s -> go s r
388 FL.Done b1 -> return $ Right (b1, r)
389 in foldStream defState yieldk single stop m1
391-- | Like 'fold' but also returns the remaining stream. The resulting stream
392-- would be 'StreamK.nil' if the stream finished before the fold.
393--
394{-# INLINE foldBreak #-}
395foldBreak :: Monad m => Fold m a b -> StreamK m a -> m (b, StreamK m a)
396foldBreak fld strm = do
397 r <- foldEither fld strm
398 case r of
399 Right res -> return res
400 Left (Fold _ initial extract) -> do
401 res <- initial
402 case res of
403 FL.Done _ -> error "foldBreak: unreachable state"
404 FL.Partial s -> do
405 b <- extract s
406 return (b, nil)
408-- XXX Array folds can be implemented using this.
409-- foldContainers? Specialized to foldArrays.
411-- | Generate streams from individual elements of a stream and fold the
412-- concatenation of those streams using the supplied fold. Return the result of
413-- the fold and residual stream.
414--
415-- For example, this can be used to efficiently fold an Array Word8 stream
416-- using Word8 folds.
417--
418-- /Internal/
419{-# INLINE foldConcat #-}
420foldConcat :: Monad m =>
421 Producer m a b -> Fold m b c -> StreamK m a -> m (c, StreamK m a)
422foldConcat
423 (Producer pstep pinject pextract)
424 (Fold fstep begin done)
425 stream = do
427 res <- begin
428 case res of
429 FL.Partial fs -> go fs stream
430 FL.Done fb -> return (fb, stream)
432 where
434 go !acc m1 = do
435 let stop = do
436 r <- done acc
437 return (r, nil)
438 single a = do
439 st <- pinject a
440 res <- go1 SPEC acc st
441 case res of
442 Left fs -> do
443 r <- done fs
444 return (r, nil)
445 Right (b, s) -> do
446 x <- pextract s
447 return (b, fromPure x)
448 yieldk a r = do
449 st <- pinject a
450 res <- go1 SPEC acc st
451 case res of
452 Left fs -> go fs r
453 Right (b, s) -> do
454 x <- pextract s
455 return (b, x `cons` r)
456 in foldStream defState yieldk single stop m1
458 {-# INLINE go1 #-}
459 go1 !_ !fs st = do
460 r <- pstep st
461 case r of
462 Stream.Yield x s -> do
463 res <- fstep fs x
464 case res of
465 FL.Done b -> return $ Right (b, s)
466 FL.Partial fs1 -> go1 SPEC fs1 s
467 Stream.Skip s -> go1 SPEC fs s
468 Stream.Stop -> return $ Left fs
470-- | Like 'foldl'' but with a monadic step function.
471{-# INLINE foldlM' #-}
472foldlM' :: Monad m => (b -> a -> m b) -> m b -> StreamK m a -> m b
473foldlM' step begin = foldlMx' step begin return
475------------------------------------------------------------------------------
476-- Specialized folds
477------------------------------------------------------------------------------
479{-# INLINE head #-}
480head :: Monad m => StreamK m a -> m (Maybe a)
481-- head = foldrM (\x _ -> return $ Just x) (return Nothing)
482head m =
483 let stop = return Nothing
484 single a = return (Just a)
485 yieldk a _ = return (Just a)
486 in foldStream defState yieldk single stop m
488{-# INLINE elem #-}
489elem :: (Monad m, Eq a) => a -> StreamK m a -> m Bool
490elem e = go
491 where
492 go m1 =
493 let stop = return False
494 single a = return (a == e)
495 yieldk a r = if a == e then return True else go r
496 in foldStream defState yieldk single stop m1
498{-# INLINE notElem #-}
499notElem :: (Monad m, Eq a) => a -> StreamK m a -> m Bool
500notElem e = go
501 where
502 go m1 =
503 let stop = return True
504 single a = return (a /= e)
505 yieldk a r = if a == e then return False else go r
506 in foldStream defState yieldk single stop m1
508{-# INLINABLE all #-}
509all :: Monad m => (a -> Bool) -> StreamK m a -> m Bool
510all p = go
511 where
512 go m1 =
513 let single a | p a = return True
514 | otherwise = return False
515 yieldk a r | p a = go r
516 | otherwise = return False
517 in foldStream defState yieldk single (return True) m1
519{-# INLINABLE any #-}
520any :: Monad m => (a -> Bool) -> StreamK m a -> m Bool
521any p = go
522 where
523 go m1 =
524 let single a | p a = return True
525 | otherwise = return False
526 yieldk a r | p a = return True
527 | otherwise = go r
528 in foldStream defState yieldk single (return False) m1
530-- | Extract the last element of the stream, if any.
531{-# INLINE last #-}
532last :: Monad m => StreamK m a -> m (Maybe a)
533last = foldlx' (\_ y -> Just y) Nothing id
535{-# INLINE minimum #-}
536minimum :: (Monad m, Ord a) => StreamK m a -> m (Maybe a)
537minimum = go Nothing
538 where
539 go Nothing m1 =
540 let stop = return Nothing
541 single a = return (Just a)
542 yieldk a r = go (Just a) r
543 in foldStream defState yieldk single stop m1
545 go (Just res) m1 =
546 let stop = return (Just res)
547 single a =
548 if res <= a
549 then return (Just res)
550 else return (Just a)
551 yieldk a r =
552 if res <= a
553 then go (Just res) r
554 else go (Just a) r
555 in foldStream defState yieldk single stop m1
557{-# INLINE minimumBy #-}
558minimumBy
559 :: (Monad m)
560 => (a -> a -> Ordering) -> StreamK m a -> m (Maybe a)
561minimumBy cmp = go Nothing
562 where
563 go Nothing m1 =
564 let stop = return Nothing
565 single a = return (Just a)
566 yieldk a r = go (Just a) r
567 in foldStream defState yieldk single stop m1
569 go (Just res) m1 =
570 let stop = return (Just res)
571 single a = case cmp res a of
572 GT -> return (Just a)
573 _ -> return (Just res)
574 yieldk a r = case cmp res a of
575 GT -> go (Just a) r
576 _ -> go (Just res) r
577 in foldStream defState yieldk single stop m1
579{-# INLINE maximum #-}
580maximum :: (Monad m, Ord a) => StreamK m a -> m (Maybe a)
581maximum = go Nothing
582 where
583 go Nothing m1 =
584 let stop = return Nothing
585 single a = return (Just a)
586 yieldk a r = go (Just a) r
587 in foldStream defState yieldk single stop m1
589 go (Just res) m1 =
590 let stop = return (Just res)
591 single a =
592 if res <= a
593 then return (Just a)
594 else return (Just res)
595 yieldk a r =
596 if res <= a
597 then go (Just a) r
598 else go (Just res) r
599 in foldStream defState yieldk single stop m1
601{-# INLINE maximumBy #-}
602maximumBy :: Monad m => (a -> a -> Ordering) -> StreamK m a -> m (Maybe a)
603maximumBy cmp = go Nothing
604 where
605 go Nothing m1 =
606 let stop = return Nothing
607 single a = return (Just a)
608 yieldk a r = go (Just a) r
609 in foldStream defState yieldk single stop m1
611 go (Just res) m1 =
612 let stop = return (Just res)
613 single a = case cmp res a of
614 GT -> return (Just res)
615 _ -> return (Just a)
616 yieldk a r = case cmp res a of
617 GT -> go (Just res) r
618 _ -> go (Just a) r
619 in foldStream defState yieldk single stop m1
621{-# INLINE (!!) #-}
622(!!) :: Monad m => StreamK m a -> Int -> m (Maybe a)
623m !! i = go i m
624 where
625 go n m1 =
626 let single a | n == 0 = return $ Just a
627 | otherwise = return Nothing
628 yieldk a x | n < 0 = return Nothing
629 | n == 0 = return $ Just a
630 | otherwise = go (n - 1) x
631 in foldStream defState yieldk single (return Nothing) m1
633{-# INLINE lookup #-}
634lookup :: (Monad m, Eq a) => a -> StreamK m (a, b) -> m (Maybe b)
635lookup e = go
636 where
637 go m1 =
638 let single (a, b) | a == e = return $ Just b
639 | otherwise = return Nothing
640 yieldk (a, b) x | a == e = return $ Just b
641 | otherwise = go x
642 in foldStream defState yieldk single (return Nothing) m1
644{-# INLINE findM #-}
645findM :: Monad m => (a -> m Bool) -> StreamK m a -> m (Maybe a)
646findM p = go
647 where
648 go m1 =
649 let single a = do
650 b <- p a
651 if b then return $ Just a else return Nothing
652 yieldk a x = do
653 b <- p a
654 if b then return $ Just a else go x
655 in foldStream defState yieldk single (return Nothing) m1
657{-# INLINE find #-}
658find :: Monad m => (a -> Bool) -> StreamK m a -> m (Maybe a)
659find p = findM (return . p)
661{-# INLINE findIndices #-}
662findIndices :: (a -> Bool) -> StreamK m a -> StreamK m Int
663findIndices p = go 0
664 where
665 go offset m1 = mkStream $ \st yld sng stp ->
666 let single a | p a = sng offset
667 | otherwise = stp
668 yieldk a x | p a = yld offset $ go (offset + 1) x
669 | otherwise = foldStream (adaptState st) yld sng stp $
670 go (offset + 1) x
671 in foldStream (adaptState st) yieldk single stp m1
673------------------------------------------------------------------------------
674-- Map and Fold
675------------------------------------------------------------------------------
677-- | Apply a monadic action to each element of the stream and discard the
678-- output of the action.
679{-# INLINE mapM_ #-}
680mapM_ :: Monad m => (a -> m b) -> StreamK m a -> m ()
681mapM_ f = go
682 where
683 go m1 =
684 let stop = return ()
685 single a = void (f a)
686 yieldk a r = f a >> go r
687 in foldStream defState yieldk single stop m1
689{-# INLINE mapM #-}
690mapM :: Monad m => (a -> m b) -> StreamK m a -> StreamK m b
691mapM = mapMWith consM
693------------------------------------------------------------------------------
694-- Converting folds
695------------------------------------------------------------------------------
697{-# INLINABLE toList #-}
698toList :: Monad m => StreamK m a -> m [a]
699toList = foldr (:) []
701-- Based on suggestions by David Feuer and Pranay Sashank
702{-# INLINE hoist #-}
703hoist :: (Monad m, Monad n)
704 => (forall x. m x -> n x) -> StreamK m a -> StreamK n a
705hoist f str =
706 mkStream $ \st yld sng stp ->
707 let single = return . sng
708 yieldk a s = return $ yld a (hoist f s)
709 stop = return stp
710 state = adaptState st
711 in join . f $ foldStreamShared state yieldk single stop str
713-------------------------------------------------------------------------------
714-- Transformation by folding (Scans)
715-------------------------------------------------------------------------------
717{-# INLINE scanlx' #-}
718scanlx' :: (x -> a -> x) -> x -> (x -> b) -> StreamK m a -> StreamK m b
719scanlx' step begin done m =
720 cons (done begin) $ go m begin
721 where
722 go m1 !acc = mkStream $ \st yld sng stp ->
723 let single a = sng (done $ step acc a)
724 yieldk a r =
725 let s = step acc a
726 in yld (done s) (go r s)
727 in foldStream (adaptState st) yieldk single stp m1
729{-# INLINE scanl' #-}
730scanl' :: (b -> a -> b) -> b -> StreamK m a -> StreamK m b
731scanl' step begin = scanlx' step begin id
733-------------------------------------------------------------------------------
734-- Filtering
735-------------------------------------------------------------------------------
737{-# INLINE filter #-}
738filter :: (a -> Bool) -> StreamK m a -> StreamK m a
739filter p = go
740 where
741 go m1 = mkStream $ \st yld sng stp ->
742 let single a | p a = sng a
743 | otherwise = stp
744 yieldk a r | p a = yld a (go r)
745 | otherwise = foldStream st yieldk single stp r
746 in foldStream st yieldk single stp m1
748{-# INLINE take #-}
749take :: Int -> StreamK m a -> StreamK m a
750take = go
751 where
752 go n1 m1 = mkStream $ \st yld sng stp ->
753 let yieldk a r = yld a (go (n1 - 1) r)
754 in if n1 <= 0
755 then stp
756 else foldStream st yieldk sng stp m1
758{-# INLINE takeWhile #-}
759takeWhile :: (a -> Bool) -> StreamK m a -> StreamK m a
760takeWhile p = go
761 where
762 go m1 = mkStream $ \st yld sng stp ->
763 let single a | p a = sng a
764 | otherwise = stp
765 yieldk a r | p a = yld a (go r)
766 | otherwise = stp
767 in foldStream st yieldk single stp m1
769{-# INLINE drop #-}
770drop :: Int -> StreamK m a -> StreamK m a
771drop n m = unShare (go n m)
772 where
773 go n1 m1 = mkStream $ \st yld sng stp ->
774 let single _ = stp
775 yieldk _ r = foldStreamShared st yld sng stp $ go (n1 - 1) r
776 -- Somehow "<=" check performs better than a ">"
777 in if n1 <= 0
778 then foldStreamShared st yld sng stp m1
779 else foldStreamShared st yieldk single stp m1
781{-# INLINE dropWhile #-}
782dropWhile :: (a -> Bool) -> StreamK m a -> StreamK m a
783dropWhile p = go
784 where
785 go m1 = mkStream $ \st yld sng stp ->
786 let single a | p a = stp
787 | otherwise = sng a
788 yieldk a r | p a = foldStream st yieldk single stp r
789 | otherwise = yld a r
790 in foldStream st yieldk single stp m1
792-------------------------------------------------------------------------------
793-- Mapping
794-------------------------------------------------------------------------------
796-- Be careful when modifying this, this uses a consM (|:) deliberately to allow
797-- other stream types to overload it.
798{-# INLINE sequence #-}
799sequence :: Monad m => StreamK m (m a) -> StreamK m a
800sequence = go
801 where
802 go m1 = mkStream $ \st yld sng stp ->
803 let single ma = ma >>= sng
804 yieldk ma r = foldStreamShared st yld sng stp $ ma `consM` go r
805 in foldStream (adaptState st) yieldk single stp m1
807-------------------------------------------------------------------------------
808-- Inserting
809-------------------------------------------------------------------------------
811{-# INLINE intersperseM #-}
812intersperseM :: Monad m => m a -> StreamK m a -> StreamK m a
813intersperseM a = prependingStart
814 where
815 prependingStart m1 = mkStream $ \st yld sng stp ->
816 let yieldk i x =
817 foldStreamShared st yld sng stp $ return i `consM` go x
818 in foldStream st yieldk sng stp m1
819 go m2 = mkStream $ \st yld sng stp ->
820 let single i = foldStreamShared st yld sng stp $ a `consM` fromPure i
821 yieldk i x =
822 foldStreamShared
823 st yld sng stp $ a `consM` return i `consM` go x
824 in foldStream st yieldk single stp m2
826{-# INLINE intersperse #-}
827intersperse :: Monad m => a -> StreamK m a -> StreamK m a
828intersperse a = intersperseM (return a)
830{-# INLINE insertBy #-}
831insertBy :: (a -> a -> Ordering) -> a -> StreamK m a -> StreamK m a
832insertBy cmp x = go
833 where
834 go m1 = mkStream $ \st yld _ _ ->
835 let single a = case cmp x a of
836 GT -> yld a (fromPure x)
837 _ -> yld x (fromPure a)
838 stop = yld x nil
839 yieldk a r = case cmp x a of
840 GT -> yld a (go r)
841 _ -> yld x (a `cons` r)
842 in foldStream st yieldk single stop m1
844------------------------------------------------------------------------------
845-- Deleting
846------------------------------------------------------------------------------
848{-# INLINE deleteBy #-}
849deleteBy :: (a -> a -> Bool) -> a -> StreamK m a -> StreamK m a
850deleteBy eq x = go
851 where
852 go m1 = mkStream $ \st yld sng stp ->
853 let single a = if eq x a then stp else sng a
854 yieldk a r = if eq x a
855 then foldStream st yld sng stp r
856 else yld a (go r)
857 in foldStream st yieldk single stp m1
859-------------------------------------------------------------------------------
860-- Map and Filter
861-------------------------------------------------------------------------------
863{-# INLINE mapMaybe #-}
864mapMaybe :: (a -> Maybe b) -> StreamK m a -> StreamK m b
865mapMaybe f = go
866 where
867 go m1 = mkStream $ \st yld sng stp ->
868 let single a = maybe stp sng (f a)
869 yieldk a r = case f a of
870 Just b -> yld b $ go r
871 Nothing -> foldStream (adaptState st) yieldk single stp r
872 in foldStream (adaptState st) yieldk single stp m1
874-------------------------------------------------------------------------------
875-- Exception Handling
876-------------------------------------------------------------------------------
878-- | Like Streamly.Data.Stream.'Streamly.Data.Stream.handle' but with one
879-- significant difference, this function observes exceptions from the consumer
880-- of the stream as well.
881--
882-- You can also convert 'StreamK' to 'Stream' and use exception handling from
883-- 'Stream' module:
884--
885-- >>> handle f s = StreamK.fromStream $ Stream.handle (\e -> StreamK.toStream (f e)) (StreamK.toStream s)
886--
887{-# INLINABLE handle #-}
888handle :: (MonadCatch m, Exception e)
889 => (e -> m (StreamK m a)) -> StreamK m a -> StreamK m a
890handle f stream = go stream
892 where
894 go m1 = mkStream $ \st yld sng stp ->
895 let yieldk a r = yld a $ go r
896 in do
897 res <- MC.try (foldStream (adaptState st) yieldk sng stp m1)
898 case res of
899 Right r -> return r
900 Left e -> do
901 r <- f e
902 foldStream (adaptState st) yld sng stp r
904-------------------------------------------------------------------------------
905-- Resource Management
906-------------------------------------------------------------------------------
908-- If we are folding the stream and we do not drain the entire stream (e.g. if
909-- the fold terminates before the stream) then the finalizer will run on GC.
910--
911-- XXX To implement a prompt cleanup, we will have to yield a cleanup function
912-- via the yield continuation. A chain of cleanup functions can be built and
913-- the entire chain can be invoked when the stream ends voluntarily or if
914-- someone decides to abandon the stream.
916-- | Like Streamly.Data.Stream.'Streamly.Data.Stream.bracketIO' but with one
917-- significant difference, this function observes exceptions from the consumer
918-- of the stream as well. Therefore, it cleans up the resource promptly when
919-- the consumer encounters an exception.
920--
921-- You can also convert 'StreamK' to 'Stream' and use resource handling from
922-- 'Stream' module:
923--
924-- >>> bracketIO bef aft bet = StreamK.fromStream $ Stream.bracketIO bef aft (StreamK.toStream . bet)
925--
926{-# INLINABLE bracketIO #-}
927bracketIO :: (MonadIO m, MonadCatch m)
928 => IO b -> (b -> IO c) -> (b -> StreamK m a) -> StreamK m a
929bracketIO bef aft bet =
930 concatEffect $ do
931 (r, ref) <- liftIO $ mask_ $ do
932 r <- bef
933 ref <- newIOFinalizer (aft r)
934 return (r, ref)
935 return $ go ref (bet r)
937 where
939 go ref m1 = mkStream $ \st yld sng stp ->
940 let
941 -- We can discard exceptions on continuations to make it equivalent
942 -- to StreamD, but it seems like a desirable behavior.
943 stop = liftIO (runIOFinalizer ref) >> stp
944 single a = liftIO (runIOFinalizer ref) >> sng a
945 yieldk a r = yld a $ go ref r
946 in do
947 -- Do not call the finalizer twice if it has already been
948 -- called via stop continuation and stop continuation itself
949 -- generated an exception. runIOFinalizer takes care of that.
950 res <- MC.try (foldStream (adaptState st) yieldk single stop m1)
951 case res of
952 Right r -> return r
953 Left (e :: MC.SomeException) ->
954 liftIO (runIOFinalizer ref) >> MC.throwM e
956------------------------------------------------------------------------------
957-- Serial Zipping
958------------------------------------------------------------------------------
960-- | Zipping of @n@ streams can be performed by combining the streams pair
961-- wise using 'mergeMapWith' with O(n * log n) time complexity. If used
962-- with 'concatMapWith' it will have O(n^2) performance.
963{-# INLINE zipWith #-}
964zipWith :: Monad m => (a -> b -> c) -> StreamK m a -> StreamK m b -> StreamK m c
965zipWith f = zipWithM (\a b -> return (f a b))
967{-# INLINE zipWithM #-}
968zipWithM :: Monad m =>
969 (a -> b -> m c) -> StreamK m a -> StreamK m b -> StreamK m c
970zipWithM f = go
972 where
974 go mx my = mkStream $ \st yld sng stp -> do
975 let merge a ra =
976 let single2 b = f a b >>= sng
977 yield2 b rb = f a b >>= \x -> yld x (go ra rb)
978 in foldStream (adaptState st) yield2 single2 stp my
979 let single1 a = merge a nil
980 yield1 = merge
981 foldStream (adaptState st) yield1 single1 stp mx
983------------------------------------------------------------------------------
984-- Merging
985------------------------------------------------------------------------------
987{-# INLINE mergeByM #-}
988mergeByM :: Monad m =>
989 (a -> a -> m Ordering) -> StreamK m a -> StreamK m a -> StreamK m a
990mergeByM cmp = go
992 where
994 go mx my = mkStream $ \st yld sng stp -> do
995 let stop = foldStream st yld sng stp my
996 single x = foldStream st yld sng stp (goX0 x my)
997 yield x rx = foldStream st yld sng stp (goX x rx my)
998 foldStream st yield single stop mx
1000 goX0 x my = mkStream $ \st yld sng _ -> do
1001 let stop = sng x
1002 single y = do
1003 r <- cmp x y
1004 case r of
1005 GT -> yld y (fromPure x)
1006 _ -> yld x (fromPure y)
1007 yield y ry = do
1008 r <- cmp x y
1009 case r of
1010 GT -> yld y (goX0 x ry)
1011 _ -> yld x (y `cons` ry)
1012 in foldStream st yield single stop my
1014 goX x mx my = mkStream $ \st yld _ _ -> do
1015 let stop = yld x mx
1016 single y = do
1017 r <- cmp x y
1018 case r of
1019 GT -> yld y (x `cons` mx)
1020 _ -> yld x (goY0 mx y)
1021 yield y ry = do
1022 r <- cmp x y
1023 case r of
1024 GT -> yld y (goX x mx ry)
1025 _ -> yld x (goY mx y ry)
1026 in foldStream st yield single stop my
1028 goY0 mx y = mkStream $ \st yld sng _ -> do
1029 let stop = sng y
1030 single x = do
1031 r <- cmp x y
1032 case r of
1033 GT -> yld y (fromPure x)
1034 _ -> yld x (fromPure y)
1035 yield x rx = do
1036 r <- cmp x y
1037 case r of
1038 GT -> yld y (x `cons` rx)
1039 _ -> yld x (goY0 rx y)
1040 in foldStream st yield single stop mx
1042 goY mx y my = mkStream $ \st yld _ _ -> do
1043 let stop = yld y my
1044 single x = do
1045 r <- cmp x y
1046 case r of
1047 GT -> yld y (goX0 x my)
1048 _ -> yld x (y `cons` my)
1049 yield x rx = do
1050 r <- cmp x y
1051 case r of
1052 GT -> yld y (goX x rx my)
1053 _ -> yld x (goY rx y my)
1054 in foldStream st yield single stop mx
1056-- | Merging of @n@ streams can be performed by combining the streams pair
1057-- wise using 'mergeMapWith' to give O(n * log n) time complexity. If used
1058-- with 'concatMapWith' it will have O(n^2) performance.
1060{-# INLINE mergeBy #-}
1061mergeBy :: (a -> a -> Ordering) -> StreamK m a -> StreamK m a -> StreamK m a
1062-- XXX GHC: This has slightly worse performance than replacing "r <- cmp x y"
1063-- with "let r = cmp x y" in the monadic version. The definition below is
1064-- exactly the same as mergeByM except this change.
1065-- mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
1066mergeBy cmp = go
1068 where
1070 go mx my = mkStream $ \st yld sng stp -> do
1071 let stop = foldStream st yld sng stp my
1072 single x = foldStream st yld sng stp (goX0 x my)
1073 yield x rx = foldStream st yld sng stp (goX x rx my)
1074 foldStream st yield single stop mx
1076 goX0 x my = mkStream $ \st yld sng _ -> do
1077 let stop = sng x
1078 single y = do
1079 case cmp x y of
1080 GT -> yld y (fromPure x)
1081 _ -> yld x (fromPure y)
1082 yield y ry = do
1083 case cmp x y of
1084 GT -> yld y (goX0 x ry)
1085 _ -> yld x (y `cons` ry)
1086 in foldStream st yield single stop my
1088 goX x mx my = mkStream $ \st yld _ _ -> do
1089 let stop = yld x mx
1090 single y = do
1091 case cmp x y of
1092 GT -> yld y (x `cons` mx)
1093 _ -> yld x (goY0 mx y)
1094 yield y ry = do
1095 case cmp x y of
1096 GT -> yld y (goX x mx ry)
1097 _ -> yld x (goY mx y ry)
1098 in foldStream st yield single stop my
1100 goY0 mx y = mkStream $ \st yld sng _ -> do
1101 let stop = sng y
1102 single x = do
1103 case cmp x y of
1104 GT -> yld y (fromPure x)
1105 _ -> yld x (fromPure y)
1106 yield x rx = do
1107 case cmp x y of
1108 GT -> yld y (x `cons` rx)
1109 _ -> yld x (goY0 rx y)
1110 in foldStream st yield single stop mx
1112 goY mx y my = mkStream $ \st yld _ _ -> do
1113 let stop = yld y my
1114 single x = do
1115 case cmp x y of
1116 GT -> yld y (goX0 x my)
1117 _ -> yld x (y `cons` my)
1118 yield x rx = do
1119 case cmp x y of
1120 GT -> yld y (goX x rx my)
1121 _ -> yld x (goY rx y my)
1122 in foldStream st yield single stop mx
1124------------------------------------------------------------------------------
1125-- Transformation comprehensions
1126------------------------------------------------------------------------------
1128{-# INLINE the #-}
1129the :: (Eq a, Monad m) => StreamK m a -> m (Maybe a)
1130the m = do
1131 r <- uncons m
1132 case r of
1133 Nothing -> return Nothing
1134 Just (h, t) -> go h t
1135 where
1136 go h m1 =
1137 let single a | h == a = return $ Just h
1138 | otherwise = return Nothing
1139 yieldk a r | h == a = go h r
1140 | otherwise = return Nothing
1141 in foldStream defState yieldk single (return $ Just h) m1
1143------------------------------------------------------------------------------
1144-- Alternative & MonadPlus
1145------------------------------------------------------------------------------
1147_alt :: StreamK m a -> StreamK m a -> StreamK m a
1148_alt m1 m2 = mkStream $ \st yld sng stp ->
1149 let stop = foldStream st yld sng stp m2
1150 in foldStream st yld sng stop m1
1152------------------------------------------------------------------------------
1153-- MonadError
1154------------------------------------------------------------------------------
1157-- XXX handle and test cross thread state transfer
1158withCatchError
1159 :: MonadError e m
1160 => StreamK m a -> (e -> StreamK m a) -> StreamK m a
1161withCatchError m h =
1162 mkStream $ \_ stp sng yld ->
1163 let run x = unStream x Nothing stp sng yieldk
1164 handle r = r `catchError` \e -> run $ h e
1165 yieldk a r = yld a (withCatchError r h)
1166 in handle $ run m
1169-------------------------------------------------------------------------------
1170-- Parsing
1171-------------------------------------------------------------------------------
1173-- Inlined definition.
1174{-# INLINE splitAt #-}
1175splitAt :: Int -> [a] -> ([a],[a])
1176splitAt n ls
1177 | n <= 0 = ([], ls)
1178 | otherwise = splitAt' n ls
1179 where
1180 splitAt' :: Int -> [a] -> ([a], [a])
1181 splitAt' _ [] = ([], [])
1182 splitAt' 1 (x:xs) = ([x], xs)
1183 splitAt' m (x:xs) = (x:xs', xs'')
1184 where
1185 (xs', xs'') = splitAt' (m - 1) xs
1187-- | Run a 'Parser' over a stream and return rest of the Stream.
1188{-# INLINE_NORMAL parseDBreak #-}
1189parseDBreak
1190 :: Monad m
1191 => PR.Parser a m b
1192 -> StreamK m a
1193 -> m (Either ParseError b, StreamK m a)
1194parseDBreak (PR.Parser pstep initial extract) stream = do
1195 res <- initial
1196 case res of
1197 PR.IPartial s -> goStream stream [] s
1198 PR.IDone b -> return (Right b, stream)
1199 PR.IError err -> return (Left (ParseError err), stream)
1201 where
1203 -- "buf" contains last few items in the stream that we may have to
1204 -- backtrack to.
1205 --
1206 -- XXX currently we are using a dumb list based approach for backtracking
1207 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
1208 -- That will allow us more efficient random back and forth movement.
1209 goStream st buf !pst =
1210 let stop = do
1211 r <- extract pst
1212 case r of
1213 PR.Error err -> do
1214 let src = Prelude.reverse buf
1215 return (Left (ParseError err), fromList src)
1216 PR.Done n b -> do
1217 assertM(n <= length buf)
1218 let src0 = Prelude.take n buf
1219 src = Prelude.reverse src0
1220 return (Right b, fromList src)
1221 PR.Partial _ _ -> error "Bug: parseBreak: Partial in extract"
1222 PR.Continue 0 s -> goStream nil buf s
1223 PR.Continue n s -> do
1224 assertM(n <= length buf)
1225 let (src0, buf1) = splitAt n buf
1226 src = Prelude.reverse src0
1227 goBuf nil buf1 src s
1228 single x = yieldk x nil
1229 yieldk x r = do
1230 res <- pstep pst x
1231 case res of
1232 PR.Partial 0 s -> goStream r [] s
1233 PR.Partial n s -> do
1234 assertM(n <= length (x:buf))
1235 let src0 = Prelude.take n (x:buf)
1236 src = Prelude.reverse src0
1237 goBuf r [] src s
1238 PR.Continue 0 s -> goStream r (x:buf) s
1239 PR.Continue n s -> do
1240 assertM(n <= length (x:buf))
1241 let (src0, buf1) = splitAt n (x:buf)
1242 src = Prelude.reverse src0
1243 goBuf r buf1 src s
1244 PR.Done 0 b -> return (Right b, r)
1245 PR.Done n b -> do
1246 assertM(n <= length (x:buf))
1247 let src0 = Prelude.take n (x:buf)
1248 src = Prelude.reverse src0
1249 return (Right b, append (fromList src) r)
1250 PR.Error err -> do
1251 let src = Prelude.reverse (x:buf)
1252 return (Left (ParseError err), append (fromList src) r)
1253 in foldStream defState yieldk single stop st
1255 goBuf st buf [] !pst = goStream st buf pst
1256 goBuf st buf (x:xs) !pst = do
1257 pRes <- pstep pst x
1258 case pRes of
1259 PR.Partial 0 s -> goBuf st [] xs s
1260 PR.Partial n s -> do
1261 assert (n <= length (x:buf)) (return ())
1262 let src0 = Prelude.take n (x:buf)
1263 src = Prelude.reverse src0 ++ xs
1264 goBuf st [] src s
1265 PR.Continue 0 s -> goBuf st (x:buf) xs s
1266 PR.Continue n s -> do
1267 assert (n <= length (x:buf)) (return ())
1268 let (src0, buf1) = splitAt n (x:buf)
1269 src = Prelude.reverse src0 ++ xs
1270 goBuf st buf1 src s
1271 PR.Done n b -> do
1272 assert (n <= length (x:buf)) (return ())
1273 let src0 = Prelude.take n (x:buf)
1274 src = Prelude.reverse src0
1275 return (Right b, append (fromList src) st)
1276 PR.Error err -> do
1277 let src = Prelude.reverse buf ++ x:xs
1278 return (Left (ParseError err), append (fromList src) st)
1280-- Using ParserD or ParserK on StreamK may not make much difference. We should
1281-- perhaps use only chunked parsing on StreamK. We can always convert a stream
1282-- to chunks before parsing. Or just have a ParserK element parser for StreamK
1283-- and convert ParserD to ParserK for element parsing using StreamK.
1284{-# INLINE parseD #-}
1285parseD :: Monad m =>
1286 Parser.Parser a m b -> StreamK m a -> m (Either ParseError b)
1287parseD f = fmap fst . parseDBreak f
1289-------------------------------------------------------------------------------
1290-- ParserK Chunked
1291-------------------------------------------------------------------------------
1293-- The backracking buffer consists of arrays in the most-recent-first order. We
1294-- want to take a total of n array elements from this buffer. Note: when we
1295-- have to take an array partially, we must take the last part of the array.
1296{-# INLINE backTrack #-}
1297backTrack :: forall m a. Unbox a =>
1298 Int
1299 -> [Array a]
1300 -> StreamK m (Array a)
1301 -> (StreamK m (Array a), [Array a])
1302backTrack = go
1304 where
1306 go _ [] stream = (stream, [])
1307 go n xs stream | n <= 0 = (stream, xs)
1308 go n (x:xs) stream =
1309 let len = Array.length x
1310 in if n > len
1311 then go (n - len) xs (cons x stream)
1312 else if n == len
1313 then (cons x stream, xs)
1314 else let !(Array contents start end) = x
1315 !start1 = end - (n * SIZE_OF(a))
1316 arr1 = Array contents start1 end
1317 arr2 = Array contents start start1
1318 in (cons arr1 stream, arr2:xs)
1320-- | A continuation to extract the result when a CPS parser is done.
1321{-# INLINE parserDone #-}
1322parserDone :: Applicative m =>
1323 ParserK.ParseResult b -> Int -> ParserK.Input a -> m (ParserK.Step a m b)
1324parserDone (ParserK.Success n b) _ _ = pure $ ParserK.Done n b
1325parserDone (ParserK.Failure n e) _ _ = pure $ ParserK.Error n e
1327-- XXX parseDBreakChunks may be faster than converting parserD to parserK and
1328-- using parseBreakChunks. We can also use parseBreak as an alternative to the
1329-- monad instance of ParserD.
1331-- | Run a 'ParserK' over a chunked 'StreamK' and return the rest of the Stream.
1332{-# INLINE_NORMAL parseBreakChunks #-}
1333parseBreakChunks
1334 :: (Monad m, Unbox a)
1335 => ParserK (Array a) m b
1336 -> StreamK m (Array a)
1337 -> m (Either ParseError b, StreamK m (Array a))
1338parseBreakChunks parser input = do
1339 let parserk = ParserK.runParser parser parserDone 0 0
1340 in go [] parserk input
1342 where
1344 {-# INLINE goStop #-}
1345 goStop backBuf parserk = do
1346 pRes <- parserk ParserK.None
1347 case pRes of
1348 -- If we stop in an alternative, it will try calling the next
1349 -- parser, the next parser may call initial returning Partial and
1350 -- then immediately we have to call extract on it.
1351 ParserK.Partial 0 cont1 ->
1352 go [] cont1 nil
1353 ParserK.Partial n cont1 -> do
1354 let n1 = negate n
1355 assertM(n1 >= 0 && n1 <= sum (Prelude.map Array.length backBuf))
1356 let (s1, backBuf1) = backTrack n1 backBuf nil
1357 in go backBuf1 cont1 s1
1358 ParserK.Continue 0 cont1 ->
1359 go backBuf cont1 nil
1360 ParserK.Continue n cont1 -> do
1361 let n1 = negate n
1362 assertM(n1 >= 0 && n1 <= sum (Prelude.map Array.length backBuf))
1363 let (s1, backBuf1) = backTrack n1 backBuf nil
1364 in go backBuf1 cont1 s1
1365 ParserK.Done 0 b ->
1366 return (Right b, nil)
1367 ParserK.Done n b -> do
1368 let n1 = negate n
1369 assertM(n1 >= 0 && n1 <= sum (Prelude.map Array.length backBuf))
1370 let (s1, _) = backTrack n1 backBuf nil
1371 in return (Right b, s1)
1372 ParserK.Error _ err -> do
1373 let (s1, _) = backTrack maxBound backBuf nil
1374 return (Left (ParseError err), s1)
1376 seekErr n len =
1377 error $ "parseBreak: Partial: forward seek not implemented n = "
1378 ++ show n ++ " len = " ++ show len
1380 yieldk backBuf parserk arr stream = do
1381 pRes <- parserk (ParserK.Chunk arr)
1382 let len = Array.length arr
1383 case pRes of
1384 ParserK.Partial n cont1 ->
1385 case compare n len of
1386 EQ -> go [] cont1 stream
1387 LT -> do
1388 if n >= 0
1389 then yieldk [] cont1 arr stream
1390 else do
1391 let n1 = negate n
1392 bufLen = sum (Prelude.map Array.length backBuf)
1393 s = cons arr stream
1394 assertM(n1 >= 0 && n1 <= bufLen)
1395 let (s1, _) = backTrack n1 backBuf s
1396 go [] cont1 s1
1397 GT -> seekErr n len
1398 ParserK.Continue n cont1 ->
1399 case compare n len of
1400 EQ -> go (arr:backBuf) cont1 stream
1401 LT -> do
1402 if n >= 0
1403 then yieldk backBuf cont1 arr stream
1404 else do
1405 let n1 = negate n
1406 bufLen = sum (Prelude.map Array.length backBuf)
1407 s = cons arr stream
1408 assertM(n1 >= 0 && n1 <= bufLen)
1409 let (s1, backBuf1) = backTrack n1 backBuf s
1410 go backBuf1 cont1 s1
1411 GT -> seekErr n len
1412 ParserK.Done n b -> do
1413 let n1 = len - n
1414 assertM(n1 <= sum (Prelude.map Array.length (arr:backBuf)))
1415 let (s1, _) = backTrack n1 (arr:backBuf) stream
1416 in return (Right b, s1)
1417 ParserK.Error _ err -> do
1418 let (s1, _) = backTrack maxBound (arr:backBuf) stream
1419 return (Left (ParseError err), s1)
1421 go backBuf parserk stream = do
1422 let stop = goStop backBuf parserk
1423 single a = yieldk backBuf parserk a nil
1424 in foldStream
1425 defState (yieldk backBuf parserk) single stop stream
1427{-# INLINE parseChunks #-}
1428parseChunks :: (Monad m, Unbox a) =>
1429 ParserK (Array a) m b -> StreamK m (Array a) -> m (Either ParseError b)
1430parseChunks f = fmap fst . parseBreakChunks f
1432-------------------------------------------------------------------------------
1433-- ParserK Singular
1434-------------------------------------------------------------------------------
1436{-# INLINE backTrackSingular #-}
1437backTrackSingular :: Int -> [a] -> StreamK m a -> (StreamK m a, [a])
1438backTrackSingular = go
1440 where
1442 go _ [] stream = (stream, [])
1443 go n xs stream | n <= 0 = (stream, xs)
1444 go n xs stream =
1445 let (appendBuf, newBTBuf) = splitAt n xs
1446 in (append (fromList (Prelude.reverse appendBuf)) stream, newBTBuf)
1449-- | Similar to 'parseBreak' but works on singular elements.
1451{-# INLINE_NORMAL parseBreak #-}
1452parseBreak
1453 :: forall m a b. Monad m
1454 => ParserK.ParserK a m b
1455 -> StreamK m a
1456 -> m (Either ParseError b, StreamK m a)
1457parseBreak parser input = do
1458 let parserk = ParserK.runParser parser parserDone 0 0
1459 in go [] parserk input
1461 where
1463 {-# INLINE goStop #-}
1464 goStop
1465 :: [a]
1466 -> (ParserK.Input a -> m (ParserK.Step a m b))
1467 -> m (Either ParseError b, StreamK m a)
1468 goStop backBuf parserk = do
1469 pRes <- parserk ParserK.None
1470 case pRes of
1471 -- If we stop in an alternative, it will try calling the next
1472 -- parser, the next parser may call initial returning Partial and
1473 -- then immediately we have to call extract on it.
1474 ParserK.Partial 0 cont1 ->
1475 go [] cont1 nil
1476 ParserK.Partial n cont1 -> do
1477 let n1 = negate n
1478 assertM(n1 >= 0 && n1 <= length backBuf)
1479 let (s1, backBuf1) = backTrackSingular n1 backBuf nil
1480 in go backBuf1 cont1 s1
1481 ParserK.Continue 0 cont1 ->
1482 go backBuf cont1 nil
1483 ParserK.Continue n cont1 -> do
1484 let n1 = negate n
1485 assertM(n1 >= 0 && n1 <= length backBuf)
1486 let (s1, backBuf1) = backTrackSingular n1 backBuf nil
1487 in go backBuf1 cont1 s1
1488 ParserK.Done 0 b ->
1489 return (Right b, nil)
1490 ParserK.Done n b -> do
1491 let n1 = negate n
1492 assertM(n1 >= 0 && n1 <= length backBuf)
1493 let (s1, _) = backTrackSingular n1 backBuf nil
1494 in return (Right b, s1)
1495 ParserK.Error _ err -> return (Left (ParseError err), nil)
1497 seekErr n =
1498 error $ "parseBreak: Partial: forward seek not implemented n = "
1499 ++ show n
1501 yieldk
1502 :: [a]
1503 -> (ParserK.Input a -> m (ParserK.Step a m b))
1504 -> a
1505 -> StreamK m a
1506 -> m (Either ParseError b, StreamK m a)
1507 yieldk backBuf parserk arr stream = do
1508 pRes <- parserk (ParserK.Chunk arr)
1509 case pRes of
1510 ParserK.Partial 1 cont1 -> go [] cont1 stream
1511 ParserK.Partial 0 cont1 -> go [] cont1 (cons arr stream)
1512 ParserK.Partial n _ | n > 1 -> seekErr n
1513 ParserK.Partial n cont1 -> do
1514 let n1 = negate n
1515 bufLen = length backBuf
1516 s = cons arr stream
1517 assertM(n1 >= 0 && n1 <= bufLen)
1518 let (s1, _) = backTrackSingular n1 backBuf s
1519 go [] cont1 s1
1520 ParserK.Continue 1 cont1 -> go (arr:backBuf) cont1 stream
1521 ParserK.Continue 0 cont1 ->
1522 go backBuf cont1 (cons arr stream)
1523 ParserK.Continue n _ | n > 1 -> seekErr n
1524 ParserK.Continue n cont1 -> do
1525 let n1 = negate n
1526 bufLen = length backBuf
1527 s = cons arr stream
1528 assertM(n1 >= 0 && n1 <= bufLen)
1529 let (s1, backBuf1) = backTrackSingular n1 backBuf s
1530 go backBuf1 cont1 s1
1531 ParserK.Done 1 b -> pure (Right b, stream)
1532 ParserK.Done 0 b -> pure (Right b, cons arr stream)
1533 ParserK.Done n _ | n > 1 -> seekErr n
1534 ParserK.Done n b -> do
1535 let n1 = negate n
1536 bufLen = length backBuf
1537 s = cons arr stream
1538 assertM(n1 >= 0 && n1 <= bufLen)
1539 let (s1, _) = backTrackSingular n1 backBuf s
1540 pure (Right b, s1)
1541 ParserK.Error _ err -> return (Left (ParseError err), nil)
1543 go
1544 :: [a]
1545 -> (ParserK.Input a -> m (ParserK.Step a m b))
1546 -> StreamK m a
1547 -> m (Either ParseError b, StreamK m a)
1548 go backBuf parserk stream = do
1549 let stop = goStop backBuf parserk
1550 single a = yieldk backBuf parserk a nil
1551 in foldStream
1552 defState (yieldk backBuf parserk) single stop stream
1554-- | Run a 'ParserK' over a 'StreamK'. Please use 'parseChunks' where possible,
1555-- for better performance.
1556{-# INLINE parse #-}
1557parse :: Monad m =>
1558 ParserK.ParserK a m b -> StreamK m a -> m (Either ParseError b)
1559parse f = fmap fst . parseBreak f
1561-------------------------------------------------------------------------------
1562-- ParserK Chunked Generic
1563-------------------------------------------------------------------------------
1565{-# INLINE backTrackGenericChunks #-}
1566backTrackGenericChunks ::
1567 Int
1568 -> [GenArr.Array a]
1569 -> StreamK m (GenArr.Array a)
1570 -> (StreamK m (GenArr.Array a), [GenArr.Array a])
1571backTrackGenericChunks = go
1573 where
1575 go _ [] stream = (stream, [])
1576 go n xs stream | n <= 0 = (stream, xs)
1577 go n (x:xs) stream =
1578 let len = GenArr.length x
1579 in if n > len
1580 then go (n - len) xs (cons x stream)
1581 else if n == len
1582 then (cons x stream, xs)
1583 else let arr1 = GenArr.getSliceUnsafe (len - n) n x
1584 arr2 = GenArr.getSliceUnsafe 0 (len - n) x
1585 in (cons arr1 stream, arr2:xs)
1587-- | Similar to 'parseBreak' but works on generic arrays
1589{-# INLINE_NORMAL parseBreakChunksGeneric #-}
1590parseBreakChunksGeneric
1591 :: forall m a b. Monad m
1592 => ParserK.ParserK (GenArr.Array a) m b
1593 -> StreamK m (GenArr.Array a)
1594 -> m (Either ParseError b, StreamK m (GenArr.Array a))
1595parseBreakChunksGeneric parser input = do
1596 let parserk = ParserK.runParser parser parserDone 0 0
1597 in go [] parserk input
1599 where
1601 {-# INLINE goStop #-}
1602 goStop
1603 :: [GenArr.Array a]
1604 -> (ParserK.Input (GenArr.Array a)
1605 -> m (ParserK.Step (GenArr.Array a) m b))
1606 -> m (Either ParseError b, StreamK m (GenArr.Array a))
1607 goStop backBuf parserk = do
1608 pRes <- parserk ParserK.None
1609 case pRes of
1610 -- If we stop in an alternative, it will try calling the next
1611 -- parser, the next parser may call initial returning Partial and
1612 -- then immediately we have to call extract on it.
1613 ParserK.Partial 0 cont1 ->
1614 go [] cont1 nil
1615 ParserK.Partial n cont1 -> do
1616 let n1 = negate n
1617 assertM(n1 >= 0 && n1 <= sum (Prelude.map GenArr.length backBuf))
1618 let (s1, backBuf1) = backTrackGenericChunks n1 backBuf nil
1619 in go backBuf1 cont1 s1
1620 ParserK.Continue 0 cont1 ->
1621 go backBuf cont1 nil
1622 ParserK.Continue n cont1 -> do
1623 let n1 = negate n
1624 assertM(n1 >= 0 && n1 <= sum (Prelude.map GenArr.length backBuf))
1625 let (s1, backBuf1) = backTrackGenericChunks n1 backBuf nil
1626 in go backBuf1 cont1 s1
1627 ParserK.Done 0 b ->
1628 return (Right b, nil)
1629 ParserK.Done n b -> do
1630 let n1 = negate n
1631 assertM(n1 >= 0 && n1 <= sum (Prelude.map GenArr.length backBuf))
1632 let (s1, _) = backTrackGenericChunks n1 backBuf nil
1633 in return (Right b, s1)
1634 ParserK.Error _ err -> return (Left (ParseError err), nil)
1636 seekErr n len =
1637 error $ "parseBreak: Partial: forward seek not implemented n = "
1638 ++ show n ++ " len = " ++ show len
1640 yieldk
1641 :: [GenArr.Array a]
1642 -> (ParserK.Input (GenArr.Array a)
1643 -> m (ParserK.Step (GenArr.Array a) m b))
1644 -> GenArr.Array a
1645 -> StreamK m (GenArr.Array a)
1646 -> m (Either ParseError b, StreamK m (GenArr.Array a))
1647 yieldk backBuf parserk arr stream = do
1648 pRes <- parserk (ParserK.Chunk arr)
1649 let len = GenArr.length arr
1650 case pRes of
1651 ParserK.Partial n cont1 ->
1652 case compare n len of
1653 EQ -> go [] cont1 stream
1654 LT -> do
1655 if n >= 0
1656 then yieldk [] cont1 arr stream
1657 else do
1658 let n1 = negate n
1659 bufLen = sum (Prelude.map GenArr.length backBuf)
1660 s = cons arr stream
1661 assertM(n1 >= 0 && n1 <= bufLen)
1662 let (s1, _) = backTrackGenericChunks n1 backBuf s
1663 go [] cont1 s1
1664 GT -> seekErr n len
1665 ParserK.Continue n cont1 ->
1666 case compare n len of
1667 EQ -> go (arr:backBuf) cont1 stream
1668 LT -> do
1669 if n >= 0
1670 then yieldk backBuf cont1 arr stream
1671 else do
1672 let n1 = negate n
1673 bufLen = sum (Prelude.map GenArr.length backBuf)
1674 s = cons arr stream
1675 assertM(n1 >= 0 && n1 <= bufLen)
1676 let (s1, backBuf1) = backTrackGenericChunks n1 backBuf s
1677 go backBuf1 cont1 s1
1678 GT -> seekErr n len
1679 ParserK.Done n b -> do
1680 let n1 = len - n
1681 assertM(n1 <= sum (Prelude.map GenArr.length (arr:backBuf)))
1682 let (s1, _) = backTrackGenericChunks n1 (arr:backBuf) stream
1683 in return (Right b, s1)
1684 ParserK.Error _ err -> return (Left (ParseError err), nil)
1686 go
1687 :: [GenArr.Array a]
1688 -> (ParserK.Input (GenArr.Array a)
1689 -> m (ParserK.Step (GenArr.Array a) m b))
1690 -> StreamK m (GenArr.Array a)
1691 -> m (Either ParseError b, StreamK m (GenArr.Array a))
1692 go backBuf parserk stream = do
1693 let stop = goStop backBuf parserk
1694 single a = yieldk backBuf parserk a nil
1695 in foldStream
1696 defState (yieldk backBuf parserk) single stop stream
1698{-# INLINE parseChunksGeneric #-}
1699parseChunksGeneric ::
1700 (Monad m)
1701 => ParserK.ParserK (GenArr.Array a) m b
1702 -> StreamK m (GenArr.Array a)
1703 -> m (Either ParseError b)
1704parseChunksGeneric f = fmap fst . parseBreakChunksGeneric f
1706-------------------------------------------------------------------------------
1707-- Sorting
1708-------------------------------------------------------------------------------
1710-- | Sort the input stream using a supplied comparison function.
1712-- Sorting can be achieved by simply:
1714-- >>> sortBy cmp = StreamK.mergeMapWith (StreamK.mergeBy cmp) StreamK.fromPure
1716-- However, this combinator uses a parser to first split the input stream into
1717-- down and up sorted segments and then merges them to optimize sorting when
1718-- pre-sorted sequences exist in the input stream.
1720-- /O(n) space/
1722{-# INLINE sortBy #-}
1723sortBy :: Monad m => (a -> a -> Ordering) -> StreamK m a -> StreamK m a
1724-- sortBy f = Stream.concatPairsWith (Stream.mergeBy f) Stream.fromPure
1725sortBy cmp =
1726 let p =
1727 Parser.groupByRollingEither
1728 (\x -> (< GT) . cmp x)
1729 FL.toStreamKRev
1730 FL.toStreamK
1731 in mergeMapWith (mergeBy cmp) id
1732 . Stream.toStreamK
1733 . Stream.catRights -- its a non-failing backtracking parser
1734 . Stream.parseMany (fmap (either id id) p)
1735 . Stream.fromStreamK