git.iko.soy/streamly
on 4e04e47 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 "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-------------------------------------------------------------------------------
94-- Generation
95-------------------------------------------------------------------------------
97-- | @arraysOf n stream@ groups the elements in the input stream into arrays of
98-- @n@ elements each.
99--
100-- > arraysOf n = Stream.chunksOf n (Array.writeN n)
101--
102-- /Pre-release/
103{-# INLINE arraysOf #-}
104arraysOf :: (IsStream t, MonadIO m, Storable a)
105 => Int -> t m a -> t m (Array a)
106arraysOf n str = fromStreamD $ A.arraysOf n (toStreamD str)
108-------------------------------------------------------------------------------
109-- Append
110-------------------------------------------------------------------------------
112-- XXX efficiently compare two streams of arrays. Two streams can have chunks
113-- of different sizes, we can handle that in the stream comparison abstraction.
114-- This could be useful e.g. to fast compare whether two files differ.
116-- | Convert a stream of arrays into a stream of their elements.
117--
118-- Same as the following but more efficient:
119--
120-- > concat = Stream.unfoldMany Array.read
121--
122-- @since 0.7.0
123{-# INLINE concat #-}
124concat :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
125-- concat m = fromStreamD $ A.flattenArrays (toStreamD m)
126-- concat m = fromStreamD $ D.concatMap A.toStreamD (toStreamD m)
127concat m = fromStreamD $ D.unfoldMany A.read (toStreamD m)
129-- | Convert a stream of arrays into a stream of their elements reversing the
130-- contents of each array before flattening.
131--
132-- > concatRev = Stream.unfoldMany Array.readRev
133--
134-- @since 0.7.0
135{-# INLINE concatRev #-}
136concatRev :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
137concatRev m = fromStreamD $ A.flattenArraysRev (toStreamD m)
139-------------------------------------------------------------------------------
140-- Intersperse and append
141-------------------------------------------------------------------------------
143-- | Flatten a stream of arrays after inserting the given element between
144-- arrays.
145--
146-- /Pre-release/
147{-# INLINE interpose #-}
148interpose :: (MonadIO m, IsStream t, Storable a) => a -> t m (Array a) -> t m a
149interpose x = S.interpose x A.read
151{-# INLINE intercalateSuffix #-}
152intercalateSuffix :: (MonadIO m, IsStream t, Storable a)
153 => Array a -> t m (Array a) -> t m a
154intercalateSuffix = S.intercalateSuffix A.read
156-- | Flatten a stream of arrays appending the given element after each
157-- array.
158--
159-- @since 0.7.0
160{-# INLINE interposeSuffix #-}
161interposeSuffix :: (MonadIO m, IsStream t, Storable a)
162 => a -> t m (Array a) -> t m a
163-- interposeSuffix x = fromStreamD . A.unlines x . toStreamD
164interposeSuffix x = S.interposeSuffix x A.read
166data FlattenState s a =
167 OuterLoop s
168 | InnerLoop s !MA.ArrayContents !(Ptr a) !(Ptr a)
170{-# INLINE_NORMAL unlines #-}
171unlines :: forall m a. (MonadIO m, Storable a)
172 => a -> D.Stream m (Array a) -> D.Stream m a
173unlines sep (D.Stream step state) = D.Stream step' (OuterLoop state)
174 where
175 {-# INLINE_LATE step' #-}
176 step' gst (OuterLoop st) = do
177 r <- step (adaptState gst) st
178 return $ case r of
179 D.Yield Array{..} s ->
180 D.Skip (InnerLoop s arrContents arrStart aEnd)
181 D.Skip s -> D.Skip (OuterLoop s)
182 D.Stop -> D.Stop
184 step' _ (InnerLoop st _ p end) | p == end =
185 return $ D.Yield sep $ OuterLoop st
187 step' _ (InnerLoop st contents p end) = do
188 x <- liftIO $ do
189 r <- peek p
190 touch contents
191 return r
192 return $ D.Yield x (InnerLoop st contents
193 (p `plusPtr` sizeOf (undefined :: a)) end)
195-------------------------------------------------------------------------------
196-- Compact
197-------------------------------------------------------------------------------
199-- XXX These would not be needed once we implement compactLEFold, see
200-- module Streamly.Internal.Data.Array.Stream.Mut.Foreign
201--
202{-# INLINE_NORMAL packArraysChunksOf #-}
203packArraysChunksOf :: (MonadIO m, Storable a)
204 => Int -> D.Stream m (Array a) -> D.Stream m (Array a)
205packArraysChunksOf n str =
206 D.map A.unsafeFreeze $ AS.packArraysChunksOf n $ D.map A.unsafeThaw str
208-- XXX instead of writing two different versions of this operation, we should
209-- write it as a pipe.
210{-# INLINE_NORMAL lpackArraysChunksOf #-}
211lpackArraysChunksOf :: (MonadIO m, Storable a)
212 => Int -> Fold m (Array a) () -> Fold m (Array a) ()
213lpackArraysChunksOf n fld =
214 FL.lmap A.unsafeThaw $ AS.lpackArraysChunksOf n (FL.lmap A.unsafeFreeze fld)
216-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
217-- maximum specified size in bytes.
218--
219-- @since 0.7.0
220{-# INLINE compact #-}
221compact :: (MonadIO m, Storable a)
222 => Int -> SerialT m (Array a) -> SerialT m (Array a)
223compact n xs = fromStreamD $ packArraysChunksOf n (toStreamD xs)
225-------------------------------------------------------------------------------
226-- Split
227-------------------------------------------------------------------------------
229data SplitState s arr
230 = Initial s
231 | Buffering s arr
232 | Splitting s arr
233 | Yielding arr (SplitState s arr)
234 | Finishing
236-- | Split a stream of arrays on a given separator byte, dropping the separator
237-- and coalescing all the arrays between two separators into a single array.
238--
239-- @since 0.7.0
240{-# INLINE_NORMAL _splitOn #-}
241_splitOn
242 :: MonadIO m
243 => Word8
244 -> D.Stream m (Array Word8)
245 -> D.Stream m (Array Word8)
246_splitOn byte (D.Stream step state) = D.Stream step' (Initial state)
248 where
250 {-# INLINE_LATE step' #-}
251 step' gst (Initial st) = do
252 r <- step gst st
253 case r of
254 D.Yield arr s -> do
255 (arr1, marr2) <- A.breakOn byte arr
256 return $ case marr2 of
257 Nothing -> D.Skip (Buffering s arr1)
258 Just arr2 -> D.Skip (Yielding arr1 (Splitting s arr2))
259 D.Skip s -> return $ D.Skip (Initial s)
260 D.Stop -> return D.Stop
262 step' gst (Buffering st buf) = do
263 r <- step gst st
264 case r of
265 D.Yield arr s -> do
266 (arr1, marr2) <- A.breakOn byte arr
267 buf' <- A.splice buf arr1
268 return $ case marr2 of
269 Nothing -> D.Skip (Buffering s buf')
270 Just x -> D.Skip (Yielding buf' (Splitting s x))
271 D.Skip s -> return $ D.Skip (Buffering s buf)
272 D.Stop -> return $
273 if A.byteLength buf == 0
274 then D.Stop
275 else D.Skip (Yielding buf Finishing)
277 step' _ (Splitting st buf) = do
278 (arr1, marr2) <- A.breakOn byte buf
279 return $ case marr2 of
280 Nothing -> D.Skip $ Buffering st arr1
281 Just arr2 -> D.Skip $ Yielding arr1 (Splitting st arr2)
283 step' _ (Yielding arr next) = return $ D.Yield arr next
284 step' _ Finishing = return D.Stop
286-- | Split a stream of arrays on a given separator byte, dropping the separator
287-- and coalescing all the arrays between two separators into a single array.
288--
289-- @since 0.7.0
290{-# INLINE splitOn #-}
291splitOn
292 :: (IsStream t, MonadIO m)
293 => Word8
294 -> t m (Array Word8)
295 -> t m (Array Word8)
296splitOn byte s =
297 fromStreamD $ D.splitInnerBy (A.breakOn byte) A.splice $ toStreamD s
299{-# INLINE splitOnSuffix #-}
300splitOnSuffix
301 :: (IsStream t, MonadIO m)
302 => Word8
303 -> t m (Array Word8)
304 -> t m (Array Word8)
305-- splitOn byte s = fromStreamD $ A.splitOn byte $ toStreamD s
306splitOnSuffix byte s =
307 fromStreamD $ D.splitInnerBySuffix (A.breakOn byte) A.splice $ toStreamD s
309-------------------------------------------------------------------------------
310-- Elimination - Running folds
311-------------------------------------------------------------------------------
313-- XXX This should be written using CPS (as foldK) if we want it to scale wrt
314-- to the number of times it can be called on the same stream.
315--
316{-# INLINE_NORMAL foldD #-}
317foldD :: forall m a b. (MonadIO m, Storable a) =>
318 Fold m a b -> D.Stream m (Array a) -> m (b, D.Stream m (Array a))
319foldD (Fold fstep initial extract) stream@(D.Stream step state) = do
320 res <- initial
321 case res of
322 FL.Partial fs -> go SPEC state fs
323 FL.Done fb -> return $! (fb, stream)
325 where
327 {-# INLINE go #-}
328 go !_ st !fs = do
329 r <- step defState st
330 case r of
331 D.Yield (Array contents start (Ptr end)) s ->
332 let fp = ForeignPtr end (arrayToFptrContents contents)
333 in goArray SPEC s fp start fs
334 D.Skip s -> go SPEC s fs
335 D.Stop -> do
336 b <- extract fs
337 return (b, D.nil)
339 goArray !_ s fp@(ForeignPtr end _) !cur !fs
340 | cur == Ptr end = do
341 liftIO $ touchForeignPtr fp
342 go SPEC s fs
343 goArray !_ st fp@(ForeignPtr end contents) !cur !fs = do
344 x <- liftIO $ peek cur
345 res <- fstep fs x
346 let elemSize = sizeOf (undefined :: a)
347 next = cur `plusPtr` elemSize
348 case res of
349 FL.Done b -> do
350 let arr = Array (fptrToArrayContents contents) next (Ptr end)
351 return $! (b, D.cons arr (D.Stream step st))
352 FL.Partial fs1 -> goArray SPEC st fp next fs1
354-- | Fold an array stream using the supplied 'Fold'. Returns the fold result
355-- and the unconsumed stream.
356--
357-- /Internal/
358--
359{-# INLINE_NORMAL fold #-}
360fold ::
361 (MonadIO m, Storable a)
362 => FL.Fold m a b
363 -> SerialT m (A.Array a)
364 -> m (b, SerialT m (A.Array a))
365fold f s = fmap fromStreamD <$> foldD f (toStreamD s)
367-------------------------------------------------------------------------------
368-- Fold to a single Array
369-------------------------------------------------------------------------------
371-- When we have to take an array partially, take the last part of the array.
372{-# INLINE takeArrayListRev #-}
373takeArrayListRev :: forall a. Storable a => Int -> [Array a] -> [Array a]
374takeArrayListRev = go
376 where
378 go _ [] = []
379 go n _ | n <= 0 = []
380 go n (x:xs) =
381 let len = Array.length x
382 in if n > len
383 then x : go (n - len) xs
384 else if n == len
385 then [x]
386 else let !(Array contents _ end) = x
387 sz = sizeOf (undefined :: a)
388 !start = end `plusPtr` negate (n * sz)
389 in [Array contents start end]
391-- When we have to take an array partially, take the last part of the array in
392-- the first split.
393{-# INLINE splitAtArrayListRev #-}
394splitAtArrayListRev :: forall a. Storable a =>
395 Int -> [Array a] -> ([Array a],[Array a])
396splitAtArrayListRev n ls
397 | n <= 0 = ([], ls)
398 | otherwise = go n ls
399 where
400 go :: Int -> [Array a] -> ([Array a], [Array a])
401 go _ [] = ([], [])
402 go m (x:xs) =
403 let len = Array.length x
404 (xs', xs'') = go (m - len) xs
405 in if m > len
406 then (x:xs', xs'')
407 else if m == len
408 then ([x],xs)
409 else let !(Array contents start end) = x
410 sz = sizeOf (undefined :: a)
411 end1 = end `plusPtr` negate (m * sz)
412 arr2 = Array contents start end1
413 arr1 = Array contents end1 end
414 in ([arr1], arr2:xs)
416-------------------------------------------------------------------------------
417-- Fold to a single Array
418-------------------------------------------------------------------------------
420-- XXX Both of these implementations of splicing seem to perform equally well.
421-- We need to perform benchmarks over a range of sizes though.
423-- CAUTION! length must more than equal to lengths of all the arrays in the
424-- stream.
425{-# INLINE spliceArraysLenUnsafe #-}
426spliceArraysLenUnsafe :: (MonadIO m, Storable a)
427 => Int -> SerialT m (MA.Array a) -> m (MA.Array a)
428spliceArraysLenUnsafe len buffered = do
429 arr <- liftIO $ MA.newArray len
430 end <- S.foldlM' writeArr (return $ MA.aEnd arr) buffered
431 return $ arr {MA.aEnd = end}
433 where
435 writeArr dst (MA.Array ac src ae _) =
436 liftIO $ do
437 let count = ae `minusPtr` src
438 memcpy (castPtr dst) (castPtr src) count
439 touch ac
440 return $ dst `plusPtr` count
442{-# INLINE _spliceArrays #-}
443_spliceArrays :: (MonadIO m, Storable a)
444 => SerialT m (Array a) -> m (Array a)
445_spliceArrays s = do
446 buffered <- S.foldr S.cons S.nil s
447 len <- S.sum (S.map Array.length buffered)
448 arr <- liftIO $ MA.newArray len
449 end <- S.foldlM' writeArr (return $ MA.aEnd arr) s
450 return $ A.unsafeFreeze $ arr {MA.aEnd = end}
452 where
454 writeArr dst (Array ac src ae) =
455 liftIO $ do
456 let count = ae `minusPtr` src
457 memcpy (castPtr dst) (castPtr src) count
458 touch ac
459 return $ dst `plusPtr` count
461{-# INLINE _spliceArraysBuffered #-}
462_spliceArraysBuffered :: (MonadIO m, Storable a)
463 => SerialT m (Array a) -> m (Array a)
464_spliceArraysBuffered s = do
465 buffered <- S.foldr S.cons S.nil s
466 len <- S.sum (S.map Array.length buffered)
467 A.unsafeFreeze <$> spliceArraysLenUnsafe len (S.map A.unsafeThaw s)
469{-# INLINE spliceArraysRealloced #-}
470spliceArraysRealloced :: forall m a. (MonadIO m, Storable a)
471 => SerialT m (Array a) -> m (Array a)
472spliceArraysRealloced s = do
473 let n = allocBytesToElemCount (undefined :: a) (4 * 1024)
474 idst = liftIO $ MA.newArray n
476 arr <- S.foldlM' MA.spliceExp idst (S.map A.unsafeThaw s)
477 liftIO $ A.unsafeFreeze <$> MA.rightSize arr
479-- XXX This should just be "fold A.write"
480--
481-- | Given a stream of arrays, splice them all together to generate a single
482-- array. The stream must be /finite/.
483--
484-- @since 0.7.0
485{-# INLINE toArray #-}
486toArray :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)
487toArray = spliceArraysRealloced
488-- spliceArrays = _spliceArraysBuffered
490-- exponentially increasing sizes of the chunks upto the max limit.
491-- XXX this will be easier to implement with parsers/terminating folds
492-- With this we should be able to reduce the number of chunks/allocations.
493-- The reallocation/copy based toArray can also be implemented using this.
494--
495{-
496{-# INLINE toArraysInRange #-}
497toArraysInRange :: (IsStream t, MonadIO m, Storable a)
498 => Int -> Int -> Fold m (Array a) b -> Fold m a b
499toArraysInRange low high (Fold step initial extract) =
500-}
502{-
503-- | Fold the input to a pure buffered stream (List) of arrays.
504{-# INLINE _toArraysOf #-}
505_toArraysOf :: (MonadIO m, Storable a)
506 => Int -> Fold m a (SerialT Identity (Array a))
507_toArraysOf n = FL.chunksOf n (A.writeNF n) FL.toStream
508-}
510-------------------------------------------------------------------------------
511-- Elimination - running element parsers
512-------------------------------------------------------------------------------
514-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
515-- to make a newtype.
516{-# ANN type List NoSpecConstr #-}
517newtype List a = List {getList :: [a]}
519-- This can be generalized to any type provided it can be unfolded to a stream
520-- and it can be combined using a semigroup operation.
521--
522-- XXX This should be written using CPS (as parseK) if we want it to scale wrt
523-- to the number of times it can be called on the same stream.
524{-# INLINE_NORMAL parseD #-}
525parseD ::
526 forall m a b. (MonadIO m, MonadThrow m, Storable a)
527 => PRD.Parser m a b
528 -> D.Stream m (Array.Array a)
529 -> m (b, D.Stream m (Array.Array a))
530parseD (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
531 res <- initial
532 case res of
533 PRD.IPartial s -> go SPEC state (List []) s
534 PRD.IDone b -> return (b, stream)
535 PRD.IError err -> throwM $ ParseError err
537 where
539 -- "backBuf" contains last few items in the stream that we may have to
540 -- backtrack to.
541 --
542 -- XXX currently we are using a dumb list based approach for backtracking
543 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
544 -- That will allow us more efficient random back and forth movement.
545 {-# INLINE go #-}
546 go !_ st backBuf !pst = do
547 r <- step defState st
548 case r of
549 D.Yield (Array contents start (Ptr end)) s ->
550 gobuf SPEC s backBuf
551 (ForeignPtr end (arrayToFptrContents contents)) start pst
552 D.Skip s -> go SPEC s backBuf pst
553 D.Stop -> do
554 b <- extract pst
555 return (b, D.nil)
557 -- Use strictness on "cur" to keep it unboxed
558 gobuf !_ s backBuf fp@(ForeignPtr end _) !cur !pst
559 | cur == Ptr end = do
560 liftIO $ touchForeignPtr fp
561 go SPEC s backBuf pst
562 gobuf !_ s backBuf fp@(ForeignPtr end contents) !cur !pst = do
563 x <- liftIO $ peek cur
564 pRes <- pstep pst x
565 let elemSize = sizeOf (undefined :: a)
566 next = cur `plusPtr` elemSize
567 case pRes of
568 PR.Partial 0 pst1 ->
569 gobuf SPEC s (List []) fp next pst1
570 PR.Partial n pst1 -> do
571 assert (n <= Prelude.length (x:getList backBuf)) (return ())
572 let src0 = Prelude.take n (x:getList backBuf)
573 arr0 = A.fromListN n (Prelude.reverse src0)
574 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
575 src = arr0 <> arr1
576 let !(Array cont1 start (Ptr end1)) = src
577 fp1 = ForeignPtr end1 (arrayToFptrContents cont1)
578 gobuf SPEC s (List []) fp1 start pst1
579 PR.Continue 0 pst1 ->
580 gobuf SPEC s (List (x:getList backBuf)) fp next pst1
581 PR.Continue n pst1 -> do
582 assert (n <= Prelude.length (x:getList backBuf)) (return ())
583 let (src0, buf1) = splitAt n (x:getList backBuf)
584 arr0 = A.fromListN n (Prelude.reverse src0)
585 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
586 src = arr0 <> arr1
587 let !(Array cont1 start (Ptr end1)) = src
588 fp1 = ForeignPtr end1 (arrayToFptrContents cont1)
589 gobuf SPEC s (List buf1) fp1 start pst1
590 PR.Done 0 b -> do
591 let arr = Array (fptrToArrayContents contents) next (Ptr end)
592 return (b, D.cons arr (D.Stream step s))
593 PR.Done n b -> do
594 assert (n <= Prelude.length (x:getList backBuf)) (return ())
595 let src0 = Prelude.take n (x:getList backBuf)
596 -- XXX create the array in reverse instead
597 arr0 = A.fromListN n (Prelude.reverse src0)
598 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
599 -- XXX Use StreamK to avoid adding arbitrary layers of
600 -- constructors every time.
601 str = D.cons arr0 (D.cons arr1 (D.Stream step s))
602 return (b, str)
603 PR.Error err -> throwM $ ParseError err
605-- | Parse an array stream using the supplied 'Parser'. Returns the parse
606-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
607--
608-- /Internal/
609--
610{-# INLINE_NORMAL parse #-}
611parse ::
612 (MonadIO m, MonadThrow m, Storable a)
613 => PRD.Parser m a b
614 -> SerialT m (A.Array a)
615 -> m (b, SerialT m (A.Array a))
616parse p s = fmap fromStreamD <$> parseD p (toStreamD s)
618-------------------------------------------------------------------------------
619-- Elimination - Running Array Folds and parsers
620-------------------------------------------------------------------------------
622{-# INLINE_NORMAL parseArrD #-}
623parseArrD ::
624 forall m a b. (MonadIO m, MonadThrow m, Storable a)
625 => PRD.Parser m (Array a) b
626 -> D.Stream m (Array.Array a)
627 -> m (b, D.Stream m (Array.Array a))
628parseArrD (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
629 res <- initial
630 case res of
631 PRD.IPartial s -> go SPEC state (List []) s
632 PRD.IDone b -> return (b, stream)
633 PRD.IError err -> throwM $ ParseError err
635 where
637 -- "backBuf" contains last few items in the stream that we may have to
638 -- backtrack to.
639 --
640 -- XXX currently we are using a dumb list based approach for backtracking
641 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
642 -- That will allow us more efficient random back and forth movement.
643 {-# INLINE go #-}
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)