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