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