git.iko.soy/streamly
on 08f82ab help sign in
streamly/src/Streamly/Internal/Data/Array/Stream/Foreign.hs
865 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 "inline.hs"
51import Data.Bifunctor (second)
52import Control.Exception (assert)
53import Control.Monad.Catch (MonadThrow, throwM)
54import Control.Monad.IO.Class (MonadIO(..))
55#if __GLASGOW_HASKELL__ < 808
56import Data.Semigroup (Semigroup(..))
57#endif
58import Data.Word (Word8)
59import Foreign.ForeignPtr (touchForeignPtr)
60import Foreign.Ptr (minusPtr, plusPtr, castPtr)
61import Foreign.Storable (Storable(..))
62import Fusion.Plugin.Types (Fuse(..))
63import GHC.Exts (SpecConstrAnnotation(..))
64import GHC.ForeignPtr (ForeignPtr(..))
65import GHC.Ptr (Ptr(..))
66import GHC.Types (SPEC(..))
67import Prelude hiding (null, last, (!!), read, concat, unlines)
69import Streamly.Internal.Data.Array.Foreign.Mut.Type
70 (arrayToFptrContents, fptrToArrayContents, touch)
71import Streamly.Internal.Data.Array.Foreign.Type (Array(..))
72import Streamly.Internal.Data.Fold.Type (Fold(..))
73import Streamly.Internal.Data.Parser (ParseError(..))
74import Streamly.Internal.Data.Stream.Serial (SerialT)
75import Streamly.Internal.Data.Stream.IsStream.Type
76 (IsStream, fromStreamD, toStreamD)
77import Streamly.Internal.Data.SVar (adaptState, defState)
78import Streamly.Internal.Data.Array.Foreign.Mut.Type
79 (memcpy, allocBytesToElemCount)
81import qualified Streamly.Internal.Data.Array.Foreign as A
82import qualified Streamly.Internal.Data.Array.Foreign as Array
83import qualified Streamly.Internal.Data.Array.Foreign.Type as A
84import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MA
85import qualified Streamly.Internal.Data.Array.Stream.Mut.Foreign as AS
86import qualified Streamly.Internal.Data.Array.Stream.Fold.Foreign as ASF
87import qualified Streamly.Internal.Data.Fold.Type as FL
88import qualified Streamly.Internal.Data.Parser as PR
89import qualified Streamly.Internal.Data.Parser.ParserD as PRD
90import qualified Streamly.Internal.Data.Stream.IsStream as S
91import qualified Streamly.Internal.Data.Stream.StreamD as D
93-- XXX Since these are immutable arrays MonadIO constraint can be removed from
94-- most places.
96-------------------------------------------------------------------------------
97-- Generation
98-------------------------------------------------------------------------------
100-- | @arraysOf n stream@ groups the elements in the input stream into arrays of
101-- @n@ elements each.
102--
103-- > arraysOf n = Stream.chunksOf n (Array.writeN n)
104--
105-- /Pre-release/
106{-# INLINE arraysOf #-}
107arraysOf :: (IsStream t, MonadIO m, Storable a)
108 => Int -> t m a -> t m (Array a)
109arraysOf n str = fromStreamD $ A.arraysOf n (toStreamD str)
111-------------------------------------------------------------------------------
112-- Append
113-------------------------------------------------------------------------------
115-- XXX efficiently compare two streams of arrays. Two streams can have chunks
116-- of different sizes, we can handle that in the stream comparison abstraction.
117-- This could be useful e.g. to fast compare whether two files differ.
119-- | Convert a stream of arrays into a stream of their elements.
120--
121-- Same as the following:
122--
123-- > concat = Stream.unfoldMany Array.read
124--
125-- @since 0.7.0
126{-# INLINE concat #-}
127concat :: (IsStream t, Monad m, Storable a) => t m (Array a) -> t m a
128-- concat m = fromStreamD $ A.flattenArrays (toStreamD m)
129-- concat m = fromStreamD $ D.concatMap A.toStreamD (toStreamD m)
130concat m = fromStreamD $ D.unfoldMany A.read (toStreamD m)
132-- | Convert a stream of arrays into a stream of their elements reversing the
133-- contents of each array before flattening.
134--
135-- > concatRev = Stream.unfoldMany Array.readRev
136--
137-- @since 0.7.0
138{-# INLINE concatRev #-}
139concatRev :: (IsStream t, Monad m, Storable a) => t m (Array a) -> t m a
140-- concatRev m = fromStreamD $ A.flattenArraysRev (toStreamD m)
141concatRev m = fromStreamD $ D.unfoldMany A.readRev (toStreamD m)
143-------------------------------------------------------------------------------
144-- Intersperse and append
145-------------------------------------------------------------------------------
147-- | Flatten a stream of arrays after inserting the given element between
148-- arrays.
149--
150-- /Pre-release/
151{-# INLINE interpose #-}
152interpose :: (Monad m, IsStream t, Storable a) => a -> t m (Array a) -> t m a
153interpose x = S.interpose x A.read
155{-# INLINE intercalateSuffix #-}
156intercalateSuffix :: (Monad m, IsStream t, Storable a)
157 => Array a -> t m (Array a) -> t m a
158intercalateSuffix = S.intercalateSuffix A.read
160-- | Flatten a stream of arrays appending the given element after each
161-- array.
162--
163-- @since 0.7.0
164{-# INLINE interposeSuffix #-}
165interposeSuffix :: (Monad m, IsStream t, Storable a)
166 => a -> t m (Array a) -> t m a
167-- interposeSuffix x = fromStreamD . A.unlines x . toStreamD
168interposeSuffix x = S.interposeSuffix x A.read
170data FlattenState s a =
171 OuterLoop s
172 | InnerLoop s !MA.ArrayContents !(Ptr a) !(Ptr a)
174-- XXX Remove monadIO constraint
175{-# INLINE_NORMAL unlines #-}
176unlines :: forall m a. (MonadIO m, Storable a)
177 => a -> D.Stream m (Array a) -> D.Stream m a
178unlines sep (D.Stream step state) = D.Stream step' (OuterLoop state)
179 where
180 {-# INLINE_LATE step' #-}
181 step' gst (OuterLoop st) = do
182 r <- step (adaptState gst) st
183 return $ case r of
184 D.Yield Array{..} s ->
185 D.Skip (InnerLoop s arrContents arrStart aEnd)
186 D.Skip s -> D.Skip (OuterLoop s)
187 D.Stop -> D.Stop
189 step' _ (InnerLoop st _ p end) | p == end =
190 return $ D.Yield sep $ OuterLoop st
192 step' _ (InnerLoop st contents p end) = do
193 x <- liftIO $ do
194 r <- peek p
195 touch contents
196 return r
197 return $ D.Yield x (InnerLoop st contents
198 (p `plusPtr` sizeOf (undefined :: 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 elemSize = sizeOf (undefined :: a)
353 next = cur `plusPtr` elemSize
354 case res of
355 FL.Done b -> do
356 let arr = Array (fptrToArrayContents contents) next (Ptr end)
357 return $! (b, D.cons arr (D.Stream step st))
358 FL.Partial fs1 -> goArray SPEC st fp next fs1
360-- | Fold an array stream using the supplied 'Fold'. Returns the fold result
361-- and the unconsumed stream.
362--
363-- /Internal/
364--
365{-# INLINE_NORMAL fold #-}
366fold ::
367 (MonadIO m, Storable a)
368 => FL.Fold m a b
369 -> SerialT m (A.Array a)
370 -> m (b, SerialT m (A.Array a))
371fold f s = fmap fromStreamD <$> foldD f (toStreamD s)
373-------------------------------------------------------------------------------
374-- Fold to a single Array
375-------------------------------------------------------------------------------
377-- When we have to take an array partially, take the last part of the array.
378{-# INLINE takeArrayListRev #-}
379takeArrayListRev :: forall a. Storable a => Int -> [Array a] -> [Array a]
380takeArrayListRev = go
382 where
384 go _ [] = []
385 go n _ | n <= 0 = []
386 go n (x:xs) =
387 let len = Array.length x
388 in if n > len
389 then x : go (n - len) xs
390 else if n == len
391 then [x]
392 else let !(Array contents _ end) = x
393 sz = sizeOf (undefined :: a)
394 !start = end `plusPtr` negate (n * sz)
395 in [Array contents start end]
397-- When we have to take an array partially, take the last part of the array in
398-- the first split.
399{-# INLINE splitAtArrayListRev #-}
400splitAtArrayListRev :: forall a. Storable a =>
401 Int -> [Array a] -> ([Array a],[Array a])
402splitAtArrayListRev n ls
403 | n <= 0 = ([], ls)
404 | otherwise = go n ls
405 where
406 go :: Int -> [Array a] -> ([Array a], [Array a])
407 go _ [] = ([], [])
408 go m (x:xs) =
409 let len = Array.length x
410 (xs', xs'') = go (m - len) xs
411 in if m > len
412 then (x:xs', xs'')
413 else if m == len
414 then ([x],xs)
415 else let !(Array contents start end) = x
416 sz = sizeOf (undefined :: a)
417 end1 = end `plusPtr` negate (m * sz)
418 arr2 = Array contents start end1
419 arr1 = Array contents end1 end
420 in ([arr1], arr2:xs)
422-------------------------------------------------------------------------------
423-- Fold to a single Array
424-------------------------------------------------------------------------------
426-- XXX Both of these implementations of splicing seem to perform equally well.
427-- We need to perform benchmarks over a range of sizes though.
429-- CAUTION! length must more than equal to lengths of all the arrays in the
430-- stream.
431{-# INLINE spliceArraysLenUnsafe #-}
432spliceArraysLenUnsafe :: (MonadIO m, Storable a)
433 => Int -> SerialT m (MA.Array a) -> m (MA.Array a)
434spliceArraysLenUnsafe len buffered = do
435 arr <- liftIO $ MA.newArray len
436 end <- S.foldlM' writeArr (return $ MA.aEnd arr) buffered
437 return $ arr {MA.aEnd = end}
439 where
441 writeArr dst (MA.Array ac src ae _) =
442 liftIO $ do
443 let count = ae `minusPtr` src
444 memcpy (castPtr dst) (castPtr src) count
445 touch ac
446 return $ dst `plusPtr` count
448{-# INLINE _spliceArrays #-}
449_spliceArrays :: (MonadIO m, Storable a)
450 => SerialT m (Array a) -> m (Array a)
451_spliceArrays s = do
452 buffered <- S.foldr S.cons S.nil s
453 len <- S.sum (S.map Array.length buffered)
454 arr <- liftIO $ MA.newArray len
455 end <- S.foldlM' writeArr (return $ MA.aEnd arr) s
456 return $ A.unsafeFreeze $ arr {MA.aEnd = end}
458 where
460 writeArr dst (Array ac src ae) =
461 liftIO $ do
462 let count = ae `minusPtr` src
463 memcpy (castPtr dst) (castPtr src) count
464 touch ac
465 return $ dst `plusPtr` count
467{-# INLINE _spliceArraysBuffered #-}
468_spliceArraysBuffered :: (MonadIO m, Storable a)
469 => SerialT m (Array a) -> m (Array a)
470_spliceArraysBuffered s = do
471 buffered <- S.foldr S.cons S.nil s
472 len <- S.sum (S.map Array.length buffered)
473 A.unsafeFreeze <$> spliceArraysLenUnsafe len (S.map A.unsafeThaw s)
475{-# INLINE spliceArraysRealloced #-}
476spliceArraysRealloced :: forall m a. (MonadIO m, Storable a)
477 => SerialT m (Array a) -> m (Array a)
478spliceArraysRealloced s = do
479 let n = allocBytesToElemCount (undefined :: a) (4 * 1024)
480 idst = liftIO $ MA.newArray n
482 arr <- S.foldlM' MA.spliceExp idst (S.map A.unsafeThaw s)
483 liftIO $ A.unsafeFreeze <$> MA.rightSize arr
485-- XXX This should just be "fold A.write"
486--
487-- | Given a stream of arrays, splice them all together to generate a single
488-- array. The stream must be /finite/.
489--
490-- @since 0.7.0
491{-# INLINE toArray #-}
492toArray :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)
493toArray = spliceArraysRealloced
494-- spliceArrays = _spliceArraysBuffered
496-- exponentially increasing sizes of the chunks upto the max limit.
497-- XXX this will be easier to implement with parsers/terminating folds
498-- With this we should be able to reduce the number of chunks/allocations.
499-- The reallocation/copy based toArray can also be implemented using this.
500--
501{-
502{-# INLINE toArraysInRange #-}
503toArraysInRange :: (IsStream t, MonadIO m, Storable a)
504 => Int -> Int -> Fold m (Array a) b -> Fold m a b
505toArraysInRange low high (Fold step initial extract) =
506-}
508{-
509-- | Fold the input to a pure buffered stream (List) of arrays.
510{-# INLINE _toArraysOf #-}
511_toArraysOf :: (MonadIO m, Storable a)
512 => Int -> Fold m a (SerialT Identity (Array a))
513_toArraysOf n = FL.chunksOf n (A.writeNF n) FL.toStream
514-}
516-------------------------------------------------------------------------------
517-- Elimination - running element parsers
518-------------------------------------------------------------------------------
520-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
521-- to make a newtype.
522{-# ANN type List NoSpecConstr #-}
523newtype List a = List {getList :: [a]}
525-- This can be generalized to any type provided it can be unfolded to a stream
526-- and it can be combined using a semigroup operation.
527--
528-- XXX This should be written using CPS (as parseK) if we want it to scale wrt
529-- to the number of times it can be called on the same stream.
530{-# INLINE_NORMAL parseD #-}
531parseD ::
532 forall m a b. (MonadIO m, MonadThrow m, Storable a)
533 => PRD.Parser m a b
534 -> D.Stream m (Array.Array a)
535 -> m (b, D.Stream m (Array.Array a))
536parseD (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
537 res <- initial
538 case res of
539 PRD.IPartial s -> go SPEC state (List []) s
540 PRD.IDone b -> return (b, stream)
541 PRD.IError err -> throwM $ ParseError err
543 where
545 -- "backBuf" contains last few items in the stream that we may have to
546 -- backtrack to.
547 --
548 -- XXX currently we are using a dumb list based approach for backtracking
549 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
550 -- That will allow us more efficient random back and forth movement.
551 {-# INLINE go #-}
552 go !_ st backBuf !pst = do
553 r <- step defState st
554 case r of
555 D.Yield (Array contents start (Ptr end)) s ->
556 gobuf SPEC s backBuf
557 (ForeignPtr end (arrayToFptrContents contents)) start pst
558 D.Skip s -> go SPEC s backBuf pst
559 D.Stop -> do
560 b <- extract pst
561 return (b, D.nil)
563 -- Use strictness on "cur" to keep it unboxed
564 gobuf !_ s backBuf fp@(ForeignPtr end _) !cur !pst
565 | cur == Ptr end = do
566 liftIO $ touchForeignPtr fp
567 go SPEC s backBuf pst
568 gobuf !_ s backBuf fp@(ForeignPtr end contents) !cur !pst = do
569 x <- liftIO $ peek cur
570 pRes <- pstep pst x
571 let elemSize = sizeOf (undefined :: a)
572 next = cur `plusPtr` elemSize
573 case pRes of
574 PR.Partial 0 pst1 ->
575 gobuf SPEC s (List []) fp next pst1
576 PR.Partial n pst1 -> do
577 assert (n <= Prelude.length (x:getList backBuf)) (return ())
578 let src0 = Prelude.take n (x:getList backBuf)
579 arr0 = A.fromListN n (Prelude.reverse src0)
580 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
581 src = arr0 <> arr1
582 let !(Array cont1 start (Ptr end1)) = src
583 fp1 = ForeignPtr end1 (arrayToFptrContents cont1)
584 gobuf SPEC s (List []) fp1 start pst1
585 PR.Continue 0 pst1 ->
586 gobuf SPEC s (List (x:getList backBuf)) fp next pst1
587 PR.Continue n pst1 -> do
588 assert (n <= Prelude.length (x:getList backBuf)) (return ())
589 let (src0, buf1) = splitAt n (x:getList backBuf)
590 arr0 = A.fromListN n (Prelude.reverse src0)
591 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
592 src = arr0 <> arr1
593 let !(Array cont1 start (Ptr end1)) = src
594 fp1 = ForeignPtr end1 (arrayToFptrContents cont1)
595 gobuf SPEC s (List buf1) fp1 start pst1
596 PR.Done 0 b -> do
597 let arr = Array (fptrToArrayContents contents) next (Ptr end)
598 return (b, D.cons arr (D.Stream step s))
599 PR.Done n b -> do
600 assert (n <= Prelude.length (x:getList backBuf)) (return ())
601 let src0 = Prelude.take n (x:getList backBuf)
602 -- XXX create the array in reverse instead
603 arr0 = A.fromListN n (Prelude.reverse src0)
604 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
605 -- XXX Use StreamK to avoid adding arbitrary layers of
606 -- constructors every time.
607 str = D.cons arr0 (D.cons arr1 (D.Stream step s))
608 return (b, str)
609 PR.Error err -> throwM $ ParseError err
611-- | Parse an array stream using the supplied 'Parser'. Returns the parse
612-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
613--
614-- /Internal/
615--
616{-# INLINE_NORMAL parse #-}
617parse ::
618 (MonadIO m, MonadThrow m, Storable a)
619 => PRD.Parser m a b
620 -> SerialT m (A.Array a)
621 -> m (b, SerialT m (A.Array a))
622parse p s = fmap fromStreamD <$> parseD p (toStreamD s)
624-------------------------------------------------------------------------------
625-- Elimination - Running Array Folds and parsers
626-------------------------------------------------------------------------------
628{-# INLINE_NORMAL parseArrD #-}
629parseArrD ::
630 forall m a b. (MonadIO m, MonadThrow m, Storable a)
631 => PRD.Parser m (Array a) b
632 -> D.Stream m (Array.Array a)
633 -> m (b, D.Stream m (Array.Array a))
634parseArrD (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
635 res <- initial
636 case res of
637 PRD.IPartial s -> go SPEC state (List []) s
638 PRD.IDone b -> return (b, stream)
639 PRD.IError err -> throwM $ ParseError err
641 where
643 -- "backBuf" contains last few items in the stream that we may have to
644 -- backtrack to.
645 --
646 -- XXX currently we are using a dumb list based approach for backtracking
647 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
648 -- That will allow us more efficient random back and forth movement.
649 {-# INLINE go #-}
650 go !_ st backBuf !pst = do
651 r <- step defState st
652 case r of
653 D.Yield x s -> gobuf SPEC [x] s backBuf pst
654 D.Skip s -> go SPEC s backBuf pst
655 D.Stop -> do
656 b <- extract pst
657 return (b, D.nil)
659 gobuf !_ [] s backBuf !pst = go SPEC s backBuf pst
660 gobuf !_ (x:xs) s backBuf !pst = do
661 pRes <- pstep pst x
662 case pRes of
663 PR.Partial 0 pst1 ->
664 gobuf SPEC xs s (List []) pst1
665 PR.Partial n pst1 -> do
666 assert
667 (n <= sum (map Array.length (x:getList backBuf)))
668 (return ())
669 let src0 = takeArrayListRev n (x:getList backBuf)
670 src = Prelude.reverse src0 ++ xs
671 gobuf SPEC src s (List []) pst1
672 PR.Continue 0 pst1 ->
673 gobuf SPEC xs s (List (x:getList backBuf)) pst1
674 PR.Continue n pst1 -> do
675 assert
676 (n <= sum (map Array.length (x:getList backBuf)))
677 (return ())
678 let (src0, buf1) = splitAtArrayListRev n (x:getList backBuf)
679 src = Prelude.reverse src0 ++ xs
680 gobuf SPEC src s (List buf1) pst1
681 PR.Done 0 b ->
682 return (b, D.Stream step s)
683 PR.Done n b -> do
684 assert
685 (n <= sum (map Array.length (x:getList backBuf)))
686 (return ())
687 let src0 = takeArrayListRev n (x:getList backBuf)
688 src = Prelude.reverse src0 ++ xs
689 return (b, D.append (D.fromList src) (D.Stream step s))
690 PR.Error err -> throwM $ ParseError err
692{-
693-- | Parse an array stream using the supplied 'Parser'. Returns the parse
694-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
695--
696-- /Internal/
697--
698{-# INLINE parseArr #-}
699parseArr ::
700 (MonadIO m, MonadThrow m, Storable a)
701 => ASF.Parser m a b
702 -> SerialT m (A.Array a)
703 -> m (b, SerialT m (A.Array a))
704parseArr p s = fmap fromStreamD <$> parseD p (toStreamD s)
705-}
707-- | Fold an array stream using the supplied array stream 'Fold'.
708--
709-- /Pre-release/
710--
711{-# INLINE foldArr #-}
712foldArr :: (MonadIO m, MonadThrow m, Storable a) =>
713 ASF.Fold m a b -> SerialT m (A.Array a) -> m b
714foldArr (ASF.Fold p) s = fst <$> parseArrD p (toStreamD s)
716-- | Like 'fold' but also returns the remaining stream.
717--
718-- /Pre-release/
719--
720{-# INLINE foldArr_ #-}
721foldArr_ :: (MonadIO m, MonadThrow m, Storable a) =>
722 ASF.Fold m a b -> SerialT m (A.Array a) -> m (b, SerialT m (A.Array a))
723foldArr_ (ASF.Fold p) s = second fromStreamD <$> parseArrD p (toStreamD s)
725{-# ANN type ParseChunksState Fuse #-}
726data ParseChunksState x inpBuf st pst =
727 ParseChunksInit inpBuf st
728 | ParseChunksInitLeftOver inpBuf
729 | ParseChunksStream st inpBuf !pst
730 | ParseChunksBuf inpBuf st inpBuf !pst
731 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
733{-# INLINE_NORMAL foldArrManyD #-}
734foldArrManyD
735 :: (MonadThrow m, Storable a)
736 => ASF.Fold m a b
737 -> D.Stream m (Array a)
738 -> D.Stream m b
739foldArrManyD (ASF.Fold (PRD.Parser pstep initial extract)) (D.Stream step state) =
740 D.Stream stepOuter (ParseChunksInit [] state)
742 where
744 {-# INLINE_LATE stepOuter #-}
745 -- Buffer is empty, get the first element from the stream, initialize the
746 -- fold and then go to stream processing loop.
747 stepOuter gst (ParseChunksInit [] st) = do
748 r <- step (adaptState gst) st
749 case r of
750 D.Yield x s -> do
751 res <- initial
752 case res of
753 PRD.IPartial ps ->
754 return $ D.Skip $ ParseChunksBuf [x] s [] ps
755 PRD.IDone pb ->
756 let next = ParseChunksInit [x] s
757 in return $ D.Skip $ ParseChunksYield pb next
758 PRD.IError err -> throwM $ ParseError err
759 D.Skip s -> return $ D.Skip $ ParseChunksInit [] s
760 D.Stop -> return D.Stop
762 -- Buffer is not empty, go to buffered processing loop
763 stepOuter _ (ParseChunksInit src st) = do
764 res <- initial
765 case res of
766 PRD.IPartial ps ->
767 return $ D.Skip $ ParseChunksBuf src st [] ps
768 PRD.IDone pb ->
769 let next = ParseChunksInit src st
770 in return $ D.Skip $ ParseChunksYield pb next
771 PRD.IError err -> throwM $ ParseError err
773 -- XXX we just discard any leftover input at the end
774 stepOuter _ (ParseChunksInitLeftOver _) = return D.Stop
776 -- Buffer is empty, process elements from the stream
777 stepOuter gst (ParseChunksStream st backBuf pst) = do
778 r <- step (adaptState gst) st
779 case r of
780 D.Yield x s -> do
781 pRes <- pstep pst x
782 case pRes of
783 PR.Partial 0 pst1 ->
784 return $ D.Skip $ ParseChunksStream s [] pst1
785 PR.Partial n pst1 -> do
786 assert
787 (n <= sum (map Array.length (x:backBuf)))
788 (return ())
789 let src0 = takeArrayListRev n (x:backBuf)
790 src = Prelude.reverse src0
791 return $ D.Skip $ ParseChunksBuf src s [] pst1
792 PR.Continue 0 pst1 ->
793 return $ D.Skip $ ParseChunksStream s (x:backBuf) pst1
794 PR.Continue n pst1 -> do
795 assert
796 (n <= sum (map Array.length (x:backBuf)))
797 (return ())
798 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
799 src = Prelude.reverse src0
800 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
801 PR.Done 0 b -> do
802 return $ D.Skip $
803 ParseChunksYield b (ParseChunksInit [] s)
804 PR.Done n b -> do
805 assert
806 (n <= sum (map Array.length (x:backBuf)))
807 (return ())
808 let src0 = takeArrayListRev n (x:backBuf)
809 let src = Prelude.reverse src0
810 return $ D.Skip $
811 ParseChunksYield b (ParseChunksInit src s)
812 PR.Error err -> throwM $ ParseError err
813 D.Skip s -> return $ D.Skip $ ParseChunksStream s backBuf pst
814 D.Stop -> do
815 b <- extract pst
816 let src = Prelude.reverse backBuf
817 return $ D.Skip $
818 ParseChunksYield b (ParseChunksInitLeftOver src)
820 -- go back to stream processing mode
821 stepOuter _ (ParseChunksBuf [] s buf pst) =
822 return $ D.Skip $ ParseChunksStream s buf pst
824 -- buffered processing loop
825 stepOuter _ (ParseChunksBuf (x:xs) s backBuf pst) = do
826 pRes <- pstep pst x
827 case pRes of
828 PR.Partial 0 pst1 ->
829 return $ D.Skip $ ParseChunksBuf xs s [] pst1
830 PR.Partial n pst1 -> do
831 assert (n <= sum (map Array.length (x:backBuf))) (return ())
832 let src0 = takeArrayListRev n (x:backBuf)
833 src = Prelude.reverse src0 ++ xs
834 return $ D.Skip $ ParseChunksBuf src s [] pst1
835 PR.Continue 0 pst1 ->
836 return $ D.Skip $ ParseChunksBuf xs s (x:backBuf) pst1
837 PR.Continue n pst1 -> do
838 assert (n <= sum (map Array.length (x:backBuf))) (return ())
839 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
840 src = Prelude.reverse src0 ++ xs
841 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
842 PR.Done 0 b ->
843 return $ D.Skip $ ParseChunksYield b (ParseChunksInit xs s)
844 PR.Done n b -> do
845 assert (n <= sum (map Array.length (x:backBuf))) (return ())
846 let src0 = takeArrayListRev n (x:backBuf)
847 src = Prelude.reverse src0 ++ xs
848 return $ D.Skip $ ParseChunksYield b (ParseChunksInit src s)
849 PR.Error err -> throwM $ ParseError err
851 stepOuter _ (ParseChunksYield a next) = return $ D.Yield a next
853-- | Apply an array stream 'Fold' repeatedly on an array stream and emit the
854-- fold outputs in the output stream.
855--
856-- See "Streamly.Prelude.foldMany" for more details.
857--
858-- /Pre-release/
859{-# INLINE foldArrMany #-}
860foldArrMany
861 :: (IsStream t, MonadThrow m, Storable a)
862 => ASF.Fold m a b
863 -> t m (Array a)
864 -> t m b
865foldArrMany p m = fromStreamD $ foldArrManyD p (toStreamD m)