git.iko.soy/streamly
on 770887d help sign in
streamly/core/src/Streamly/Internal/Data/Array/Unboxed/Stream.hs
1158 lines 45 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Array.Unboxed.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.Unboxed.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 (Unboxed, 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.Unboxed.Type (Array(..))
67import Streamly.Internal.Data.Array.Unboxed.Stream.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.Unboxed.Mut.Type
73 (allocBytesToElemCount)
74import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
76import qualified Streamly.Data.Fold as FL
77import qualified Streamly.Internal.Data.Array.Unboxed as A
78import qualified Streamly.Internal.Data.Array.Unboxed as Array
79import qualified Streamly.Internal.Data.Array.Unboxed.Type as A
80import qualified Streamly.Internal.Data.Array.Unboxed.Mut.Type as MA
81import qualified Streamly.Internal.Data.Array.Unboxed.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, Unboxed 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, Unboxed 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.read (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.readRev
139--
140-- @since 0.7.0
141{-# INLINE concatRev #-}
142concatRev :: (Monad m, Unboxed a) => Stream m (Array a) -> Stream m a
143-- concatRev m = fromStreamD $ A.flattenArraysRev (toStreamD m)
144concatRev m = fromStreamD $ D.unfoldMany A.readRev (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, Unboxed a) => a -> Stream m (Array a) -> Stream m a
156interpose x = S.interpose x A.read
158{-# INLINE intercalateSuffix #-}
159intercalateSuffix :: (Monad m, Unboxed a)
160 => Array a -> Stream m (Array a) -> Stream m a
161intercalateSuffix = S.intercalateSuffix A.read
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, Unboxed a)
169 => a -> Stream m (Array a) -> Stream m a
170-- interposeSuffix x = fromStreamD . A.unlines x . toStreamD
171interposeSuffix x = S.interposeSuffix x A.read
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, Unboxed 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.Unboxed.Mut.Stream
205--
206{-# INLINE_NORMAL packArraysChunksOf #-}
207packArraysChunksOf :: (MonadIO m, Unboxed 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, Unboxed 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, Unboxed 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-- XXX This should be written using CPS (as foldK) if we want it to scale wrt
319-- to the number of times it can be called on the same stream.
320--
321{-# INLINE_NORMAL foldBreakD #-}
322foldBreakD :: forall m a b. (MonadIO m, Unboxed a) =>
323 Fold m a b -> D.Stream m (Array a) -> m (b, D.Stream m (Array a))
324foldBreakD (FL.Fold fstep initial extract) stream@(D.Stream step state) = do
325 res <- initial
326 case res of
327 FL.Partial fs -> go SPEC state fs
328 FL.Done fb -> return $! (fb, stream)
330 where
332 {-# INLINE go #-}
333 go !_ st !fs = do
334 r <- step defState st
335 case r of
336 D.Yield (Array contents start end) s ->
337 let fp = Tuple' end contents
338 in goArray SPEC s fp start fs
339 D.Skip s -> go SPEC s fs
340 D.Stop -> do
341 b <- extract fs
342 return (b, D.nil)
344 goArray !_ s (Tuple' end _) !cur !fs
345 | cur == end = do
346 go SPEC s fs
347 goArray !_ st fp@(Tuple' end contents) !cur !fs = do
348 x <- liftIO $ peekWith contents cur
349 res <- fstep fs x
350 let next = INDEX_NEXT(cur,a)
351 case res of
352 FL.Done b -> do
353 let arr = Array contents next end
354 return $! (b, D.cons arr (D.Stream step st))
355 FL.Partial fs1 -> goArray SPEC st fp next fs1
357{-# INLINE_NORMAL foldBreakK #-}
358foldBreakK :: forall m a b. (MonadIO m, Unboxed a) =>
359 Fold m a b -> K.Stream m (Array a) -> m (b, K.Stream m (Array a))
360foldBreakK (FL.Fold fstep initial extract) stream = do
361 res <- initial
362 case res of
363 FL.Partial fs -> go fs stream
364 FL.Done fb -> return (fb, stream)
366 where
368 {-# INLINE go #-}
369 go !fs st = do
370 let stop = (, K.nil) <$> extract fs
371 single a = yieldk a K.nil
372 yieldk (Array contents start end) r =
373 let fp = Tuple' end contents
374 in goArray fs r fp start
375 in K.foldStream defState yieldk single stop st
377 goArray !fs st (Tuple' end _) !cur
378 | cur == end = do
379 go fs st
380 goArray !fs st fp@(Tuple' end contents) !cur = do
381 x <- liftIO $ peekWith contents cur
382 res <- fstep fs x
383 let next = INDEX_NEXT(cur,a)
384 case res of
385 FL.Done b -> do
386 let arr = Array contents next end
387 return $! (b, K.cons arr st)
388 FL.Partial fs1 -> goArray fs1 st fp next
390-- | Fold an array stream using the supplied 'Fold'. Returns the fold result
391-- and the unconsumed stream.
392--
393-- > foldBreak f = runArrayFoldBreak (ArrayFold.fromFold f)
394--
395-- /Internal/
396--
397{-# INLINE_NORMAL foldBreak #-}
398foldBreak ::
399 (MonadIO m, Unboxed a)
400 => Fold m a b
401 -> Stream m (A.Array a)
402 -> m (b, Stream m (A.Array a))
403-- foldBreak f s = fmap fromStreamD <$> foldBreakD f (toStreamD s)
404foldBreak f =
405 fmap (fmap fromStreamK)
406 . foldBreakK f
407 . toStreamK
408-- If foldBreak performs better than runArrayFoldBreak we can use a rewrite
409-- rule to rewrite runArrayFoldBreak to fold.
410-- foldBreak f = runArrayFoldBreak (ArrayFold.fromFold f)
412-------------------------------------------------------------------------------
413-- Fold to a single Array
414-------------------------------------------------------------------------------
416-- When we have to take an array partially, take the last part of the array.
417{-# INLINE takeArrayListRev #-}
418takeArrayListRev :: forall a. Unboxed a => Int -> [Array a] -> [Array a]
419takeArrayListRev = go
421 where
423 go _ [] = []
424 go n _ | n <= 0 = []
425 go n (x:xs) =
426 let len = Array.length x
427 in if n > len
428 then x : go (n - len) xs
429 else if n == len
430 then [x]
431 else let !(Array contents _ end) = x
432 !start = end - (n * SIZE_OF(a))
433 in [Array contents start end]
435-- When we have to take an array partially, take the last part of the array in
436-- the first split.
437{-# INLINE splitAtArrayListRev #-}
438splitAtArrayListRev ::
439 forall a. Unboxed a => Int -> [Array a] -> ([Array a],[Array a])
440splitAtArrayListRev n ls
441 | n <= 0 = ([], ls)
442 | otherwise = go n ls
443 where
444 go :: Int -> [Array a] -> ([Array a], [Array a])
445 go _ [] = ([], [])
446 go m (x:xs) =
447 let len = Array.length x
448 (xs', xs'') = go (m - len) xs
449 in if m > len
450 then (x:xs', xs'')
451 else if m == len
452 then ([x],xs)
453 else let !(Array contents start end) = x
454 end1 = end - (m * SIZE_OF(a))
455 arr2 = Array contents start end1
456 arr1 = Array contents end1 end
457 in ([arr1], arr2:xs)
459-------------------------------------------------------------------------------
460-- Fold to a single Array
461-------------------------------------------------------------------------------
463{-# INLINE foldlM' #-}
464foldlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> m b
465foldlM' step begin = D.foldlM' step begin . S.toStreamD
467-- XXX Both of these implementations of splicing seem to perform equally well.
468-- We need to perform benchmarks over a range of sizes though.
470-- CAUTION! length must more than equal to lengths of all the arrays in the
471-- stream.
472{-# INLINE spliceArraysLenUnsafe #-}
473spliceArraysLenUnsafe :: (MonadIO m, Unboxed a)
474 => Int -> Stream m (MA.Array a) -> m (MA.Array a)
475spliceArraysLenUnsafe len buffered = do
476 arr <- liftIO $ MA.newPinned len
477 foldlM' MA.spliceUnsafe (return arr) buffered
479{-# INLINE _spliceArrays #-}
480_spliceArrays :: (MonadIO m, Unboxed a)
481 => Stream m (Array a) -> m (Array a)
482_spliceArrays s = do
483 buffered <- S.foldr S.cons S.nil s
484 len <- S.fold FL.sum (fmap Array.length buffered)
485 arr <- liftIO $ MA.newPinned len
486 final <- foldlM' writeArr (return arr) s
487 return $ A.unsafeFreeze final
489 where
491 writeArr dst arr = MA.spliceUnsafe dst (A.unsafeThaw arr)
493{-# INLINE _spliceArraysBuffered #-}
494_spliceArraysBuffered :: (MonadIO m, Unboxed a)
495 => Stream m (Array a) -> m (Array a)
496_spliceArraysBuffered s = do
497 buffered <- S.foldr S.cons S.nil s
498 len <- S.fold FL.sum (fmap Array.length buffered)
499 A.unsafeFreeze <$> spliceArraysLenUnsafe len (fmap A.unsafeThaw s)
501{-# INLINE spliceArraysRealloced #-}
502spliceArraysRealloced :: forall m a. (MonadIO m, Unboxed a)
503 => Stream m (Array a) -> m (Array a)
504spliceArraysRealloced s = do
505 let n = allocBytesToElemCount (undefined :: a) (4 * 1024)
506 idst = liftIO $ MA.newPinned n
508 arr <- foldlM' MA.spliceExp idst (fmap A.unsafeThaw s)
509 liftIO $ A.unsafeFreeze <$> MA.rightSize arr
511-- XXX This should just be "fold A.write"
512--
513-- | Given a stream of arrays, splice them all together to generate a single
514-- array. The stream must be /finite/.
515--
516-- @since 0.7.0
517{-# INLINE toArray #-}
518toArray :: (MonadIO m, Unboxed a) => Stream m (Array a) -> m (Array a)
519toArray = spliceArraysRealloced
520-- spliceArrays = _spliceArraysBuffered
522-- exponentially increasing sizes of the chunks upto the max limit.
523-- XXX this will be easier to implement with parsers/terminating folds
524-- With this we should be able to reduce the number of chunks/allocations.
525-- The reallocation/copy based toArray can also be implemented using this.
526--
527{-
528{-# INLINE toArraysInRange #-}
529toArraysInRange :: (MonadIO m, Unboxed a)
530 => Int -> Int -> Fold m (Array a) b -> Fold m a b
531toArraysInRange low high (Fold step initial extract) =
532-}
534{-
535-- | Fold the input to a pure buffered stream (List) of arrays.
536{-# INLINE _toArraysOf #-}
537_toArraysOf :: (MonadIO m, Unboxed a)
538 => Int -> Fold m a (Stream Identity (Array a))
539_toArraysOf n = FL.chunksOf n (A.writeNF n) FL.toStream
540-}
542-------------------------------------------------------------------------------
543-- Elimination - running element parsers
544-------------------------------------------------------------------------------
546-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
547-- to make a newtype.
548{-# ANN type List NoSpecConstr #-}
549newtype List a = List {getList :: [a]}
551{-
552-- This can be generalized to any type provided it can be unfolded to a stream
553-- and it can be combined using a semigroup operation.
554--
555-- XXX This should be written using CPS (as parseK) if we want it to scale wrt
556-- to the number of times it can be called on the same stream.
557{-# INLINE_NORMAL parseBreakD #-}
558parseBreakD ::
559 forall m a b. (MonadIO m, MonadThrow m, Unboxed a)
560 => PRD.Parser m a b
561 -> D.Stream m (Array.Array a)
562 -> m (b, D.Stream m (Array.Array a))
563parseBreakD
564 (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
566 res <- initial
567 case res of
568 PRD.IPartial s -> go SPEC state (List []) s
569 PRD.IDone b -> return (b, stream)
570 PRD.IError err -> throwM $ ParseError err
572 where
574 -- "backBuf" contains last few items in the stream that we may have to
575 -- backtrack to.
576 --
577 -- XXX currently we are using a dumb list based approach for backtracking
578 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
579 -- That will allow us more efficient random back and forth movement.
580 go !_ st backBuf !pst = do
581 r <- step defState st
582 case r of
583 D.Yield (Array contents start end) s ->
584 gobuf SPEC s backBuf
585 (Tuple' end contents) start pst
586 D.Skip s -> go SPEC s backBuf pst
587 D.Stop -> do
588 b <- extract pst
589 return (b, D.nil)
591 -- Use strictness on "cur" to keep it unboxed
592 gobuf !_ s backBuf (Tuple' end _) !cur !pst
593 | cur == end = do
594 go SPEC s backBuf pst
595 gobuf !_ s backBuf fp@(Tuple' end contents) !cur !pst = do
596 x <- liftIO $ peekWith contents cur
597 pRes <- pstep pst x
598 let next = INDEX_NEXT(cur,a)
599 case pRes of
600 PR.Partial 0 pst1 ->
601 gobuf SPEC s (List []) fp next pst1
602 PR.Partial n pst1 -> do
603 assert (n <= Prelude.length (x:getList backBuf)) (return ())
604 let src0 = Prelude.take n (x:getList backBuf)
605 arr0 = A.fromListN n (Prelude.reverse src0)
606 arr1 = Array contents next end
607 src = arr0 <> arr1
608 let !(Array cont1 start end1) = src
609 fp1 = Tuple' end1 cont1
610 gobuf SPEC s (List []) fp1 start pst1
611 PR.Continue 0 pst1 ->
612 gobuf SPEC s (List (x:getList backBuf)) fp next pst1
613 PR.Continue n pst1 -> do
614 assert (n <= Prelude.length (x:getList backBuf)) (return ())
615 let (src0, buf1) = splitAt n (x:getList backBuf)
616 arr0 = A.fromListN n (Prelude.reverse src0)
617 arr1 = Array contents next end
618 src = arr0 <> arr1
619 let !(Array cont1 start end1) = src
620 fp1 = Tuple' end1 cont1
621 gobuf SPEC s (List buf1) fp1 start pst1
622 PR.Done 0 b -> do
623 let arr = Array contents next end
624 return (b, D.cons arr (D.Stream step s))
625 PR.Done n b -> do
626 assert (n <= Prelude.length (x:getList backBuf)) (return ())
627 let src0 = Prelude.take n (x:getList backBuf)
628 -- XXX create the array in reverse instead
629 arr0 = A.fromListN n (Prelude.reverse src0)
630 arr1 = Array contents next end
631 -- XXX Use StreamK to avoid adding arbitrary layers of
632 -- constructors every time.
633 str = D.cons arr0 (D.cons arr1 (D.Stream step s))
634 return (b, str)
635 PR.Error err -> throwM $ ParseError err
636-}
638{-# INLINE_NORMAL parseBreakK #-}
639parseBreakK ::
640 forall m a b. (MonadIO m, MonadThrow m, Unboxed a)
641 => PRD.Parser m a b
642 -> K.Stream m (Array.Array a)
643 -> m (b, K.Stream m (Array.Array a))
644parseBreakK (PRD.Parser pstep initial extract) stream = do
645 res <- initial
646 case res of
647 PRD.IPartial s -> go s stream []
648 PRD.IDone b -> return (b, stream)
649 PRD.IError err -> throwM $ ParseError err
651 where
653 -- "backBuf" contains last few items in the stream that we may have to
654 -- backtrack to.
655 --
656 -- XXX currently we are using a dumb list based approach for backtracking
657 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
658 -- That will allow us more efficient random back and forth movement.
659 go !pst st backBuf = do
660 let stop = goStop pst backBuf -- (, K.nil) <$> extract pst
661 single a = yieldk a K.nil
662 yieldk arr r = goArray pst backBuf r arr
663 in K.foldStream defState yieldk single stop st
665 -- Use strictness on "cur" to keep it unboxed
666 goArray !pst backBuf st (Array _ cur end) | cur == end = go pst st backBuf
667 goArray !pst backBuf st (Array contents cur end) = do
668 x <- liftIO $ peekWith contents cur
669 pRes <- pstep pst x
670 let next = INDEX_NEXT(cur,a)
671 case pRes of
672 PR.Partial 0 s ->
673 goArray s [] st (Array contents next end)
674 PR.Partial n s -> do
675 assert (n <= Prelude.length (x:backBuf)) (return ())
676 let src0 = Prelude.take n (x:backBuf)
677 arr0 = A.fromListN n (Prelude.reverse src0)
678 arr1 = Array contents next end
679 src = arr0 <> arr1
680 goArray s [] st src
681 PR.Continue 0 s ->
682 goArray s (x:backBuf) st (Array contents next end)
683 PR.Continue n s -> do
684 assert (n <= Prelude.length (x:backBuf)) (return ())
685 let (src0, buf1) = splitAt n (x:backBuf)
686 arr0 = A.fromListN n (Prelude.reverse src0)
687 arr1 = Array contents next end
688 src = arr0 <> arr1
689 goArray s buf1 st src
690 PR.Done 0 b -> do
691 let arr = Array contents next end
692 return (b, K.cons arr st)
693 PR.Done n b -> do
694 assert (n <= Prelude.length (x:backBuf)) (return ())
695 let src0 = Prelude.take n (x:backBuf)
696 -- XXX Use fromListRevN once implemented
697 -- arr0 = A.fromListRevN n src0
698 arr0 = A.fromListN n (Prelude.reverse src0)
699 arr1 = Array contents next end
700 str = K.cons arr0 (K.cons arr1 st)
701 return (b, str)
702 PR.Error err -> throwM $ ParseError err
704 -- This is a simplified goArray
705 goExtract !pst backBuf (Array _ cur end)
706 | cur == end = goStop pst backBuf
707 goExtract !pst backBuf (Array contents cur end) = do
708 x <- liftIO $ peekWith contents cur
709 pRes <- pstep pst x
710 let next = INDEX_NEXT(cur,a)
711 case pRes of
712 PR.Partial 0 s ->
713 goExtract s [] (Array contents next end)
714 PR.Partial n s -> do
715 assert (n <= Prelude.length (x:backBuf)) (return ())
716 let src0 = Prelude.take n (x:backBuf)
717 arr0 = A.fromListN n (Prelude.reverse src0)
718 arr1 = Array contents next end
719 src = arr0 <> arr1
720 goExtract s [] src
721 PR.Continue 0 s ->
722 goExtract s backBuf (Array contents next end)
723 PR.Continue n s -> do
724 assert (n <= Prelude.length (x:backBuf)) (return ())
725 let (src0, buf1) = splitAt n (x:backBuf)
726 arr0 = A.fromListN n (Prelude.reverse src0)
727 arr1 = Array contents next end
728 src = arr0 <> arr1
729 goExtract s buf1 src
730 PR.Done 0 b -> do
731 let arr = Array contents next end
732 return (b, K.fromPure arr)
733 PR.Done n b -> do
734 assert (n <= Prelude.length backBuf) (return ())
735 let src0 = Prelude.take n backBuf
736 -- XXX Use fromListRevN once implemented
737 -- arr0 = A.fromListRevN n src0
738 arr0 = A.fromListN n (Prelude.reverse src0)
739 arr1 = Array contents next end
740 str = K.cons arr0 (K.fromPure arr1)
741 return (b, str)
742 PR.Error err -> throwM $ ParseError err
744 -- This is a simplified goExtract
745 {-# INLINE goStop #-}
746 goStop !pst backBuf = do
747 pRes <- extract pst
748 case pRes of
749 PR.Partial _ _ -> error "Bug: parseBreak: Partial in extract"
750 PR.Continue 0 _ ->
751 error "parseBreak: extract, Continue 0 creates infinite loop"
752 PR.Continue n s -> do
753 assert (n <= Prelude.length backBuf) (return ())
754 let (src0, buf1) = splitAt n backBuf
755 arr = A.fromListN n (Prelude.reverse src0)
756 goExtract s buf1 arr
757 PR.Done 0 b ->
758 return (b, K.nil)
759 PR.Done n b -> do
760 assert (n <= Prelude.length backBuf) (return ())
761 let src0 = Prelude.take n backBuf
762 -- XXX Use fromListRevN once implemented
763 -- arr0 = A.fromListRevN n src0
764 arr0 = A.fromListN n (Prelude.reverse src0)
765 return (b, K.fromPure arr0)
766 PR.Error err -> throwM $ ParseError err
768-- | Parse an array stream using the supplied 'Parser'. Returns the parse
769-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
770--
771-- /Internal/
772--
773{-# INLINE_NORMAL parseBreak #-}
774parseBreak ::
775 (MonadIO m, MonadThrow m, Unboxed a)
776 => PR.Parser m a b
777 -> Stream m (A.Array a)
778 -> m (b, Stream m (A.Array a))
779{-
780parseBreak p s =
781 fmap fromStreamD <$> parseBreakD (PRD.fromParserK p) (toStreamD s)
782-}
783parseBreak p =
784 fmap (fmap fromStreamK)
785 . parseBreakK (PRD.fromParserK p)
786 . toStreamK
788-------------------------------------------------------------------------------
789-- Elimination - Running Array Folds and parsers
790-------------------------------------------------------------------------------
792-- | Note that this is not the same as using a @Parser m (Array a) b@ with the
793-- regular "Streamly.Internal.Data.IsStream.parse" function. The regular parse
794-- would consume the input arrays as single unit. This parser parses in the way
795-- as described in the ArrayFold module. The input arrays are treated as @n@
796-- element units and can be consumed partially. The remaining elements are
797-- inserted in the source stream as an array.
798--
799{-# INLINE_NORMAL runArrayParserDBreak #-}
800runArrayParserDBreak ::
801 forall m a b. (MonadIO m, MonadThrow m, Unboxed a)
802 => PRD.Parser m (Array a) b
803 -> D.Stream m (Array.Array a)
804 -> m (b, D.Stream m (Array.Array a))
805runArrayParserDBreak
806 (PRD.Parser pstep initial extract)
807 stream@(D.Stream step state) = do
809 res <- initial
810 case res of
811 PRD.IPartial s -> go SPEC state (List []) s
812 PRD.IDone b -> return (b, stream)
813 PRD.IError err -> throwM $ ParseError err
815 where
817 -- "backBuf" contains last few items in the stream that we may have to
818 -- backtrack to.
819 --
820 -- XXX currently we are using a dumb list based approach for backtracking
821 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
822 -- That will allow us more efficient random back and forth movement.
823 go _ st backBuf !pst = do
824 r <- step defState st
825 case r of
826 D.Yield x s -> gobuf SPEC [x] s backBuf pst
827 D.Skip s -> go SPEC s backBuf pst
828 D.Stop -> goStop backBuf pst
830 gobuf !_ [] s backBuf !pst = go SPEC s backBuf pst
831 gobuf !_ (x:xs) s backBuf !pst = do
832 pRes <- pstep pst x
833 case pRes of
834 PR.Partial 0 pst1 ->
835 gobuf SPEC xs s (List []) pst1
836 PR.Partial n pst1 -> do
837 assert
838 (n <= sum (map Array.length (x:getList backBuf)))
839 (return ())
840 let src0 = takeArrayListRev n (x:getList backBuf)
841 src = Prelude.reverse src0 ++ xs
842 gobuf SPEC src s (List []) pst1
843 PR.Continue 0 pst1 ->
844 gobuf SPEC xs s (List (x:getList backBuf)) pst1
845 PR.Continue n pst1 -> do
846 assert
847 (n <= sum (map Array.length (x:getList backBuf)))
848 (return ())
849 let (src0, buf1) = splitAtArrayListRev n (x:getList backBuf)
850 src = Prelude.reverse src0 ++ xs
851 gobuf SPEC src s (List buf1) pst1
852 PR.Done 0 b ->
853 return (b, D.Stream step s)
854 PR.Done n b -> do
855 assert
856 (n <= sum (map Array.length (x:getList backBuf)))
857 (return ())
858 let src0 = takeArrayListRev n (x:getList backBuf)
859 src = Prelude.reverse src0 ++ xs
860 return (b, D.append (D.fromList src) (D.Stream step s))
861 PR.Error err -> throwM $ ParseError err
863 -- This is a simplified gobuf
864 goExtract _ [] backBuf !pst = goStop backBuf pst
865 goExtract _ (x:xs) backBuf !pst = do
866 pRes <- pstep pst x
867 case pRes of
868 PR.Partial 0 pst1 ->
869 goExtract SPEC xs (List []) pst1
870 PR.Partial n pst1 -> do
871 assert
872 (n <= sum (map Array.length (x:getList backBuf)))
873 (return ())
874 let src0 = takeArrayListRev n (x:getList backBuf)
875 src = Prelude.reverse src0 ++ xs
876 goExtract SPEC src (List []) pst1
877 PR.Continue 0 pst1 ->
878 goExtract SPEC xs (List (x:getList backBuf)) pst1
879 PR.Continue n pst1 -> do
880 assert
881 (n <= sum (map Array.length (x:getList backBuf)))
882 (return ())
883 let (src0, buf1) = splitAtArrayListRev n (x:getList backBuf)
884 src = Prelude.reverse src0 ++ xs
885 goExtract SPEC src (List buf1) pst1
886 PR.Done 0 b ->
887 return (b, D.nil)
888 PR.Done n b -> do
889 assert
890 (n <= sum (map Array.length (x:getList backBuf)))
891 (return ())
892 let src0 = takeArrayListRev n (x:getList backBuf)
893 src = Prelude.reverse src0 ++ xs
894 return (b, D.fromList src)
895 PR.Error err -> throwM $ ParseError err
897 -- This is a simplified goExtract
898 {-# INLINE goStop #-}
899 goStop backBuf pst = do
900 pRes <- extract pst
901 case pRes of
902 PR.Partial _ _ -> error "Bug: runArrayParserDBreak: Partial in extract"
903 PR.Continue 0 _ ->
904 error "runArrayParserDBreak: extract, Continue 0 creates infinite loop"
905 PR.Continue n pst1 -> do
906 assert
907 (n <= sum (map Array.length (getList backBuf)))
908 (return ())
909 let (src0, buf1) = splitAtArrayListRev n (getList backBuf)
910 src = Prelude.reverse src0
911 goExtract SPEC src (List buf1) pst1
912 PR.Done 0 b -> return (b, D.nil)
913 PR.Done n b -> do
914 assert
915 (n <= sum (map Array.length (getList backBuf)))
916 (return ())
917 let src0 = takeArrayListRev n (getList backBuf)
918 src = Prelude.reverse src0
919 return (b, D.fromList src)
920 PR.Error err -> throwM $ ParseError err
922{-
923-- | Parse an array stream using the supplied 'Parser'. Returns the parse
924-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
925--
926-- /Internal/
927--
928{-# INLINE parseArr #-}
929parseArr ::
930 (MonadIO m, MonadThrow m, Unboxed a)
931 => ASF.Parser m a b
932 -> Stream m (A.Array a)
933 -> m (b, Stream m (A.Array a))
934parseArr p s = fmap fromStreamD <$> parseBreakD p (toStreamD s)
935-}
937-- | Fold an array stream using the supplied array stream 'Fold'.
938--
939-- /Pre-release/
940--
941{-# INLINE runArrayFold #-}
942runArrayFold :: (MonadIO m, MonadThrow m, Unboxed a) =>
943 ArrayFold m a b -> Stream m (A.Array a) -> m b
944runArrayFold (ArrayFold p) s = fst <$> runArrayParserDBreak p (toStreamD s)
946-- | Like 'fold' but also returns the remaining stream.
947--
948-- /Pre-release/
949--
950{-# INLINE runArrayFoldBreak #-}
951runArrayFoldBreak :: (MonadIO m, MonadThrow m, Unboxed a) =>
952 ArrayFold m a b -> Stream m (A.Array a) -> m (b, Stream m (A.Array a))
953runArrayFoldBreak (ArrayFold p) s =
954 second fromStreamD <$> runArrayParserDBreak p (toStreamD s)
956{-# ANN type ParseChunksState Fuse #-}
957data ParseChunksState x inpBuf st pst =
958 ParseChunksInit inpBuf st
959 | ParseChunksInitBuf inpBuf
960 | ParseChunksInitLeftOver inpBuf
961 | ParseChunksStream st inpBuf !pst
962 | ParseChunksStop inpBuf !pst
963 | ParseChunksBuf inpBuf st inpBuf !pst
964 | ParseChunksExtract inpBuf inpBuf !pst
965 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
967{-# INLINE_NORMAL runArrayFoldManyD #-}
968runArrayFoldManyD
969 :: (MonadThrow m, Unboxed a)
970 => ArrayFold m a b
971 -> D.Stream m (Array a)
972 -> D.Stream m b
973runArrayFoldManyD
974 (ArrayFold (PRD.Parser pstep initial extract)) (D.Stream step state) =
976 D.Stream stepOuter (ParseChunksInit [] state)
978 where
980 {-# INLINE_LATE stepOuter #-}
981 -- Buffer is empty, get the first element from the stream, initialize the
982 -- fold and then go to stream processing loop.
983 stepOuter gst (ParseChunksInit [] st) = do
984 r <- step (adaptState gst) st
985 case r of
986 D.Yield x s -> do
987 res <- initial
988 case res of
989 PRD.IPartial ps ->
990 return $ D.Skip $ ParseChunksBuf [x] s [] ps
991 PRD.IDone pb ->
992 let next = ParseChunksInit [x] s
993 in return $ D.Skip $ ParseChunksYield pb next
994 PRD.IError err -> throwM $ ParseError err
995 D.Skip s -> return $ D.Skip $ ParseChunksInit [] s
996 D.Stop -> return D.Stop
998 -- Buffer is not empty, go to buffered processing loop
999 stepOuter _ (ParseChunksInit src st) = do
1000 res <- initial
1001 case res of
1002 PRD.IPartial ps ->
1003 return $ D.Skip $ ParseChunksBuf src st [] ps
1004 PRD.IDone pb ->
1005 let next = ParseChunksInit src st
1006 in return $ D.Skip $ ParseChunksYield pb next
1007 PRD.IError err -> throwM $ ParseError err
1009 -- This is a simplified ParseChunksInit
1010 stepOuter _ (ParseChunksInitBuf src) = do
1011 res <- initial
1012 case res of
1013 PRD.IPartial ps ->
1014 return $ D.Skip $ ParseChunksExtract src [] ps
1015 PRD.IDone pb ->
1016 let next = ParseChunksInitBuf src
1017 in return $ D.Skip $ ParseChunksYield pb next
1018 PRD.IError err -> throwM $ ParseError err
1020 -- XXX we just discard any leftover input at the end
1021 stepOuter _ (ParseChunksInitLeftOver _) = return D.Stop
1023 -- Buffer is empty, process elements from the stream
1024 stepOuter gst (ParseChunksStream st backBuf pst) = do
1025 r <- step (adaptState gst) st
1026 case r of
1027 D.Yield x s -> do
1028 pRes <- pstep pst x
1029 case pRes of
1030 PR.Partial 0 pst1 ->
1031 return $ D.Skip $ ParseChunksStream s [] pst1
1032 PR.Partial n pst1 -> do
1033 assert
1034 (n <= sum (map Array.length (x:backBuf)))
1035 (return ())
1036 let src0 = takeArrayListRev n (x:backBuf)
1037 src = Prelude.reverse src0
1038 return $ D.Skip $ ParseChunksBuf src s [] pst1
1039 PR.Continue 0 pst1 ->
1040 return $ D.Skip $ ParseChunksStream s (x:backBuf) pst1
1041 PR.Continue n pst1 -> do
1042 assert
1043 (n <= sum (map Array.length (x:backBuf)))
1044 (return ())
1045 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
1046 src = Prelude.reverse src0
1047 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
1048 PR.Done 0 b -> do
1049 return $ D.Skip $
1050 ParseChunksYield b (ParseChunksInit [] s)
1051 PR.Done n b -> do
1052 assert
1053 (n <= sum (map Array.length (x:backBuf)))
1054 (return ())
1055 let src0 = takeArrayListRev n (x:backBuf)
1056 let src = Prelude.reverse src0
1057 return $ D.Skip $
1058 ParseChunksYield b (ParseChunksInit src s)
1059 PR.Error err -> throwM $ ParseError err
1060 D.Skip s -> return $ D.Skip $ ParseChunksStream s backBuf pst
1061 D.Stop -> return $ D.Skip $ ParseChunksStop backBuf pst
1063 -- go back to stream processing mode
1064 stepOuter _ (ParseChunksBuf [] s buf pst) =
1065 return $ D.Skip $ ParseChunksStream s buf pst
1067 -- buffered processing loop
1068 stepOuter _ (ParseChunksBuf (x:xs) s backBuf pst) = do
1069 pRes <- pstep pst x
1070 case pRes of
1071 PR.Partial 0 pst1 ->
1072 return $ D.Skip $ ParseChunksBuf xs s [] pst1
1073 PR.Partial n pst1 -> do
1074 assert (n <= sum (map Array.length (x:backBuf))) (return ())
1075 let src0 = takeArrayListRev n (x:backBuf)
1076 src = Prelude.reverse src0 ++ xs
1077 return $ D.Skip $ ParseChunksBuf src s [] pst1
1078 PR.Continue 0 pst1 ->
1079 return $ D.Skip $ ParseChunksBuf xs s (x:backBuf) pst1
1080 PR.Continue n pst1 -> do
1081 assert (n <= sum (map Array.length (x:backBuf))) (return ())
1082 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
1083 src = Prelude.reverse src0 ++ xs
1084 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
1085 PR.Done 0 b ->
1086 return $ D.Skip $ ParseChunksYield b (ParseChunksInit xs s)
1087 PR.Done n b -> do
1088 assert (n <= sum (map Array.length (x:backBuf))) (return ())
1089 let src0 = takeArrayListRev n (x:backBuf)
1090 src = Prelude.reverse src0 ++ xs
1091 return $ D.Skip $ ParseChunksYield b (ParseChunksInit src s)
1092 PR.Error err -> throwM $ ParseError err
1094 -- This is a simplified ParseChunksBuf
1095 stepOuter _ (ParseChunksExtract [] buf pst) =
1096 return $ D.Skip $ ParseChunksStop buf pst
1098 stepOuter _ (ParseChunksExtract (x:xs) backBuf pst) = do
1099 pRes <- pstep pst x
1100 case pRes of
1101 PR.Partial 0 pst1 ->
1102 return $ D.Skip $ ParseChunksExtract xs [] pst1
1103 PR.Partial n pst1 -> do
1104 assert (n <= sum (map Array.length (x:backBuf))) (return ())
1105 let src0 = takeArrayListRev n (x:backBuf)
1106 src = Prelude.reverse src0 ++ xs
1107 return $ D.Skip $ ParseChunksExtract src [] pst1
1108 PR.Continue 0 pst1 ->
1109 return $ D.Skip $ ParseChunksExtract xs (x:backBuf) pst1
1110 PR.Continue n pst1 -> do
1111 assert (n <= sum (map Array.length (x:backBuf))) (return ())
1112 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
1113 src = Prelude.reverse src0 ++ xs
1114 return $ D.Skip $ ParseChunksExtract src buf1 pst1
1115 PR.Done 0 b ->
1116 return $ D.Skip $ ParseChunksYield b (ParseChunksInitBuf xs)
1117 PR.Done n b -> do
1118 assert (n <= sum (map Array.length (x:backBuf))) (return ())
1119 let src0 = takeArrayListRev n (x:backBuf)
1120 src = Prelude.reverse src0 ++ xs
1121 return $ D.Skip $ ParseChunksYield b (ParseChunksInitBuf src)
1122 PR.Error err -> throwM $ ParseError err
1124 -- This is a simplified ParseChunksExtract
1125 stepOuter _ (ParseChunksStop backBuf pst) = do
1126 pRes <- extract pst
1127 case pRes of
1128 PR.Partial _ _ -> error "runArrayFoldManyD: Partial in extract"
1129 PR.Continue 0 _ -> error "runArrayFoldManyD: Continue 0 in extract"
1130 PR.Continue n pst1 -> do
1131 assert (n <= sum (map Array.length backBuf)) (return ())
1132 let (src0, buf1) = splitAtArrayListRev n backBuf
1133 src = Prelude.reverse src0
1134 return $ D.Skip $ ParseChunksExtract src buf1 pst1
1135 PR.Done 0 b ->
1136 return $ D.Skip $ ParseChunksYield b (ParseChunksInitLeftOver [])
1137 PR.Done n b -> do
1138 assert (n <= sum (map Array.length backBuf)) (return ())
1139 let src0 = takeArrayListRev n backBuf
1140 src = Prelude.reverse src0
1141 return $ D.Skip $ ParseChunksYield b (ParseChunksInitBuf src)
1142 PR.Error err -> throwM $ ParseError err
1144 stepOuter _ (ParseChunksYield a next) = return $ D.Yield a next
1146-- | Apply an 'ArrayFold' repeatedly on an array stream and emit the
1147-- fold outputs in the output stream.
1149-- See "Streamly.Prelude.foldMany" for more details.
1151-- /Pre-release/
1152{-# INLINE runArrayFoldMany #-}
1153runArrayFoldMany
1154 :: (MonadThrow m, Unboxed a)
1155 => ArrayFold m a b
1156 -> Stream m (Array a)
1157 -> Stream m b
1158runArrayFoldMany p m = fromStreamD $ runArrayFoldManyD p (toStreamD m)