git.iko.soy/streamly
on ca6cfde help sign in
streamly/src/Streamly/Internal/Data/Array/Stream/Foreign.hs
882 lines 33 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 , parseBreak
28 , parseBreakD
30 -- ** Array Folds
31 , runArrayFold
32 , runArrayFoldBreak
33 -- , parseArr
34 , runArrayParserDBreak
35 , runArrayFoldMany
37 , toArray
39 -- * Compaction
40 , lpackArraysChunksOf
41 , compact
43 -- * Splitting
44 , splitOn
45 , splitOnSuffix
46 )
47where
49#include "ArrayMacros.h"
50#include "inline.hs"
52import Data.Bifunctor (second)
53import Control.Exception (assert)
54import Control.Monad.Catch (MonadThrow, throwM)
55import Control.Monad.IO.Class (MonadIO(..))
56#if __GLASGOW_HASKELL__ < 808
57import Data.Semigroup (Semigroup(..))
58#endif
59import Data.Word (Word8)
60import Foreign.ForeignPtr (touchForeignPtr)
61import Foreign.Ptr (minusPtr, plusPtr, castPtr)
62import Foreign.Storable (Storable(..))
63import Fusion.Plugin.Types (Fuse(..))
64import GHC.Exts (SpecConstrAnnotation(..))
65import GHC.ForeignPtr (ForeignPtr(..))
66import GHC.Ptr (Ptr(..))
67import GHC.Types (SPEC(..))
68import Prelude hiding (null, last, (!!), read, concat, unlines)
70import Streamly.Internal.Data.Array.Foreign.Mut.Type
71 (arrayToFptrContents, fptrToArrayContents, touch)
72import Streamly.Internal.Data.Array.Foreign.Type (Array(..))
73import Streamly.Internal.Data.Array.Stream.Fold.Foreign (ArrayFold(..))
74import Streamly.Internal.Data.Fold.Type (Fold(..))
75import Streamly.Internal.Data.Parser (ParseError(..))
76import Streamly.Internal.Data.Stream.Serial (SerialT)
77import Streamly.Internal.Data.Stream.IsStream.Type
78 (IsStream, fromStreamD, toStreamD)
79import Streamly.Internal.Data.SVar (adaptState, defState)
80import Streamly.Internal.Data.Array.Foreign.Mut.Type
81 (memcpy, allocBytesToElemCount)
83import qualified Streamly.Internal.Data.Array.Foreign as A
84import qualified Streamly.Internal.Data.Array.Foreign as Array
85-- import qualified Streamly.Internal.Data.Array.Stream.Fold.Foreign
86-- as ArrayFold
87import qualified Streamly.Internal.Data.Array.Foreign.Type as A
88import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MA
89import qualified Streamly.Internal.Data.Array.Stream.Mut.Foreign as AS
90import qualified Streamly.Internal.Data.Fold.Type as FL
91import qualified Streamly.Internal.Data.Parser as PR
92import qualified Streamly.Internal.Data.Parser.ParserD as PRD
93import qualified Streamly.Internal.Data.Stream.IsStream as S
94import qualified Streamly.Internal.Data.Stream.StreamD as D
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 :: (IsStream t, MonadIO m, Storable a)
111 => Int -> t m a -> t 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 :: (IsStream t, Monad m, Storable a) => t m (Array a) -> t 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 :: (IsStream t, Monad m, Storable a) => t m (Array a) -> t 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, IsStream t, Storable a) => a -> t m (Array a) -> t m a
156interpose x = S.interpose x A.read
158{-# INLINE intercalateSuffix #-}
159intercalateSuffix :: (Monad m, IsStream t, Storable a)
160 => Array a -> t m (Array a) -> t 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, IsStream t, Storable a)
169 => a -> t m (Array a) -> t 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.ArrayContents !(Ptr a) !(Ptr a)
177-- XXX Remove monadIO constraint
178{-# INLINE_NORMAL unlines #-}
179unlines :: forall m a. (MonadIO m, Storable 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 aEnd)
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 $ do
197 r <- peek p
198 touch contents
199 return r
200 return $ D.Yield x (InnerLoop st contents (PTR_NEXT(p,a)) end)
202-------------------------------------------------------------------------------
203-- Compact
204-------------------------------------------------------------------------------
206-- XXX These would not be needed once we implement compactLEFold, see
207-- module Streamly.Internal.Data.Array.Stream.Mut.Foreign
208--
209{-# INLINE_NORMAL packArraysChunksOf #-}
210packArraysChunksOf :: (MonadIO m, Storable a)
211 => Int -> D.Stream m (Array a) -> D.Stream m (Array a)
212packArraysChunksOf n str =
213 D.map A.unsafeFreeze $ AS.packArraysChunksOf n $ D.map A.unsafeThaw str
215-- XXX instead of writing two different versions of this operation, we should
216-- write it as a pipe.
217{-# INLINE_NORMAL lpackArraysChunksOf #-}
218lpackArraysChunksOf :: (MonadIO m, Storable a)
219 => Int -> Fold m (Array a) () -> Fold m (Array a) ()
220lpackArraysChunksOf n fld =
221 FL.lmap A.unsafeThaw $ AS.lpackArraysChunksOf n (FL.lmap A.unsafeFreeze fld)
223-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
224-- maximum specified size in bytes.
225--
226-- @since 0.7.0
227{-# INLINE compact #-}
228compact :: (MonadIO m, Storable a)
229 => Int -> SerialT m (Array a) -> SerialT m (Array a)
230compact n xs = fromStreamD $ packArraysChunksOf n (toStreamD xs)
232-------------------------------------------------------------------------------
233-- Split
234-------------------------------------------------------------------------------
236data SplitState s arr
237 = Initial s
238 | Buffering s arr
239 | Splitting s arr
240 | Yielding arr (SplitState s arr)
241 | Finishing
243-- | Split a stream of arrays on a given separator byte, dropping the separator
244-- and coalescing all the arrays between two separators into a single array.
245--
246-- @since 0.7.0
247{-# INLINE_NORMAL _splitOn #-}
248_splitOn
249 :: MonadIO m
250 => Word8
251 -> D.Stream m (Array Word8)
252 -> D.Stream m (Array Word8)
253_splitOn byte (D.Stream step state) = D.Stream step' (Initial state)
255 where
257 {-# INLINE_LATE step' #-}
258 step' gst (Initial st) = do
259 r <- step gst st
260 case r of
261 D.Yield arr s -> do
262 (arr1, marr2) <- A.breakOn byte arr
263 return $ case marr2 of
264 Nothing -> D.Skip (Buffering s arr1)
265 Just arr2 -> D.Skip (Yielding arr1 (Splitting s arr2))
266 D.Skip s -> return $ D.Skip (Initial s)
267 D.Stop -> return D.Stop
269 step' gst (Buffering st buf) = do
270 r <- step gst st
271 case r of
272 D.Yield arr s -> do
273 (arr1, marr2) <- A.breakOn byte arr
274 buf' <- A.splice buf arr1
275 return $ case marr2 of
276 Nothing -> D.Skip (Buffering s buf')
277 Just x -> D.Skip (Yielding buf' (Splitting s x))
278 D.Skip s -> return $ D.Skip (Buffering s buf)
279 D.Stop -> return $
280 if A.byteLength buf == 0
281 then D.Stop
282 else D.Skip (Yielding buf Finishing)
284 step' _ (Splitting st buf) = do
285 (arr1, marr2) <- A.breakOn byte buf
286 return $ case marr2 of
287 Nothing -> D.Skip $ Buffering st arr1
288 Just arr2 -> D.Skip $ Yielding arr1 (Splitting st arr2)
290 step' _ (Yielding arr next) = return $ D.Yield arr next
291 step' _ Finishing = return D.Stop
293-- XXX Remove MonadIO constraint.
294-- | Split a stream of arrays on a given separator byte, dropping the separator
295-- and coalescing all the arrays between two separators into a single array.
296--
297-- @since 0.7.0
298{-# INLINE splitOn #-}
299splitOn
300 :: (IsStream t, MonadIO m)
301 => Word8
302 -> t m (Array Word8)
303 -> t m (Array Word8)
304splitOn byte s =
305 fromStreamD $ D.splitInnerBy (A.breakOn byte) A.splice $ toStreamD s
307{-# INLINE splitOnSuffix #-}
308splitOnSuffix
309 :: (IsStream t, MonadIO m)
310 => Word8
311 -> t m (Array Word8)
312 -> t m (Array Word8)
313-- splitOn byte s = fromStreamD $ A.splitOn byte $ toStreamD s
314splitOnSuffix byte s =
315 fromStreamD $ D.splitInnerBySuffix (A.breakOn byte) A.splice $ toStreamD s
317-------------------------------------------------------------------------------
318-- Elimination - Running folds
319-------------------------------------------------------------------------------
321-- XXX This should be written using CPS (as foldK) if we want it to scale wrt
322-- to the number of times it can be called on the same stream.
323--
324{-# INLINE_NORMAL foldD #-}
325foldD :: forall m a b. (MonadIO m, Storable a) =>
326 Fold m a b -> D.Stream m (Array a) -> m (b, D.Stream m (Array a))
327foldD (Fold fstep initial extract) stream@(D.Stream step state) = do
328 res <- initial
329 case res of
330 FL.Partial fs -> go SPEC state fs
331 FL.Done fb -> return $! (fb, stream)
333 where
335 {-# INLINE go #-}
336 go !_ st !fs = do
337 r <- step defState st
338 case r of
339 D.Yield (Array contents start (Ptr end)) s ->
340 let fp = ForeignPtr end (arrayToFptrContents contents)
341 in goArray SPEC s fp start fs
342 D.Skip s -> go SPEC s fs
343 D.Stop -> do
344 b <- extract fs
345 return (b, D.nil)
347 goArray !_ s fp@(ForeignPtr end _) !cur !fs
348 | cur == Ptr end = do
349 liftIO $ touchForeignPtr fp
350 go SPEC s fs
351 goArray !_ st fp@(ForeignPtr end contents) !cur !fs = do
352 x <- liftIO $ peek cur
353 res <- fstep fs x
354 let next = PTR_NEXT(cur,a)
355 case res of
356 FL.Done b -> do
357 let arr = Array (fptrToArrayContents contents) next (Ptr end)
358 return $! (b, D.cons arr (D.Stream step st))
359 FL.Partial fs1 -> goArray SPEC st fp next fs1
361-- | Fold an array stream using the supplied 'Fold'. Returns the fold result
362-- and the unconsumed stream.
363--
364-- > foldBreak f = runArrayFoldBreak (ArrayFold.fromFold f)
365--
366-- /Internal/
367--
368{-# INLINE_NORMAL foldBreak #-}
369foldBreak ::
370 (MonadIO m, Storable a)
371 => FL.Fold m a b
372 -> SerialT m (A.Array a)
373 -> m (b, SerialT m (A.Array a))
374foldBreak f s = fmap fromStreamD <$> foldD f (toStreamD s)
375-- If fold performs better than runArrayFoldBreak we can rewrite runArrayFoldBreak to
376-- fold.
377-- fold f = runArrayFoldBreak (ArrayFold.fromFold f)
379-------------------------------------------------------------------------------
380-- Fold to a single Array
381-------------------------------------------------------------------------------
383-- When we have to take an array partially, take the last part of the array.
384{-# INLINE takeArrayListRev #-}
385takeArrayListRev :: forall a. Storable a => Int -> [Array a] -> [Array a]
386takeArrayListRev = go
388 where
390 go _ [] = []
391 go n _ | n <= 0 = []
392 go n (x:xs) =
393 let len = Array.length x
394 in if n > len
395 then x : go (n - len) xs
396 else if n == len
397 then [x]
398 else let !(Array contents _ end) = x
399 !start = end `plusPtr` negate (n * SIZE_OF(a))
400 in [Array contents start end]
402-- When we have to take an array partially, take the last part of the array in
403-- the first split.
404{-# INLINE splitAtArrayListRev #-}
405splitAtArrayListRev :: forall a. Storable a =>
406 Int -> [Array a] -> ([Array a],[Array a])
407splitAtArrayListRev n ls
408 | n <= 0 = ([], ls)
409 | otherwise = go n ls
410 where
411 go :: Int -> [Array a] -> ([Array a], [Array a])
412 go _ [] = ([], [])
413 go m (x:xs) =
414 let len = Array.length x
415 (xs', xs'') = go (m - len) xs
416 in if m > len
417 then (x:xs', xs'')
418 else if m == len
419 then ([x],xs)
420 else let !(Array contents start end) = x
421 end1 = end `plusPtr` negate (m * SIZE_OF(a))
422 arr2 = Array contents start end1
423 arr1 = Array contents end1 end
424 in ([arr1], arr2:xs)
426-------------------------------------------------------------------------------
427-- Fold to a single Array
428-------------------------------------------------------------------------------
430-- XXX Both of these implementations of splicing seem to perform equally well.
431-- We need to perform benchmarks over a range of sizes though.
433-- CAUTION! length must more than equal to lengths of all the arrays in the
434-- stream.
435{-# INLINE spliceArraysLenUnsafe #-}
436spliceArraysLenUnsafe :: (MonadIO m, Storable a)
437 => Int -> SerialT m (MA.Array a) -> m (MA.Array a)
438spliceArraysLenUnsafe len buffered = do
439 arr <- liftIO $ MA.newArray len
440 end <- S.foldlM' writeArr (return $ MA.aEnd arr) buffered
441 return $ arr {MA.aEnd = end}
443 where
445 writeArr dst (MA.Array ac src ae _) =
446 liftIO $ do
447 let count = ae `minusPtr` src
448 memcpy (castPtr dst) (castPtr src) count
449 touch ac
450 return $ dst `plusPtr` count
452{-# INLINE _spliceArrays #-}
453_spliceArrays :: (MonadIO m, Storable a)
454 => SerialT m (Array a) -> m (Array a)
455_spliceArrays s = do
456 buffered <- S.foldr S.cons S.nil s
457 len <- S.sum (S.map Array.length buffered)
458 arr <- liftIO $ MA.newArray len
459 end <- S.foldlM' writeArr (return $ MA.aEnd arr) s
460 return $ A.unsafeFreeze $ arr {MA.aEnd = end}
462 where
464 writeArr dst (Array ac src ae) =
465 liftIO $ do
466 let count = ae `minusPtr` src
467 memcpy (castPtr dst) (castPtr src) count
468 touch ac
469 return $ dst `plusPtr` count
471{-# INLINE _spliceArraysBuffered #-}
472_spliceArraysBuffered :: (MonadIO m, Storable a)
473 => SerialT m (Array a) -> m (Array a)
474_spliceArraysBuffered s = do
475 buffered <- S.foldr S.cons S.nil s
476 len <- S.sum (S.map Array.length buffered)
477 A.unsafeFreeze <$> spliceArraysLenUnsafe len (S.map A.unsafeThaw s)
479{-# INLINE spliceArraysRealloced #-}
480spliceArraysRealloced :: forall m a. (MonadIO m, Storable a)
481 => SerialT m (Array a) -> m (Array a)
482spliceArraysRealloced s = do
483 let n = allocBytesToElemCount (undefined :: a) (4 * 1024)
484 idst = liftIO $ MA.newArray n
486 arr <- S.foldlM' MA.spliceExp idst (S.map A.unsafeThaw s)
487 liftIO $ A.unsafeFreeze <$> MA.rightSize arr
489-- XXX This should just be "fold A.write"
490--
491-- | Given a stream of arrays, splice them all together to generate a single
492-- array. The stream must be /finite/.
493--
494-- @since 0.7.0
495{-# INLINE toArray #-}
496toArray :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)
497toArray = spliceArraysRealloced
498-- spliceArrays = _spliceArraysBuffered
500-- exponentially increasing sizes of the chunks upto the max limit.
501-- XXX this will be easier to implement with parsers/terminating folds
502-- With this we should be able to reduce the number of chunks/allocations.
503-- The reallocation/copy based toArray can also be implemented using this.
504--
505{-
506{-# INLINE toArraysInRange #-}
507toArraysInRange :: (IsStream t, MonadIO m, Storable a)
508 => Int -> Int -> Fold m (Array a) b -> Fold m a b
509toArraysInRange low high (Fold step initial extract) =
510-}
512{-
513-- | Fold the input to a pure buffered stream (List) of arrays.
514{-# INLINE _toArraysOf #-}
515_toArraysOf :: (MonadIO m, Storable a)
516 => Int -> Fold m a (SerialT Identity (Array a))
517_toArraysOf n = FL.chunksOf n (A.writeNF n) FL.toStream
518-}
520-------------------------------------------------------------------------------
521-- Elimination - running element parsers
522-------------------------------------------------------------------------------
524-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
525-- to make a newtype.
526{-# ANN type List NoSpecConstr #-}
527newtype List a = List {getList :: [a]}
529-- This can be generalized to any type provided it can be unfolded to a stream
530-- and it can be combined using a semigroup operation.
531--
532-- XXX This should be written using CPS (as parseK) if we want it to scale wrt
533-- to the number of times it can be called on the same stream.
534{-# INLINE_NORMAL parseBreakD #-}
535parseBreakD ::
536 forall m a b. (MonadIO m, MonadThrow m, Storable a)
537 => PRD.Parser m a b
538 -> D.Stream m (Array.Array a)
539 -> m (b, D.Stream m (Array.Array a))
540parseBreakD
541 (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
543 res <- initial
544 case res of
545 PRD.IPartial s -> go SPEC state (List []) s
546 PRD.IDone b -> return (b, stream)
547 PRD.IError err -> throwM $ ParseError err
549 where
551 -- "backBuf" contains last few items in the stream that we may have to
552 -- backtrack to.
553 --
554 -- XXX currently we are using a dumb list based approach for backtracking
555 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
556 -- That will allow us more efficient random back and forth movement.
557 go !_ st backBuf !pst = do
558 r <- step defState st
559 case r of
560 D.Yield (Array contents start (Ptr end)) s ->
561 gobuf SPEC s backBuf
562 (ForeignPtr end (arrayToFptrContents contents)) start pst
563 D.Skip s -> go SPEC s backBuf pst
564 D.Stop -> do
565 b <- extract pst
566 return (b, D.nil)
568 -- Use strictness on "cur" to keep it unboxed
569 gobuf !_ s backBuf fp@(ForeignPtr end _) !cur !pst
570 | cur == Ptr end = do
571 liftIO $ touchForeignPtr fp
572 go SPEC s backBuf pst
573 gobuf !_ s backBuf fp@(ForeignPtr end contents) !cur !pst = do
574 x <- liftIO $ peek cur
575 pRes <- pstep pst x
576 let next = PTR_NEXT(cur,a)
577 case pRes of
578 PR.Partial 0 pst1 ->
579 gobuf SPEC s (List []) fp next pst1
580 PR.Partial n pst1 -> do
581 assert (n <= Prelude.length (x:getList backBuf)) (return ())
582 let src0 = Prelude.take n (x:getList backBuf)
583 arr0 = A.fromListN n (Prelude.reverse src0)
584 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
585 src = arr0 <> arr1
586 let !(Array cont1 start (Ptr end1)) = src
587 fp1 = ForeignPtr end1 (arrayToFptrContents cont1)
588 gobuf SPEC s (List []) fp1 start pst1
589 PR.Continue 0 pst1 ->
590 gobuf SPEC s (List (x:getList backBuf)) fp next pst1
591 PR.Continue n pst1 -> do
592 assert (n <= Prelude.length (x:getList backBuf)) (return ())
593 let (src0, buf1) = splitAt n (x:getList backBuf)
594 arr0 = A.fromListN n (Prelude.reverse src0)
595 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
596 src = arr0 <> arr1
597 let !(Array cont1 start (Ptr end1)) = src
598 fp1 = ForeignPtr end1 (arrayToFptrContents cont1)
599 gobuf SPEC s (List buf1) fp1 start pst1
600 PR.Done 0 b -> do
601 let arr = Array (fptrToArrayContents contents) next (Ptr end)
602 return (b, D.cons arr (D.Stream step s))
603 PR.Done n b -> do
604 assert (n <= Prelude.length (x:getList backBuf)) (return ())
605 let src0 = Prelude.take n (x:getList backBuf)
606 -- XXX create the array in reverse instead
607 arr0 = A.fromListN n (Prelude.reverse src0)
608 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
609 -- XXX Use StreamK to avoid adding arbitrary layers of
610 -- constructors every time.
611 str = D.cons arr0 (D.cons arr1 (D.Stream step s))
612 return (b, str)
613 PR.Error err -> throwM $ ParseError err
615-- | Parse an array stream using the supplied 'Parser'. Returns the parse
616-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
617--
618-- /Internal/
619--
620{-# INLINE_NORMAL parseBreak #-}
621parseBreak ::
622 (MonadIO m, MonadThrow m, Storable a)
623 => PR.Parser m a b
624 -> SerialT m (A.Array a)
625 -> m (b, SerialT m (A.Array a))
626parseBreak p s =
627 fmap fromStreamD <$> parseBreakD (PRD.fromParserK p) (toStreamD s)
629-------------------------------------------------------------------------------
630-- Elimination - Running Array Folds and parsers
631-------------------------------------------------------------------------------
633-- | Note that this is not the same as using a @Parser m (Array a) b@ with the
634-- regular "Streamly.Internal.Data.IsStream.parse" function. The regular parse
635-- would consume the input arrays as single unit. This parser parses in the way
636-- as described in the ArrayFold module. The input arrays are treated as @n@
637-- element units and can be consumed partially. The remaining elements are
638-- inserted in the source stream as an array.
639--
640{-# INLINE_NORMAL runArrayParserDBreak #-}
641runArrayParserDBreak ::
642 forall m a b. (MonadIO m, MonadThrow m, Storable a)
643 => PRD.Parser m (Array a) b
644 -> D.Stream m (Array.Array a)
645 -> m (b, D.Stream m (Array.Array a))
646runArrayParserDBreak
647 (PRD.Parser pstep initial extract)
648 stream@(D.Stream step state) = do
650 res <- initial
651 case res of
652 PRD.IPartial s -> go SPEC state (List []) s
653 PRD.IDone b -> return (b, stream)
654 PRD.IError err -> throwM $ ParseError err
656 where
658 -- "backBuf" contains last few items in the stream that we may have to
659 -- backtrack to.
660 --
661 -- XXX currently we are using a dumb list based approach for backtracking
662 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
663 -- That will allow us more efficient random back and forth movement.
664 go !_ st backBuf !pst = do
665 r <- step defState st
666 case r of
667 D.Yield x s -> gobuf SPEC [x] s backBuf pst
668 D.Skip s -> go SPEC s backBuf pst
669 D.Stop -> do
670 b <- extract pst
671 return (b, D.nil)
673 gobuf !_ [] s backBuf !pst = go SPEC s backBuf pst
674 gobuf !_ (x:xs) s backBuf !pst = do
675 pRes <- pstep pst x
676 case pRes of
677 PR.Partial 0 pst1 ->
678 gobuf SPEC xs s (List []) pst1
679 PR.Partial n pst1 -> do
680 assert
681 (n <= sum (map Array.length (x:getList backBuf)))
682 (return ())
683 let src0 = takeArrayListRev n (x:getList backBuf)
684 src = Prelude.reverse src0 ++ xs
685 gobuf SPEC src s (List []) pst1
686 PR.Continue 0 pst1 ->
687 gobuf SPEC xs s (List (x:getList backBuf)) pst1
688 PR.Continue n pst1 -> do
689 assert
690 (n <= sum (map Array.length (x:getList backBuf)))
691 (return ())
692 let (src0, buf1) = splitAtArrayListRev n (x:getList backBuf)
693 src = Prelude.reverse src0 ++ xs
694 gobuf SPEC src s (List buf1) pst1
695 PR.Done 0 b ->
696 return (b, D.Stream step s)
697 PR.Done n b -> do
698 assert
699 (n <= sum (map Array.length (x:getList backBuf)))
700 (return ())
701 let src0 = takeArrayListRev n (x:getList backBuf)
702 src = Prelude.reverse src0 ++ xs
703 return (b, D.append (D.fromList src) (D.Stream step s))
704 PR.Error err -> throwM $ ParseError err
706{-
707-- | Parse an array stream using the supplied 'Parser'. Returns the parse
708-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
709--
710-- /Internal/
711--
712{-# INLINE parseArr #-}
713parseArr ::
714 (MonadIO m, MonadThrow m, Storable a)
715 => ASF.Parser m a b
716 -> SerialT m (A.Array a)
717 -> m (b, SerialT m (A.Array a))
718parseArr p s = fmap fromStreamD <$> parseBreakD p (toStreamD s)
719-}
721-- | Fold an array stream using the supplied array stream 'Fold'.
722--
723-- /Pre-release/
724--
725{-# INLINE runArrayFold #-}
726runArrayFold :: (MonadIO m, MonadThrow m, Storable a) =>
727 ArrayFold m a b -> SerialT m (A.Array a) -> m b
728runArrayFold (ArrayFold p) s = fst <$> runArrayParserDBreak p (toStreamD s)
730-- | Like 'fold' but also returns the remaining stream.
731--
732-- /Pre-release/
733--
734{-# INLINE runArrayFoldBreak #-}
735runArrayFoldBreak :: (MonadIO m, MonadThrow m, Storable a) =>
736 ArrayFold m a b -> SerialT m (A.Array a) -> m (b, SerialT m (A.Array a))
737runArrayFoldBreak (ArrayFold p) s =
738 second fromStreamD <$> runArrayParserDBreak p (toStreamD s)
740{-# ANN type ParseChunksState Fuse #-}
741data ParseChunksState x inpBuf st pst =
742 ParseChunksInit inpBuf st
743 | ParseChunksInitLeftOver inpBuf
744 | ParseChunksStream st inpBuf !pst
745 | ParseChunksBuf inpBuf st inpBuf !pst
746 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
748{-# INLINE_NORMAL runArrayFoldManyD #-}
749runArrayFoldManyD
750 :: (MonadThrow m, Storable a)
751 => ArrayFold m a b
752 -> D.Stream m (Array a)
753 -> D.Stream m b
754runArrayFoldManyD
755 (ArrayFold (PRD.Parser pstep initial extract)) (D.Stream step state) =
757 D.Stream stepOuter (ParseChunksInit [] state)
759 where
761 {-# INLINE_LATE stepOuter #-}
762 -- Buffer is empty, get the first element from the stream, initialize the
763 -- fold and then go to stream processing loop.
764 stepOuter gst (ParseChunksInit [] st) = do
765 r <- step (adaptState gst) st
766 case r of
767 D.Yield x s -> do
768 res <- initial
769 case res of
770 PRD.IPartial ps ->
771 return $ D.Skip $ ParseChunksBuf [x] s [] ps
772 PRD.IDone pb ->
773 let next = ParseChunksInit [x] s
774 in return $ D.Skip $ ParseChunksYield pb next
775 PRD.IError err -> throwM $ ParseError err
776 D.Skip s -> return $ D.Skip $ ParseChunksInit [] s
777 D.Stop -> return D.Stop
779 -- Buffer is not empty, go to buffered processing loop
780 stepOuter _ (ParseChunksInit src st) = do
781 res <- initial
782 case res of
783 PRD.IPartial ps ->
784 return $ D.Skip $ ParseChunksBuf src st [] ps
785 PRD.IDone pb ->
786 let next = ParseChunksInit src st
787 in return $ D.Skip $ ParseChunksYield pb next
788 PRD.IError err -> throwM $ ParseError err
790 -- XXX we just discard any leftover input at the end
791 stepOuter _ (ParseChunksInitLeftOver _) = return D.Stop
793 -- Buffer is empty, process elements from the stream
794 stepOuter gst (ParseChunksStream st backBuf pst) = do
795 r <- step (adaptState gst) st
796 case r of
797 D.Yield x s -> do
798 pRes <- pstep pst x
799 case pRes of
800 PR.Partial 0 pst1 ->
801 return $ D.Skip $ ParseChunksStream s [] pst1
802 PR.Partial n pst1 -> do
803 assert
804 (n <= sum (map Array.length (x:backBuf)))
805 (return ())
806 let src0 = takeArrayListRev n (x:backBuf)
807 src = Prelude.reverse src0
808 return $ D.Skip $ ParseChunksBuf src s [] pst1
809 PR.Continue 0 pst1 ->
810 return $ D.Skip $ ParseChunksStream s (x:backBuf) pst1
811 PR.Continue n pst1 -> do
812 assert
813 (n <= sum (map Array.length (x:backBuf)))
814 (return ())
815 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
816 src = Prelude.reverse src0
817 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
818 PR.Done 0 b -> do
819 return $ D.Skip $
820 ParseChunksYield b (ParseChunksInit [] s)
821 PR.Done n b -> do
822 assert
823 (n <= sum (map Array.length (x:backBuf)))
824 (return ())
825 let src0 = takeArrayListRev n (x:backBuf)
826 let src = Prelude.reverse src0
827 return $ D.Skip $
828 ParseChunksYield b (ParseChunksInit src s)
829 PR.Error err -> throwM $ ParseError err
830 D.Skip s -> return $ D.Skip $ ParseChunksStream s backBuf pst
831 D.Stop -> do
832 b <- extract pst
833 let src = Prelude.reverse backBuf
834 return $ D.Skip $
835 ParseChunksYield b (ParseChunksInitLeftOver src)
837 -- go back to stream processing mode
838 stepOuter _ (ParseChunksBuf [] s buf pst) =
839 return $ D.Skip $ ParseChunksStream s buf pst
841 -- buffered processing loop
842 stepOuter _ (ParseChunksBuf (x:xs) s backBuf pst) = do
843 pRes <- pstep pst x
844 case pRes of
845 PR.Partial 0 pst1 ->
846 return $ D.Skip $ ParseChunksBuf xs s [] pst1
847 PR.Partial n pst1 -> do
848 assert (n <= sum (map Array.length (x:backBuf))) (return ())
849 let src0 = takeArrayListRev n (x:backBuf)
850 src = Prelude.reverse src0 ++ xs
851 return $ D.Skip $ ParseChunksBuf src s [] pst1
852 PR.Continue 0 pst1 ->
853 return $ D.Skip $ ParseChunksBuf xs s (x:backBuf) pst1
854 PR.Continue n pst1 -> do
855 assert (n <= sum (map Array.length (x:backBuf))) (return ())
856 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
857 src = Prelude.reverse src0 ++ xs
858 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
859 PR.Done 0 b ->
860 return $ D.Skip $ ParseChunksYield b (ParseChunksInit xs s)
861 PR.Done n b -> do
862 assert (n <= sum (map Array.length (x:backBuf))) (return ())
863 let src0 = takeArrayListRev n (x:backBuf)
864 src = Prelude.reverse src0 ++ xs
865 return $ D.Skip $ ParseChunksYield b (ParseChunksInit src s)
866 PR.Error err -> throwM $ ParseError err
868 stepOuter _ (ParseChunksYield a next) = return $ D.Yield a next
870-- | Apply an 'ArrayFold' repeatedly on an array stream and emit the
871-- fold outputs in the output stream.
872--
873-- See "Streamly.Prelude.foldMany" for more details.
874--
875-- /Pre-release/
876{-# INLINE runArrayFoldMany #-}
877runArrayFoldMany
878 :: (IsStream t, MonadThrow m, Storable a)
879 => ArrayFold m a b
880 -> t m (Array a)
881 -> t m b
882runArrayFoldMany p m = fromStreamD $ runArrayFoldManyD p (toStreamD m)