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