git.iko.soy/streamly
on d2aaef0 help sign in
streamly/src/Streamly/Internal/Data/Array/Stream/Foreign.hs
977 lines 37 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 , foldBreak
27 , foldBreakD
28 , parseBreak
29 , parseBreakD
31 -- ** Array Folds
32 , runArrayFold
33 , runArrayFoldBreak
34 -- , parseArr
35 , runArrayParserDBreak
36 , runArrayFoldMany
38 , toArray
40 -- * Compaction
41 , lpackArraysChunksOf
42 , compact
44 -- * Splitting
45 , splitOn
46 , splitOnSuffix
47 )
48where
50#include "ArrayMacros.h"
51#include "inline.hs"
53import Data.Bifunctor (second)
54import Control.Exception (assert)
55import Control.Monad.Catch (MonadThrow, throwM)
56import Control.Monad.IO.Class (MonadIO(..))
57#if __GLASGOW_HASKELL__ < 808
58import Data.Semigroup (Semigroup(..))
59#endif
60import Data.Word (Word8)
61import Streamly.Internal.Data.Unboxed (Unboxed, peekWith, sizeOf)
62import Fusion.Plugin.Types (Fuse(..))
63import GHC.Exts (SpecConstrAnnotation(..))
64import GHC.Types (SPEC(..))
65import Prelude hiding (null, last, (!!), read, concat, unlines)
67import Streamly.Internal.Data.Array.Unboxed.Type (Array(..))
68import Streamly.Internal.Data.Array.Stream.Fold.Foreign (ArrayFold(..))
69import Streamly.Internal.Data.Fold.Type (Fold(..))
70import Streamly.Internal.Data.Parser (ParseError(..))
71import Streamly.Internal.Data.Stream.Serial (SerialT)
72import Streamly.Internal.Data.Stream.IsStream.Type
73 (IsStream, fromStreamD, toStreamD)
74import Streamly.Internal.Data.SVar (adaptState, defState)
75import Streamly.Internal.Data.Array.Unboxed.Mut.Type
76 (allocBytesToElemCount)
77import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
79import qualified Streamly.Internal.Data.Array.Unboxed as A
80import qualified Streamly.Internal.Data.Array.Unboxed as Array
81-- import qualified Streamly.Internal.Data.Array.Stream.Fold.Foreign
82-- as ArrayFold
83import qualified Streamly.Internal.Data.Array.Unboxed.Type as A
84import qualified Streamly.Internal.Data.Array.Unboxed.Mut.Type as MA
85import qualified Streamly.Internal.Data.Array.Stream.Mut.Foreign as AS
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
91import qualified Streamly.Internal.Data.Stream.StreamK as K
92import qualified Streamly.Internal.Data.Stream.Type as Stream
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, Unboxed 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, Unboxed 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, Unboxed 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, Unboxed a) => a -> t m (Array a) -> t m a
154interpose x = S.interpose x A.read
156{-# INLINE intercalateSuffix #-}
157intercalateSuffix :: (Monad m, IsStream t, Unboxed 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, Unboxed 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 a) !Int !Int
175-- XXX Remove monadIO constraint
176{-# INLINE_NORMAL unlines #-}
177unlines :: forall m a. (MonadIO m, Unboxed 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 $ peekWith contents p
195 return $ D.Yield x (InnerLoop st contents (INDEX_NEXT(p,a)) end)
197-------------------------------------------------------------------------------
198-- Compact
199-------------------------------------------------------------------------------
201-- XXX These would not be needed once we implement compactLEFold, see
202-- module Streamly.Internal.Data.Array.Stream.Mut.Foreign
203--
204{-# INLINE_NORMAL packArraysChunksOf #-}
205packArraysChunksOf :: (MonadIO m, Unboxed a)
206 => Int -> D.Stream m (Array a) -> D.Stream m (Array a)
207packArraysChunksOf n str =
208 D.map A.unsafeFreeze $ AS.packArraysChunksOf n $ D.map A.unsafeThaw str
210-- XXX instead of writing two different versions of this operation, we should
211-- write it as a pipe.
212{-# INLINE_NORMAL lpackArraysChunksOf #-}
213lpackArraysChunksOf :: (MonadIO m, Unboxed a)
214 => Int -> Fold m (Array a) () -> Fold m (Array a) ()
215lpackArraysChunksOf n fld =
216 FL.lmap A.unsafeThaw $ AS.lpackArraysChunksOf n (FL.lmap A.unsafeFreeze fld)
218-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
219-- maximum specified size in bytes.
220--
221-- @since 0.7.0
222{-# INLINE compact #-}
223compact :: (MonadIO m, Unboxed a)
224 => Int -> SerialT m (Array a) -> SerialT m (Array a)
225compact n xs = fromStreamD $ packArraysChunksOf n (toStreamD xs)
227-------------------------------------------------------------------------------
228-- Split
229-------------------------------------------------------------------------------
231data SplitState s arr
232 = Initial s
233 | Buffering s arr
234 | Splitting s arr
235 | Yielding arr (SplitState s arr)
236 | Finishing
238-- | Split a stream of arrays on a given separator byte, dropping the separator
239-- and coalescing all the arrays between two separators into a single array.
240--
241-- @since 0.7.0
242{-# INLINE_NORMAL _splitOn #-}
243_splitOn
244 :: MonadIO m
245 => Word8
246 -> D.Stream m (Array Word8)
247 -> D.Stream m (Array Word8)
248_splitOn byte (D.Stream step state) = D.Stream step' (Initial state)
250 where
252 {-# INLINE_LATE step' #-}
253 step' gst (Initial st) = do
254 r <- step gst st
255 case r of
256 D.Yield arr s -> do
257 (arr1, marr2) <- A.breakOn byte arr
258 return $ case marr2 of
259 Nothing -> D.Skip (Buffering s arr1)
260 Just arr2 -> D.Skip (Yielding arr1 (Splitting s arr2))
261 D.Skip s -> return $ D.Skip (Initial s)
262 D.Stop -> return D.Stop
264 step' gst (Buffering st buf) = do
265 r <- step gst st
266 case r of
267 D.Yield arr s -> do
268 (arr1, marr2) <- A.breakOn byte arr
269 buf' <- A.splice buf arr1
270 return $ case marr2 of
271 Nothing -> D.Skip (Buffering s buf')
272 Just x -> D.Skip (Yielding buf' (Splitting s x))
273 D.Skip s -> return $ D.Skip (Buffering s buf)
274 D.Stop -> return $
275 if A.byteLength buf == 0
276 then D.Stop
277 else D.Skip (Yielding buf Finishing)
279 step' _ (Splitting st buf) = do
280 (arr1, marr2) <- A.breakOn byte buf
281 return $ case marr2 of
282 Nothing -> D.Skip $ Buffering st arr1
283 Just arr2 -> D.Skip $ Yielding arr1 (Splitting st arr2)
285 step' _ (Yielding arr next) = return $ D.Yield arr next
286 step' _ Finishing = return D.Stop
288-- XXX Remove MonadIO constraint.
289-- | Split a stream of arrays on a given separator byte, dropping the separator
290-- and coalescing all the arrays between two separators into a single array.
291--
292-- @since 0.7.0
293{-# INLINE splitOn #-}
294splitOn
295 :: (IsStream t, MonadIO m)
296 => Word8
297 -> t m (Array Word8)
298 -> t m (Array Word8)
299splitOn byte s =
300 fromStreamD $ D.splitInnerBy (A.breakOn byte) A.splice $ toStreamD s
302{-# INLINE splitOnSuffix #-}
303splitOnSuffix
304 :: (IsStream t, MonadIO m)
305 => Word8
306 -> t m (Array Word8)
307 -> t m (Array Word8)
308-- splitOn byte s = fromStreamD $ A.splitOn byte $ toStreamD s
309splitOnSuffix byte s =
310 fromStreamD $ D.splitInnerBySuffix (A.breakOn byte) A.splice $ toStreamD s
312-------------------------------------------------------------------------------
313-- Elimination - Running folds
314-------------------------------------------------------------------------------
316-- XXX This should be written using CPS (as foldK) if we want it to scale wrt
317-- to the number of times it can be called on the same stream.
318--
319{-# INLINE_NORMAL foldBreakD #-}
320foldBreakD :: forall m a b. (MonadIO m, Unboxed a) =>
321 Fold m a b -> D.Stream m (Array a) -> m (b, D.Stream m (Array a))
322foldBreakD (Fold fstep initial extract) stream@(D.Stream step state) = do
323 res <- initial
324 case res of
325 FL.Partial fs -> go SPEC state fs
326 FL.Done fb -> return $! (fb, stream)
328 where
330 {-# INLINE go #-}
331 go !_ st !fs = do
332 r <- step defState st
333 case r of
334 D.Yield (Array contents start end) s ->
335 let fp = Tuple' end contents
336 in goArray SPEC s fp start fs
337 D.Skip s -> go SPEC s fs
338 D.Stop -> do
339 b <- extract fs
340 return (b, D.nil)
342 goArray !_ s (Tuple' end _) !cur !fs
343 | cur == end = do
344 go SPEC s fs
345 goArray !_ st fp@(Tuple' end contents) !cur !fs = do
346 x <- liftIO $ peekWith contents cur
347 res <- fstep fs x
348 let next = INDEX_NEXT(cur,a)
349 case res of
350 FL.Done b -> do
351 let arr = Array contents next end
352 return $! (b, D.cons arr (D.Stream step st))
353 FL.Partial fs1 -> goArray SPEC st fp next fs1
355{-# INLINE_NORMAL foldBreakK #-}
356foldBreakK :: forall m a b. (MonadIO m, Unboxed a) =>
357 Fold m a b -> K.Stream m (Array a) -> m (b, K.Stream m (Array a))
358foldBreakK (Fold fstep initial extract) stream = do
359 res <- initial
360 case res of
361 FL.Partial fs -> go fs stream
362 FL.Done fb -> return (fb, stream)
364 where
366 {-# INLINE go #-}
367 go !fs st = do
368 let stop = (, K.nil) <$> extract fs
369 single a = yieldk a K.nil
370 yieldk (Array contents start end) r =
371 let fp = Tuple' end contents
372 in goArray fs r fp start
373 in K.foldStream defState yieldk single stop st
375 goArray !fs st (Tuple' end _) !cur
376 | cur == end = do
377 go fs st
378 goArray !fs st fp@(Tuple' end contents) !cur = do
379 x <- liftIO $ peekWith contents cur
380 res <- fstep fs x
381 let next = INDEX_NEXT(cur,a)
382 case res of
383 FL.Done b -> do
384 let arr = Array contents next end
385 return $! (b, K.cons arr st)
386 FL.Partial fs1 -> goArray fs1 st fp next
388-- | Fold an array stream using the supplied 'Fold'. Returns the fold result
389-- and the unconsumed stream.
390--
391-- > foldBreak f = runArrayFoldBreak (ArrayFold.fromFold f)
392--
393-- /Internal/
394--
395{-# INLINE_NORMAL foldBreak #-}
396foldBreak ::
397 (MonadIO m, Unboxed a)
398 => FL.Fold m a b
399 -> SerialT m (A.Array a)
400 -> m (b, SerialT m (A.Array a))
401-- foldBreak f s = fmap fromStreamD <$> foldBreakD f (toStreamD s)
402foldBreak f =
403 fmap (fmap Stream.fromStreamK)
404 . foldBreakK f
405 . Stream.toStreamK
406-- If foldBreak performs better than runArrayFoldBreak we can use a rewrite
407-- rule to rewrite runArrayFoldBreak to fold.
408-- foldBreak f = runArrayFoldBreak (ArrayFold.fromFold f)
410-------------------------------------------------------------------------------
411-- Fold to a single Array
412-------------------------------------------------------------------------------
414-- When we have to take an array partially, take the last part of the array.
415{-# INLINE takeArrayListRev #-}
416takeArrayListRev :: forall a. Unboxed a => Int -> [Array a] -> [Array a]
417takeArrayListRev = go
419 where
421 go _ [] = []
422 go n _ | n <= 0 = []
423 go n (x:xs) =
424 let len = Array.length x
425 in if n > len
426 then x : go (n - len) xs
427 else if n == len
428 then [x]
429 else let !(Array contents _ end) = x
430 !start = end - (n * SIZE_OF(a))
431 in [Array contents start end]
433-- When we have to take an array partially, take the last part of the array in
434-- the first split.
435{-# INLINE splitAtArrayListRev #-}
436splitAtArrayListRev ::
437 forall a. Unboxed a => Int -> [Array a] -> ([Array a],[Array a])
438splitAtArrayListRev n ls
439 | n <= 0 = ([], ls)
440 | otherwise = go n ls
441 where
442 go :: Int -> [Array a] -> ([Array a], [Array a])
443 go _ [] = ([], [])
444 go m (x:xs) =
445 let len = Array.length x
446 (xs', xs'') = go (m - len) xs
447 in if m > len
448 then (x:xs', xs'')
449 else if m == len
450 then ([x],xs)
451 else let !(Array contents start end) = x
452 end1 = end - (m * SIZE_OF(a))
453 arr2 = Array contents start end1
454 arr1 = Array contents end1 end
455 in ([arr1], arr2:xs)
457-------------------------------------------------------------------------------
458-- Fold to a single Array
459-------------------------------------------------------------------------------
461-- XXX Both of these implementations of splicing seem to perform equally well.
462-- We need to perform benchmarks over a range of sizes though.
464-- CAUTION! length must more than equal to lengths of all the arrays in the
465-- stream.
466{-# INLINE spliceArraysLenUnsafe #-}
467spliceArraysLenUnsafe :: (MonadIO m, Unboxed a)
468 => Int -> SerialT m (MA.Array a) -> m (MA.Array a)
469spliceArraysLenUnsafe len buffered = do
470 arr <- liftIO $ MA.newArray len
471 S.foldlM' MA.spliceUnsafe (return arr) buffered
473{-# INLINE _spliceArrays #-}
474_spliceArrays :: (MonadIO m, Unboxed a)
475 => SerialT m (Array a) -> m (Array a)
476_spliceArrays s = do
477 buffered <- S.foldr S.cons S.nil s
478 len <- S.sum (S.map Array.length buffered)
479 arr <- liftIO $ MA.newArray len
480 final <- S.foldlM' writeArr (return arr) s
481 return $ A.unsafeFreeze final
483 where
485 writeArr dst arr = MA.spliceUnsafe dst (A.unsafeThaw arr)
487{-# INLINE _spliceArraysBuffered #-}
488_spliceArraysBuffered :: (MonadIO m, Unboxed a)
489 => SerialT m (Array a) -> m (Array a)
490_spliceArraysBuffered s = do
491 buffered <- S.foldr S.cons S.nil s
492 len <- S.sum (S.map Array.length buffered)
493 A.unsafeFreeze <$> spliceArraysLenUnsafe len (S.map A.unsafeThaw s)
495{-# INLINE spliceArraysRealloced #-}
496spliceArraysRealloced :: forall m a. (MonadIO m, Unboxed a)
497 => SerialT m (Array a) -> m (Array a)
498spliceArraysRealloced s = do
499 let n = allocBytesToElemCount (undefined :: a) (4 * 1024)
500 idst = liftIO $ MA.newArray n
502 arr <- S.foldlM' MA.spliceExp idst (S.map A.unsafeThaw s)
503 liftIO $ A.unsafeFreeze <$> MA.rightSize arr
505-- XXX This should just be "fold A.write"
506--
507-- | Given a stream of arrays, splice them all together to generate a single
508-- array. The stream must be /finite/.
509--
510-- @since 0.7.0
511{-# INLINE toArray #-}
512toArray :: (MonadIO m, Unboxed a) => SerialT m (Array a) -> m (Array a)
513toArray = spliceArraysRealloced
514-- spliceArrays = _spliceArraysBuffered
516-- exponentially increasing sizes of the chunks upto the max limit.
517-- XXX this will be easier to implement with parsers/terminating folds
518-- With this we should be able to reduce the number of chunks/allocations.
519-- The reallocation/copy based toArray can also be implemented using this.
520--
521{-
522{-# INLINE toArraysInRange #-}
523toArraysInRange :: (IsStream t, MonadIO m, Unboxed a)
524 => Int -> Int -> Fold m (Array a) b -> Fold m a b
525toArraysInRange low high (Fold step initial extract) =
526-}
528{-
529-- | Fold the input to a pure buffered stream (List) of arrays.
530{-# INLINE _toArraysOf #-}
531_toArraysOf :: (MonadIO m, Unboxed a)
532 => Int -> Fold m a (SerialT Identity (Array a))
533_toArraysOf n = FL.chunksOf n (A.writeNF n) FL.toStream
534-}
536-------------------------------------------------------------------------------
537-- Elimination - running element parsers
538-------------------------------------------------------------------------------
540-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
541-- to make a newtype.
542{-# ANN type List NoSpecConstr #-}
543newtype List a = List {getList :: [a]}
545-- This can be generalized to any type provided it can be unfolded to a stream
546-- and it can be combined using a semigroup operation.
547--
548-- XXX This should be written using CPS (as parseK) if we want it to scale wrt
549-- to the number of times it can be called on the same stream.
550{-# INLINE_NORMAL parseBreakD #-}
551parseBreakD ::
552 forall m a b. (MonadIO m, MonadThrow m, Unboxed a)
553 => PRD.Parser m a b
554 -> D.Stream m (Array.Array a)
555 -> m (b, D.Stream m (Array.Array a))
556parseBreakD
557 (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
559 res <- initial
560 case res of
561 PRD.IPartial s -> go SPEC state (List []) s
562 PRD.IDone b -> return (b, stream)
563 PRD.IError err -> throwM $ ParseError err
565 where
567 -- "backBuf" contains last few items in the stream that we may have to
568 -- backtrack to.
569 --
570 -- XXX currently we are using a dumb list based approach for backtracking
571 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
572 -- That will allow us more efficient random back and forth movement.
573 go !_ st backBuf !pst = do
574 r <- step defState st
575 case r of
576 D.Yield (Array contents start end) s ->
577 gobuf SPEC s backBuf
578 (Tuple' end contents) start pst
579 D.Skip s -> go SPEC s backBuf pst
580 D.Stop -> do
581 b <- extract pst
582 return (b, D.nil)
584 -- Use strictness on "cur" to keep it unboxed
585 gobuf !_ s backBuf (Tuple' end _) !cur !pst
586 | cur == end = do
587 go SPEC s backBuf pst
588 gobuf !_ s backBuf fp@(Tuple' end contents) !cur !pst = do
589 x <- liftIO $ peekWith contents cur
590 pRes <- pstep pst x
591 let next = INDEX_NEXT(cur,a)
592 case pRes of
593 PR.Partial 0 pst1 ->
594 gobuf SPEC s (List []) fp next pst1
595 PR.Partial n pst1 -> do
596 assert (n <= Prelude.length (x:getList backBuf)) (return ())
597 let src0 = Prelude.take n (x:getList backBuf)
598 arr0 = A.fromListN n (Prelude.reverse src0)
599 arr1 = Array contents next end
600 src = arr0 <> arr1
601 let !(Array cont1 start end1) = src
602 fp1 = Tuple' end1 cont1
603 gobuf SPEC s (List []) fp1 start pst1
604 PR.Continue 0 pst1 ->
605 gobuf SPEC s (List (x:getList backBuf)) fp next pst1
606 PR.Continue n pst1 -> do
607 assert (n <= Prelude.length (x:getList backBuf)) (return ())
608 let (src0, buf1) = splitAt n (x:getList backBuf)
609 arr0 = A.fromListN n (Prelude.reverse src0)
610 arr1 = Array contents next end
611 src = arr0 <> arr1
612 let !(Array cont1 start end1) = src
613 fp1 = Tuple' end1 cont1
614 gobuf SPEC s (List buf1) fp1 start pst1
615 PR.Done 0 b -> do
616 let arr = Array contents next end
617 return (b, D.cons arr (D.Stream step s))
618 PR.Done n b -> do
619 assert (n <= Prelude.length (x:getList backBuf)) (return ())
620 let src0 = Prelude.take n (x:getList backBuf)
621 -- XXX create the array in reverse instead
622 arr0 = A.fromListN n (Prelude.reverse src0)
623 arr1 = Array contents next end
624 -- XXX Use StreamK to avoid adding arbitrary layers of
625 -- constructors every time.
626 str = D.cons arr0 (D.cons arr1 (D.Stream step s))
627 return (b, str)
628 PR.Error err -> throwM $ ParseError err
630{-# INLINE_NORMAL parseBreakK #-}
631parseBreakK ::
632 forall m a b. (MonadIO m, MonadThrow m, Unboxed a)
633 => PRD.Parser m a b
634 -> K.Stream m (Array.Array a)
635 -> m (b, K.Stream m (Array.Array a))
636parseBreakK (PRD.Parser pstep initial extract) stream = do
637 res <- initial
638 case res of
639 PRD.IPartial s -> go s stream []
640 PRD.IDone b -> return (b, stream)
641 PRD.IError err -> throwM $ ParseError err
643 where
645 -- "backBuf" contains last few items in the stream that we may have to
646 -- backtrack to.
647 --
648 -- XXX currently we are using a dumb list based approach for backtracking
649 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
650 -- That will allow us more efficient random back and forth movement.
651 go !pst st backBuf = do
652 let stop = (, K.nil) <$> extract pst
653 single a = yieldk a K.nil
654 yieldk (Array contents start end) r =
655 let fp = Tuple' end contents
656 in goArray pst backBuf r fp start
657 in K.foldStream defState yieldk single stop st
659 -- Use strictness on "cur" to keep it unboxed
660 goArray !pst backBuf st (Tuple' end _) !cur
661 | cur == end = do
662 go pst st backBuf
663 goArray !pst backBuf st fp@(Tuple' end contents) !cur = do
664 x <- liftIO $ peekWith contents cur
665 pRes <- pstep pst x
666 let next = INDEX_NEXT(cur,a)
667 case pRes of
668 PR.Partial 0 s ->
669 goArray s [] st fp next
670 PR.Partial n s -> do
671 assert (n <= Prelude.length (x:backBuf)) (return ())
672 let src0 = Prelude.take n (x:backBuf)
673 arr0 = A.fromListN n (Prelude.reverse src0)
674 arr1 = Array contents next end
675 src = arr0 <> arr1
676 let !(Array cont1 start end1) = src
677 fp1 = Tuple' end1 cont1
678 goArray s [] st fp1 start
679 PR.Continue 0 s ->
680 goArray s (x:backBuf) st fp next
681 PR.Continue n s -> do
682 assert (n <= Prelude.length (x:backBuf)) (return ())
683 let (src0, buf1) = splitAt n (x:backBuf)
684 arr0 = A.fromListN n (Prelude.reverse src0)
685 arr1 = Array contents next end
686 src = arr0 <> arr1
687 let !(Array cont1 start end1) = src
688 fp1 = Tuple' end1 cont1
689 goArray s buf1 st fp1 start
690 PR.Done 0 b -> do
691 let arr = Array contents next end
692 return (b, K.cons arr st)
693 PR.Done n b -> do
694 assert (n <= Prelude.length (x:backBuf)) (return ())
695 let src0 = Prelude.take n (x:backBuf)
696 -- XXX Use fromListRevN once implemented
697 -- arr0 = A.fromListRevN n src0
698 arr0 = A.fromListN n (Prelude.reverse src0)
699 arr1 = Array contents next end
700 str = K.cons arr0 (K.cons arr1 st)
701 return (b, str)
702 PR.Error err -> throwM $ ParseError err
704-- | Parse an array stream using the supplied 'Parser'. Returns the parse
705-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
706--
707-- /Internal/
708--
709{-# INLINE_NORMAL parseBreak #-}
710parseBreak ::
711 (MonadIO m, MonadThrow m, Unboxed a)
712 => PR.Parser m a b
713 -> SerialT m (A.Array a)
714 -> m (b, SerialT m (A.Array a))
715{-
716parseBreak p s =
717 fmap fromStreamD <$> parseBreakD (PRD.fromParserK p) (toStreamD s)
718-}
719parseBreak p =
720 fmap (fmap Stream.fromStreamK)
721 . parseBreakK (PRD.fromParserK p)
722 . Stream.toStreamK
724-------------------------------------------------------------------------------
725-- Elimination - Running Array Folds and parsers
726-------------------------------------------------------------------------------
728-- | Note that this is not the same as using a @Parser m (Array a) b@ with the
729-- regular "Streamly.Internal.Data.IsStream.parse" function. The regular parse
730-- would consume the input arrays as single unit. This parser parses in the way
731-- as described in the ArrayFold module. The input arrays are treated as @n@
732-- element units and can be consumed partially. The remaining elements are
733-- inserted in the source stream as an array.
734--
735{-# INLINE_NORMAL runArrayParserDBreak #-}
736runArrayParserDBreak ::
737 forall m a b. (MonadIO m, MonadThrow m, Unboxed a)
738 => PRD.Parser m (Array a) b
739 -> D.Stream m (Array.Array a)
740 -> m (b, D.Stream m (Array.Array a))
741runArrayParserDBreak
742 (PRD.Parser pstep initial extract)
743 stream@(D.Stream step state) = do
745 res <- initial
746 case res of
747 PRD.IPartial s -> go SPEC state (List []) s
748 PRD.IDone b -> return (b, stream)
749 PRD.IError err -> throwM $ ParseError err
751 where
753 -- "backBuf" contains last few items in the stream that we may have to
754 -- backtrack to.
755 --
756 -- XXX currently we are using a dumb list based approach for backtracking
757 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
758 -- That will allow us more efficient random back and forth movement.
759 go !_ st backBuf !pst = do
760 r <- step defState st
761 case r of
762 D.Yield x s -> gobuf SPEC [x] s backBuf pst
763 D.Skip s -> go SPEC s backBuf pst
764 D.Stop -> do
765 b <- extract pst
766 return (b, D.nil)
768 gobuf !_ [] s backBuf !pst = go SPEC s backBuf pst
769 gobuf !_ (x:xs) s backBuf !pst = do
770 pRes <- pstep pst x
771 case pRes of
772 PR.Partial 0 pst1 ->
773 gobuf SPEC xs s (List []) pst1
774 PR.Partial n pst1 -> do
775 assert
776 (n <= sum (map Array.length (x:getList backBuf)))
777 (return ())
778 let src0 = takeArrayListRev n (x:getList backBuf)
779 src = Prelude.reverse src0 ++ xs
780 gobuf SPEC src s (List []) pst1
781 PR.Continue 0 pst1 ->
782 gobuf SPEC xs s (List (x:getList backBuf)) pst1
783 PR.Continue n pst1 -> do
784 assert
785 (n <= sum (map Array.length (x:getList backBuf)))
786 (return ())
787 let (src0, buf1) = splitAtArrayListRev n (x:getList backBuf)
788 src = Prelude.reverse src0 ++ xs
789 gobuf SPEC src s (List buf1) pst1
790 PR.Done 0 b ->
791 return (b, D.Stream step s)
792 PR.Done n b -> do
793 assert
794 (n <= sum (map Array.length (x:getList backBuf)))
795 (return ())
796 let src0 = takeArrayListRev n (x:getList backBuf)
797 src = Prelude.reverse src0 ++ xs
798 return (b, D.append (D.fromList src) (D.Stream step s))
799 PR.Error err -> throwM $ ParseError err
801{-
802-- | Parse an array stream using the supplied 'Parser'. Returns the parse
803-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
804--
805-- /Internal/
806--
807{-# INLINE parseArr #-}
808parseArr ::
809 (MonadIO m, MonadThrow m, Unboxed a)
810 => ASF.Parser m a b
811 -> SerialT m (A.Array a)
812 -> m (b, SerialT m (A.Array a))
813parseArr p s = fmap fromStreamD <$> parseBreakD p (toStreamD s)
814-}
816-- | Fold an array stream using the supplied array stream 'Fold'.
817--
818-- /Pre-release/
819--
820{-# INLINE runArrayFold #-}
821runArrayFold :: (MonadIO m, MonadThrow m, Unboxed a) =>
822 ArrayFold m a b -> SerialT m (A.Array a) -> m b
823runArrayFold (ArrayFold p) s = fst <$> runArrayParserDBreak p (toStreamD s)
825-- | Like 'fold' but also returns the remaining stream.
826--
827-- /Pre-release/
828--
829{-# INLINE runArrayFoldBreak #-}
830runArrayFoldBreak :: (MonadIO m, MonadThrow m, Unboxed a) =>
831 ArrayFold m a b -> SerialT m (A.Array a) -> m (b, SerialT m (A.Array a))
832runArrayFoldBreak (ArrayFold p) s =
833 second fromStreamD <$> runArrayParserDBreak p (toStreamD s)
835{-# ANN type ParseChunksState Fuse #-}
836data ParseChunksState x inpBuf st pst =
837 ParseChunksInit inpBuf st
838 | ParseChunksInitLeftOver inpBuf
839 | ParseChunksStream st inpBuf !pst
840 | ParseChunksBuf inpBuf st inpBuf !pst
841 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
843{-# INLINE_NORMAL runArrayFoldManyD #-}
844runArrayFoldManyD
845 :: (MonadThrow m, Unboxed a)
846 => ArrayFold m a b
847 -> D.Stream m (Array a)
848 -> D.Stream m b
849runArrayFoldManyD
850 (ArrayFold (PRD.Parser pstep initial extract)) (D.Stream step state) =
852 D.Stream stepOuter (ParseChunksInit [] state)
854 where
856 {-# INLINE_LATE stepOuter #-}
857 -- Buffer is empty, get the first element from the stream, initialize the
858 -- fold and then go to stream processing loop.
859 stepOuter gst (ParseChunksInit [] st) = do
860 r <- step (adaptState gst) st
861 case r of
862 D.Yield x s -> do
863 res <- initial
864 case res of
865 PRD.IPartial ps ->
866 return $ D.Skip $ ParseChunksBuf [x] s [] ps
867 PRD.IDone pb ->
868 let next = ParseChunksInit [x] s
869 in return $ D.Skip $ ParseChunksYield pb next
870 PRD.IError err -> throwM $ ParseError err
871 D.Skip s -> return $ D.Skip $ ParseChunksInit [] s
872 D.Stop -> return D.Stop
874 -- Buffer is not empty, go to buffered processing loop
875 stepOuter _ (ParseChunksInit src st) = do
876 res <- initial
877 case res of
878 PRD.IPartial ps ->
879 return $ D.Skip $ ParseChunksBuf src st [] ps
880 PRD.IDone pb ->
881 let next = ParseChunksInit src st
882 in return $ D.Skip $ ParseChunksYield pb next
883 PRD.IError err -> throwM $ ParseError err
885 -- XXX we just discard any leftover input at the end
886 stepOuter _ (ParseChunksInitLeftOver _) = return D.Stop
888 -- Buffer is empty, process elements from the stream
889 stepOuter gst (ParseChunksStream st backBuf pst) = do
890 r <- step (adaptState gst) st
891 case r of
892 D.Yield x s -> do
893 pRes <- pstep pst x
894 case pRes of
895 PR.Partial 0 pst1 ->
896 return $ D.Skip $ ParseChunksStream s [] pst1
897 PR.Partial n pst1 -> do
898 assert
899 (n <= sum (map Array.length (x:backBuf)))
900 (return ())
901 let src0 = takeArrayListRev n (x:backBuf)
902 src = Prelude.reverse src0
903 return $ D.Skip $ ParseChunksBuf src s [] pst1
904 PR.Continue 0 pst1 ->
905 return $ D.Skip $ ParseChunksStream s (x:backBuf) pst1
906 PR.Continue n pst1 -> do
907 assert
908 (n <= sum (map Array.length (x:backBuf)))
909 (return ())
910 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
911 src = Prelude.reverse src0
912 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
913 PR.Done 0 b -> do
914 return $ D.Skip $
915 ParseChunksYield b (ParseChunksInit [] s)
916 PR.Done n b -> do
917 assert
918 (n <= sum (map Array.length (x:backBuf)))
919 (return ())
920 let src0 = takeArrayListRev n (x:backBuf)
921 let src = Prelude.reverse src0
922 return $ D.Skip $
923 ParseChunksYield b (ParseChunksInit src s)
924 PR.Error err -> throwM $ ParseError err
925 D.Skip s -> return $ D.Skip $ ParseChunksStream s backBuf pst
926 D.Stop -> do
927 b <- extract pst
928 let src = Prelude.reverse backBuf
929 return $ D.Skip $
930 ParseChunksYield b (ParseChunksInitLeftOver src)
932 -- go back to stream processing mode
933 stepOuter _ (ParseChunksBuf [] s buf pst) =
934 return $ D.Skip $ ParseChunksStream s buf pst
936 -- buffered processing loop
937 stepOuter _ (ParseChunksBuf (x:xs) s backBuf pst) = do
938 pRes <- pstep pst x
939 case pRes of
940 PR.Partial 0 pst1 ->
941 return $ D.Skip $ ParseChunksBuf xs s [] pst1
942 PR.Partial n pst1 -> do
943 assert (n <= sum (map Array.length (x:backBuf))) (return ())
944 let src0 = takeArrayListRev n (x:backBuf)
945 src = Prelude.reverse src0 ++ xs
946 return $ D.Skip $ ParseChunksBuf src s [] pst1
947 PR.Continue 0 pst1 ->
948 return $ D.Skip $ ParseChunksBuf xs s (x:backBuf) pst1
949 PR.Continue n pst1 -> do
950 assert (n <= sum (map Array.length (x:backBuf))) (return ())
951 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
952 src = Prelude.reverse src0 ++ xs
953 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
954 PR.Done 0 b ->
955 return $ D.Skip $ ParseChunksYield b (ParseChunksInit xs s)
956 PR.Done n b -> do
957 assert (n <= sum (map Array.length (x:backBuf))) (return ())
958 let src0 = takeArrayListRev n (x:backBuf)
959 src = Prelude.reverse src0 ++ xs
960 return $ D.Skip $ ParseChunksYield b (ParseChunksInit src s)
961 PR.Error err -> throwM $ ParseError err
963 stepOuter _ (ParseChunksYield a next) = return $ D.Yield a next
965-- | Apply an 'ArrayFold' repeatedly on an array stream and emit the
966-- fold outputs in the output stream.
967--
968-- See "Streamly.Prelude.foldMany" for more details.
969--
970-- /Pre-release/
971{-# INLINE runArrayFoldMany #-}
972runArrayFoldMany
973 :: (IsStream t, MonadThrow m, Unboxed a)
974 => ArrayFold m a b
975 -> t m (Array a)
976 -> t m b
977runArrayFoldMany p m = fromStreamD $ runArrayFoldManyD p (toStreamD m)