git.iko.soy/streamly
on 2ae9759 help sign in
streamly/src/Streamly/Internal/Data/Array/Stream/Foreign.hs
1152 lines 45 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 (Stream)
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 arrEnd)
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 -> Stream m (Array a) -> Stream 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 -> Stream m (A.Array a)
400 -> m (b, Stream 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 -> Stream 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 => Stream 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 => Stream 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 => Stream 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) => Stream 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 (Stream 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{-
546-- This can be generalized to any type provided it can be unfolded to a stream
547-- and it can be combined using a semigroup operation.
548--
549-- XXX This should be written using CPS (as parseK) if we want it to scale wrt
550-- to the number of times it can be called on the same stream.
551{-# INLINE_NORMAL parseBreakD #-}
552parseBreakD ::
553 forall m a b. (MonadIO m, MonadThrow m, Unboxed a)
554 => PRD.Parser m a b
555 -> D.Stream m (Array.Array a)
556 -> m (b, D.Stream m (Array.Array a))
557parseBreakD
558 (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
560 res <- initial
561 case res of
562 PRD.IPartial s -> go SPEC state (List []) s
563 PRD.IDone b -> return (b, stream)
564 PRD.IError err -> throwM $ ParseError err
566 where
568 -- "backBuf" contains last few items in the stream that we may have to
569 -- backtrack to.
570 --
571 -- XXX currently we are using a dumb list based approach for backtracking
572 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
573 -- That will allow us more efficient random back and forth movement.
574 go !_ st backBuf !pst = do
575 r <- step defState st
576 case r of
577 D.Yield (Array contents start end) s ->
578 gobuf SPEC s backBuf
579 (Tuple' end contents) start pst
580 D.Skip s -> go SPEC s backBuf pst
581 D.Stop -> do
582 b <- extract pst
583 return (b, D.nil)
585 -- Use strictness on "cur" to keep it unboxed
586 gobuf !_ s backBuf (Tuple' end _) !cur !pst
587 | cur == end = do
588 go SPEC s backBuf pst
589 gobuf !_ s backBuf fp@(Tuple' end contents) !cur !pst = do
590 x <- liftIO $ peekWith contents cur
591 pRes <- pstep pst x
592 let next = INDEX_NEXT(cur,a)
593 case pRes of
594 PR.Partial 0 pst1 ->
595 gobuf SPEC s (List []) fp next pst1
596 PR.Partial n pst1 -> do
597 assert (n <= Prelude.length (x:getList backBuf)) (return ())
598 let src0 = Prelude.take n (x:getList backBuf)
599 arr0 = A.fromListN n (Prelude.reverse src0)
600 arr1 = Array contents next end
601 src = arr0 <> arr1
602 let !(Array cont1 start end1) = src
603 fp1 = Tuple' end1 cont1
604 gobuf SPEC s (List []) fp1 start pst1
605 PR.Continue 0 pst1 ->
606 gobuf SPEC s (List (x:getList backBuf)) fp next pst1
607 PR.Continue n pst1 -> do
608 assert (n <= Prelude.length (x:getList backBuf)) (return ())
609 let (src0, buf1) = splitAt n (x:getList backBuf)
610 arr0 = A.fromListN n (Prelude.reverse src0)
611 arr1 = Array contents next end
612 src = arr0 <> arr1
613 let !(Array cont1 start end1) = src
614 fp1 = Tuple' end1 cont1
615 gobuf SPEC s (List buf1) fp1 start pst1
616 PR.Done 0 b -> do
617 let arr = Array contents next end
618 return (b, D.cons arr (D.Stream step s))
619 PR.Done n b -> do
620 assert (n <= Prelude.length (x:getList backBuf)) (return ())
621 let src0 = Prelude.take n (x:getList backBuf)
622 -- XXX create the array in reverse instead
623 arr0 = A.fromListN n (Prelude.reverse src0)
624 arr1 = Array contents next end
625 -- XXX Use StreamK to avoid adding arbitrary layers of
626 -- constructors every time.
627 str = D.cons arr0 (D.cons arr1 (D.Stream step s))
628 return (b, str)
629 PR.Error err -> throwM $ ParseError err
630-}
632{-# INLINE_NORMAL parseBreakK #-}
633parseBreakK ::
634 forall m a b. (MonadIO m, MonadThrow m, Unboxed a)
635 => PRD.Parser m a b
636 -> K.Stream m (Array.Array a)
637 -> m (b, K.Stream m (Array.Array a))
638parseBreakK (PRD.Parser pstep initial extract) stream = do
639 res <- initial
640 case res of
641 PRD.IPartial s -> go s stream []
642 PRD.IDone b -> return (b, stream)
643 PRD.IError err -> throwM $ ParseError err
645 where
647 -- "backBuf" contains last few items in the stream that we may have to
648 -- backtrack to.
649 --
650 -- XXX currently we are using a dumb list based approach for backtracking
651 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
652 -- That will allow us more efficient random back and forth movement.
653 go !pst st backBuf = do
654 let stop = goStop pst backBuf -- (, K.nil) <$> extract pst
655 single a = yieldk a K.nil
656 yieldk arr r = goArray pst backBuf r arr
657 in K.foldStream defState yieldk single stop st
659 -- Use strictness on "cur" to keep it unboxed
660 goArray !pst backBuf st (Array _ cur end) | cur == end = go pst st backBuf
661 goArray !pst backBuf st (Array contents cur end) = do
662 x <- liftIO $ peekWith contents cur
663 pRes <- pstep pst x
664 let next = INDEX_NEXT(cur,a)
665 case pRes of
666 PR.Partial 0 s ->
667 goArray s [] st (Array contents next end)
668 PR.Partial n s -> do
669 assert (n <= Prelude.length (x:backBuf)) (return ())
670 let src0 = Prelude.take n (x:backBuf)
671 arr0 = A.fromListN n (Prelude.reverse src0)
672 arr1 = Array contents next end
673 src = arr0 <> arr1
674 goArray s [] st src
675 PR.Continue 0 s ->
676 goArray s (x:backBuf) st (Array contents next end)
677 PR.Continue n s -> do
678 assert (n <= Prelude.length (x:backBuf)) (return ())
679 let (src0, buf1) = splitAt n (x:backBuf)
680 arr0 = A.fromListN n (Prelude.reverse src0)
681 arr1 = Array contents next end
682 src = arr0 <> arr1
683 goArray s buf1 st src
684 PR.Done 0 b -> do
685 let arr = Array contents next end
686 return (b, K.cons arr st)
687 PR.Done n b -> do
688 assert (n <= Prelude.length (x:backBuf)) (return ())
689 let src0 = Prelude.take n (x:backBuf)
690 -- XXX Use fromListRevN once implemented
691 -- arr0 = A.fromListRevN n src0
692 arr0 = A.fromListN n (Prelude.reverse src0)
693 arr1 = Array contents next end
694 str = K.cons arr0 (K.cons arr1 st)
695 return (b, str)
696 PR.Error err -> throwM $ ParseError err
698 -- This is a simplified goArray
699 goExtract !pst backBuf (Array _ cur end)
700 | cur == end = goStop pst backBuf
701 goExtract !pst backBuf (Array contents cur end) = do
702 x <- liftIO $ peekWith contents cur
703 pRes <- pstep pst x
704 let next = INDEX_NEXT(cur,a)
705 case pRes of
706 PR.Partial 0 s ->
707 goExtract s [] (Array contents next end)
708 PR.Partial n s -> do
709 assert (n <= Prelude.length (x:backBuf)) (return ())
710 let src0 = Prelude.take n (x:backBuf)
711 arr0 = A.fromListN n (Prelude.reverse src0)
712 arr1 = Array contents next end
713 src = arr0 <> arr1
714 goExtract s [] src
715 PR.Continue 0 s ->
716 goExtract s backBuf (Array contents next end)
717 PR.Continue n s -> do
718 assert (n <= Prelude.length (x:backBuf)) (return ())
719 let (src0, buf1) = splitAt n (x:backBuf)
720 arr0 = A.fromListN n (Prelude.reverse src0)
721 arr1 = Array contents next end
722 src = arr0 <> arr1
723 goExtract s buf1 src
724 PR.Done 0 b -> do
725 let arr = Array contents next end
726 return (b, K.fromPure arr)
727 PR.Done n b -> do
728 assert (n <= Prelude.length backBuf) (return ())
729 let src0 = Prelude.take n backBuf
730 -- XXX Use fromListRevN once implemented
731 -- arr0 = A.fromListRevN n src0
732 arr0 = A.fromListN n (Prelude.reverse src0)
733 arr1 = Array contents next end
734 str = K.cons arr0 (K.fromPure arr1)
735 return (b, str)
736 PR.Error err -> throwM $ ParseError err
738 -- This is a simplified goExtract
739 {-# INLINE goStop #-}
740 goStop !pst backBuf = do
741 pRes <- extract pst
742 case pRes of
743 PR.Partial _ _ -> error "Bug: parseBreak: Partial in extract"
744 PR.Continue 0 _ ->
745 error "parseBreak: extract, Continue 0 creates infinite loop"
746 PR.Continue n s -> do
747 assert (n <= Prelude.length backBuf) (return ())
748 let (src0, buf1) = splitAt n backBuf
749 arr = A.fromListN n (Prelude.reverse src0)
750 goExtract s buf1 arr
751 PR.Done 0 b ->
752 return (b, K.nil)
753 PR.Done n b -> do
754 assert (n <= Prelude.length backBuf) (return ())
755 let src0 = Prelude.take n backBuf
756 -- XXX Use fromListRevN once implemented
757 -- arr0 = A.fromListRevN n src0
758 arr0 = A.fromListN n (Prelude.reverse src0)
759 return (b, K.fromPure arr0)
760 PR.Error err -> throwM $ ParseError err
762-- | Parse an array stream using the supplied 'Parser'. Returns the parse
763-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
764--
765-- /Internal/
766--
767{-# INLINE_NORMAL parseBreak #-}
768parseBreak ::
769 (MonadIO m, MonadThrow m, Unboxed a)
770 => PR.Parser m a b
771 -> Stream m (A.Array a)
772 -> m (b, Stream m (A.Array a))
773{-
774parseBreak p s =
775 fmap fromStreamD <$> parseBreakD (PRD.fromParserK p) (toStreamD s)
776-}
777parseBreak p =
778 fmap (fmap Stream.fromStreamK)
779 . parseBreakK (PRD.fromParserK p)
780 . Stream.toStreamK
782-------------------------------------------------------------------------------
783-- Elimination - Running Array Folds and parsers
784-------------------------------------------------------------------------------
786-- | Note that this is not the same as using a @Parser m (Array a) b@ with the
787-- regular "Streamly.Internal.Data.IsStream.parse" function. The regular parse
788-- would consume the input arrays as single unit. This parser parses in the way
789-- as described in the ArrayFold module. The input arrays are treated as @n@
790-- element units and can be consumed partially. The remaining elements are
791-- inserted in the source stream as an array.
792--
793{-# INLINE_NORMAL runArrayParserDBreak #-}
794runArrayParserDBreak ::
795 forall m a b. (MonadIO m, MonadThrow m, Unboxed a)
796 => PRD.Parser m (Array a) b
797 -> D.Stream m (Array.Array a)
798 -> m (b, D.Stream m (Array.Array a))
799runArrayParserDBreak
800 (PRD.Parser pstep initial extract)
801 stream@(D.Stream step state) = do
803 res <- initial
804 case res of
805 PRD.IPartial s -> go SPEC state (List []) s
806 PRD.IDone b -> return (b, stream)
807 PRD.IError err -> throwM $ ParseError err
809 where
811 -- "backBuf" contains last few items in the stream that we may have to
812 -- backtrack to.
813 --
814 -- XXX currently we are using a dumb list based approach for backtracking
815 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
816 -- That will allow us more efficient random back and forth movement.
817 go _ st backBuf !pst = do
818 r <- step defState st
819 case r of
820 D.Yield x s -> gobuf SPEC [x] s backBuf pst
821 D.Skip s -> go SPEC s backBuf pst
822 D.Stop -> goStop backBuf pst
824 gobuf !_ [] s backBuf !pst = go SPEC s backBuf pst
825 gobuf !_ (x:xs) s backBuf !pst = do
826 pRes <- pstep pst x
827 case pRes of
828 PR.Partial 0 pst1 ->
829 gobuf SPEC xs s (List []) pst1
830 PR.Partial n pst1 -> do
831 assert
832 (n <= sum (map Array.length (x:getList backBuf)))
833 (return ())
834 let src0 = takeArrayListRev n (x:getList backBuf)
835 src = Prelude.reverse src0 ++ xs
836 gobuf SPEC src s (List []) pst1
837 PR.Continue 0 pst1 ->
838 gobuf SPEC xs s (List (x:getList backBuf)) pst1
839 PR.Continue n pst1 -> do
840 assert
841 (n <= sum (map Array.length (x:getList backBuf)))
842 (return ())
843 let (src0, buf1) = splitAtArrayListRev n (x:getList backBuf)
844 src = Prelude.reverse src0 ++ xs
845 gobuf SPEC src s (List buf1) pst1
846 PR.Done 0 b ->
847 return (b, D.Stream step s)
848 PR.Done n b -> do
849 assert
850 (n <= sum (map Array.length (x:getList backBuf)))
851 (return ())
852 let src0 = takeArrayListRev n (x:getList backBuf)
853 src = Prelude.reverse src0 ++ xs
854 return (b, D.append (D.fromList src) (D.Stream step s))
855 PR.Error err -> throwM $ ParseError err
857 -- This is a simplified gobuf
858 goExtract _ [] backBuf !pst = goStop backBuf pst
859 goExtract _ (x:xs) backBuf !pst = do
860 pRes <- pstep pst x
861 case pRes of
862 PR.Partial 0 pst1 ->
863 goExtract SPEC xs (List []) pst1
864 PR.Partial n pst1 -> do
865 assert
866 (n <= sum (map Array.length (x:getList backBuf)))
867 (return ())
868 let src0 = takeArrayListRev n (x:getList backBuf)
869 src = Prelude.reverse src0 ++ xs
870 goExtract SPEC src (List []) pst1
871 PR.Continue 0 pst1 ->
872 goExtract SPEC xs (List (x:getList backBuf)) pst1
873 PR.Continue n pst1 -> do
874 assert
875 (n <= sum (map Array.length (x:getList backBuf)))
876 (return ())
877 let (src0, buf1) = splitAtArrayListRev n (x:getList backBuf)
878 src = Prelude.reverse src0 ++ xs
879 goExtract SPEC src (List buf1) pst1
880 PR.Done 0 b ->
881 return (b, D.nil)
882 PR.Done n b -> do
883 assert
884 (n <= sum (map Array.length (x:getList backBuf)))
885 (return ())
886 let src0 = takeArrayListRev n (x:getList backBuf)
887 src = Prelude.reverse src0 ++ xs
888 return (b, D.fromList src)
889 PR.Error err -> throwM $ ParseError err
891 -- This is a simplified goExtract
892 {-# INLINE goStop #-}
893 goStop backBuf pst = do
894 pRes <- extract pst
895 case pRes of
896 PR.Partial _ _ -> error "Bug: runArrayParserDBreak: Partial in extract"
897 PR.Continue 0 _ ->
898 error "runArrayParserDBreak: extract, Continue 0 creates infinite loop"
899 PR.Continue n pst1 -> do
900 assert
901 (n <= sum (map Array.length (getList backBuf)))
902 (return ())
903 let (src0, buf1) = splitAtArrayListRev n (getList backBuf)
904 src = Prelude.reverse src0
905 goExtract SPEC src (List buf1) pst1
906 PR.Done 0 b -> return (b, D.nil)
907 PR.Done n b -> do
908 assert
909 (n <= sum (map Array.length (getList backBuf)))
910 (return ())
911 let src0 = takeArrayListRev n (getList backBuf)
912 src = Prelude.reverse src0
913 return (b, D.fromList src)
914 PR.Error err -> throwM $ ParseError err
916{-
917-- | Parse an array stream using the supplied 'Parser'. Returns the parse
918-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
919--
920-- /Internal/
921--
922{-# INLINE parseArr #-}
923parseArr ::
924 (MonadIO m, MonadThrow m, Unboxed a)
925 => ASF.Parser m a b
926 -> Stream m (A.Array a)
927 -> m (b, Stream m (A.Array a))
928parseArr p s = fmap fromStreamD <$> parseBreakD p (toStreamD s)
929-}
931-- | Fold an array stream using the supplied array stream 'Fold'.
932--
933-- /Pre-release/
934--
935{-# INLINE runArrayFold #-}
936runArrayFold :: (MonadIO m, MonadThrow m, Unboxed a) =>
937 ArrayFold m a b -> Stream m (A.Array a) -> m b
938runArrayFold (ArrayFold p) s = fst <$> runArrayParserDBreak p (toStreamD s)
940-- | Like 'fold' but also returns the remaining stream.
941--
942-- /Pre-release/
943--
944{-# INLINE runArrayFoldBreak #-}
945runArrayFoldBreak :: (MonadIO m, MonadThrow m, Unboxed a) =>
946 ArrayFold m a b -> Stream m (A.Array a) -> m (b, Stream m (A.Array a))
947runArrayFoldBreak (ArrayFold p) s =
948 second fromStreamD <$> runArrayParserDBreak p (toStreamD s)
950{-# ANN type ParseChunksState Fuse #-}
951data ParseChunksState x inpBuf st pst =
952 ParseChunksInit inpBuf st
953 | ParseChunksInitBuf inpBuf
954 | ParseChunksInitLeftOver inpBuf
955 | ParseChunksStream st inpBuf !pst
956 | ParseChunksStop inpBuf !pst
957 | ParseChunksBuf inpBuf st inpBuf !pst
958 | ParseChunksExtract inpBuf inpBuf !pst
959 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
961{-# INLINE_NORMAL runArrayFoldManyD #-}
962runArrayFoldManyD
963 :: (MonadThrow m, Unboxed a)
964 => ArrayFold m a b
965 -> D.Stream m (Array a)
966 -> D.Stream m b
967runArrayFoldManyD
968 (ArrayFold (PRD.Parser pstep initial extract)) (D.Stream step state) =
970 D.Stream stepOuter (ParseChunksInit [] state)
972 where
974 {-# INLINE_LATE stepOuter #-}
975 -- Buffer is empty, get the first element from the stream, initialize the
976 -- fold and then go to stream processing loop.
977 stepOuter gst (ParseChunksInit [] st) = do
978 r <- step (adaptState gst) st
979 case r of
980 D.Yield x s -> do
981 res <- initial
982 case res of
983 PRD.IPartial ps ->
984 return $ D.Skip $ ParseChunksBuf [x] s [] ps
985 PRD.IDone pb ->
986 let next = ParseChunksInit [x] s
987 in return $ D.Skip $ ParseChunksYield pb next
988 PRD.IError err -> throwM $ ParseError err
989 D.Skip s -> return $ D.Skip $ ParseChunksInit [] s
990 D.Stop -> return D.Stop
992 -- Buffer is not empty, go to buffered processing loop
993 stepOuter _ (ParseChunksInit src st) = do
994 res <- initial
995 case res of
996 PRD.IPartial ps ->
997 return $ D.Skip $ ParseChunksBuf src st [] ps
998 PRD.IDone pb ->
999 let next = ParseChunksInit src st
1000 in return $ D.Skip $ ParseChunksYield pb next
1001 PRD.IError err -> throwM $ ParseError err
1003 -- This is a simplified ParseChunksInit
1004 stepOuter _ (ParseChunksInitBuf src) = do
1005 res <- initial
1006 case res of
1007 PRD.IPartial ps ->
1008 return $ D.Skip $ ParseChunksExtract src [] ps
1009 PRD.IDone pb ->
1010 let next = ParseChunksInitBuf src
1011 in return $ D.Skip $ ParseChunksYield pb next
1012 PRD.IError err -> throwM $ ParseError err
1014 -- XXX we just discard any leftover input at the end
1015 stepOuter _ (ParseChunksInitLeftOver _) = return D.Stop
1017 -- Buffer is empty, process elements from the stream
1018 stepOuter gst (ParseChunksStream st backBuf pst) = do
1019 r <- step (adaptState gst) st
1020 case r of
1021 D.Yield x s -> do
1022 pRes <- pstep pst x
1023 case pRes of
1024 PR.Partial 0 pst1 ->
1025 return $ D.Skip $ ParseChunksStream s [] pst1
1026 PR.Partial n pst1 -> do
1027 assert
1028 (n <= sum (map Array.length (x:backBuf)))
1029 (return ())
1030 let src0 = takeArrayListRev n (x:backBuf)
1031 src = Prelude.reverse src0
1032 return $ D.Skip $ ParseChunksBuf src s [] pst1
1033 PR.Continue 0 pst1 ->
1034 return $ D.Skip $ ParseChunksStream s (x:backBuf) pst1
1035 PR.Continue n pst1 -> do
1036 assert
1037 (n <= sum (map Array.length (x:backBuf)))
1038 (return ())
1039 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
1040 src = Prelude.reverse src0
1041 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
1042 PR.Done 0 b -> do
1043 return $ D.Skip $
1044 ParseChunksYield b (ParseChunksInit [] s)
1045 PR.Done n b -> do
1046 assert
1047 (n <= sum (map Array.length (x:backBuf)))
1048 (return ())
1049 let src0 = takeArrayListRev n (x:backBuf)
1050 let src = Prelude.reverse src0
1051 return $ D.Skip $
1052 ParseChunksYield b (ParseChunksInit src s)
1053 PR.Error err -> throwM $ ParseError err
1054 D.Skip s -> return $ D.Skip $ ParseChunksStream s backBuf pst
1055 D.Stop -> return $ D.Skip $ ParseChunksStop backBuf pst
1057 -- go back to stream processing mode
1058 stepOuter _ (ParseChunksBuf [] s buf pst) =
1059 return $ D.Skip $ ParseChunksStream s buf pst
1061 -- buffered processing loop
1062 stepOuter _ (ParseChunksBuf (x:xs) s backBuf pst) = do
1063 pRes <- pstep pst x
1064 case pRes of
1065 PR.Partial 0 pst1 ->
1066 return $ D.Skip $ ParseChunksBuf xs s [] pst1
1067 PR.Partial n pst1 -> do
1068 assert (n <= sum (map Array.length (x:backBuf))) (return ())
1069 let src0 = takeArrayListRev n (x:backBuf)
1070 src = Prelude.reverse src0 ++ xs
1071 return $ D.Skip $ ParseChunksBuf src s [] pst1
1072 PR.Continue 0 pst1 ->
1073 return $ D.Skip $ ParseChunksBuf xs s (x:backBuf) pst1
1074 PR.Continue n pst1 -> do
1075 assert (n <= sum (map Array.length (x:backBuf))) (return ())
1076 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
1077 src = Prelude.reverse src0 ++ xs
1078 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
1079 PR.Done 0 b ->
1080 return $ D.Skip $ ParseChunksYield b (ParseChunksInit xs s)
1081 PR.Done n b -> do
1082 assert (n <= sum (map Array.length (x:backBuf))) (return ())
1083 let src0 = takeArrayListRev n (x:backBuf)
1084 src = Prelude.reverse src0 ++ xs
1085 return $ D.Skip $ ParseChunksYield b (ParseChunksInit src s)
1086 PR.Error err -> throwM $ ParseError err
1088 -- This is a simplified ParseChunksBuf
1089 stepOuter _ (ParseChunksExtract [] buf pst) =
1090 return $ D.Skip $ ParseChunksStop buf pst
1092 stepOuter _ (ParseChunksExtract (x:xs) backBuf pst) = do
1093 pRes <- pstep pst x
1094 case pRes of
1095 PR.Partial 0 pst1 ->
1096 return $ D.Skip $ ParseChunksExtract xs [] pst1
1097 PR.Partial n pst1 -> do
1098 assert (n <= sum (map Array.length (x:backBuf))) (return ())
1099 let src0 = takeArrayListRev n (x:backBuf)
1100 src = Prelude.reverse src0 ++ xs
1101 return $ D.Skip $ ParseChunksExtract src [] pst1
1102 PR.Continue 0 pst1 ->
1103 return $ D.Skip $ ParseChunksExtract xs (x:backBuf) pst1
1104 PR.Continue n pst1 -> do
1105 assert (n <= sum (map Array.length (x:backBuf))) (return ())
1106 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
1107 src = Prelude.reverse src0 ++ xs
1108 return $ D.Skip $ ParseChunksExtract src buf1 pst1
1109 PR.Done 0 b ->
1110 return $ D.Skip $ ParseChunksYield b (ParseChunksInitBuf xs)
1111 PR.Done n b -> do
1112 assert (n <= sum (map Array.length (x:backBuf))) (return ())
1113 let src0 = takeArrayListRev n (x:backBuf)
1114 src = Prelude.reverse src0 ++ xs
1115 return $ D.Skip $ ParseChunksYield b (ParseChunksInitBuf src)
1116 PR.Error err -> throwM $ ParseError err
1118 -- This is a simplified ParseChunksExtract
1119 stepOuter _ (ParseChunksStop backBuf pst) = do
1120 pRes <- extract pst
1121 case pRes of
1122 PR.Partial _ _ -> error "runArrayFoldManyD: Partial in extract"
1123 PR.Continue 0 _ -> error "runArrayFoldManyD: Continue 0 in extract"
1124 PR.Continue n pst1 -> do
1125 assert (n <= sum (map Array.length backBuf)) (return ())
1126 let (src0, buf1) = splitAtArrayListRev n backBuf
1127 src = Prelude.reverse src0
1128 return $ D.Skip $ ParseChunksExtract src buf1 pst1
1129 PR.Done 0 b ->
1130 return $ D.Skip $ ParseChunksYield b (ParseChunksInitLeftOver [])
1131 PR.Done n b -> do
1132 assert (n <= sum (map Array.length backBuf)) (return ())
1133 let src0 = takeArrayListRev n backBuf
1134 src = Prelude.reverse src0
1135 return $ D.Skip $ ParseChunksYield b (ParseChunksInitBuf src)
1136 PR.Error err -> throwM $ ParseError err
1138 stepOuter _ (ParseChunksYield a next) = return $ D.Yield a next
1140-- | Apply an 'ArrayFold' repeatedly on an array stream and emit the
1141-- fold outputs in the output stream.
1143-- See "Streamly.Prelude.foldMany" for more details.
1145-- /Pre-release/
1146{-# INLINE runArrayFoldMany #-}
1147runArrayFoldMany
1148 :: (IsStream t, MonadThrow m, Unboxed a)
1149 => ArrayFold m a b
1150 -> t m (Array a)
1151 -> t m b
1152runArrayFoldMany p m = fromStreamD $ runArrayFoldManyD p (toStreamD m)