git.iko.soy/streamly
on 4e604ef help sign in
streamly/src/Streamly/Internal/Data/Array/Stream/Foreign.hs
859 lines 32 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 , fold
27 , parse
28 , parseD
30 -- ** Array Folds
31 , foldArr
32 , foldArr_
33 -- , parseArr
34 , parseArrD
35 , foldArrMany
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.Fold.Type (Fold(..))
74import Streamly.Internal.Data.Parser (ParseError(..))
75import Streamly.Internal.Data.Stream.Serial (SerialT)
76import Streamly.Internal.Data.Stream.IsStream.Type
77 (IsStream, fromStreamD, toStreamD)
78import Streamly.Internal.Data.SVar (adaptState, defState)
79import Streamly.Internal.Data.Array.Foreign.Mut.Type
80 (memcpy, allocBytesToElemCount)
82import qualified Streamly.Internal.Data.Array.Foreign as A
83import qualified Streamly.Internal.Data.Array.Foreign as Array
84import qualified Streamly.Internal.Data.Array.Foreign.Type as A
85import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MA
86import qualified Streamly.Internal.Data.Array.Stream.Mut.Foreign as AS
87import qualified Streamly.Internal.Data.Array.Stream.Fold.Foreign as ASF
88import qualified Streamly.Internal.Data.Fold.Type as FL
89import qualified Streamly.Internal.Data.Parser as PR
90import qualified Streamly.Internal.Data.Parser.ParserD as PRD
91import qualified Streamly.Internal.Data.Stream.IsStream as S
92import qualified Streamly.Internal.Data.Stream.StreamD as D
94-- XXX Since these are immutable arrays MonadIO constraint can be removed from
95-- most places.
97-------------------------------------------------------------------------------
98-- Generation
99-------------------------------------------------------------------------------
101-- | @arraysOf n stream@ groups the elements in the input stream into arrays of
102-- @n@ elements each.
103--
104-- > arraysOf n = Stream.chunksOf n (Array.writeN n)
105--
106-- /Pre-release/
107{-# INLINE arraysOf #-}
108arraysOf :: (IsStream t, MonadIO m, Storable a)
109 => Int -> t m a -> t m (Array a)
110arraysOf n str = fromStreamD $ A.arraysOf n (toStreamD str)
112-------------------------------------------------------------------------------
113-- Append
114-------------------------------------------------------------------------------
116-- XXX efficiently compare two streams of arrays. Two streams can have chunks
117-- of different sizes, we can handle that in the stream comparison abstraction.
118-- This could be useful e.g. to fast compare whether two files differ.
120-- | Convert a stream of arrays into a stream of their elements.
121--
122-- Same as the following:
123--
124-- > concat = Stream.unfoldMany Array.read
125--
126-- @since 0.7.0
127{-# INLINE concat #-}
128concat :: (IsStream t, Monad m, Storable a) => t m (Array a) -> t m a
129-- concat m = fromStreamD $ A.flattenArrays (toStreamD m)
130-- concat m = fromStreamD $ D.concatMap A.toStreamD (toStreamD m)
131concat m = fromStreamD $ D.unfoldMany A.read (toStreamD m)
133-- | Convert a stream of arrays into a stream of their elements reversing the
134-- contents of each array before flattening.
135--
136-- > concatRev = Stream.unfoldMany Array.readRev
137--
138-- @since 0.7.0
139{-# INLINE concatRev #-}
140concatRev :: (IsStream t, Monad m, Storable a) => t m (Array a) -> t m a
141-- concatRev m = fromStreamD $ A.flattenArraysRev (toStreamD m)
142concatRev m = fromStreamD $ D.unfoldMany A.readRev (toStreamD m)
144-------------------------------------------------------------------------------
145-- Intersperse and append
146-------------------------------------------------------------------------------
148-- | Flatten a stream of arrays after inserting the given element between
149-- arrays.
150--
151-- /Pre-release/
152{-# INLINE interpose #-}
153interpose :: (Monad m, IsStream t, Storable a) => a -> t m (Array a) -> t m a
154interpose x = S.interpose x A.read
156{-# INLINE intercalateSuffix #-}
157intercalateSuffix :: (Monad m, IsStream t, Storable a)
158 => Array a -> t m (Array a) -> t m a
159intercalateSuffix = S.intercalateSuffix A.read
161-- | Flatten a stream of arrays appending the given element after each
162-- array.
163--
164-- @since 0.7.0
165{-# INLINE interposeSuffix #-}
166interposeSuffix :: (Monad m, IsStream t, Storable a)
167 => a -> t m (Array a) -> t m a
168-- interposeSuffix x = fromStreamD . A.unlines x . toStreamD
169interposeSuffix x = S.interposeSuffix x A.read
171data FlattenState s a =
172 OuterLoop s
173 | InnerLoop s !MA.ArrayContents !(Ptr a) !(Ptr a)
175-- XXX Remove monadIO constraint
176{-# INLINE_NORMAL unlines #-}
177unlines :: forall m a. (MonadIO m, Storable a)
178 => a -> D.Stream m (Array a) -> D.Stream m a
179unlines sep (D.Stream step state) = D.Stream step' (OuterLoop state)
180 where
181 {-# INLINE_LATE step' #-}
182 step' gst (OuterLoop st) = do
183 r <- step (adaptState gst) st
184 return $ case r of
185 D.Yield Array{..} s ->
186 D.Skip (InnerLoop s arrContents arrStart aEnd)
187 D.Skip s -> D.Skip (OuterLoop s)
188 D.Stop -> D.Stop
190 step' _ (InnerLoop st _ p end) | p == end =
191 return $ D.Yield sep $ OuterLoop st
193 step' _ (InnerLoop st contents p end) = do
194 x <- liftIO $ do
195 r <- peek p
196 touch contents
197 return r
198 return $ D.Yield x (InnerLoop st contents (PTR_NEXT(p,a)) end)
200-------------------------------------------------------------------------------
201-- Compact
202-------------------------------------------------------------------------------
204-- XXX These would not be needed once we implement compactLEFold, see
205-- module Streamly.Internal.Data.Array.Stream.Mut.Foreign
206--
207{-# INLINE_NORMAL packArraysChunksOf #-}
208packArraysChunksOf :: (MonadIO m, Storable a)
209 => Int -> D.Stream m (Array a) -> D.Stream m (Array a)
210packArraysChunksOf n str =
211 D.map A.unsafeFreeze $ AS.packArraysChunksOf n $ D.map A.unsafeThaw str
213-- XXX instead of writing two different versions of this operation, we should
214-- write it as a pipe.
215{-# INLINE_NORMAL lpackArraysChunksOf #-}
216lpackArraysChunksOf :: (MonadIO m, Storable a)
217 => Int -> Fold m (Array a) () -> Fold m (Array a) ()
218lpackArraysChunksOf n fld =
219 FL.lmap A.unsafeThaw $ AS.lpackArraysChunksOf n (FL.lmap A.unsafeFreeze fld)
221-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
222-- maximum specified size in bytes.
223--
224-- @since 0.7.0
225{-# INLINE compact #-}
226compact :: (MonadIO m, Storable a)
227 => Int -> SerialT m (Array a) -> SerialT m (Array a)
228compact n xs = fromStreamD $ packArraysChunksOf n (toStreamD xs)
230-------------------------------------------------------------------------------
231-- Split
232-------------------------------------------------------------------------------
234data SplitState s arr
235 = Initial s
236 | Buffering s arr
237 | Splitting s arr
238 | Yielding arr (SplitState s arr)
239 | Finishing
241-- | Split a stream of arrays on a given separator byte, dropping the separator
242-- and coalescing all the arrays between two separators into a single array.
243--
244-- @since 0.7.0
245{-# INLINE_NORMAL _splitOn #-}
246_splitOn
247 :: MonadIO m
248 => Word8
249 -> D.Stream m (Array Word8)
250 -> D.Stream m (Array Word8)
251_splitOn byte (D.Stream step state) = D.Stream step' (Initial state)
253 where
255 {-# INLINE_LATE step' #-}
256 step' gst (Initial st) = do
257 r <- step gst st
258 case r of
259 D.Yield arr s -> do
260 (arr1, marr2) <- A.breakOn byte arr
261 return $ case marr2 of
262 Nothing -> D.Skip (Buffering s arr1)
263 Just arr2 -> D.Skip (Yielding arr1 (Splitting s arr2))
264 D.Skip s -> return $ D.Skip (Initial s)
265 D.Stop -> return D.Stop
267 step' gst (Buffering st buf) = do
268 r <- step gst st
269 case r of
270 D.Yield arr s -> do
271 (arr1, marr2) <- A.breakOn byte arr
272 buf' <- A.splice buf arr1
273 return $ case marr2 of
274 Nothing -> D.Skip (Buffering s buf')
275 Just x -> D.Skip (Yielding buf' (Splitting s x))
276 D.Skip s -> return $ D.Skip (Buffering s buf)
277 D.Stop -> return $
278 if A.byteLength buf == 0
279 then D.Stop
280 else D.Skip (Yielding buf Finishing)
282 step' _ (Splitting st buf) = do
283 (arr1, marr2) <- A.breakOn byte buf
284 return $ case marr2 of
285 Nothing -> D.Skip $ Buffering st arr1
286 Just arr2 -> D.Skip $ Yielding arr1 (Splitting st arr2)
288 step' _ (Yielding arr next) = return $ D.Yield arr next
289 step' _ Finishing = return D.Stop
291-- XXX Remove MonadIO constraint.
292-- | Split a stream of arrays on a given separator byte, dropping the separator
293-- and coalescing all the arrays between two separators into a single array.
294--
295-- @since 0.7.0
296{-# INLINE splitOn #-}
297splitOn
298 :: (IsStream t, MonadIO m)
299 => Word8
300 -> t m (Array Word8)
301 -> t m (Array Word8)
302splitOn byte s =
303 fromStreamD $ D.splitInnerBy (A.breakOn byte) A.splice $ toStreamD s
305{-# INLINE splitOnSuffix #-}
306splitOnSuffix
307 :: (IsStream t, MonadIO m)
308 => Word8
309 -> t m (Array Word8)
310 -> t m (Array Word8)
311-- splitOn byte s = fromStreamD $ A.splitOn byte $ toStreamD s
312splitOnSuffix byte s =
313 fromStreamD $ D.splitInnerBySuffix (A.breakOn byte) A.splice $ toStreamD s
315-------------------------------------------------------------------------------
316-- Elimination - Running folds
317-------------------------------------------------------------------------------
319-- XXX This should be written using CPS (as foldK) if we want it to scale wrt
320-- to the number of times it can be called on the same stream.
321--
322{-# INLINE_NORMAL foldD #-}
323foldD :: forall m a b. (MonadIO m, Storable a) =>
324 Fold m a b -> D.Stream m (Array a) -> m (b, D.Stream m (Array a))
325foldD (Fold fstep initial extract) stream@(D.Stream step state) = do
326 res <- initial
327 case res of
328 FL.Partial fs -> go SPEC state fs
329 FL.Done fb -> return $! (fb, stream)
331 where
333 {-# INLINE go #-}
334 go !_ st !fs = do
335 r <- step defState st
336 case r of
337 D.Yield (Array contents start (Ptr end)) s ->
338 let fp = ForeignPtr end (arrayToFptrContents contents)
339 in goArray SPEC s fp start fs
340 D.Skip s -> go SPEC s fs
341 D.Stop -> do
342 b <- extract fs
343 return (b, D.nil)
345 goArray !_ s fp@(ForeignPtr end _) !cur !fs
346 | cur == Ptr end = do
347 liftIO $ touchForeignPtr fp
348 go SPEC s fs
349 goArray !_ st fp@(ForeignPtr end contents) !cur !fs = do
350 x <- liftIO $ peek cur
351 res <- fstep fs x
352 let next = PTR_NEXT(cur,a)
353 case res of
354 FL.Done b -> do
355 let arr = Array (fptrToArrayContents contents) next (Ptr end)
356 return $! (b, D.cons arr (D.Stream step st))
357 FL.Partial fs1 -> goArray SPEC st fp next fs1
359-- | Fold an array stream using the supplied 'Fold'. Returns the fold result
360-- and the unconsumed stream.
361--
362-- /Internal/
363--
364{-# INLINE_NORMAL fold #-}
365fold ::
366 (MonadIO m, Storable a)
367 => FL.Fold m a b
368 -> SerialT m (A.Array a)
369 -> m (b, SerialT m (A.Array a))
370fold f s = fmap fromStreamD <$> foldD f (toStreamD s)
372-------------------------------------------------------------------------------
373-- Fold to a single Array
374-------------------------------------------------------------------------------
376-- When we have to take an array partially, take the last part of the array.
377{-# INLINE takeArrayListRev #-}
378takeArrayListRev :: forall a. Storable a => Int -> [Array a] -> [Array a]
379takeArrayListRev = go
381 where
383 go _ [] = []
384 go n _ | n <= 0 = []
385 go n (x:xs) =
386 let len = Array.length x
387 in if n > len
388 then x : go (n - len) xs
389 else if n == len
390 then [x]
391 else let !(Array contents _ end) = x
392 !start = end `plusPtr` negate (n * SIZE_OF(a))
393 in [Array contents start end]
395-- When we have to take an array partially, take the last part of the array in
396-- the first split.
397{-# INLINE splitAtArrayListRev #-}
398splitAtArrayListRev :: forall a. Storable a =>
399 Int -> [Array a] -> ([Array a],[Array a])
400splitAtArrayListRev n ls
401 | n <= 0 = ([], ls)
402 | otherwise = go n ls
403 where
404 go :: Int -> [Array a] -> ([Array a], [Array a])
405 go _ [] = ([], [])
406 go m (x:xs) =
407 let len = Array.length x
408 (xs', xs'') = go (m - len) xs
409 in if m > len
410 then (x:xs', xs'')
411 else if m == len
412 then ([x],xs)
413 else let !(Array contents start end) = x
414 end1 = end `plusPtr` negate (m * SIZE_OF(a))
415 arr2 = Array contents start end1
416 arr1 = Array contents end1 end
417 in ([arr1], arr2:xs)
419-------------------------------------------------------------------------------
420-- Fold to a single Array
421-------------------------------------------------------------------------------
423-- XXX Both of these implementations of splicing seem to perform equally well.
424-- We need to perform benchmarks over a range of sizes though.
426-- CAUTION! length must more than equal to lengths of all the arrays in the
427-- stream.
428{-# INLINE spliceArraysLenUnsafe #-}
429spliceArraysLenUnsafe :: (MonadIO m, Storable a)
430 => Int -> SerialT m (MA.Array a) -> m (MA.Array a)
431spliceArraysLenUnsafe len buffered = do
432 arr <- liftIO $ MA.newArray len
433 end <- S.foldlM' writeArr (return $ MA.aEnd arr) buffered
434 return $ arr {MA.aEnd = end}
436 where
438 writeArr dst (MA.Array ac src ae _) =
439 liftIO $ do
440 let count = ae `minusPtr` src
441 memcpy (castPtr dst) (castPtr src) count
442 touch ac
443 return $ dst `plusPtr` count
445{-# INLINE _spliceArrays #-}
446_spliceArrays :: (MonadIO m, Storable a)
447 => SerialT m (Array a) -> m (Array a)
448_spliceArrays s = do
449 buffered <- S.foldr S.cons S.nil s
450 len <- S.sum (S.map Array.length buffered)
451 arr <- liftIO $ MA.newArray len
452 end <- S.foldlM' writeArr (return $ MA.aEnd arr) s
453 return $ A.unsafeFreeze $ arr {MA.aEnd = end}
455 where
457 writeArr dst (Array ac src ae) =
458 liftIO $ do
459 let count = ae `minusPtr` src
460 memcpy (castPtr dst) (castPtr src) count
461 touch ac
462 return $ dst `plusPtr` count
464{-# INLINE _spliceArraysBuffered #-}
465_spliceArraysBuffered :: (MonadIO m, Storable a)
466 => SerialT m (Array a) -> m (Array a)
467_spliceArraysBuffered s = do
468 buffered <- S.foldr S.cons S.nil s
469 len <- S.sum (S.map Array.length buffered)
470 A.unsafeFreeze <$> spliceArraysLenUnsafe len (S.map A.unsafeThaw s)
472{-# INLINE spliceArraysRealloced #-}
473spliceArraysRealloced :: forall m a. (MonadIO m, Storable a)
474 => SerialT m (Array a) -> m (Array a)
475spliceArraysRealloced s = do
476 let n = allocBytesToElemCount (undefined :: a) (4 * 1024)
477 idst = liftIO $ MA.newArray n
479 arr <- S.foldlM' MA.spliceExp idst (S.map A.unsafeThaw s)
480 liftIO $ A.unsafeFreeze <$> MA.rightSize arr
482-- XXX This should just be "fold A.write"
483--
484-- | Given a stream of arrays, splice them all together to generate a single
485-- array. The stream must be /finite/.
486--
487-- @since 0.7.0
488{-# INLINE toArray #-}
489toArray :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)
490toArray = spliceArraysRealloced
491-- spliceArrays = _spliceArraysBuffered
493-- exponentially increasing sizes of the chunks upto the max limit.
494-- XXX this will be easier to implement with parsers/terminating folds
495-- With this we should be able to reduce the number of chunks/allocations.
496-- The reallocation/copy based toArray can also be implemented using this.
497--
498{-
499{-# INLINE toArraysInRange #-}
500toArraysInRange :: (IsStream t, MonadIO m, Storable a)
501 => Int -> Int -> Fold m (Array a) b -> Fold m a b
502toArraysInRange low high (Fold step initial extract) =
503-}
505{-
506-- | Fold the input to a pure buffered stream (List) of arrays.
507{-# INLINE _toArraysOf #-}
508_toArraysOf :: (MonadIO m, Storable a)
509 => Int -> Fold m a (SerialT Identity (Array a))
510_toArraysOf n = FL.chunksOf n (A.writeNF n) FL.toStream
511-}
513-------------------------------------------------------------------------------
514-- Elimination - running element parsers
515-------------------------------------------------------------------------------
517-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
518-- to make a newtype.
519{-# ANN type List NoSpecConstr #-}
520newtype List a = List {getList :: [a]}
522-- This can be generalized to any type provided it can be unfolded to a stream
523-- and it can be combined using a semigroup operation.
524--
525-- XXX This should be written using CPS (as parseK) if we want it to scale wrt
526-- to the number of times it can be called on the same stream.
527{-# INLINE_NORMAL parseD #-}
528parseD ::
529 forall m a b. (MonadIO m, MonadThrow m, Storable a)
530 => PRD.Parser m a b
531 -> D.Stream m (Array.Array a)
532 -> m (b, D.Stream m (Array.Array a))
533parseD (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
534 res <- initial
535 case res of
536 PRD.IPartial s -> go SPEC state (List []) s
537 PRD.IDone b -> return (b, stream)
538 PRD.IError err -> throwM $ ParseError err
540 where
542 -- "backBuf" contains last few items in the stream that we may have to
543 -- backtrack to.
544 --
545 -- XXX currently we are using a dumb list based approach for backtracking
546 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
547 -- That will allow us more efficient random back and forth movement.
548 go !_ st backBuf !pst = do
549 r <- step defState st
550 case r of
551 D.Yield (Array contents start (Ptr end)) s ->
552 gobuf SPEC s backBuf
553 (ForeignPtr end (arrayToFptrContents contents)) start pst
554 D.Skip s -> go SPEC s backBuf pst
555 D.Stop -> do
556 b <- extract pst
557 return (b, D.nil)
559 -- Use strictness on "cur" to keep it unboxed
560 gobuf !_ s backBuf fp@(ForeignPtr end _) !cur !pst
561 | cur == Ptr end = do
562 liftIO $ touchForeignPtr fp
563 go SPEC s backBuf pst
564 gobuf !_ s backBuf fp@(ForeignPtr end contents) !cur !pst = do
565 x <- liftIO $ peek cur
566 pRes <- pstep pst x
567 let next = PTR_NEXT(cur,a)
568 case pRes of
569 PR.Partial 0 pst1 ->
570 gobuf SPEC s (List []) fp next pst1
571 PR.Partial n pst1 -> do
572 assert (n <= Prelude.length (x:getList backBuf)) (return ())
573 let src0 = Prelude.take n (x:getList backBuf)
574 arr0 = A.fromListN n (Prelude.reverse src0)
575 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
576 src = arr0 <> arr1
577 let !(Array cont1 start (Ptr end1)) = src
578 fp1 = ForeignPtr end1 (arrayToFptrContents cont1)
579 gobuf SPEC s (List []) fp1 start pst1
580 PR.Continue 0 pst1 ->
581 gobuf SPEC s (List (x:getList backBuf)) fp next pst1
582 PR.Continue n pst1 -> do
583 assert (n <= Prelude.length (x:getList backBuf)) (return ())
584 let (src0, buf1) = splitAt n (x:getList backBuf)
585 arr0 = A.fromListN n (Prelude.reverse src0)
586 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
587 src = arr0 <> arr1
588 let !(Array cont1 start (Ptr end1)) = src
589 fp1 = ForeignPtr end1 (arrayToFptrContents cont1)
590 gobuf SPEC s (List buf1) fp1 start pst1
591 PR.Done 0 b -> do
592 let arr = Array (fptrToArrayContents contents) next (Ptr end)
593 return (b, D.cons arr (D.Stream step s))
594 PR.Done n b -> do
595 assert (n <= Prelude.length (x:getList backBuf)) (return ())
596 let src0 = Prelude.take n (x:getList backBuf)
597 -- XXX create the array in reverse instead
598 arr0 = A.fromListN n (Prelude.reverse src0)
599 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
600 -- XXX Use StreamK to avoid adding arbitrary layers of
601 -- constructors every time.
602 str = D.cons arr0 (D.cons arr1 (D.Stream step s))
603 return (b, str)
604 PR.Error err -> throwM $ ParseError err
606-- | Parse an array stream using the supplied 'Parser'. Returns the parse
607-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
608--
609-- /Internal/
610--
611{-# INLINE_NORMAL parse #-}
612parse ::
613 (MonadIO m, MonadThrow m, Storable a)
614 => PRD.Parser m a b
615 -> SerialT m (A.Array a)
616 -> m (b, SerialT m (A.Array a))
617parse p s = fmap fromStreamD <$> parseD p (toStreamD s)
619-------------------------------------------------------------------------------
620-- Elimination - Running Array Folds and parsers
621-------------------------------------------------------------------------------
623{-# INLINE_NORMAL parseArrD #-}
624parseArrD ::
625 forall m a b. (MonadIO m, MonadThrow m, Storable a)
626 => PRD.Parser m (Array a) b
627 -> D.Stream m (Array.Array a)
628 -> m (b, D.Stream m (Array.Array a))
629parseArrD (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
630 res <- initial
631 case res of
632 PRD.IPartial s -> go SPEC state (List []) s
633 PRD.IDone b -> return (b, stream)
634 PRD.IError err -> throwM $ ParseError err
636 where
638 -- "backBuf" contains last few items in the stream that we may have to
639 -- backtrack to.
640 --
641 -- XXX currently we are using a dumb list based approach for backtracking
642 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
643 -- That will allow us more efficient random back and forth movement.
644 go !_ st backBuf !pst = do
645 r <- step defState st
646 case r of
647 D.Yield x s -> gobuf SPEC [x] s backBuf pst
648 D.Skip s -> go SPEC s backBuf pst
649 D.Stop -> do
650 b <- extract pst
651 return (b, D.nil)
653 gobuf !_ [] s backBuf !pst = go SPEC s backBuf pst
654 gobuf !_ (x:xs) s backBuf !pst = do
655 pRes <- pstep pst x
656 case pRes of
657 PR.Partial 0 pst1 ->
658 gobuf SPEC xs s (List []) pst1
659 PR.Partial n pst1 -> do
660 assert
661 (n <= sum (map Array.length (x:getList backBuf)))
662 (return ())
663 let src0 = takeArrayListRev n (x:getList backBuf)
664 src = Prelude.reverse src0 ++ xs
665 gobuf SPEC src s (List []) pst1
666 PR.Continue 0 pst1 ->
667 gobuf SPEC xs s (List (x:getList backBuf)) pst1
668 PR.Continue n pst1 -> do
669 assert
670 (n <= sum (map Array.length (x:getList backBuf)))
671 (return ())
672 let (src0, buf1) = splitAtArrayListRev n (x:getList backBuf)
673 src = Prelude.reverse src0 ++ xs
674 gobuf SPEC src s (List buf1) pst1
675 PR.Done 0 b ->
676 return (b, D.Stream step s)
677 PR.Done n b -> do
678 assert
679 (n <= sum (map Array.length (x:getList backBuf)))
680 (return ())
681 let src0 = takeArrayListRev n (x:getList backBuf)
682 src = Prelude.reverse src0 ++ xs
683 return (b, D.append (D.fromList src) (D.Stream step s))
684 PR.Error err -> throwM $ ParseError err
686{-
687-- | Parse an array stream using the supplied 'Parser'. Returns the parse
688-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
689--
690-- /Internal/
691--
692{-# INLINE parseArr #-}
693parseArr ::
694 (MonadIO m, MonadThrow m, Storable a)
695 => ASF.Parser m a b
696 -> SerialT m (A.Array a)
697 -> m (b, SerialT m (A.Array a))
698parseArr p s = fmap fromStreamD <$> parseD p (toStreamD s)
699-}
701-- | Fold an array stream using the supplied array stream 'Fold'.
702--
703-- /Pre-release/
704--
705{-# INLINE foldArr #-}
706foldArr :: (MonadIO m, MonadThrow m, Storable a) =>
707 ASF.Fold m a b -> SerialT m (A.Array a) -> m b
708foldArr (ASF.Fold p) s = fst <$> parseArrD p (toStreamD s)
710-- | Like 'fold' but also returns the remaining stream.
711--
712-- /Pre-release/
713--
714{-# INLINE foldArr_ #-}
715foldArr_ :: (MonadIO m, MonadThrow m, Storable a) =>
716 ASF.Fold m a b -> SerialT m (A.Array a) -> m (b, SerialT m (A.Array a))
717foldArr_ (ASF.Fold p) s = second fromStreamD <$> parseArrD p (toStreamD s)
719{-# ANN type ParseChunksState Fuse #-}
720data ParseChunksState x inpBuf st pst =
721 ParseChunksInit inpBuf st
722 | ParseChunksInitLeftOver inpBuf
723 | ParseChunksStream st inpBuf !pst
724 | ParseChunksBuf inpBuf st inpBuf !pst
725 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
727{-# INLINE_NORMAL foldArrManyD #-}
728foldArrManyD
729 :: (MonadThrow m, Storable a)
730 => ASF.Fold m a b
731 -> D.Stream m (Array a)
732 -> D.Stream m b
733foldArrManyD (ASF.Fold (PRD.Parser pstep initial extract)) (D.Stream step state) =
734 D.Stream stepOuter (ParseChunksInit [] state)
736 where
738 {-# INLINE_LATE stepOuter #-}
739 -- Buffer is empty, get the first element from the stream, initialize the
740 -- fold and then go to stream processing loop.
741 stepOuter gst (ParseChunksInit [] st) = do
742 r <- step (adaptState gst) st
743 case r of
744 D.Yield x s -> do
745 res <- initial
746 case res of
747 PRD.IPartial ps ->
748 return $ D.Skip $ ParseChunksBuf [x] s [] ps
749 PRD.IDone pb ->
750 let next = ParseChunksInit [x] s
751 in return $ D.Skip $ ParseChunksYield pb next
752 PRD.IError err -> throwM $ ParseError err
753 D.Skip s -> return $ D.Skip $ ParseChunksInit [] s
754 D.Stop -> return D.Stop
756 -- Buffer is not empty, go to buffered processing loop
757 stepOuter _ (ParseChunksInit src st) = do
758 res <- initial
759 case res of
760 PRD.IPartial ps ->
761 return $ D.Skip $ ParseChunksBuf src st [] ps
762 PRD.IDone pb ->
763 let next = ParseChunksInit src st
764 in return $ D.Skip $ ParseChunksYield pb next
765 PRD.IError err -> throwM $ ParseError err
767 -- XXX we just discard any leftover input at the end
768 stepOuter _ (ParseChunksInitLeftOver _) = return D.Stop
770 -- Buffer is empty, process elements from the stream
771 stepOuter gst (ParseChunksStream st backBuf pst) = do
772 r <- step (adaptState gst) st
773 case r of
774 D.Yield x s -> do
775 pRes <- pstep pst x
776 case pRes of
777 PR.Partial 0 pst1 ->
778 return $ D.Skip $ ParseChunksStream s [] pst1
779 PR.Partial n pst1 -> do
780 assert
781 (n <= sum (map Array.length (x:backBuf)))
782 (return ())
783 let src0 = takeArrayListRev n (x:backBuf)
784 src = Prelude.reverse src0
785 return $ D.Skip $ ParseChunksBuf src s [] pst1
786 PR.Continue 0 pst1 ->
787 return $ D.Skip $ ParseChunksStream s (x:backBuf) pst1
788 PR.Continue n pst1 -> do
789 assert
790 (n <= sum (map Array.length (x:backBuf)))
791 (return ())
792 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
793 src = Prelude.reverse src0
794 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
795 PR.Done 0 b -> do
796 return $ D.Skip $
797 ParseChunksYield b (ParseChunksInit [] s)
798 PR.Done n b -> do
799 assert
800 (n <= sum (map Array.length (x:backBuf)))
801 (return ())
802 let src0 = takeArrayListRev n (x:backBuf)
803 let src = Prelude.reverse src0
804 return $ D.Skip $
805 ParseChunksYield b (ParseChunksInit src s)
806 PR.Error err -> throwM $ ParseError err
807 D.Skip s -> return $ D.Skip $ ParseChunksStream s backBuf pst
808 D.Stop -> do
809 b <- extract pst
810 let src = Prelude.reverse backBuf
811 return $ D.Skip $
812 ParseChunksYield b (ParseChunksInitLeftOver src)
814 -- go back to stream processing mode
815 stepOuter _ (ParseChunksBuf [] s buf pst) =
816 return $ D.Skip $ ParseChunksStream s buf pst
818 -- buffered processing loop
819 stepOuter _ (ParseChunksBuf (x:xs) s backBuf pst) = do
820 pRes <- pstep pst x
821 case pRes of
822 PR.Partial 0 pst1 ->
823 return $ D.Skip $ ParseChunksBuf xs s [] pst1
824 PR.Partial n pst1 -> do
825 assert (n <= sum (map Array.length (x:backBuf))) (return ())
826 let src0 = takeArrayListRev n (x:backBuf)
827 src = Prelude.reverse src0 ++ xs
828 return $ D.Skip $ ParseChunksBuf src s [] pst1
829 PR.Continue 0 pst1 ->
830 return $ D.Skip $ ParseChunksBuf xs s (x:backBuf) pst1
831 PR.Continue n pst1 -> do
832 assert (n <= sum (map Array.length (x:backBuf))) (return ())
833 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
834 src = Prelude.reverse src0 ++ xs
835 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
836 PR.Done 0 b ->
837 return $ D.Skip $ ParseChunksYield b (ParseChunksInit xs s)
838 PR.Done n b -> do
839 assert (n <= sum (map Array.length (x:backBuf))) (return ())
840 let src0 = takeArrayListRev n (x:backBuf)
841 src = Prelude.reverse src0 ++ xs
842 return $ D.Skip $ ParseChunksYield b (ParseChunksInit src s)
843 PR.Error err -> throwM $ ParseError err
845 stepOuter _ (ParseChunksYield a next) = return $ D.Yield a next
847-- | Apply an array stream 'Fold' repeatedly on an array stream and emit the
848-- fold outputs in the output stream.
849--
850-- See "Streamly.Prelude.foldMany" for more details.
851--
852-- /Pre-release/
853{-# INLINE foldArrMany #-}
854foldArrMany
855 :: (IsStream t, MonadThrow m, Storable a)
856 => ASF.Fold m a b
857 -> t m (Array a)
858 -> t m b
859foldArrMany p m = fromStreamD $ foldArrManyD p (toStreamD m)