git.iko.soy/streamly
on eecc006 help sign in
streamly/core/src/Streamly/Internal/Data/Array/Stream.hs
1156 lines 45 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Array.Stream
3-- Copyright : (c) 2019 Composewell Technologies
4-- License : BSD3-3-Clause
5-- Maintainer : streamly@composewell.com
6-- Stability : pre-release
7-- Portability : GHC
8--
9-- Combinators to efficiently manipulate streams of immutable arrays.
10--
11module Streamly.Internal.Data.Array.Stream
12 (
13 -- * Creation
14 arraysOf
16 -- * Flattening to elements
17 , concat
18 , concatRev
19 , interpose
20 , interposeSuffix
21 , intercalateSuffix
22 , unlines
24 -- * Elimination
25 -- ** Element Folds
26 , foldBreak
27 , foldBreakD
28 , parseBreak
29 -- , parseBreakD
31 -- ** Array Folds
32 , runArrayFold
33 , runArrayFoldBreak
34 -- , parseArr
35 , runArrayParserDBreak
36 , runArrayFoldMany
38 , toArray
40 -- * Compaction
41 , lpackArraysChunksOf
42 , compact
44 -- * Splitting
45 , splitOn
46 , splitOnSuffix
47 )
48where
50#include "ArrayMacros.h"
51#include "inline.hs"
53import Data.Bifunctor (second)
54import Control.Exception (assert)
55import Control.Monad.Catch (MonadThrow, throwM)
56import Control.Monad.IO.Class (MonadIO(..))
57import Data.Word (Word8)
58import Streamly.Internal.Data.Unboxed (Unbox, peekWith, sizeOf)
59import Fusion.Plugin.Types (Fuse(..))
60import GHC.Exts (SpecConstrAnnotation(..))
61import GHC.Types (SPEC(..))
62import Prelude hiding (null, last, (!!), read, concat, unlines)
64import Streamly.Data.Fold (Fold)
65import Streamly.Data.Stream (Stream)
66import Streamly.Internal.Data.Array.Type (Array(..))
67import Streamly.Internal.Data.Array.Fold (ArrayFold(..))
68import Streamly.Internal.Data.Parser (ParseError(..))
69import Streamly.Internal.Data.Stream
70 (fromStreamD, fromStreamK, toStreamD, toStreamK)
71import Streamly.Internal.Data.SVar.Type (adaptState, defState)
72import Streamly.Internal.Data.Array.Mut.Type
73 (allocBytesToElemCount)
74import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
76import qualified Streamly.Data.Fold as FL
77import qualified Streamly.Internal.Data.Array as A
78import qualified Streamly.Internal.Data.Array as Array
79import qualified Streamly.Internal.Data.Array.Type as A
80import qualified Streamly.Internal.Data.Array.Mut.Type as MA
81import qualified Streamly.Internal.Data.Array.Mut.Stream as AS
82import qualified Streamly.Internal.Data.Fold.Type as FL (Fold(..), Step(..))
83import qualified Streamly.Internal.Data.Parser as PR
84import qualified Streamly.Internal.Data.Parser.ParserD as PRD
85 (Parser(..), Initial(..), fromParserK)
86import qualified Streamly.Internal.Data.Stream as S
87import qualified Streamly.Internal.Data.Stream.StreamD as D
88 ( fromList, nil, cons, map
89 , unfoldMany, append, splitInnerBy, splitInnerBySuffix, foldlM'
90 )
91import qualified Streamly.Internal.Data.Stream.StreamD.Type as D
92 (Step(Yield, Stop, Skip), Stream(Stream))
93import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
94 (Stream, cons, nil, fromPure, foldStream)
96-- XXX Since these are immutable arrays MonadIO constraint can be removed from
97-- most places.
99-------------------------------------------------------------------------------
100-- Generation
101-------------------------------------------------------------------------------
103-- | @arraysOf n stream@ groups the elements in the input stream into arrays of
104-- @n@ elements each.
105--
106-- > arraysOf n = Stream.chunksOf n (Array.writeN n)
107--
108-- /Pre-release/
109{-# INLINE arraysOf #-}
110arraysOf :: (MonadIO m, Unbox a)
111 => Int -> Stream m a -> Stream m (Array a)
112arraysOf n str = fromStreamD $ A.arraysOf n (toStreamD str)
114-------------------------------------------------------------------------------
115-- Append
116-------------------------------------------------------------------------------
118-- XXX efficiently compare two streams of arrays. Two streams can have chunks
119-- of different sizes, we can handle that in the stream comparison abstraction.
120-- This could be useful e.g. to fast compare whether two files differ.
122-- | Convert a stream of arrays into a stream of their elements.
123--
124-- Same as the following:
125--
126-- > concat = Stream.unfoldMany Array.read
127--
128-- @since 0.7.0
129{-# INLINE concat #-}
130concat :: (Monad m, Unbox a) => Stream m (Array a) -> Stream m a
131-- concat m = fromStreamD $ A.flattenArrays (toStreamD m)
132-- concat m = fromStreamD $ D.concatMap A.toStreamD (toStreamD m)
133concat m = fromStreamD $ D.unfoldMany A.reader (toStreamD m)
135-- | Convert a stream of arrays into a stream of their elements reversing the
136-- contents of each array before flattening.
137--
138-- > concatRev = Stream.unfoldMany Array.readerRev
139--
140-- @since 0.7.0
141{-# INLINE concatRev #-}
142concatRev :: (Monad m, Unbox a) => Stream m (Array a) -> Stream m a
143-- concatRev m = fromStreamD $ A.flattenArraysRev (toStreamD m)
144concatRev m = fromStreamD $ D.unfoldMany A.readerRev (toStreamD m)
146-------------------------------------------------------------------------------
147-- Intersperse and append
148-------------------------------------------------------------------------------
150-- | Flatten a stream of arrays after inserting the given element between
151-- arrays.
152--
153-- /Pre-release/
154{-# INLINE interpose #-}
155interpose :: (Monad m, Unbox a) => a -> Stream m (Array a) -> Stream m a
156interpose x = S.interpose x A.reader
158{-# INLINE intercalateSuffix #-}
159intercalateSuffix :: (Monad m, Unbox a)
160 => Array a -> Stream m (Array a) -> Stream m a
161intercalateSuffix = S.intercalateSuffix A.reader
163-- | Flatten a stream of arrays appending the given element after each
164-- array.
165--
166-- @since 0.7.0
167{-# INLINE interposeSuffix #-}
168interposeSuffix :: (Monad m, Unbox a)
169 => a -> Stream m (Array a) -> Stream m a
170-- interposeSuffix x = fromStreamD . A.unlines x . toStreamD
171interposeSuffix x = S.interposeSuffix x A.reader
173data FlattenState s a =
174 OuterLoop s
175 | InnerLoop s !(MA.MutableByteArray a) !Int !Int
177-- XXX Remove monadIO constraint
178{-# INLINE_NORMAL unlines #-}
179unlines :: forall m a. (MonadIO m, Unbox a)
180 => a -> D.Stream m (Array a) -> D.Stream m a
181unlines sep (D.Stream step state) = D.Stream step' (OuterLoop state)
182 where
183 {-# INLINE_LATE step' #-}
184 step' gst (OuterLoop st) = do
185 r <- step (adaptState gst) st
186 return $ case r of
187 D.Yield Array{..} s ->
188 D.Skip (InnerLoop s arrContents arrStart arrEnd)
189 D.Skip s -> D.Skip (OuterLoop s)
190 D.Stop -> D.Stop
192 step' _ (InnerLoop st _ p end) | p == end =
193 return $ D.Yield sep $ OuterLoop st
195 step' _ (InnerLoop st contents p end) = do
196 x <- liftIO $ peekWith contents p
197 return $ D.Yield x (InnerLoop st contents (INDEX_NEXT(p,a)) end)
199-------------------------------------------------------------------------------
200-- Compact
201-------------------------------------------------------------------------------
203-- XXX These would not be needed once we implement compactLEFold, see
204-- module Streamly.Internal.Data.Array.Mut.Stream
205--
206{-# INLINE_NORMAL packArraysChunksOf #-}
207packArraysChunksOf :: (MonadIO m, Unbox a)
208 => Int -> D.Stream m (Array a) -> D.Stream m (Array a)
209packArraysChunksOf n str =
210 D.map A.unsafeFreeze $ AS.packArraysChunksOf n $ D.map A.unsafeThaw str
212-- XXX instead of writing two different versions of this operation, we should
213-- write it as a pipe.
214{-# INLINE_NORMAL lpackArraysChunksOf #-}
215lpackArraysChunksOf :: (MonadIO m, Unbox a)
216 => Int -> Fold m (Array a) () -> Fold m (Array a) ()
217lpackArraysChunksOf n fld =
218 FL.lmap A.unsafeThaw $ AS.lpackArraysChunksOf n (FL.lmap A.unsafeFreeze fld)
220-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
221-- maximum specified size in bytes.
222--
223-- @since 0.7.0
224{-# INLINE compact #-}
225compact :: (MonadIO m, Unbox a)
226 => Int -> Stream m (Array a) -> Stream m (Array a)
227compact n xs = fromStreamD $ packArraysChunksOf n (toStreamD xs)
229-------------------------------------------------------------------------------
230-- Split
231-------------------------------------------------------------------------------
233data SplitState s arr
234 = Initial s
235 | Buffering s arr
236 | Splitting s arr
237 | Yielding arr (SplitState s arr)
238 | Finishing
240-- | Split a stream of arrays on a given separator byte, dropping the separator
241-- and coalescing all the arrays between two separators into a single array.
242--
243-- @since 0.7.0
244{-# INLINE_NORMAL _splitOn #-}
245_splitOn
246 :: MonadIO m
247 => Word8
248 -> D.Stream m (Array Word8)
249 -> D.Stream m (Array Word8)
250_splitOn byte (D.Stream step state) = D.Stream step' (Initial state)
252 where
254 {-# INLINE_LATE step' #-}
255 step' gst (Initial st) = do
256 r <- step gst st
257 case r of
258 D.Yield arr s -> do
259 (arr1, marr2) <- A.breakOn byte arr
260 return $ case marr2 of
261 Nothing -> D.Skip (Buffering s arr1)
262 Just arr2 -> D.Skip (Yielding arr1 (Splitting s arr2))
263 D.Skip s -> return $ D.Skip (Initial s)
264 D.Stop -> return D.Stop
266 step' gst (Buffering st buf) = do
267 r <- step gst st
268 case r of
269 D.Yield arr s -> do
270 (arr1, marr2) <- A.breakOn byte arr
271 buf' <- A.splice buf arr1
272 return $ case marr2 of
273 Nothing -> D.Skip (Buffering s buf')
274 Just x -> D.Skip (Yielding buf' (Splitting s x))
275 D.Skip s -> return $ D.Skip (Buffering s buf)
276 D.Stop -> return $
277 if A.byteLength buf == 0
278 then D.Stop
279 else D.Skip (Yielding buf Finishing)
281 step' _ (Splitting st buf) = do
282 (arr1, marr2) <- A.breakOn byte buf
283 return $ case marr2 of
284 Nothing -> D.Skip $ Buffering st arr1
285 Just arr2 -> D.Skip $ Yielding arr1 (Splitting st arr2)
287 step' _ (Yielding arr next) = return $ D.Yield arr next
288 step' _ Finishing = return D.Stop
290-- XXX Remove MonadIO constraint.
291-- | Split a stream of arrays on a given separator byte, dropping the separator
292-- and coalescing all the arrays between two separators into a single array.
293--
294-- @since 0.7.0
295{-# INLINE splitOn #-}
296splitOn
297 :: (MonadIO m)
298 => Word8
299 -> Stream m (Array Word8)
300 -> Stream m (Array Word8)
301splitOn byte s =
302 fromStreamD $ D.splitInnerBy (A.breakOn byte) A.splice $ toStreamD s
304{-# INLINE splitOnSuffix #-}
305splitOnSuffix
306 :: (MonadIO m)
307 => Word8
308 -> Stream m (Array Word8)
309 -> Stream m (Array Word8)
310-- splitOn byte s = fromStreamD $ A.splitOn byte $ toStreamD s
311splitOnSuffix byte s =
312 fromStreamD $ D.splitInnerBySuffix (A.breakOn byte) A.splice $ toStreamD s
314-------------------------------------------------------------------------------
315-- Elimination - Running folds
316-------------------------------------------------------------------------------
318{-# INLINE_NORMAL foldBreakD #-}
319foldBreakD :: forall m a b. (MonadIO m, Unbox a) =>
320 Fold m a b -> D.Stream m (Array a) -> m (b, D.Stream m (Array a))
321foldBreakD (FL.Fold fstep initial extract) stream@(D.Stream step state) = do
322 res <- initial
323 case res of
324 FL.Partial fs -> go SPEC state fs
325 FL.Done fb -> return $! (fb, stream)
327 where
329 {-# INLINE go #-}
330 go !_ st !fs = do
331 r <- step defState st
332 case r of
333 D.Yield (Array contents start end) s ->
334 let fp = Tuple' end contents
335 in goArray SPEC s fp start fs
336 D.Skip s -> go SPEC s fs
337 D.Stop -> do
338 b <- extract fs
339 return (b, D.nil)
341 goArray !_ s (Tuple' end _) !cur !fs
342 | cur == end = do
343 go SPEC s fs
344 goArray !_ st fp@(Tuple' end contents) !cur !fs = do
345 x <- liftIO $ peekWith contents cur
346 res <- fstep fs x
347 let next = INDEX_NEXT(cur,a)
348 case res of
349 FL.Done b -> do
350 let arr = Array contents next end
351 return $! (b, D.cons arr (D.Stream step st))
352 FL.Partial fs1 -> goArray SPEC st fp next fs1
354{-# INLINE_NORMAL foldBreakK #-}
355foldBreakK :: forall m a b. (MonadIO m, Unbox a) =>
356 Fold m a b -> K.Stream m (Array a) -> m (b, K.Stream m (Array a))
357foldBreakK (FL.Fold fstep initial extract) stream = do
358 res <- initial
359 case res of
360 FL.Partial fs -> go fs stream
361 FL.Done fb -> return (fb, stream)
363 where
365 {-# INLINE go #-}
366 go !fs st = do
367 let stop = (, K.nil) <$> extract fs
368 single a = yieldk a K.nil
369 yieldk (Array contents start end) r =
370 let fp = Tuple' end contents
371 in goArray fs r fp start
372 in K.foldStream defState yieldk single stop st
374 goArray !fs st (Tuple' end _) !cur
375 | cur == end = do
376 go fs st
377 goArray !fs st fp@(Tuple' end contents) !cur = do
378 x <- liftIO $ peekWith contents cur
379 res <- fstep fs x
380 let next = INDEX_NEXT(cur,a)
381 case res of
382 FL.Done b -> do
383 let arr = Array contents next end
384 return $! (b, K.cons arr st)
385 FL.Partial fs1 -> goArray fs1 st fp next
387-- | Fold an array stream using the supplied 'Fold'. Returns the fold result
388-- and the unconsumed stream.
389--
390-- > foldBreak f = runArrayFoldBreak (ArrayFold.fromFold f)
391--
392-- /Internal/
393--
394{-# INLINE_NORMAL foldBreak #-}
395foldBreak ::
396 (MonadIO m, Unbox a)
397 => Fold m a b
398 -> Stream m (A.Array a)
399 -> m (b, Stream m (A.Array a))
400-- foldBreak f s = fmap fromStreamD <$> foldBreakD f (toStreamD s)
401foldBreak f =
402 fmap (fmap fromStreamK)
403 . foldBreakK f
404 . toStreamK
405-- If foldBreak performs better than runArrayFoldBreak we can use a rewrite
406-- rule to rewrite runArrayFoldBreak to fold.
407-- foldBreak f = runArrayFoldBreak (ArrayFold.fromFold f)
409-------------------------------------------------------------------------------
410-- Fold to a single Array
411-------------------------------------------------------------------------------
413-- When we have to take an array partially, take the last part of the array.
414{-# INLINE takeArrayListRev #-}
415takeArrayListRev :: forall a. Unbox a => Int -> [Array a] -> [Array a]
416takeArrayListRev = go
418 where
420 go _ [] = []
421 go n _ | n <= 0 = []
422 go n (x:xs) =
423 let len = Array.length x
424 in if n > len
425 then x : go (n - len) xs
426 else if n == len
427 then [x]
428 else let !(Array contents _ end) = x
429 !start = end - (n * SIZE_OF(a))
430 in [Array contents start end]
432-- When we have to take an array partially, take the last part of the array in
433-- the first split.
434{-# INLINE splitAtArrayListRev #-}
435splitAtArrayListRev ::
436 forall a. Unbox a => Int -> [Array a] -> ([Array a],[Array a])
437splitAtArrayListRev n ls
438 | n <= 0 = ([], ls)
439 | otherwise = go n ls
440 where
441 go :: Int -> [Array a] -> ([Array a], [Array a])
442 go _ [] = ([], [])
443 go m (x:xs) =
444 let len = Array.length x
445 (xs', xs'') = go (m - len) xs
446 in if m > len
447 then (x:xs', xs'')
448 else if m == len
449 then ([x],xs)
450 else let !(Array contents start end) = x
451 end1 = end - (m * SIZE_OF(a))
452 arr2 = Array contents start end1
453 arr1 = Array contents end1 end
454 in ([arr1], arr2:xs)
456-------------------------------------------------------------------------------
457-- Fold to a single Array
458-------------------------------------------------------------------------------
460{-# INLINE foldlM' #-}
461foldlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> m b
462foldlM' step begin = D.foldlM' step begin . S.toStreamD
464-- XXX Both of these implementations of splicing seem to perform equally well.
465-- We need to perform benchmarks over a range of sizes though.
467-- CAUTION! length must more than equal to lengths of all the arrays in the
468-- stream.
469{-# INLINE spliceArraysLenUnsafe #-}
470spliceArraysLenUnsafe :: (MonadIO m, Unbox a)
471 => Int -> Stream m (MA.Array a) -> m (MA.Array a)
472spliceArraysLenUnsafe len buffered = do
473 arr <- liftIO $ MA.newPinned len
474 foldlM' MA.spliceUnsafe (return arr) buffered
476{-# INLINE _spliceArrays #-}
477_spliceArrays :: (MonadIO m, Unbox a)
478 => Stream m (Array a) -> m (Array a)
479_spliceArrays s = do
480 buffered <- S.foldr S.cons S.nil s
481 len <- S.fold FL.sum (fmap Array.length buffered)
482 arr <- liftIO $ MA.newPinned len
483 final <- foldlM' writeArr (return arr) s
484 return $ A.unsafeFreeze final
486 where
488 writeArr dst arr = MA.spliceUnsafe dst (A.unsafeThaw arr)
490{-# INLINE _spliceArraysBuffered #-}
491_spliceArraysBuffered :: (MonadIO m, Unbox a)
492 => Stream m (Array a) -> m (Array a)
493_spliceArraysBuffered s = do
494 buffered <- S.foldr S.cons S.nil s
495 len <- S.fold FL.sum (fmap Array.length buffered)
496 A.unsafeFreeze <$> spliceArraysLenUnsafe len (fmap A.unsafeThaw s)
498{-# INLINE spliceArraysRealloced #-}
499spliceArraysRealloced :: forall m a. (MonadIO m, Unbox a)
500 => Stream m (Array a) -> m (Array a)
501spliceArraysRealloced s = do
502 let n = allocBytesToElemCount (undefined :: a) (4 * 1024)
503 idst = liftIO $ MA.newPinned n
505 arr <- foldlM' MA.spliceExp idst (fmap A.unsafeThaw s)
506 liftIO $ A.unsafeFreeze <$> MA.rightSize arr
508-- XXX This should just be "fold A.write"
509--
510-- | Given a stream of arrays, splice them all together to generate a single
511-- array. The stream must be /finite/.
512--
513-- @since 0.7.0
514{-# INLINE toArray #-}
515toArray :: (MonadIO m, Unbox a) => Stream m (Array a) -> m (Array a)
516toArray = spliceArraysRealloced
517-- spliceArrays = _spliceArraysBuffered
519-- exponentially increasing sizes of the chunks upto the max limit.
520-- XXX this will be easier to implement with parsers/terminating folds
521-- With this we should be able to reduce the number of chunks/allocations.
522-- The reallocation/copy based toArray can also be implemented using this.
523--
524{-
525{-# INLINE toArraysInRange #-}
526toArraysInRange :: (MonadIO m, Unbox a)
527 => Int -> Int -> Fold m (Array a) b -> Fold m a b
528toArraysInRange low high (Fold step initial extract) =
529-}
531{-
532-- | Fold the input to a pure buffered stream (List) of arrays.
533{-# INLINE _toArraysOf #-}
534_toArraysOf :: (MonadIO m, Unbox a)
535 => Int -> Fold m a (Stream Identity (Array a))
536_toArraysOf n = FL.chunksOf n (A.writeNF n) FL.toStream
537-}
539-------------------------------------------------------------------------------
540-- Elimination - running element parsers
541-------------------------------------------------------------------------------
543-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
544-- to make a newtype.
545{-# ANN type List NoSpecConstr #-}
546newtype List a = List {getList :: [a]}
548{-
549-- This can be generalized to any type provided it can be unfolded to a stream
550-- and it can be combined using a semigroup operation.
551--
552-- XXX This should be written using CPS (as parseK) if we want it to scale wrt
553-- to the number of times it can be called on the same stream.
554{-# INLINE_NORMAL parseBreakD #-}
555parseBreakD ::
556 forall m a b. (MonadIO m, MonadThrow m, Unbox a)
557 => PRD.Parser a m b
558 -> D.Stream m (Array.Array a)
559 -> m (b, D.Stream m (Array.Array a))
560parseBreakD
561 (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
563 res <- initial
564 case res of
565 PRD.IPartial s -> go SPEC state (List []) s
566 PRD.IDone b -> return (b, stream)
567 PRD.IError err -> throwM $ ParseError err
569 where
571 -- "backBuf" contains last few items in the stream that we may have to
572 -- backtrack to.
573 --
574 -- XXX currently we are using a dumb list based approach for backtracking
575 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
576 -- That will allow us more efficient random back and forth movement.
577 go !_ st backBuf !pst = do
578 r <- step defState st
579 case r of
580 D.Yield (Array contents start end) s ->
581 gobuf SPEC s backBuf
582 (Tuple' end contents) start pst
583 D.Skip s -> go SPEC s backBuf pst
584 D.Stop -> do
585 b <- extract pst
586 return (b, D.nil)
588 -- Use strictness on "cur" to keep it unboxed
589 gobuf !_ s backBuf (Tuple' end _) !cur !pst
590 | cur == end = do
591 go SPEC s backBuf pst
592 gobuf !_ s backBuf fp@(Tuple' end contents) !cur !pst = do
593 x <- liftIO $ peekWith contents cur
594 pRes <- pstep pst x
595 let next = INDEX_NEXT(cur,a)
596 case pRes of
597 PR.Partial 0 pst1 ->
598 gobuf SPEC s (List []) fp next pst1
599 PR.Partial n pst1 -> do
600 assert (n <= Prelude.length (x:getList backBuf)) (return ())
601 let src0 = Prelude.take n (x:getList backBuf)
602 arr0 = A.fromListN n (Prelude.reverse src0)
603 arr1 = Array contents next end
604 src = arr0 <> arr1
605 let !(Array cont1 start end1) = src
606 fp1 = Tuple' end1 cont1
607 gobuf SPEC s (List []) fp1 start pst1
608 PR.Continue 0 pst1 ->
609 gobuf SPEC s (List (x:getList backBuf)) fp next pst1
610 PR.Continue n pst1 -> do
611 assert (n <= Prelude.length (x:getList backBuf)) (return ())
612 let (src0, buf1) = splitAt n (x:getList backBuf)
613 arr0 = A.fromListN n (Prelude.reverse src0)
614 arr1 = Array contents next end
615 src = arr0 <> arr1
616 let !(Array cont1 start end1) = src
617 fp1 = Tuple' end1 cont1
618 gobuf SPEC s (List buf1) fp1 start pst1
619 PR.Done 0 b -> do
620 let arr = Array contents next end
621 return (b, D.cons arr (D.Stream step s))
622 PR.Done n b -> do
623 assert (n <= Prelude.length (x:getList backBuf)) (return ())
624 let src0 = Prelude.take n (x:getList backBuf)
625 -- XXX create the array in reverse instead
626 arr0 = A.fromListN n (Prelude.reverse src0)
627 arr1 = Array contents next end
628 -- XXX Use StreamK to avoid adding arbitrary layers of
629 -- constructors every time.
630 str = D.cons arr0 (D.cons arr1 (D.Stream step s))
631 return (b, str)
632 PR.Error err -> throwM $ ParseError err
633-}
635{-# INLINE_NORMAL parseBreakK #-}
636parseBreakK ::
637 forall m a b. (MonadIO m, MonadThrow m, Unbox a)
638 => PRD.Parser a m b
639 -> K.Stream m (Array.Array a)
640 -> m (b, K.Stream m (Array.Array a))
641parseBreakK (PRD.Parser pstep initial extract) stream = do
642 res <- initial
643 case res of
644 PRD.IPartial s -> go s stream []
645 PRD.IDone b -> return (b, stream)
646 PRD.IError err -> throwM $ ParseError err
648 where
650 -- "backBuf" contains last few items in the stream that we may have to
651 -- backtrack to.
652 --
653 -- XXX currently we are using a dumb list based approach for backtracking
654 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
655 -- That will allow us more efficient random back and forth movement.
656 go !pst st backBuf = do
657 let stop = goStop pst backBuf -- (, K.nil) <$> extract pst
658 single a = yieldk a K.nil
659 yieldk arr r = goArray pst backBuf r arr
660 in K.foldStream defState yieldk single stop st
662 -- Use strictness on "cur" to keep it unboxed
663 goArray !pst backBuf st (Array _ cur end) | cur == end = go pst st backBuf
664 goArray !pst backBuf st (Array contents cur end) = do
665 x <- liftIO $ peekWith contents cur
666 pRes <- pstep pst x
667 let next = INDEX_NEXT(cur,a)
668 case pRes of
669 PR.Partial 0 s ->
670 goArray s [] st (Array contents next end)
671 PR.Partial n s -> do
672 assert (n <= Prelude.length (x:backBuf)) (return ())
673 let src0 = Prelude.take n (x:backBuf)
674 arr0 = A.fromListN n (Prelude.reverse src0)
675 arr1 = Array contents next end
676 src = arr0 <> arr1
677 goArray s [] st src
678 PR.Continue 0 s ->
679 goArray s (x:backBuf) st (Array contents next end)
680 PR.Continue n s -> do
681 assert (n <= Prelude.length (x:backBuf)) (return ())
682 let (src0, buf1) = splitAt n (x:backBuf)
683 arr0 = A.fromListN n (Prelude.reverse src0)
684 arr1 = Array contents next end
685 src = arr0 <> arr1
686 goArray s buf1 st src
687 PR.Done 0 b -> do
688 let arr = Array contents next end
689 return (b, K.cons arr st)
690 PR.Done n b -> do
691 assert (n <= Prelude.length (x:backBuf)) (return ())
692 let src0 = Prelude.take n (x:backBuf)
693 -- XXX Use fromListRevN once implemented
694 -- arr0 = A.fromListRevN n src0
695 arr0 = A.fromListN n (Prelude.reverse src0)
696 arr1 = Array contents next end
697 str = K.cons arr0 (K.cons arr1 st)
698 return (b, str)
699 PR.Error err -> throwM $ ParseError err
701 -- This is a simplified goArray
702 goExtract !pst backBuf (Array _ cur end)
703 | cur == end = goStop pst backBuf
704 goExtract !pst backBuf (Array contents cur end) = do
705 x <- liftIO $ peekWith contents cur
706 pRes <- pstep pst x
707 let next = INDEX_NEXT(cur,a)
708 case pRes of
709 PR.Partial 0 s ->
710 goExtract s [] (Array contents next end)
711 PR.Partial n s -> do
712 assert (n <= Prelude.length (x:backBuf)) (return ())
713 let src0 = Prelude.take n (x:backBuf)
714 arr0 = A.fromListN n (Prelude.reverse src0)
715 arr1 = Array contents next end
716 src = arr0 <> arr1
717 goExtract s [] src
718 PR.Continue 0 s ->
719 goExtract s backBuf (Array contents next end)
720 PR.Continue n s -> do
721 assert (n <= Prelude.length (x:backBuf)) (return ())
722 let (src0, buf1) = splitAt n (x:backBuf)
723 arr0 = A.fromListN n (Prelude.reverse src0)
724 arr1 = Array contents next end
725 src = arr0 <> arr1
726 goExtract s buf1 src
727 PR.Done 0 b -> do
728 let arr = Array contents next end
729 return (b, K.fromPure arr)
730 PR.Done n b -> do
731 assert (n <= Prelude.length backBuf) (return ())
732 let src0 = Prelude.take n backBuf
733 -- XXX Use fromListRevN once implemented
734 -- arr0 = A.fromListRevN n src0
735 arr0 = A.fromListN n (Prelude.reverse src0)
736 arr1 = Array contents next end
737 str = K.cons arr0 (K.fromPure arr1)
738 return (b, str)
739 PR.Error err -> throwM $ ParseError err
741 -- This is a simplified goExtract
742 {-# INLINE goStop #-}
743 goStop !pst backBuf = do
744 pRes <- extract pst
745 case pRes of
746 PR.Partial _ _ -> error "Bug: parseBreak: Partial in extract"
747 PR.Continue 0 s ->
748 goStop s backBuf
749 PR.Continue n s -> do
750 assert (n <= Prelude.length backBuf) (return ())
751 let (src0, buf1) = splitAt n backBuf
752 arr = A.fromListN n (Prelude.reverse src0)
753 goExtract s buf1 arr
754 PR.Done 0 b ->
755 return (b, K.nil)
756 PR.Done n b -> do
757 assert (n <= Prelude.length backBuf) (return ())
758 let src0 = Prelude.take n backBuf
759 -- XXX Use fromListRevN once implemented
760 -- arr0 = A.fromListRevN n src0
761 arr0 = A.fromListN n (Prelude.reverse src0)
762 return (b, K.fromPure arr0)
763 PR.Error err -> throwM $ ParseError err
765-- | Parse an array stream using the supplied 'Parser'. Returns the parse
766-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
767--
768-- /Internal/
769--
770{-# INLINE_NORMAL parseBreak #-}
771parseBreak ::
772 (MonadIO m, MonadThrow m, Unbox a)
773 => PR.Parser a m b
774 -> Stream m (A.Array a)
775 -> m (b, Stream m (A.Array a))
776{-
777parseBreak p s =
778 fmap fromStreamD <$> parseBreakD (PRD.fromParserK p) (toStreamD s)
779-}
780parseBreak p =
781 fmap (fmap fromStreamK)
782 . parseBreakK (PRD.fromParserK p)
783 . toStreamK
785-------------------------------------------------------------------------------
786-- Elimination - Running Array Folds and parsers
787-------------------------------------------------------------------------------
789-- | Note that this is not the same as using a @Parser (Array a) m b@ with the
790-- regular "Streamly.Internal.Data.IsStream.parse" function. The regular parse
791-- would consume the input arrays as single unit. This parser parses in the way
792-- as described in the ArrayFold module. The input arrays are treated as @n@
793-- element units and can be consumed partially. The remaining elements are
794-- inserted in the source stream as an array.
795--
796{-# INLINE_NORMAL runArrayParserDBreak #-}
797runArrayParserDBreak ::
798 forall m a b. (MonadIO m, MonadThrow m, Unbox a)
799 => PRD.Parser (Array a) m b
800 -> D.Stream m (Array.Array a)
801 -> m (b, D.Stream m (Array.Array a))
802runArrayParserDBreak
803 (PRD.Parser pstep initial extract)
804 stream@(D.Stream step state) = do
806 res <- initial
807 case res of
808 PRD.IPartial s -> go SPEC state (List []) s
809 PRD.IDone b -> return (b, stream)
810 PRD.IError err -> throwM $ ParseError err
812 where
814 -- "backBuf" contains last few items in the stream that we may have to
815 -- backtrack to.
816 --
817 -- XXX currently we are using a dumb list based approach for backtracking
818 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
819 -- That will allow us more efficient random back and forth movement.
820 go _ st backBuf !pst = do
821 r <- step defState st
822 case r of
823 D.Yield x s -> gobuf SPEC [x] s backBuf pst
824 D.Skip s -> go SPEC s backBuf pst
825 D.Stop -> goStop backBuf pst
827 gobuf !_ [] s backBuf !pst = go SPEC s backBuf pst
828 gobuf !_ (x:xs) s backBuf !pst = do
829 pRes <- pstep pst x
830 case pRes of
831 PR.Partial 0 pst1 ->
832 gobuf SPEC xs s (List []) pst1
833 PR.Partial n pst1 -> do
834 assert
835 (n <= sum (map Array.length (x:getList backBuf)))
836 (return ())
837 let src0 = takeArrayListRev n (x:getList backBuf)
838 src = Prelude.reverse src0 ++ xs
839 gobuf SPEC src s (List []) pst1
840 PR.Continue 0 pst1 ->
841 gobuf SPEC xs s (List (x:getList backBuf)) pst1
842 PR.Continue n pst1 -> do
843 assert
844 (n <= sum (map Array.length (x:getList backBuf)))
845 (return ())
846 let (src0, buf1) = splitAtArrayListRev n (x:getList backBuf)
847 src = Prelude.reverse src0 ++ xs
848 gobuf SPEC src s (List buf1) pst1
849 PR.Done 0 b ->
850 return (b, D.Stream step s)
851 PR.Done n b -> do
852 assert
853 (n <= sum (map Array.length (x:getList backBuf)))
854 (return ())
855 let src0 = takeArrayListRev n (x:getList backBuf)
856 src = Prelude.reverse src0 ++ xs
857 return (b, D.append (D.fromList src) (D.Stream step s))
858 PR.Error err -> throwM $ ParseError err
860 -- This is a simplified gobuf
861 goExtract _ [] backBuf !pst = goStop backBuf pst
862 goExtract _ (x:xs) backBuf !pst = do
863 pRes <- pstep pst x
864 case pRes of
865 PR.Partial 0 pst1 ->
866 goExtract SPEC xs (List []) pst1
867 PR.Partial n pst1 -> do
868 assert
869 (n <= sum (map Array.length (x:getList backBuf)))
870 (return ())
871 let src0 = takeArrayListRev n (x:getList backBuf)
872 src = Prelude.reverse src0 ++ xs
873 goExtract SPEC src (List []) pst1
874 PR.Continue 0 pst1 ->
875 goExtract SPEC xs (List (x:getList backBuf)) pst1
876 PR.Continue n pst1 -> do
877 assert
878 (n <= sum (map Array.length (x:getList backBuf)))
879 (return ())
880 let (src0, buf1) = splitAtArrayListRev n (x:getList backBuf)
881 src = Prelude.reverse src0 ++ xs
882 goExtract SPEC src (List buf1) pst1
883 PR.Done 0 b ->
884 return (b, D.nil)
885 PR.Done n b -> do
886 assert
887 (n <= sum (map Array.length (x:getList backBuf)))
888 (return ())
889 let src0 = takeArrayListRev n (x:getList backBuf)
890 src = Prelude.reverse src0 ++ xs
891 return (b, D.fromList src)
892 PR.Error err -> throwM $ ParseError err
894 -- This is a simplified goExtract
895 {-# INLINE goStop #-}
896 goStop backBuf pst = do
897 pRes <- extract pst
898 case pRes of
899 PR.Partial _ _ -> error "Bug: runArrayParserDBreak: Partial in extract"
900 PR.Continue 0 pst1 ->
901 goStop backBuf pst1
902 PR.Continue n pst1 -> do
903 assert
904 (n <= sum (map Array.length (getList backBuf)))
905 (return ())
906 let (src0, buf1) = splitAtArrayListRev n (getList backBuf)
907 src = Prelude.reverse src0
908 goExtract SPEC src (List buf1) pst1
909 PR.Done 0 b -> return (b, D.nil)
910 PR.Done n b -> do
911 assert
912 (n <= sum (map Array.length (getList backBuf)))
913 (return ())
914 let src0 = takeArrayListRev n (getList backBuf)
915 src = Prelude.reverse src0
916 return (b, D.fromList src)
917 PR.Error err -> throwM $ ParseError err
919{-
920-- | Parse an array stream using the supplied 'Parser'. Returns the parse
921-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
922--
923-- /Internal/
924--
925{-# INLINE parseArr #-}
926parseArr ::
927 (MonadIO m, MonadThrow m, Unbox a)
928 => ASF.Parser a m b
929 -> Stream m (A.Array a)
930 -> m (b, Stream m (A.Array a))
931parseArr p s = fmap fromStreamD <$> parseBreakD p (toStreamD s)
932-}
934-- | Fold an array stream using the supplied array stream 'Fold'.
935--
936-- /Pre-release/
937--
938{-# INLINE runArrayFold #-}
939runArrayFold :: (MonadIO m, MonadThrow m, Unbox a) =>
940 ArrayFold m a b -> Stream m (A.Array a) -> m b
941runArrayFold (ArrayFold p) s = fst <$> runArrayParserDBreak p (toStreamD s)
943-- | Like 'fold' but also returns the remaining stream.
944--
945-- /Pre-release/
946--
947{-# INLINE runArrayFoldBreak #-}
948runArrayFoldBreak :: (MonadIO m, MonadThrow m, Unbox a) =>
949 ArrayFold m a b -> Stream m (A.Array a) -> m (b, Stream m (A.Array a))
950runArrayFoldBreak (ArrayFold p) s =
951 second fromStreamD <$> runArrayParserDBreak p (toStreamD s)
953{-# ANN type ParseChunksState Fuse #-}
954data ParseChunksState x inpBuf st pst =
955 ParseChunksInit inpBuf st
956 | ParseChunksInitBuf inpBuf
957 | ParseChunksInitLeftOver inpBuf
958 | ParseChunksStream st inpBuf !pst
959 | ParseChunksStop inpBuf !pst
960 | ParseChunksBuf inpBuf st inpBuf !pst
961 | ParseChunksExtract inpBuf inpBuf !pst
962 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
964{-# INLINE_NORMAL runArrayFoldManyD #-}
965runArrayFoldManyD
966 :: (MonadThrow m, Unbox a)
967 => ArrayFold m a b
968 -> D.Stream m (Array a)
969 -> D.Stream m b
970runArrayFoldManyD
971 (ArrayFold (PRD.Parser pstep initial extract)) (D.Stream step state) =
973 D.Stream stepOuter (ParseChunksInit [] state)
975 where
977 {-# INLINE_LATE stepOuter #-}
978 -- Buffer is empty, get the first element from the stream, initialize the
979 -- fold and then go to stream processing loop.
980 stepOuter gst (ParseChunksInit [] st) = do
981 r <- step (adaptState gst) st
982 case r of
983 D.Yield x s -> do
984 res <- initial
985 case res of
986 PRD.IPartial ps ->
987 return $ D.Skip $ ParseChunksBuf [x] s [] ps
988 PRD.IDone pb ->
989 let next = ParseChunksInit [x] s
990 in return $ D.Skip $ ParseChunksYield pb next
991 PRD.IError err -> throwM $ ParseError err
992 D.Skip s -> return $ D.Skip $ ParseChunksInit [] s
993 D.Stop -> return D.Stop
995 -- Buffer is not empty, go to buffered processing loop
996 stepOuter _ (ParseChunksInit src st) = do
997 res <- initial
998 case res of
999 PRD.IPartial ps ->
1000 return $ D.Skip $ ParseChunksBuf src st [] ps
1001 PRD.IDone pb ->
1002 let next = ParseChunksInit src st
1003 in return $ D.Skip $ ParseChunksYield pb next
1004 PRD.IError err -> throwM $ ParseError err
1006 -- This is a simplified ParseChunksInit
1007 stepOuter _ (ParseChunksInitBuf src) = do
1008 res <- initial
1009 case res of
1010 PRD.IPartial ps ->
1011 return $ D.Skip $ ParseChunksExtract src [] ps
1012 PRD.IDone pb ->
1013 let next = ParseChunksInitBuf src
1014 in return $ D.Skip $ ParseChunksYield pb next
1015 PRD.IError err -> throwM $ ParseError err
1017 -- XXX we just discard any leftover input at the end
1018 stepOuter _ (ParseChunksInitLeftOver _) = return D.Stop
1020 -- Buffer is empty, process elements from the stream
1021 stepOuter gst (ParseChunksStream st backBuf pst) = do
1022 r <- step (adaptState gst) st
1023 case r of
1024 D.Yield x s -> do
1025 pRes <- pstep pst x
1026 case pRes of
1027 PR.Partial 0 pst1 ->
1028 return $ D.Skip $ ParseChunksStream s [] pst1
1029 PR.Partial n pst1 -> do
1030 assert
1031 (n <= sum (map Array.length (x:backBuf)))
1032 (return ())
1033 let src0 = takeArrayListRev n (x:backBuf)
1034 src = Prelude.reverse src0
1035 return $ D.Skip $ ParseChunksBuf src s [] pst1
1036 PR.Continue 0 pst1 ->
1037 return $ D.Skip $ ParseChunksStream s (x:backBuf) pst1
1038 PR.Continue n pst1 -> do
1039 assert
1040 (n <= sum (map Array.length (x:backBuf)))
1041 (return ())
1042 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
1043 src = Prelude.reverse src0
1044 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
1045 PR.Done 0 b -> do
1046 return $ D.Skip $
1047 ParseChunksYield b (ParseChunksInit [] s)
1048 PR.Done n b -> do
1049 assert
1050 (n <= sum (map Array.length (x:backBuf)))
1051 (return ())
1052 let src0 = takeArrayListRev n (x:backBuf)
1053 let src = Prelude.reverse src0
1054 return $ D.Skip $
1055 ParseChunksYield b (ParseChunksInit src s)
1056 PR.Error err -> throwM $ ParseError err
1057 D.Skip s -> return $ D.Skip $ ParseChunksStream s backBuf pst
1058 D.Stop -> return $ D.Skip $ ParseChunksStop backBuf pst
1060 -- go back to stream processing mode
1061 stepOuter _ (ParseChunksBuf [] s buf pst) =
1062 return $ D.Skip $ ParseChunksStream s buf pst
1064 -- buffered processing loop
1065 stepOuter _ (ParseChunksBuf (x:xs) s backBuf pst) = do
1066 pRes <- pstep pst x
1067 case pRes of
1068 PR.Partial 0 pst1 ->
1069 return $ D.Skip $ ParseChunksBuf xs s [] pst1
1070 PR.Partial n pst1 -> do
1071 assert (n <= sum (map Array.length (x:backBuf))) (return ())
1072 let src0 = takeArrayListRev n (x:backBuf)
1073 src = Prelude.reverse src0 ++ xs
1074 return $ D.Skip $ ParseChunksBuf src s [] pst1
1075 PR.Continue 0 pst1 ->
1076 return $ D.Skip $ ParseChunksBuf xs s (x:backBuf) pst1
1077 PR.Continue n pst1 -> do
1078 assert (n <= sum (map Array.length (x:backBuf))) (return ())
1079 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
1080 src = Prelude.reverse src0 ++ xs
1081 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
1082 PR.Done 0 b ->
1083 return $ D.Skip $ ParseChunksYield b (ParseChunksInit xs s)
1084 PR.Done n b -> do
1085 assert (n <= sum (map Array.length (x:backBuf))) (return ())
1086 let src0 = takeArrayListRev n (x:backBuf)
1087 src = Prelude.reverse src0 ++ xs
1088 return $ D.Skip $ ParseChunksYield b (ParseChunksInit src s)
1089 PR.Error err -> throwM $ ParseError err
1091 -- This is a simplified ParseChunksBuf
1092 stepOuter _ (ParseChunksExtract [] buf pst) =
1093 return $ D.Skip $ ParseChunksStop buf pst
1095 stepOuter _ (ParseChunksExtract (x:xs) backBuf pst) = do
1096 pRes <- pstep pst x
1097 case pRes of
1098 PR.Partial 0 pst1 ->
1099 return $ D.Skip $ ParseChunksExtract xs [] pst1
1100 PR.Partial n pst1 -> do
1101 assert (n <= sum (map Array.length (x:backBuf))) (return ())
1102 let src0 = takeArrayListRev n (x:backBuf)
1103 src = Prelude.reverse src0 ++ xs
1104 return $ D.Skip $ ParseChunksExtract src [] pst1
1105 PR.Continue 0 pst1 ->
1106 return $ D.Skip $ ParseChunksExtract xs (x:backBuf) pst1
1107 PR.Continue n pst1 -> do
1108 assert (n <= sum (map Array.length (x:backBuf))) (return ())
1109 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
1110 src = Prelude.reverse src0 ++ xs
1111 return $ D.Skip $ ParseChunksExtract src buf1 pst1
1112 PR.Done 0 b ->
1113 return $ D.Skip $ ParseChunksYield b (ParseChunksInitBuf xs)
1114 PR.Done n b -> do
1115 assert (n <= sum (map Array.length (x:backBuf))) (return ())
1116 let src0 = takeArrayListRev n (x:backBuf)
1117 src = Prelude.reverse src0 ++ xs
1118 return $ D.Skip $ ParseChunksYield b (ParseChunksInitBuf src)
1119 PR.Error err -> throwM $ ParseError err
1121 -- This is a simplified ParseChunksExtract
1122 stepOuter _ (ParseChunksStop backBuf pst) = do
1123 pRes <- extract pst
1124 case pRes of
1125 PR.Partial _ _ -> error "runArrayFoldManyD: Partial in extract"
1126 PR.Continue 0 pst1 ->
1127 return $ D.Skip $ ParseChunksStop backBuf pst1
1128 PR.Continue n pst1 -> do
1129 assert (n <= sum (map Array.length backBuf)) (return ())
1130 let (src0, buf1) = splitAtArrayListRev n backBuf
1131 src = Prelude.reverse src0
1132 return $ D.Skip $ ParseChunksExtract src buf1 pst1
1133 PR.Done 0 b ->
1134 return $ D.Skip $ ParseChunksYield b (ParseChunksInitLeftOver [])
1135 PR.Done n b -> do
1136 assert (n <= sum (map Array.length backBuf)) (return ())
1137 let src0 = takeArrayListRev n backBuf
1138 src = Prelude.reverse src0
1139 return $ D.Skip $ ParseChunksYield b (ParseChunksInitBuf src)
1140 PR.Error err -> throwM $ ParseError err
1142 stepOuter _ (ParseChunksYield a next) = return $ D.Yield a next
1144-- | Apply an 'ArrayFold' repeatedly on an array stream and emit the
1145-- fold outputs in the output stream.
1147-- See "Streamly.Data.Stream.foldMany" for more details.
1149-- /Pre-release/
1150{-# INLINE runArrayFoldMany #-}
1151runArrayFoldMany
1152 :: (MonadThrow m, Unbox a)
1153 => ArrayFold m a b
1154 -> Stream m (Array a)
1155 -> Stream m b
1156runArrayFoldMany p m = fromStreamD $ runArrayFoldManyD p (toStreamD m)