git.iko.soy/streamly
on a02f2d0 help sign in
streamly/src/Streamly/Internal/Data/Array/Stream/Foreign.hs
859 lines 33 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Array.Stream.Foreign
3-- Copyright : (c) 2019 Composewell Technologies
4-- License : BSD3-3-Clause
5-- Maintainer : streamly@composewell.com
6-- Stability : pre-release
7-- Portability : GHC
8--
9-- Combinators to efficiently manipulate streams of immutable arrays.
10--
11module Streamly.Internal.Data.Array.Stream.Foreign
12 (
13 -- * Creation
14 arraysOf
16 -- * Flattening to elements
17 , concat
18 , concatRev
19 , interpose
20 , interposeSuffix
21 , intercalateSuffix
22 , unlines
24 -- * Elimination
25 -- ** Element Folds
26 , 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.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
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.BaseCompat
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.StreamK.Type (IsStream)
76import Streamly.Internal.Data.SVar (adaptState, defState)
77import Streamly.Internal.Data.Array.Foreign.Mut.Type (memcpy, bytesToElemCount)
78import Streamly.Internal.System.IO (mkChunkSizeKB)
80import qualified Streamly.Internal.Data.Array.Foreign as A
81import qualified Streamly.Internal.Data.Array.Foreign as Array
82import qualified Streamly.Internal.Data.Array.Foreign.Type as A
83import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MA
84import qualified Streamly.Internal.Data.Array.Stream.Mut.Foreign as AS
85import qualified Streamly.Internal.Data.Array.Stream.Fold.Foreign as ASF
86import qualified Streamly.Internal.Data.Fold.Type as FL
87import qualified Streamly.Internal.Data.Parser as PR
88import qualified Streamly.Internal.Data.Parser.ParserD as PRD
89import qualified Streamly.Internal.Data.Stream.IsStream as S
90import qualified Streamly.Internal.Data.Stream.StreamD as D
92-------------------------------------------------------------------------------
93-- Generation
94-------------------------------------------------------------------------------
96-- | @arraysOf n stream@ groups the elements in the input stream into arrays of
97-- @n@ elements each.
98--
99-- > arraysOf n = Stream.chunksOf n (Array.writeN n)
100--
101-- /Pre-release/
102{-# INLINE arraysOf #-}
103arraysOf :: (IsStream t, MonadIO m, Storable a)
104 => Int -> t m a -> t m (Array a)
105arraysOf n str = D.fromStreamD $ A.arraysOf n (D.toStreamD str)
107-------------------------------------------------------------------------------
108-- Append
109-------------------------------------------------------------------------------
111-- XXX efficiently compare two streams of arrays. Two streams can have chunks
112-- of different sizes, we can handle that in the stream comparison abstraction.
113-- This could be useful e.g. to fast compare whether two files differ.
115-- | Convert a stream of arrays into a stream of their elements.
116--
117-- Same as the following but more efficient:
118--
119-- > concat = Stream.unfoldMany Array.read
120--
121-- @since 0.7.0
122{-# INLINE concat #-}
123concat :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
124-- concat m = D.fromStreamD $ A.flattenArrays (D.toStreamD m)
125-- concat m = D.fromStreamD $ D.concatMap A.toStreamD (D.toStreamD m)
126concat m = D.fromStreamD $ D.unfoldMany A.read (D.toStreamD m)
128-- | Convert a stream of arrays into a stream of their elements reversing the
129-- contents of each array before flattening.
130--
131-- > concatRev = Stream.unfoldMany Array.readRev
132--
133-- @since 0.7.0
134{-# INLINE concatRev #-}
135concatRev :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
136concatRev m = D.fromStreamD $ A.flattenArraysRev (D.toStreamD m)
138-------------------------------------------------------------------------------
139-- Intersperse and append
140-------------------------------------------------------------------------------
142-- | Flatten a stream of arrays after inserting the given element between
143-- arrays.
144--
145-- /Pre-release/
146{-# INLINE interpose #-}
147interpose :: (MonadIO m, IsStream t, Storable a) => a -> t m (Array a) -> t m a
148interpose x = S.interpose x A.read
150{-# INLINE intercalateSuffix #-}
151intercalateSuffix :: (MonadIO m, IsStream t, Storable a)
152 => Array a -> t m (Array a) -> t m a
153intercalateSuffix = S.intercalateSuffix A.read
155-- | Flatten a stream of arrays appending the given element after each
156-- array.
157--
158-- @since 0.7.0
159{-# INLINE interposeSuffix #-}
160interposeSuffix :: (MonadIO m, IsStream t, Storable a)
161 => a -> t m (Array a) -> t m a
162-- interposeSuffix x = D.fromStreamD . A.unlines x . D.toStreamD
163interposeSuffix x = S.interposeSuffix x A.read
165data FlattenState s a =
166 OuterLoop s
167 | InnerLoop s !(ForeignPtr a) !(Ptr a) !(Ptr a)
169{-# INLINE_NORMAL unlines #-}
170unlines :: forall m a. (MonadIO m, Storable a)
171 => a -> D.Stream m (Array a) -> D.Stream m a
172unlines sep (D.Stream step state) = D.Stream step' (OuterLoop state)
173 where
174 {-# INLINE_LATE step' #-}
175 step' gst (OuterLoop st) = do
176 r <- step (adaptState gst) st
177 return $ case r of
178 D.Yield Array{..} s ->
179 let p = unsafeForeignPtrToPtr aStart
180 in D.Skip (InnerLoop s aStart p 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 startf p end) = do
188 x <- liftIO $ do
189 r <- peek p
190 touchForeignPtr startf
191 return r
192 return $ D.Yield x (InnerLoop st startf
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.map A.unsafeThaw $ AS.lpackArraysChunksOf n (FL.map 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 = D.fromStreamD $ packArraysChunksOf n (D.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.spliceTwo 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 D.fromStreamD $ D.splitInnerBy (A.breakOn byte) A.spliceTwo $ D.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 = D.fromStreamD $ A.splitOn byte $ D.toStreamD s
306splitOnSuffix byte s =
307 D.fromStreamD $ D.splitInnerBySuffix (A.breakOn byte) A.spliceTwo $ D.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 (ForeignPtr start contents) (Ptr end)) s ->
332 goArray SPEC s (ForeignPtr end contents) (Ptr start) fs
333 D.Skip s -> go SPEC s fs
334 D.Stop -> do
335 b <- extract fs
336 return (b, D.nil)
338 goArray !_ s fp@(ForeignPtr end _) !cur !fs
339 | cur == Ptr end = do
340 liftIO $ touchForeignPtr fp
341 go SPEC s fs
342 goArray !_ st fp@(ForeignPtr end contents) !cur !fs = do
343 x <- liftIO $ peek cur
344 res <- fstep fs x
345 let elemSize = sizeOf (undefined :: a)
346 next = cur `plusPtr` elemSize
347 case res of
348 FL.Done b -> do
349 let !(Ptr curAddr) = next
350 arr = Array (ForeignPtr curAddr contents) (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 D.fromStreamD <$> foldD f (D.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 (ForeignPtr _ contents) end) = x
387 sz = sizeOf (undefined :: a)
388 !(Ptr start) = end `plusPtr` negate (n * sz)
389 in [Array (ForeignPtr start contents) 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 (ForeignPtr start contents) end) = x
410 sz = sizeOf (undefined :: a)
411 end1 = end `plusPtr` negate (m * sz)
412 arr2 = Array (ForeignPtr start contents) end1
413 !(Ptr addrEnd1) = end1
414 arr1 = Array (ForeignPtr addrEnd1 contents) end
415 in ([arr1], arr2:xs)
417-------------------------------------------------------------------------------
418-- Fold to a single Array
419-------------------------------------------------------------------------------
421-- XXX Both of these implementations of splicing seem to perform equally well.
422-- We need to perform benchmarks over a range of sizes though.
424-- CAUTION! length must more than equal to lengths of all the arrays in the
425-- stream.
426{-# INLINE spliceArraysLenUnsafe #-}
427spliceArraysLenUnsafe :: (MonadIO m, Storable a)
428 => Int -> SerialT m (MA.Array a) -> m (MA.Array a)
429spliceArraysLenUnsafe len buffered = do
430 arr <- liftIO $ MA.newArray len
431 end <- S.foldlM' writeArr (return $ MA.aEnd arr) buffered
432 return $ arr {MA.aEnd = end}
434 where
436 writeArr dst (MA.Array as ae _) =
437 liftIO $ unsafeWithForeignPtr as $ \src -> do
438 let count = ae `minusPtr` src
439 memcpy (castPtr dst) (castPtr src) count
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 as ae) =
455 liftIO $ unsafeWithForeignPtr as $ \src -> do
456 let count = ae `minusPtr` src
457 memcpy (castPtr dst) (castPtr src) count
458 return $ dst `plusPtr` count
460{-# INLINE _spliceArraysBuffered #-}
461_spliceArraysBuffered :: (MonadIO m, Storable a)
462 => SerialT m (Array a) -> m (Array a)
463_spliceArraysBuffered s = do
464 buffered <- S.foldr S.cons S.nil s
465 len <- S.sum (S.map Array.length buffered)
466 A.unsafeFreeze <$> spliceArraysLenUnsafe len (S.map A.unsafeThaw s)
468{-# INLINE spliceArraysRealloced #-}
469spliceArraysRealloced :: forall m a. (MonadIO m, Storable a)
470 => SerialT m (Array a) -> m (Array a)
471spliceArraysRealloced s = do
472 let idst = liftIO $ MA.newArray (bytesToElemCount (undefined :: a)
473 (mkChunkSizeKB 4))
475 arr <- S.foldlM' MA.spliceWithDoubling idst (S.map A.unsafeThaw s)
476 liftIO $ A.unsafeFreeze <$> MA.shrinkToFit arr
478-- XXX This should just be "fold A.write"
479--
480-- | Given a stream of arrays, splice them all together to generate a single
481-- array. The stream must be /finite/.
482--
483-- @since 0.7.0
484{-# INLINE toArray #-}
485toArray :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)
486toArray = spliceArraysRealloced
487-- spliceArrays = _spliceArraysBuffered
489-- exponentially increasing sizes of the chunks upto the max limit.
490-- XXX this will be easier to implement with parsers/terminating folds
491-- With this we should be able to reduce the number of chunks/allocations.
492-- The reallocation/copy based toArray can also be implemented using this.
493--
494{-
495{-# INLINE toArraysInRange #-}
496toArraysInRange :: (IsStream t, MonadIO m, Storable a)
497 => Int -> Int -> Fold m (Array a) b -> Fold m a b
498toArraysInRange low high (Fold step initial extract) =
499-}
501{-
502-- | Fold the input to a pure buffered stream (List) of arrays.
503{-# INLINE _toArraysOf #-}
504_toArraysOf :: (MonadIO m, Storable a)
505 => Int -> Fold m a (SerialT Identity (Array a))
506_toArraysOf n = FL.chunksOf n (A.writeNF n) FL.toStream
507-}
509-------------------------------------------------------------------------------
510-- Elimination - running element parsers
511-------------------------------------------------------------------------------
513-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
514-- to make a newtype.
515{-# ANN type List NoSpecConstr #-}
516newtype List a = List {getList :: [a]}
518-- This can be generalized to any type provided it can be unfolded to a stream
519-- and it can be combined using a semigroup operation.
520--
521-- XXX This should be written using CPS (as parseK) if we want it to scale wrt
522-- to the number of times it can be called on the same stream.
523{-# INLINE_NORMAL parseD #-}
524parseD ::
525 forall m a b. (MonadIO m, MonadThrow m, Storable a)
526 => PRD.Parser m a b
527 -> D.Stream m (Array.Array a)
528 -> m (b, D.Stream m (Array.Array a))
529parseD (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
530 res <- initial
531 case res of
532 PRD.IPartial s -> go SPEC state (List []) s
533 PRD.IDone b -> return (b, stream)
534 PRD.IError err -> throwM $ ParseError err
536 where
538 -- "backBuf" contains last few items in the stream that we may have to
539 -- backtrack to.
540 --
541 -- XXX currently we are using a dumb list based approach for backtracking
542 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
543 -- That will allow us more efficient random back and forth movement.
544 {-# INLINE go #-}
545 go !_ st backBuf !pst = do
546 r <- step defState st
547 case r of
548 D.Yield (Array (ForeignPtr start contents) (Ptr end)) s ->
549 gobuf SPEC s backBuf (ForeignPtr end contents) (Ptr start) pst
550 D.Skip s -> go SPEC s backBuf pst
551 D.Stop -> do
552 b <- extract pst
553 return (b, D.nil)
555 -- Use strictness on "cur" to keep it unboxed
556 gobuf !_ s backBuf fp@(ForeignPtr end _) !cur !pst
557 | cur == Ptr end = do
558 liftIO $ touchForeignPtr fp
559 go SPEC s backBuf pst
560 gobuf !_ s backBuf fp@(ForeignPtr end contents) !cur !pst = do
561 x <- liftIO $ peek cur
562 pRes <- pstep pst x
563 let elemSize = sizeOf (undefined :: a)
564 next = cur `plusPtr` elemSize
565 case pRes of
566 PR.Partial 0 pst1 ->
567 gobuf SPEC s (List []) fp next pst1
568 PR.Partial n pst1 -> do
569 assert (n <= Prelude.length (x:getList backBuf)) (return ())
570 let src0 = Prelude.take n (x:getList backBuf)
571 arr0 = A.fromListN n (Prelude.reverse src0)
572 !(Ptr curAddr) = next
573 arr1 = Array (ForeignPtr curAddr contents) (Ptr end)
574 src = arr0 <> arr1
575 let !(Array (ForeignPtr start cont1) (Ptr end1)) = src
576 gobuf SPEC s (List []) (ForeignPtr end1 cont1) (Ptr start) pst1
577 PR.Continue 0 pst1 ->
578 gobuf SPEC s (List (x:getList backBuf)) fp next pst1
579 PR.Continue n pst1 -> do
580 assert (n <= Prelude.length (x:getList backBuf)) (return ())
581 let (src0, buf1) = splitAt n (x:getList backBuf)
582 arr0 = A.fromListN n (Prelude.reverse src0)
583 !(Ptr curAddr) = next
584 arr1 = Array (ForeignPtr curAddr contents) (Ptr end)
585 src = arr0 <> arr1
586 let !(Array (ForeignPtr start cont1) (Ptr end1)) = src
587 gobuf SPEC s (List buf1) (ForeignPtr end1 cont1) (Ptr start) pst1
588 PR.Done 0 b -> do
589 let !(Ptr curAddr) = next
590 arr = Array (ForeignPtr curAddr contents) (Ptr end)
591 return (b, D.cons arr (D.Stream step s))
592 PR.Done n b -> do
593 assert (n <= Prelude.length (x:getList backBuf)) (return ())
594 let src0 = Prelude.take n (x:getList backBuf)
595 -- XXX create the array in reverse instead
596 arr0 = A.fromListN n (Prelude.reverse src0)
597 !(Ptr curAddr) = next
598 arr1 = Array (ForeignPtr curAddr contents) (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 D.fromStreamD <$> parseD p (D.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 D.fromStreamD <$> parseD p (D.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 (D.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 D.fromStreamD <$> parseArrD p (D.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 = D.fromStreamD $ foldArrManyD p (D.toStreamD m)