git.iko.soy/streamly
on 01643fe help sign in
streamly/src/Streamly/Internal/Data/Array/Stream/Foreign.hs
860 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.IsStream.Type
76 (IsStream, fromStreamD, toStreamD)
77import Streamly.Internal.Data.SVar (adaptState, defState)
78import Streamly.Internal.Data.Array.Foreign.Mut.Type (memcpy, bytesToElemCount)
79import Streamly.Internal.System.IO (mkChunkSizeKB)
81import qualified Streamly.Internal.Data.Array.Foreign as A
82import qualified Streamly.Internal.Data.Array.Foreign as Array
83import qualified Streamly.Internal.Data.Array.Foreign.Type as A
84import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MA
85import qualified Streamly.Internal.Data.Array.Stream.Mut.Foreign as AS
86import qualified Streamly.Internal.Data.Array.Stream.Fold.Foreign as ASF
87import qualified Streamly.Internal.Data.Fold.Type as FL
88import qualified Streamly.Internal.Data.Parser as PR
89import qualified Streamly.Internal.Data.Parser.ParserD as PRD
90import qualified Streamly.Internal.Data.Stream.IsStream as S
91import qualified Streamly.Internal.Data.Stream.StreamD as D
93-------------------------------------------------------------------------------
94-- Generation
95-------------------------------------------------------------------------------
97-- | @arraysOf n stream@ groups the elements in the input stream into arrays of
98-- @n@ elements each.
99--
100-- > arraysOf n = Stream.chunksOf n (Array.writeN n)
101--
102-- /Pre-release/
103{-# INLINE arraysOf #-}
104arraysOf :: (IsStream t, MonadIO m, Storable a)
105 => Int -> t m a -> t m (Array a)
106arraysOf n str = fromStreamD $ A.arraysOf n (toStreamD str)
108-------------------------------------------------------------------------------
109-- Append
110-------------------------------------------------------------------------------
112-- XXX efficiently compare two streams of arrays. Two streams can have chunks
113-- of different sizes, we can handle that in the stream comparison abstraction.
114-- This could be useful e.g. to fast compare whether two files differ.
116-- | Convert a stream of arrays into a stream of their elements.
117--
118-- Same as the following but more efficient:
119--
120-- > concat = Stream.unfoldMany Array.read
121--
122-- @since 0.7.0
123{-# INLINE concat #-}
124concat :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
125-- concat m = fromStreamD $ A.flattenArrays (toStreamD m)
126-- concat m = fromStreamD $ D.concatMap A.toStreamD (toStreamD m)
127concat m = fromStreamD $ D.unfoldMany A.read (toStreamD m)
129-- | Convert a stream of arrays into a stream of their elements reversing the
130-- contents of each array before flattening.
131--
132-- > concatRev = Stream.unfoldMany Array.readRev
133--
134-- @since 0.7.0
135{-# INLINE concatRev #-}
136concatRev :: (IsStream t, MonadIO m, Storable a) => t m (Array a) -> t m a
137concatRev m = fromStreamD $ A.flattenArraysRev (toStreamD m)
139-------------------------------------------------------------------------------
140-- Intersperse and append
141-------------------------------------------------------------------------------
143-- | Flatten a stream of arrays after inserting the given element between
144-- arrays.
145--
146-- /Pre-release/
147{-# INLINE interpose #-}
148interpose :: (MonadIO m, IsStream t, Storable a) => a -> t m (Array a) -> t m a
149interpose x = S.interpose x A.read
151{-# INLINE intercalateSuffix #-}
152intercalateSuffix :: (MonadIO m, IsStream t, Storable a)
153 => Array a -> t m (Array a) -> t m a
154intercalateSuffix = S.intercalateSuffix A.read
156-- | Flatten a stream of arrays appending the given element after each
157-- array.
158--
159-- @since 0.7.0
160{-# INLINE interposeSuffix #-}
161interposeSuffix :: (MonadIO m, IsStream t, Storable a)
162 => a -> t m (Array a) -> t m a
163-- interposeSuffix x = fromStreamD . A.unlines x . toStreamD
164interposeSuffix x = S.interposeSuffix x A.read
166data FlattenState s a =
167 OuterLoop s
168 | InnerLoop s !(ForeignPtr a) !(Ptr a) !(Ptr a)
170{-# INLINE_NORMAL unlines #-}
171unlines :: forall m a. (MonadIO m, Storable a)
172 => a -> D.Stream m (Array a) -> D.Stream m a
173unlines sep (D.Stream step state) = D.Stream step' (OuterLoop state)
174 where
175 {-# INLINE_LATE step' #-}
176 step' gst (OuterLoop st) = do
177 r <- step (adaptState gst) st
178 return $ case r of
179 D.Yield Array{..} s ->
180 let p = unsafeForeignPtrToPtr aStart
181 in D.Skip (InnerLoop s aStart p aEnd)
182 D.Skip s -> D.Skip (OuterLoop s)
183 D.Stop -> D.Stop
185 step' _ (InnerLoop st _ p end) | p == end =
186 return $ D.Yield sep $ OuterLoop st
188 step' _ (InnerLoop st startf p end) = do
189 x <- liftIO $ do
190 r <- peek p
191 touchForeignPtr startf
192 return r
193 return $ D.Yield x (InnerLoop st startf
194 (p `plusPtr` sizeOf (undefined :: a)) end)
196-------------------------------------------------------------------------------
197-- Compact
198-------------------------------------------------------------------------------
200-- XXX These would not be needed once we implement compactLEFold, see
201-- module Streamly.Internal.Data.Array.Stream.Mut.Foreign
202--
203{-# INLINE_NORMAL packArraysChunksOf #-}
204packArraysChunksOf :: (MonadIO m, Storable a)
205 => Int -> D.Stream m (Array a) -> D.Stream m (Array a)
206packArraysChunksOf n str =
207 D.map A.unsafeFreeze $ AS.packArraysChunksOf n $ D.map A.unsafeThaw str
209-- XXX instead of writing two different versions of this operation, we should
210-- write it as a pipe.
211{-# INLINE_NORMAL lpackArraysChunksOf #-}
212lpackArraysChunksOf :: (MonadIO m, Storable a)
213 => Int -> Fold m (Array a) () -> Fold m (Array a) ()
214lpackArraysChunksOf n fld =
215 FL.lmap A.unsafeThaw $ AS.lpackArraysChunksOf n (FL.lmap A.unsafeFreeze fld)
217-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
218-- maximum specified size in bytes.
219--
220-- @since 0.7.0
221{-# INLINE compact #-}
222compact :: (MonadIO m, Storable a)
223 => Int -> SerialT m (Array a) -> SerialT m (Array a)
224compact n xs = fromStreamD $ packArraysChunksOf n (toStreamD xs)
226-------------------------------------------------------------------------------
227-- Split
228-------------------------------------------------------------------------------
230data SplitState s arr
231 = Initial s
232 | Buffering s arr
233 | Splitting s arr
234 | Yielding arr (SplitState s arr)
235 | Finishing
237-- | Split a stream of arrays on a given separator byte, dropping the separator
238-- and coalescing all the arrays between two separators into a single array.
239--
240-- @since 0.7.0
241{-# INLINE_NORMAL _splitOn #-}
242_splitOn
243 :: MonadIO m
244 => Word8
245 -> D.Stream m (Array Word8)
246 -> D.Stream m (Array Word8)
247_splitOn byte (D.Stream step state) = D.Stream step' (Initial state)
249 where
251 {-# INLINE_LATE step' #-}
252 step' gst (Initial st) = do
253 r <- step gst st
254 case r of
255 D.Yield arr s -> do
256 (arr1, marr2) <- A.breakOn byte arr
257 return $ case marr2 of
258 Nothing -> D.Skip (Buffering s arr1)
259 Just arr2 -> D.Skip (Yielding arr1 (Splitting s arr2))
260 D.Skip s -> return $ D.Skip (Initial s)
261 D.Stop -> return D.Stop
263 step' gst (Buffering st buf) = do
264 r <- step gst st
265 case r of
266 D.Yield arr s -> do
267 (arr1, marr2) <- A.breakOn byte arr
268 buf' <- A.spliceTwo buf arr1
269 return $ case marr2 of
270 Nothing -> D.Skip (Buffering s buf')
271 Just x -> D.Skip (Yielding buf' (Splitting s x))
272 D.Skip s -> return $ D.Skip (Buffering s buf)
273 D.Stop -> return $
274 if A.byteLength buf == 0
275 then D.Stop
276 else D.Skip (Yielding buf Finishing)
278 step' _ (Splitting st buf) = do
279 (arr1, marr2) <- A.breakOn byte buf
280 return $ case marr2 of
281 Nothing -> D.Skip $ Buffering st arr1
282 Just arr2 -> D.Skip $ Yielding arr1 (Splitting st arr2)
284 step' _ (Yielding arr next) = return $ D.Yield arr next
285 step' _ Finishing = return D.Stop
287-- | Split a stream of arrays on a given separator byte, dropping the separator
288-- and coalescing all the arrays between two separators into a single array.
289--
290-- @since 0.7.0
291{-# INLINE splitOn #-}
292splitOn
293 :: (IsStream t, MonadIO m)
294 => Word8
295 -> t m (Array Word8)
296 -> t m (Array Word8)
297splitOn byte s =
298 fromStreamD $ D.splitInnerBy (A.breakOn byte) A.spliceTwo $ toStreamD s
300{-# INLINE splitOnSuffix #-}
301splitOnSuffix
302 :: (IsStream t, MonadIO m)
303 => Word8
304 -> t m (Array Word8)
305 -> t m (Array Word8)
306-- splitOn byte s = fromStreamD $ A.splitOn byte $ toStreamD s
307splitOnSuffix byte s =
308 fromStreamD $ D.splitInnerBySuffix (A.breakOn byte) A.spliceTwo $ toStreamD s
310-------------------------------------------------------------------------------
311-- Elimination - Running folds
312-------------------------------------------------------------------------------
314-- XXX This should be written using CPS (as foldK) if we want it to scale wrt
315-- to the number of times it can be called on the same stream.
316--
317{-# INLINE_NORMAL foldD #-}
318foldD :: forall m a b. (MonadIO m, Storable a) =>
319 Fold m a b -> D.Stream m (Array a) -> m (b, D.Stream m (Array a))
320foldD (Fold fstep initial extract) stream@(D.Stream step state) = do
321 res <- initial
322 case res of
323 FL.Partial fs -> go SPEC state fs
324 FL.Done fb -> return $! (fb, stream)
326 where
328 {-# INLINE go #-}
329 go !_ st !fs = do
330 r <- step defState st
331 case r of
332 D.Yield (Array (ForeignPtr start contents) (Ptr end)) s ->
333 goArray SPEC s (ForeignPtr end contents) (Ptr start) fs
334 D.Skip s -> go SPEC s fs
335 D.Stop -> do
336 b <- extract fs
337 return (b, D.nil)
339 goArray !_ s fp@(ForeignPtr end _) !cur !fs
340 | cur == Ptr end = do
341 liftIO $ touchForeignPtr fp
342 go SPEC s fs
343 goArray !_ st fp@(ForeignPtr end contents) !cur !fs = do
344 x <- liftIO $ peek cur
345 res <- fstep fs x
346 let elemSize = sizeOf (undefined :: a)
347 next = cur `plusPtr` elemSize
348 case res of
349 FL.Done b -> do
350 let !(Ptr curAddr) = next
351 arr = Array (ForeignPtr curAddr contents) (Ptr end)
352 return $! (b, D.cons arr (D.Stream step st))
353 FL.Partial fs1 -> goArray SPEC st fp next fs1
355-- | Fold an array stream using the supplied 'Fold'. Returns the fold result
356-- and the unconsumed stream.
357--
358-- /Internal/
359--
360{-# INLINE_NORMAL fold #-}
361fold ::
362 (MonadIO m, Storable a)
363 => FL.Fold m a b
364 -> SerialT m (A.Array a)
365 -> m (b, SerialT m (A.Array a))
366fold f s = fmap fromStreamD <$> foldD f (toStreamD s)
368-------------------------------------------------------------------------------
369-- Fold to a single Array
370-------------------------------------------------------------------------------
372-- When we have to take an array partially, take the last part of the array.
373{-# INLINE takeArrayListRev #-}
374takeArrayListRev :: forall a. Storable a => Int -> [Array a] -> [Array a]
375takeArrayListRev = go
377 where
379 go _ [] = []
380 go n _ | n <= 0 = []
381 go n (x:xs) =
382 let len = Array.length x
383 in if n > len
384 then x : go (n - len) xs
385 else if n == len
386 then [x]
387 else let !(Array (ForeignPtr _ contents) end) = x
388 sz = sizeOf (undefined :: a)
389 !(Ptr start) = end `plusPtr` negate (n * sz)
390 in [Array (ForeignPtr start contents) end]
392-- When we have to take an array partially, take the last part of the array in
393-- the first split.
394{-# INLINE splitAtArrayListRev #-}
395splitAtArrayListRev :: forall a. Storable a =>
396 Int -> [Array a] -> ([Array a],[Array a])
397splitAtArrayListRev n ls
398 | n <= 0 = ([], ls)
399 | otherwise = go n ls
400 where
401 go :: Int -> [Array a] -> ([Array a], [Array a])
402 go _ [] = ([], [])
403 go m (x:xs) =
404 let len = Array.length x
405 (xs', xs'') = go (m - len) xs
406 in if m > len
407 then (x:xs', xs'')
408 else if m == len
409 then ([x],xs)
410 else let !(Array (ForeignPtr start contents) end) = x
411 sz = sizeOf (undefined :: a)
412 end1 = end `plusPtr` negate (m * sz)
413 arr2 = Array (ForeignPtr start contents) end1
414 !(Ptr addrEnd1) = end1
415 arr1 = Array (ForeignPtr addrEnd1 contents) end
416 in ([arr1], arr2:xs)
418-------------------------------------------------------------------------------
419-- Fold to a single Array
420-------------------------------------------------------------------------------
422-- XXX Both of these implementations of splicing seem to perform equally well.
423-- We need to perform benchmarks over a range of sizes though.
425-- CAUTION! length must more than equal to lengths of all the arrays in the
426-- stream.
427{-# INLINE spliceArraysLenUnsafe #-}
428spliceArraysLenUnsafe :: (MonadIO m, Storable a)
429 => Int -> SerialT m (MA.Array a) -> m (MA.Array a)
430spliceArraysLenUnsafe len buffered = do
431 arr <- liftIO $ MA.newArray len
432 end <- S.foldlM' writeArr (return $ MA.aEnd arr) buffered
433 return $ arr {MA.aEnd = end}
435 where
437 writeArr dst (MA.Array as ae _) =
438 liftIO $ unsafeWithForeignPtr as $ \src -> do
439 let count = ae `minusPtr` src
440 memcpy (castPtr dst) (castPtr src) count
441 return $ dst `plusPtr` count
443{-# INLINE _spliceArrays #-}
444_spliceArrays :: (MonadIO m, Storable a)
445 => SerialT m (Array a) -> m (Array a)
446_spliceArrays s = do
447 buffered <- S.foldr S.cons S.nil s
448 len <- S.sum (S.map Array.length buffered)
449 arr <- liftIO $ MA.newArray len
450 end <- S.foldlM' writeArr (return $ MA.aEnd arr) s
451 return $ A.unsafeFreeze $ arr {MA.aEnd = end}
453 where
455 writeArr dst (Array as ae) =
456 liftIO $ unsafeWithForeignPtr as $ \src -> do
457 let count = ae `minusPtr` src
458 memcpy (castPtr dst) (castPtr src) count
459 return $ dst `plusPtr` count
461{-# INLINE _spliceArraysBuffered #-}
462_spliceArraysBuffered :: (MonadIO m, Storable a)
463 => SerialT m (Array a) -> m (Array a)
464_spliceArraysBuffered s = do
465 buffered <- S.foldr S.cons S.nil s
466 len <- S.sum (S.map Array.length buffered)
467 A.unsafeFreeze <$> spliceArraysLenUnsafe len (S.map A.unsafeThaw s)
469{-# INLINE spliceArraysRealloced #-}
470spliceArraysRealloced :: forall m a. (MonadIO m, Storable a)
471 => SerialT m (Array a) -> m (Array a)
472spliceArraysRealloced s = do
473 let idst = liftIO $ MA.newArray (bytesToElemCount (undefined :: a)
474 (mkChunkSizeKB 4))
476 arr <- S.foldlM' MA.spliceWithDoubling idst (S.map A.unsafeThaw s)
477 liftIO $ A.unsafeFreeze <$> MA.shrinkToFit arr
479-- XXX This should just be "fold A.write"
480--
481-- | Given a stream of arrays, splice them all together to generate a single
482-- array. The stream must be /finite/.
483--
484-- @since 0.7.0
485{-# INLINE toArray #-}
486toArray :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)
487toArray = spliceArraysRealloced
488-- spliceArrays = _spliceArraysBuffered
490-- exponentially increasing sizes of the chunks upto the max limit.
491-- XXX this will be easier to implement with parsers/terminating folds
492-- With this we should be able to reduce the number of chunks/allocations.
493-- The reallocation/copy based toArray can also be implemented using this.
494--
495{-
496{-# INLINE toArraysInRange #-}
497toArraysInRange :: (IsStream t, MonadIO m, Storable a)
498 => Int -> Int -> Fold m (Array a) b -> Fold m a b
499toArraysInRange low high (Fold step initial extract) =
500-}
502{-
503-- | Fold the input to a pure buffered stream (List) of arrays.
504{-# INLINE _toArraysOf #-}
505_toArraysOf :: (MonadIO m, Storable a)
506 => Int -> Fold m a (SerialT Identity (Array a))
507_toArraysOf n = FL.chunksOf n (A.writeNF n) FL.toStream
508-}
510-------------------------------------------------------------------------------
511-- Elimination - running element parsers
512-------------------------------------------------------------------------------
514-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
515-- to make a newtype.
516{-# ANN type List NoSpecConstr #-}
517newtype List a = List {getList :: [a]}
519-- This can be generalized to any type provided it can be unfolded to a stream
520-- and it can be combined using a semigroup operation.
521--
522-- XXX This should be written using CPS (as parseK) if we want it to scale wrt
523-- to the number of times it can be called on the same stream.
524{-# INLINE_NORMAL parseD #-}
525parseD ::
526 forall m a b. (MonadIO m, MonadThrow m, Storable a)
527 => PRD.Parser m a b
528 -> D.Stream m (Array.Array a)
529 -> m (b, D.Stream m (Array.Array a))
530parseD (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
531 res <- initial
532 case res of
533 PRD.IPartial s -> go SPEC state (List []) s
534 PRD.IDone b -> return (b, stream)
535 PRD.IError err -> throwM $ ParseError err
537 where
539 -- "backBuf" contains last few items in the stream that we may have to
540 -- backtrack to.
541 --
542 -- XXX currently we are using a dumb list based approach for backtracking
543 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
544 -- That will allow us more efficient random back and forth movement.
545 {-# INLINE go #-}
546 go !_ st backBuf !pst = do
547 r <- step defState st
548 case r of
549 D.Yield (Array (ForeignPtr start contents) (Ptr end)) s ->
550 gobuf SPEC s backBuf (ForeignPtr end contents) (Ptr start) pst
551 D.Skip s -> go SPEC s backBuf pst
552 D.Stop -> do
553 b <- extract pst
554 return (b, D.nil)
556 -- Use strictness on "cur" to keep it unboxed
557 gobuf !_ s backBuf fp@(ForeignPtr end _) !cur !pst
558 | cur == Ptr end = do
559 liftIO $ touchForeignPtr fp
560 go SPEC s backBuf pst
561 gobuf !_ s backBuf fp@(ForeignPtr end contents) !cur !pst = do
562 x <- liftIO $ peek cur
563 pRes <- pstep pst x
564 let elemSize = sizeOf (undefined :: a)
565 next = cur `plusPtr` elemSize
566 case pRes of
567 PR.Partial 0 pst1 ->
568 gobuf SPEC s (List []) fp next pst1
569 PR.Partial n pst1 -> do
570 assert (n <= Prelude.length (x:getList backBuf)) (return ())
571 let src0 = Prelude.take n (x:getList backBuf)
572 arr0 = A.fromListN n (Prelude.reverse src0)
573 !(Ptr curAddr) = next
574 arr1 = Array (ForeignPtr curAddr contents) (Ptr end)
575 src = arr0 <> arr1
576 let !(Array (ForeignPtr start cont1) (Ptr end1)) = src
577 gobuf SPEC s (List []) (ForeignPtr end1 cont1) (Ptr start) pst1
578 PR.Continue 0 pst1 ->
579 gobuf SPEC s (List (x:getList backBuf)) fp next pst1
580 PR.Continue n pst1 -> do
581 assert (n <= Prelude.length (x:getList backBuf)) (return ())
582 let (src0, buf1) = splitAt n (x:getList backBuf)
583 arr0 = A.fromListN n (Prelude.reverse src0)
584 !(Ptr curAddr) = next
585 arr1 = Array (ForeignPtr curAddr contents) (Ptr end)
586 src = arr0 <> arr1
587 let !(Array (ForeignPtr start cont1) (Ptr end1)) = src
588 gobuf SPEC s (List buf1) (ForeignPtr end1 cont1) (Ptr start) pst1
589 PR.Done 0 b -> do
590 let !(Ptr curAddr) = next
591 arr = Array (ForeignPtr curAddr contents) (Ptr end)
592 return (b, D.cons arr (D.Stream step s))
593 PR.Done n b -> do
594 assert (n <= Prelude.length (x:getList backBuf)) (return ())
595 let src0 = Prelude.take n (x:getList backBuf)
596 -- XXX create the array in reverse instead
597 arr0 = A.fromListN n (Prelude.reverse src0)
598 !(Ptr curAddr) = next
599 arr1 = Array (ForeignPtr curAddr contents) (Ptr end)
600 -- XXX Use StreamK to avoid adding arbitrary layers of
601 -- constructors every time.
602 str = D.cons arr0 (D.cons arr1 (D.Stream step s))
603 return (b, str)
604 PR.Error err -> throwM $ ParseError err
606-- | Parse an array stream using the supplied 'Parser'. Returns the parse
607-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
608--
609-- /Internal/
610--
611{-# INLINE_NORMAL parse #-}
612parse ::
613 (MonadIO m, MonadThrow m, Storable a)
614 => PRD.Parser m a b
615 -> SerialT m (A.Array a)
616 -> m (b, SerialT m (A.Array a))
617parse p s = fmap fromStreamD <$> parseD p (toStreamD s)
619-------------------------------------------------------------------------------
620-- Elimination - Running Array Folds and parsers
621-------------------------------------------------------------------------------
623{-# INLINE_NORMAL parseArrD #-}
624parseArrD ::
625 forall m a b. (MonadIO m, MonadThrow m, Storable a)
626 => PRD.Parser m (Array a) b
627 -> D.Stream m (Array.Array a)
628 -> m (b, D.Stream m (Array.Array a))
629parseArrD (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
630 res <- initial
631 case res of
632 PRD.IPartial s -> go SPEC state (List []) s
633 PRD.IDone b -> return (b, stream)
634 PRD.IError err -> throwM $ ParseError err
636 where
638 -- "backBuf" contains last few items in the stream that we may have to
639 -- backtrack to.
640 --
641 -- XXX currently we are using a dumb list based approach for backtracking
642 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
643 -- That will allow us more efficient random back and forth movement.
644 {-# INLINE go #-}
645 go !_ st backBuf !pst = do
646 r <- step defState st
647 case r of
648 D.Yield x s -> gobuf SPEC [x] s backBuf pst
649 D.Skip s -> go SPEC s backBuf pst
650 D.Stop -> do
651 b <- extract pst
652 return (b, D.nil)
654 gobuf !_ [] s backBuf !pst = go SPEC s backBuf pst
655 gobuf !_ (x:xs) s backBuf !pst = do
656 pRes <- pstep pst x
657 case pRes of
658 PR.Partial 0 pst1 ->
659 gobuf SPEC xs s (List []) pst1
660 PR.Partial n pst1 -> do
661 assert
662 (n <= sum (map Array.length (x:getList backBuf)))
663 (return ())
664 let src0 = takeArrayListRev n (x:getList backBuf)
665 src = Prelude.reverse src0 ++ xs
666 gobuf SPEC src s (List []) pst1
667 PR.Continue 0 pst1 ->
668 gobuf SPEC xs s (List (x:getList backBuf)) pst1
669 PR.Continue n pst1 -> do
670 assert
671 (n <= sum (map Array.length (x:getList backBuf)))
672 (return ())
673 let (src0, buf1) = splitAtArrayListRev n (x:getList backBuf)
674 src = Prelude.reverse src0 ++ xs
675 gobuf SPEC src s (List buf1) pst1
676 PR.Done 0 b ->
677 return (b, D.Stream step s)
678 PR.Done n b -> do
679 assert
680 (n <= sum (map Array.length (x:getList backBuf)))
681 (return ())
682 let src0 = takeArrayListRev n (x:getList backBuf)
683 src = Prelude.reverse src0 ++ xs
684 return (b, D.append (D.fromList src) (D.Stream step s))
685 PR.Error err -> throwM $ ParseError err
687{-
688-- | Parse an array stream using the supplied 'Parser'. Returns the parse
689-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
690--
691-- /Internal/
692--
693{-# INLINE parseArr #-}
694parseArr ::
695 (MonadIO m, MonadThrow m, Storable a)
696 => ASF.Parser m a b
697 -> SerialT m (A.Array a)
698 -> m (b, SerialT m (A.Array a))
699parseArr p s = fmap fromStreamD <$> parseD p (toStreamD s)
700-}
702-- | Fold an array stream using the supplied array stream 'Fold'.
703--
704-- /Pre-release/
705--
706{-# INLINE foldArr #-}
707foldArr :: (MonadIO m, MonadThrow m, Storable a) =>
708 ASF.Fold m a b -> SerialT m (A.Array a) -> m b
709foldArr (ASF.Fold p) s = fst <$> parseArrD p (toStreamD s)
711-- | Like 'fold' but also returns the remaining stream.
712--
713-- /Pre-release/
714--
715{-# INLINE foldArr_ #-}
716foldArr_ :: (MonadIO m, MonadThrow m, Storable a) =>
717 ASF.Fold m a b -> SerialT m (A.Array a) -> m (b, SerialT m (A.Array a))
718foldArr_ (ASF.Fold p) s = second fromStreamD <$> parseArrD p (toStreamD s)
720{-# ANN type ParseChunksState Fuse #-}
721data ParseChunksState x inpBuf st pst =
722 ParseChunksInit inpBuf st
723 | ParseChunksInitLeftOver inpBuf
724 | ParseChunksStream st inpBuf !pst
725 | ParseChunksBuf inpBuf st inpBuf !pst
726 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
728{-# INLINE_NORMAL foldArrManyD #-}
729foldArrManyD
730 :: (MonadThrow m, Storable a)
731 => ASF.Fold m a b
732 -> D.Stream m (Array a)
733 -> D.Stream m b
734foldArrManyD (ASF.Fold (PRD.Parser pstep initial extract)) (D.Stream step state) =
735 D.Stream stepOuter (ParseChunksInit [] state)
737 where
739 {-# INLINE_LATE stepOuter #-}
740 -- Buffer is empty, get the first element from the stream, initialize the
741 -- fold and then go to stream processing loop.
742 stepOuter gst (ParseChunksInit [] st) = do
743 r <- step (adaptState gst) st
744 case r of
745 D.Yield x s -> do
746 res <- initial
747 case res of
748 PRD.IPartial ps ->
749 return $ D.Skip $ ParseChunksBuf [x] s [] ps
750 PRD.IDone pb ->
751 let next = ParseChunksInit [x] s
752 in return $ D.Skip $ ParseChunksYield pb next
753 PRD.IError err -> throwM $ ParseError err
754 D.Skip s -> return $ D.Skip $ ParseChunksInit [] s
755 D.Stop -> return D.Stop
757 -- Buffer is not empty, go to buffered processing loop
758 stepOuter _ (ParseChunksInit src st) = do
759 res <- initial
760 case res of
761 PRD.IPartial ps ->
762 return $ D.Skip $ ParseChunksBuf src st [] ps
763 PRD.IDone pb ->
764 let next = ParseChunksInit src st
765 in return $ D.Skip $ ParseChunksYield pb next
766 PRD.IError err -> throwM $ ParseError err
768 -- XXX we just discard any leftover input at the end
769 stepOuter _ (ParseChunksInitLeftOver _) = return D.Stop
771 -- Buffer is empty, process elements from the stream
772 stepOuter gst (ParseChunksStream st backBuf pst) = do
773 r <- step (adaptState gst) st
774 case r of
775 D.Yield x s -> do
776 pRes <- pstep pst x
777 case pRes of
778 PR.Partial 0 pst1 ->
779 return $ D.Skip $ ParseChunksStream s [] pst1
780 PR.Partial n pst1 -> do
781 assert
782 (n <= sum (map Array.length (x:backBuf)))
783 (return ())
784 let src0 = takeArrayListRev n (x:backBuf)
785 src = Prelude.reverse src0
786 return $ D.Skip $ ParseChunksBuf src s [] pst1
787 PR.Continue 0 pst1 ->
788 return $ D.Skip $ ParseChunksStream s (x:backBuf) pst1
789 PR.Continue n pst1 -> do
790 assert
791 (n <= sum (map Array.length (x:backBuf)))
792 (return ())
793 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
794 src = Prelude.reverse src0
795 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
796 PR.Done 0 b -> do
797 return $ D.Skip $
798 ParseChunksYield b (ParseChunksInit [] s)
799 PR.Done n b -> do
800 assert
801 (n <= sum (map Array.length (x:backBuf)))
802 (return ())
803 let src0 = takeArrayListRev n (x:backBuf)
804 let src = Prelude.reverse src0
805 return $ D.Skip $
806 ParseChunksYield b (ParseChunksInit src s)
807 PR.Error err -> throwM $ ParseError err
808 D.Skip s -> return $ D.Skip $ ParseChunksStream s backBuf pst
809 D.Stop -> do
810 b <- extract pst
811 let src = Prelude.reverse backBuf
812 return $ D.Skip $
813 ParseChunksYield b (ParseChunksInitLeftOver src)
815 -- go back to stream processing mode
816 stepOuter _ (ParseChunksBuf [] s buf pst) =
817 return $ D.Skip $ ParseChunksStream s buf pst
819 -- buffered processing loop
820 stepOuter _ (ParseChunksBuf (x:xs) s backBuf pst) = do
821 pRes <- pstep pst x
822 case pRes of
823 PR.Partial 0 pst1 ->
824 return $ D.Skip $ ParseChunksBuf xs s [] pst1
825 PR.Partial n pst1 -> do
826 assert (n <= sum (map Array.length (x:backBuf))) (return ())
827 let src0 = takeArrayListRev n (x:backBuf)
828 src = Prelude.reverse src0 ++ xs
829 return $ D.Skip $ ParseChunksBuf src s [] pst1
830 PR.Continue 0 pst1 ->
831 return $ D.Skip $ ParseChunksBuf xs s (x:backBuf) pst1
832 PR.Continue n pst1 -> do
833 assert (n <= sum (map Array.length (x:backBuf))) (return ())
834 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
835 src = Prelude.reverse src0 ++ xs
836 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
837 PR.Done 0 b ->
838 return $ D.Skip $ ParseChunksYield b (ParseChunksInit xs s)
839 PR.Done n b -> do
840 assert (n <= sum (map Array.length (x:backBuf))) (return ())
841 let src0 = takeArrayListRev n (x:backBuf)
842 src = Prelude.reverse src0 ++ xs
843 return $ D.Skip $ ParseChunksYield b (ParseChunksInit src s)
844 PR.Error err -> throwM $ ParseError err
846 stepOuter _ (ParseChunksYield a next) = return $ D.Yield a next
848-- | Apply an array stream 'Fold' repeatedly on an array stream and emit the
849-- fold outputs in the output stream.
850--
851-- See "Streamly.Prelude.foldMany" for more details.
852--
853-- /Pre-release/
854{-# INLINE foldArrMany #-}
855foldArrMany
856 :: (IsStream t, MonadThrow m, Storable a)
857 => ASF.Fold m a b
858 -> t m (Array a)
859 -> t m b
860foldArrMany p m = fromStreamD $ foldArrManyD p (toStreamD m)