git.iko.soy/streamly
on 46314a2 help sign in
streamly/src/Streamly/Internal/Data/Array/Stream/Foreign.hs
861 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 {-# INLINE go #-}
549 go !_ st backBuf !pst = do
550 r <- step defState st
551 case r of
552 D.Yield (Array contents start (Ptr end)) s ->
553 gobuf SPEC s backBuf
554 (ForeignPtr end (arrayToFptrContents contents)) start pst
555 D.Skip s -> go SPEC s backBuf pst
556 D.Stop -> do
557 b <- extract pst
558 return (b, D.nil)
560 -- Use strictness on "cur" to keep it unboxed
561 gobuf !_ s backBuf fp@(ForeignPtr end _) !cur !pst
562 | cur == Ptr end = do
563 liftIO $ touchForeignPtr fp
564 go SPEC s backBuf pst
565 gobuf !_ s backBuf fp@(ForeignPtr end contents) !cur !pst = do
566 x <- liftIO $ peek cur
567 pRes <- pstep pst x
568 let next = PTR_NEXT(cur,a)
569 case pRes of
570 PR.Partial 0 pst1 ->
571 gobuf SPEC s (List []) fp next pst1
572 PR.Partial n pst1 -> do
573 assert (n <= Prelude.length (x:getList backBuf)) (return ())
574 let src0 = Prelude.take n (x:getList backBuf)
575 arr0 = A.fromListN n (Prelude.reverse src0)
576 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
577 src = arr0 <> arr1
578 let !(Array cont1 start (Ptr end1)) = src
579 fp1 = ForeignPtr end1 (arrayToFptrContents cont1)
580 gobuf SPEC s (List []) fp1 start pst1
581 PR.Continue 0 pst1 ->
582 gobuf SPEC s (List (x:getList backBuf)) fp next pst1
583 PR.Continue n pst1 -> do
584 assert (n <= Prelude.length (x:getList backBuf)) (return ())
585 let (src0, buf1) = splitAt n (x:getList backBuf)
586 arr0 = A.fromListN n (Prelude.reverse src0)
587 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
588 src = arr0 <> arr1
589 let !(Array cont1 start (Ptr end1)) = src
590 fp1 = ForeignPtr end1 (arrayToFptrContents cont1)
591 gobuf SPEC s (List buf1) fp1 start pst1
592 PR.Done 0 b -> do
593 let arr = Array (fptrToArrayContents contents) next (Ptr end)
594 return (b, D.cons arr (D.Stream step s))
595 PR.Done n b -> do
596 assert (n <= Prelude.length (x:getList backBuf)) (return ())
597 let src0 = Prelude.take n (x:getList backBuf)
598 -- XXX create the array in reverse instead
599 arr0 = A.fromListN n (Prelude.reverse src0)
600 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
601 -- XXX Use StreamK to avoid adding arbitrary layers of
602 -- constructors every time.
603 str = D.cons arr0 (D.cons arr1 (D.Stream step s))
604 return (b, str)
605 PR.Error err -> throwM $ ParseError err
607-- | Parse an array stream using the supplied 'Parser'. Returns the parse
608-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
609--
610-- /Internal/
611--
612{-# INLINE_NORMAL parse #-}
613parse ::
614 (MonadIO m, MonadThrow m, Storable a)
615 => PRD.Parser m a b
616 -> SerialT m (A.Array a)
617 -> m (b, SerialT m (A.Array a))
618parse p s = fmap fromStreamD <$> parseD p (toStreamD s)
620-------------------------------------------------------------------------------
621-- Elimination - Running Array Folds and parsers
622-------------------------------------------------------------------------------
624{-# INLINE_NORMAL parseArrD #-}
625parseArrD ::
626 forall m a b. (MonadIO m, MonadThrow m, Storable a)
627 => PRD.Parser m (Array a) b
628 -> D.Stream m (Array.Array a)
629 -> m (b, D.Stream m (Array.Array a))
630parseArrD (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
631 res <- initial
632 case res of
633 PRD.IPartial s -> go SPEC state (List []) s
634 PRD.IDone b -> return (b, stream)
635 PRD.IError err -> throwM $ ParseError err
637 where
639 -- "backBuf" contains last few items in the stream that we may have to
640 -- backtrack to.
641 --
642 -- XXX currently we are using a dumb list based approach for backtracking
643 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
644 -- That will allow us more efficient random back and forth movement.
645 {-# INLINE go #-}
646 go !_ st backBuf !pst = do
647 r <- step defState st
648 case r of
649 D.Yield x s -> gobuf SPEC [x] s backBuf pst
650 D.Skip s -> go SPEC s backBuf pst
651 D.Stop -> do
652 b <- extract pst
653 return (b, D.nil)
655 gobuf !_ [] s backBuf !pst = go SPEC s backBuf pst
656 gobuf !_ (x:xs) s backBuf !pst = do
657 pRes <- pstep pst x
658 case pRes of
659 PR.Partial 0 pst1 ->
660 gobuf SPEC xs s (List []) pst1
661 PR.Partial n pst1 -> do
662 assert
663 (n <= sum (map Array.length (x:getList backBuf)))
664 (return ())
665 let src0 = takeArrayListRev n (x:getList backBuf)
666 src = Prelude.reverse src0 ++ xs
667 gobuf SPEC src s (List []) pst1
668 PR.Continue 0 pst1 ->
669 gobuf SPEC xs s (List (x:getList backBuf)) pst1
670 PR.Continue n pst1 -> do
671 assert
672 (n <= sum (map Array.length (x:getList backBuf)))
673 (return ())
674 let (src0, buf1) = splitAtArrayListRev n (x:getList backBuf)
675 src = Prelude.reverse src0 ++ xs
676 gobuf SPEC src s (List buf1) pst1
677 PR.Done 0 b ->
678 return (b, D.Stream step s)
679 PR.Done n b -> 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 return (b, D.append (D.fromList src) (D.Stream step s))
686 PR.Error err -> throwM $ ParseError err
688{-
689-- | Parse an array stream using the supplied 'Parser'. Returns the parse
690-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
691--
692-- /Internal/
693--
694{-# INLINE parseArr #-}
695parseArr ::
696 (MonadIO m, MonadThrow m, Storable a)
697 => ASF.Parser m a b
698 -> SerialT m (A.Array a)
699 -> m (b, SerialT m (A.Array a))
700parseArr p s = fmap fromStreamD <$> parseD p (toStreamD s)
701-}
703-- | Fold an array stream using the supplied array stream 'Fold'.
704--
705-- /Pre-release/
706--
707{-# INLINE foldArr #-}
708foldArr :: (MonadIO m, MonadThrow m, Storable a) =>
709 ASF.Fold m a b -> SerialT m (A.Array a) -> m b
710foldArr (ASF.Fold p) s = fst <$> parseArrD p (toStreamD s)
712-- | Like 'fold' but also returns the remaining stream.
713--
714-- /Pre-release/
715--
716{-# INLINE foldArr_ #-}
717foldArr_ :: (MonadIO m, MonadThrow m, Storable a) =>
718 ASF.Fold m a b -> SerialT m (A.Array a) -> m (b, SerialT m (A.Array a))
719foldArr_ (ASF.Fold p) s = second fromStreamD <$> parseArrD p (toStreamD s)
721{-# ANN type ParseChunksState Fuse #-}
722data ParseChunksState x inpBuf st pst =
723 ParseChunksInit inpBuf st
724 | ParseChunksInitLeftOver inpBuf
725 | ParseChunksStream st inpBuf !pst
726 | ParseChunksBuf inpBuf st inpBuf !pst
727 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
729{-# INLINE_NORMAL foldArrManyD #-}
730foldArrManyD
731 :: (MonadThrow m, Storable a)
732 => ASF.Fold m a b
733 -> D.Stream m (Array a)
734 -> D.Stream m b
735foldArrManyD (ASF.Fold (PRD.Parser pstep initial extract)) (D.Stream step state) =
736 D.Stream stepOuter (ParseChunksInit [] state)
738 where
740 {-# INLINE_LATE stepOuter #-}
741 -- Buffer is empty, get the first element from the stream, initialize the
742 -- fold and then go to stream processing loop.
743 stepOuter gst (ParseChunksInit [] st) = do
744 r <- step (adaptState gst) st
745 case r of
746 D.Yield x s -> do
747 res <- initial
748 case res of
749 PRD.IPartial ps ->
750 return $ D.Skip $ ParseChunksBuf [x] s [] ps
751 PRD.IDone pb ->
752 let next = ParseChunksInit [x] s
753 in return $ D.Skip $ ParseChunksYield pb next
754 PRD.IError err -> throwM $ ParseError err
755 D.Skip s -> return $ D.Skip $ ParseChunksInit [] s
756 D.Stop -> return D.Stop
758 -- Buffer is not empty, go to buffered processing loop
759 stepOuter _ (ParseChunksInit src st) = do
760 res <- initial
761 case res of
762 PRD.IPartial ps ->
763 return $ D.Skip $ ParseChunksBuf src st [] ps
764 PRD.IDone pb ->
765 let next = ParseChunksInit src st
766 in return $ D.Skip $ ParseChunksYield pb next
767 PRD.IError err -> throwM $ ParseError err
769 -- XXX we just discard any leftover input at the end
770 stepOuter _ (ParseChunksInitLeftOver _) = return D.Stop
772 -- Buffer is empty, process elements from the stream
773 stepOuter gst (ParseChunksStream st backBuf pst) = do
774 r <- step (adaptState gst) st
775 case r of
776 D.Yield x s -> do
777 pRes <- pstep pst x
778 case pRes of
779 PR.Partial 0 pst1 ->
780 return $ D.Skip $ ParseChunksStream s [] pst1
781 PR.Partial n pst1 -> do
782 assert
783 (n <= sum (map Array.length (x:backBuf)))
784 (return ())
785 let src0 = takeArrayListRev n (x:backBuf)
786 src = Prelude.reverse src0
787 return $ D.Skip $ ParseChunksBuf src s [] pst1
788 PR.Continue 0 pst1 ->
789 return $ D.Skip $ ParseChunksStream s (x:backBuf) pst1
790 PR.Continue n pst1 -> do
791 assert
792 (n <= sum (map Array.length (x:backBuf)))
793 (return ())
794 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
795 src = Prelude.reverse src0
796 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
797 PR.Done 0 b -> do
798 return $ D.Skip $
799 ParseChunksYield b (ParseChunksInit [] s)
800 PR.Done n b -> do
801 assert
802 (n <= sum (map Array.length (x:backBuf)))
803 (return ())
804 let src0 = takeArrayListRev n (x:backBuf)
805 let src = Prelude.reverse src0
806 return $ D.Skip $
807 ParseChunksYield b (ParseChunksInit src s)
808 PR.Error err -> throwM $ ParseError err
809 D.Skip s -> return $ D.Skip $ ParseChunksStream s backBuf pst
810 D.Stop -> do
811 b <- extract pst
812 let src = Prelude.reverse backBuf
813 return $ D.Skip $
814 ParseChunksYield b (ParseChunksInitLeftOver src)
816 -- go back to stream processing mode
817 stepOuter _ (ParseChunksBuf [] s buf pst) =
818 return $ D.Skip $ ParseChunksStream s buf pst
820 -- buffered processing loop
821 stepOuter _ (ParseChunksBuf (x:xs) s backBuf pst) = do
822 pRes <- pstep pst x
823 case pRes of
824 PR.Partial 0 pst1 ->
825 return $ D.Skip $ ParseChunksBuf xs s [] pst1
826 PR.Partial n pst1 -> do
827 assert (n <= sum (map Array.length (x:backBuf))) (return ())
828 let src0 = takeArrayListRev n (x:backBuf)
829 src = Prelude.reverse src0 ++ xs
830 return $ D.Skip $ ParseChunksBuf src s [] pst1
831 PR.Continue 0 pst1 ->
832 return $ D.Skip $ ParseChunksBuf xs s (x:backBuf) pst1
833 PR.Continue n pst1 -> do
834 assert (n <= sum (map Array.length (x:backBuf))) (return ())
835 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
836 src = Prelude.reverse src0 ++ xs
837 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
838 PR.Done 0 b ->
839 return $ D.Skip $ ParseChunksYield b (ParseChunksInit xs s)
840 PR.Done n b -> do
841 assert (n <= sum (map Array.length (x:backBuf))) (return ())
842 let src0 = takeArrayListRev n (x:backBuf)
843 src = Prelude.reverse src0 ++ xs
844 return $ D.Skip $ ParseChunksYield b (ParseChunksInit src s)
845 PR.Error err -> throwM $ ParseError err
847 stepOuter _ (ParseChunksYield a next) = return $ D.Yield a next
849-- | Apply an array stream 'Fold' repeatedly on an array stream and emit the
850-- fold outputs in the output stream.
851--
852-- See "Streamly.Prelude.foldMany" for more details.
853--
854-- /Pre-release/
855{-# INLINE foldArrMany #-}
856foldArrMany
857 :: (IsStream t, MonadThrow m, Storable a)
858 => ASF.Fold m a b
859 -> t m (Array a)
860 -> t m b
861foldArrMany p m = fromStreamD $ foldArrManyD p (toStreamD m)