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