git.iko.soy/streamly
on a322c07 help sign in
streamly/src/Streamly/Internal/Data/Array/Stream/Foreign.hs
879 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{-# INLINE_NORMAL packArraysChunksOf #-}
207packArraysChunksOf :: (MonadIO m, Storable a)
208 => Int -> D.Stream m (Array a) -> D.Stream m (Array a)
209packArraysChunksOf n str =
210 D.map A.unsafeFreeze $ AS.packArraysChunksOf n $ D.map A.unsafeThaw str
212-- XXX instead of writing two different versions of this operation, we should
213-- write it as a pipe.
214{-# INLINE_NORMAL lpackArraysChunksOf #-}
215lpackArraysChunksOf :: (MonadIO m, Storable a)
216 => Int -> Fold m (Array a) () -> Fold m (Array a) ()
217lpackArraysChunksOf n fld =
218 FL.map A.unsafeThaw $ AS.lpackArraysChunksOf n (FL.map A.unsafeFreeze fld)
220#if !defined(mingw32_HOST_OS)
222-- | @groupIOVecsOf maxBytes maxEntries@ groups arrays in the incoming stream
223-- to create a stream of 'IOVec' arrays with a maximum of @maxBytes@ bytes in
224-- each array and a maximum of @maxEntries@ entries in each array.
225--
226-- @since 0.7.0
227{-# INLINE_NORMAL groupIOVecsOf #-}
228groupIOVecsOf :: MonadIO m
229 => Int -> Int -> D.Stream m (Array a) -> D.Stream m (Array IOVec)
230groupIOVecsOf n maxIOVLen str =
231 D.map A.unsafeFreeze
232 $ AS.groupIOVecsOf n maxIOVLen
233 $ D.map A.unsafeThaw str
234#endif
236-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
237-- maximum specified size in bytes.
238--
239-- @since 0.7.0
240{-# INLINE compact #-}
241compact :: (MonadIO m, Storable a)
242 => Int -> SerialT m (Array a) -> SerialT m (Array a)
243compact n xs = D.fromStreamD $ packArraysChunksOf n (D.toStreamD xs)
245-------------------------------------------------------------------------------
246-- Split
247-------------------------------------------------------------------------------
249data SplitState s arr
250 = Initial s
251 | Buffering s arr
252 | Splitting s arr
253 | Yielding arr (SplitState s arr)
254 | Finishing
256-- | Split a stream of arrays on a given separator byte, dropping the separator
257-- and coalescing all the arrays between two separators into a single array.
258--
259-- @since 0.7.0
260{-# INLINE_NORMAL _splitOn #-}
261_splitOn
262 :: MonadIO m
263 => Word8
264 -> D.Stream m (Array Word8)
265 -> D.Stream m (Array Word8)
266_splitOn byte (D.Stream step state) = D.Stream step' (Initial state)
268 where
270 {-# INLINE_LATE step' #-}
271 step' gst (Initial st) = do
272 r <- step gst st
273 case r of
274 D.Yield arr s -> do
275 (arr1, marr2) <- A.breakOn byte arr
276 return $ case marr2 of
277 Nothing -> D.Skip (Buffering s arr1)
278 Just arr2 -> D.Skip (Yielding arr1 (Splitting s arr2))
279 D.Skip s -> return $ D.Skip (Initial s)
280 D.Stop -> return D.Stop
282 step' gst (Buffering st buf) = do
283 r <- step gst st
284 case r of
285 D.Yield arr s -> do
286 (arr1, marr2) <- A.breakOn byte arr
287 buf' <- A.spliceTwo buf arr1
288 return $ case marr2 of
289 Nothing -> D.Skip (Buffering s buf')
290 Just x -> D.Skip (Yielding buf' (Splitting s x))
291 D.Skip s -> return $ D.Skip (Buffering s buf)
292 D.Stop -> return $
293 if A.byteLength buf == 0
294 then D.Stop
295 else D.Skip (Yielding buf Finishing)
297 step' _ (Splitting st buf) = do
298 (arr1, marr2) <- A.breakOn byte buf
299 return $ case marr2 of
300 Nothing -> D.Skip $ Buffering st arr1
301 Just arr2 -> D.Skip $ Yielding arr1 (Splitting st arr2)
303 step' _ (Yielding arr next) = return $ D.Yield arr next
304 step' _ Finishing = return D.Stop
306-- | Split a stream of arrays on a given separator byte, dropping the separator
307-- and coalescing all the arrays between two separators into a single array.
308--
309-- @since 0.7.0
310{-# INLINE splitOn #-}
311splitOn
312 :: (IsStream t, MonadIO m)
313 => Word8
314 -> t m (Array Word8)
315 -> t m (Array Word8)
316splitOn byte s =
317 D.fromStreamD $ D.splitInnerBy (A.breakOn byte) A.spliceTwo $ D.toStreamD s
319{-# INLINE splitOnSuffix #-}
320splitOnSuffix
321 :: (IsStream t, MonadIO m)
322 => Word8
323 -> t m (Array Word8)
324 -> t m (Array Word8)
325-- splitOn byte s = D.fromStreamD $ A.splitOn byte $ D.toStreamD s
326splitOnSuffix byte s =
327 D.fromStreamD $ D.splitInnerBySuffix (A.breakOn byte) A.spliceTwo $ D.toStreamD s
329-------------------------------------------------------------------------------
330-- Elimination - Running folds
331-------------------------------------------------------------------------------
333-- XXX This should be written using CPS (as foldK) if we want it to scale wrt
334-- to the number of times it can be called on the same stream.
335--
336{-# INLINE_NORMAL foldD #-}
337foldD :: forall m a b. (MonadIO m, Storable a) =>
338 Fold m a b -> D.Stream m (Array a) -> m (b, D.Stream m (Array a))
339foldD (Fold fstep initial extract) stream@(D.Stream step state) = do
340 res <- initial
341 case res of
342 FL.Partial fs -> go SPEC state fs
343 FL.Done fb -> return $! (fb, stream)
345 where
347 {-# INLINE go #-}
348 go !_ st !fs = do
349 r <- step defState st
350 case r of
351 D.Yield (Array (ForeignPtr start contents) (Ptr end)) s ->
352 goArray SPEC s (ForeignPtr end contents) (Ptr start) fs
353 D.Skip s -> go SPEC s fs
354 D.Stop -> do
355 b <- extract fs
356 return (b, D.nil)
358 goArray !_ s fp@(ForeignPtr end _) !cur !fs
359 | cur == Ptr end = do
360 liftIO $ touchForeignPtr fp
361 go SPEC s fs
362 goArray !_ st fp@(ForeignPtr end contents) !cur !fs = do
363 x <- liftIO $ peek cur
364 res <- fstep fs x
365 let elemSize = sizeOf (undefined :: a)
366 next = cur `plusPtr` elemSize
367 case res of
368 FL.Done b -> do
369 let !(Ptr curAddr) = next
370 arr = Array (ForeignPtr curAddr contents) (Ptr end)
371 return $! (b, D.cons arr (D.Stream step st))
372 FL.Partial fs1 -> goArray SPEC st fp next fs1
374-- | Fold an array stream using the supplied 'Fold'. Returns the fold result
375-- and the unconsumed stream.
376--
377-- /Internal/
378--
379{-# INLINE_NORMAL fold #-}
380fold ::
381 (MonadIO m, Storable a)
382 => FL.Fold m a b
383 -> SerialT m (A.Array a)
384 -> m (b, SerialT m (A.Array a))
385fold f s = fmap D.fromStreamD <$> foldD f (D.toStreamD s)
387-------------------------------------------------------------------------------
388-- Fold to a single Array
389-------------------------------------------------------------------------------
391-- When we have to take an array partially, take the last part of the array.
392{-# INLINE takeArrayListRev #-}
393takeArrayListRev :: forall a. Storable a => Int -> [Array a] -> [Array a]
394takeArrayListRev = go
396 where
398 go _ [] = []
399 go n _ | n <= 0 = []
400 go n (x:xs) =
401 let len = Array.length x
402 in if n > len
403 then x : go (n - len) xs
404 else if n == len
405 then [x]
406 else let !(Array (ForeignPtr _ contents) end) = x
407 sz = sizeOf (undefined :: a)
408 !(Ptr start) = end `plusPtr` negate (n * sz)
409 in [Array (ForeignPtr start contents) end]
411-- When we have to take an array partially, take the last part of the array in
412-- the first split.
413{-# INLINE splitAtArrayListRev #-}
414splitAtArrayListRev :: forall a. Storable a =>
415 Int -> [Array a] -> ([Array a],[Array a])
416splitAtArrayListRev n ls
417 | n <= 0 = ([], ls)
418 | otherwise = go n ls
419 where
420 go :: Int -> [Array a] -> ([Array a], [Array a])
421 go _ [] = ([], [])
422 go m (x:xs) =
423 let len = Array.length x
424 (xs', xs'') = go (m - len) xs
425 in if m > len
426 then (x:xs', xs'')
427 else if m == len
428 then ([x],xs)
429 else let !(Array (ForeignPtr start contents) end) = x
430 sz = sizeOf (undefined :: a)
431 end1 = end `plusPtr` negate (m * sz)
432 arr2 = Array (ForeignPtr start contents) end1
433 !(Ptr addrEnd1) = end1
434 arr1 = Array (ForeignPtr addrEnd1 contents) end
435 in ([arr1], arr2:xs)
437-------------------------------------------------------------------------------
438-- Fold to a single Array
439-------------------------------------------------------------------------------
441-- XXX Both of these implementations of splicing seem to perform equally well.
442-- We need to perform benchmarks over a range of sizes though.
444-- CAUTION! length must more than equal to lengths of all the arrays in the
445-- stream.
446{-# INLINE spliceArraysLenUnsafe #-}
447spliceArraysLenUnsafe :: (MonadIO m, Storable a)
448 => Int -> SerialT m (MA.Array a) -> m (MA.Array a)
449spliceArraysLenUnsafe len buffered = do
450 arr <- liftIO $ MA.newArray len
451 end <- S.foldlM' writeArr (return $ MA.aEnd arr) buffered
452 return $ arr {MA.aEnd = end}
454 where
456 writeArr dst (MA.Array as ae _) =
457 liftIO $ unsafeWithForeignPtr as $ \src -> do
458 let count = ae `minusPtr` src
459 memcpy (castPtr dst) (castPtr src) count
460 return $ dst `plusPtr` count
462{-# INLINE _spliceArrays #-}
463_spliceArrays :: (MonadIO m, Storable a)
464 => SerialT m (Array a) -> m (Array a)
465_spliceArrays s = do
466 buffered <- S.foldr S.cons S.nil s
467 len <- S.sum (S.map Array.length buffered)
468 arr <- liftIO $ MA.newArray len
469 end <- S.foldlM' writeArr (return $ MA.aEnd arr) s
470 return $ A.unsafeFreeze $ arr {MA.aEnd = end}
472 where
474 writeArr dst (Array as ae) =
475 liftIO $ unsafeWithForeignPtr as $ \src -> do
476 let count = ae `minusPtr` src
477 memcpy (castPtr dst) (castPtr src) count
478 return $ dst `plusPtr` count
480{-# INLINE _spliceArraysBuffered #-}
481_spliceArraysBuffered :: (MonadIO m, Storable a)
482 => SerialT m (Array a) -> m (Array a)
483_spliceArraysBuffered s = do
484 buffered <- S.foldr S.cons S.nil s
485 len <- S.sum (S.map Array.length buffered)
486 A.unsafeFreeze <$> spliceArraysLenUnsafe len (S.map A.unsafeThaw s)
488{-# INLINE spliceArraysRealloced #-}
489spliceArraysRealloced :: forall m a. (MonadIO m, Storable a)
490 => SerialT m (Array a) -> m (Array a)
491spliceArraysRealloced s = do
492 let idst = liftIO $ MA.newArray (bytesToElemCount (undefined :: a)
493 (mkChunkSizeKB 4))
495 arr <- S.foldlM' MA.spliceWithDoubling idst (S.map A.unsafeThaw s)
496 liftIO $ A.unsafeFreeze <$> MA.shrinkToFit arr
498-- XXX This should just be "fold A.write"
499--
500-- | Given a stream of arrays, splice them all together to generate a single
501-- array. The stream must be /finite/.
502--
503-- @since 0.7.0
504{-# INLINE toArray #-}
505toArray :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)
506toArray = spliceArraysRealloced
507-- spliceArrays = _spliceArraysBuffered
509-- exponentially increasing sizes of the chunks upto the max limit.
510-- XXX this will be easier to implement with parsers/terminating folds
511-- With this we should be able to reduce the number of chunks/allocations.
512-- The reallocation/copy based toArray can also be implemented using this.
513--
514{-
515{-# INLINE toArraysInRange #-}
516toArraysInRange :: (IsStream t, MonadIO m, Storable a)
517 => Int -> Int -> Fold m (Array a) b -> Fold m a b
518toArraysInRange low high (Fold step initial extract) =
519-}
521{-
522-- | Fold the input to a pure buffered stream (List) of arrays.
523{-# INLINE _toArraysOf #-}
524_toArraysOf :: (MonadIO m, Storable a)
525 => Int -> Fold m a (SerialT Identity (Array a))
526_toArraysOf n = FL.chunksOf n (A.writeNF n) FL.toStream
527-}
529-------------------------------------------------------------------------------
530-- Elimination - running element parsers
531-------------------------------------------------------------------------------
533-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
534-- to make a newtype.
535{-# ANN type List NoSpecConstr #-}
536newtype List a = List {getList :: [a]}
538-- This can be generalized to any type provided it can be unfolded to a stream
539-- and it can be combined using a semigroup operation.
540--
541-- XXX This should be written using CPS (as parseK) if we want it to scale wrt
542-- to the number of times it can be called on the same stream.
543{-# INLINE_NORMAL parseD #-}
544parseD ::
545 forall m a b. (MonadIO m, MonadThrow m, Storable a)
546 => PRD.Parser m a b
547 -> D.Stream m (Array.Array a)
548 -> m (b, D.Stream m (Array.Array a))
549parseD (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
550 res <- initial
551 case res of
552 PRD.IPartial s -> go SPEC state (List []) s
553 PRD.IDone b -> return (b, stream)
554 PRD.IError err -> throwM $ ParseError err
556 where
558 -- "backBuf" contains last few items in the stream that we may have to
559 -- backtrack to.
560 --
561 -- XXX currently we are using a dumb list based approach for backtracking
562 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
563 -- That will allow us more efficient random back and forth movement.
564 {-# INLINE go #-}
565 go !_ st backBuf !pst = do
566 r <- step defState st
567 case r of
568 D.Yield (Array (ForeignPtr start contents) (Ptr end)) s ->
569 gobuf SPEC s backBuf (ForeignPtr end contents) (Ptr start) pst
570 D.Skip s -> go SPEC s backBuf pst
571 D.Stop -> do
572 b <- extract pst
573 return (b, D.nil)
575 -- Use strictness on "cur" to keep it unboxed
576 gobuf !_ s backBuf fp@(ForeignPtr end _) !cur !pst
577 | cur == Ptr end = do
578 liftIO $ touchForeignPtr fp
579 go SPEC s backBuf pst
580 gobuf !_ s backBuf fp@(ForeignPtr end contents) !cur !pst = do
581 x <- liftIO $ peek cur
582 pRes <- pstep pst x
583 let elemSize = sizeOf (undefined :: a)
584 next = cur `plusPtr` elemSize
585 case pRes of
586 PR.Partial 0 pst1 ->
587 gobuf SPEC s (List []) fp next pst1
588 PR.Partial n pst1 -> do
589 assert (n <= Prelude.length (x:getList backBuf)) (return ())
590 let src0 = Prelude.take n (x:getList backBuf)
591 arr0 = A.fromListN n (Prelude.reverse src0)
592 !(Ptr curAddr) = next
593 arr1 = Array (ForeignPtr curAddr contents) (Ptr end)
594 src = arr0 <> arr1
595 let !(Array (ForeignPtr start cont1) (Ptr end1)) = src
596 gobuf SPEC s (List []) (ForeignPtr end1 cont1) (Ptr start) pst1
597 PR.Continue 0 pst1 ->
598 gobuf SPEC s (List (x:getList backBuf)) fp next pst1
599 PR.Continue n pst1 -> do
600 assert (n <= Prelude.length (x:getList backBuf)) (return ())
601 let (src0, buf1) = splitAt n (x:getList backBuf)
602 arr0 = A.fromListN n (Prelude.reverse src0)
603 !(Ptr curAddr) = next
604 arr1 = Array (ForeignPtr curAddr contents) (Ptr end)
605 src = arr0 <> arr1
606 let !(Array (ForeignPtr start cont1) (Ptr end1)) = src
607 gobuf SPEC s (List buf1) (ForeignPtr end1 cont1) (Ptr start) pst1
608 PR.Done 0 b -> do
609 let !(Ptr curAddr) = next
610 arr = Array (ForeignPtr curAddr contents) (Ptr end)
611 return (b, D.cons arr (D.Stream step s))
612 PR.Done n b -> do
613 assert (n <= Prelude.length (x:getList backBuf)) (return ())
614 let src0 = Prelude.take n (x:getList backBuf)
615 -- XXX create the array in reverse instead
616 arr0 = A.fromListN n (Prelude.reverse src0)
617 !(Ptr curAddr) = next
618 arr1 = Array (ForeignPtr curAddr contents) (Ptr end)
619 -- XXX Use StreamK to avoid adding arbitrary layers of
620 -- constructors every time.
621 str = D.cons arr0 (D.cons arr1 (D.Stream step s))
622 return (b, str)
623 PR.Error err -> throwM $ ParseError err
625-- | Parse an array stream using the supplied 'Parser'. Returns the parse
626-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
627--
628-- /Internal/
629--
630{-# INLINE_NORMAL parse #-}
631parse ::
632 (MonadIO m, MonadThrow m, Storable a)
633 => PRD.Parser m a b
634 -> SerialT m (A.Array a)
635 -> m (b, SerialT m (A.Array a))
636parse p s = fmap D.fromStreamD <$> parseD p (D.toStreamD s)
638-------------------------------------------------------------------------------
639-- Elimination - Running Array Folds and parsers
640-------------------------------------------------------------------------------
642{-# INLINE_NORMAL parseArrD #-}
643parseArrD ::
644 forall m a b. (MonadIO m, MonadThrow m, Storable a)
645 => PRD.Parser m (Array a) b
646 -> D.Stream m (Array.Array a)
647 -> m (b, D.Stream m (Array.Array a))
648parseArrD (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
649 res <- initial
650 case res of
651 PRD.IPartial s -> go SPEC state (List []) s
652 PRD.IDone b -> return (b, stream)
653 PRD.IError err -> throwM $ ParseError err
655 where
657 -- "backBuf" contains last few items in the stream that we may have to
658 -- backtrack to.
659 --
660 -- XXX currently we are using a dumb list based approach for backtracking
661 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
662 -- That will allow us more efficient random back and forth movement.
663 {-# INLINE go #-}
664 go !_ st backBuf !pst = do
665 r <- step defState st
666 case r of
667 D.Yield x s -> gobuf SPEC [x] s backBuf pst
668 D.Skip s -> go SPEC s backBuf pst
669 D.Stop -> do
670 b <- extract pst
671 return (b, D.nil)
673 gobuf !_ [] s backBuf !pst = go SPEC s backBuf pst
674 gobuf !_ (x:xs) s backBuf !pst = do
675 pRes <- pstep pst x
676 case pRes of
677 PR.Partial 0 pst1 ->
678 gobuf SPEC xs s (List []) pst1
679 PR.Partial n pst1 -> do
680 assert
681 (n <= sum (map Array.length (x:getList backBuf)))
682 (return ())
683 let src0 = takeArrayListRev n (x:getList backBuf)
684 src = Prelude.reverse src0 ++ xs
685 gobuf SPEC src s (List []) pst1
686 PR.Continue 0 pst1 ->
687 gobuf SPEC xs s (List (x:getList backBuf)) pst1
688 PR.Continue n pst1 -> do
689 assert
690 (n <= sum (map Array.length (x:getList backBuf)))
691 (return ())
692 let (src0, buf1) = splitAtArrayListRev n (x:getList backBuf)
693 src = Prelude.reverse src0 ++ xs
694 gobuf SPEC src s (List buf1) pst1
695 PR.Done 0 b ->
696 return (b, D.Stream step s)
697 PR.Done n b -> do
698 assert
699 (n <= sum (map Array.length (x:getList backBuf)))
700 (return ())
701 let src0 = takeArrayListRev n (x:getList backBuf)
702 src = Prelude.reverse src0 ++ xs
703 return (b, D.append (D.fromList src) (D.Stream step s))
704 PR.Error err -> throwM $ ParseError err
706{-
707-- | Parse an array stream using the supplied 'Parser'. Returns the parse
708-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
709--
710-- /Internal/
711--
712{-# INLINE parseArr #-}
713parseArr ::
714 (MonadIO m, MonadThrow m, Storable a)
715 => ASF.Parser m a b
716 -> SerialT m (A.Array a)
717 -> m (b, SerialT m (A.Array a))
718parseArr p s = fmap D.fromStreamD <$> parseD p (D.toStreamD s)
719-}
721-- | Fold an array stream using the supplied array stream 'Fold'.
722--
723-- /Pre-release/
724--
725{-# INLINE foldArr #-}
726foldArr :: (MonadIO m, MonadThrow m, Storable a) =>
727 ASF.Fold m a b -> SerialT m (A.Array a) -> m b
728foldArr (ASF.Fold p) s = fst <$> parseArrD p (D.toStreamD s)
730-- | Like 'fold' but also returns the remaining stream.
731--
732-- /Pre-release/
733--
734{-# INLINE foldArr_ #-}
735foldArr_ :: (MonadIO m, MonadThrow m, Storable a) =>
736 ASF.Fold m a b -> SerialT m (A.Array a) -> m (b, SerialT m (A.Array a))
737foldArr_ (ASF.Fold p) s = second D.fromStreamD <$> parseArrD p (D.toStreamD s)
739{-# ANN type ParseChunksState Fuse #-}
740data ParseChunksState x inpBuf st pst =
741 ParseChunksInit inpBuf st
742 | ParseChunksInitLeftOver inpBuf
743 | ParseChunksStream st inpBuf !pst
744 | ParseChunksBuf inpBuf st inpBuf !pst
745 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
747{-# INLINE_NORMAL foldArrManyD #-}
748foldArrManyD
749 :: (MonadThrow m, Storable a)
750 => ASF.Fold m a b
751 -> D.Stream m (Array a)
752 -> D.Stream m b
753foldArrManyD (ASF.Fold (PRD.Parser pstep initial extract)) (D.Stream step state) =
754 D.Stream stepOuter (ParseChunksInit [] state)
756 where
758 {-# INLINE_LATE stepOuter #-}
759 -- Buffer is empty, get the first element from the stream, initialize the
760 -- fold and then go to stream processing loop.
761 stepOuter gst (ParseChunksInit [] st) = do
762 r <- step (adaptState gst) st
763 case r of
764 D.Yield x s -> do
765 res <- initial
766 case res of
767 PRD.IPartial ps ->
768 return $ D.Skip $ ParseChunksBuf [x] s [] ps
769 PRD.IDone pb ->
770 let next = ParseChunksInit [x] s
771 in return $ D.Skip $ ParseChunksYield pb next
772 PRD.IError err -> throwM $ ParseError err
773 D.Skip s -> return $ D.Skip $ ParseChunksInit [] s
774 D.Stop -> return D.Stop
776 -- Buffer is not empty, go to buffered processing loop
777 stepOuter _ (ParseChunksInit src st) = do
778 res <- initial
779 case res of
780 PRD.IPartial ps ->
781 return $ D.Skip $ ParseChunksBuf src st [] ps
782 PRD.IDone pb ->
783 let next = ParseChunksInit src st
784 in return $ D.Skip $ ParseChunksYield pb next
785 PRD.IError err -> throwM $ ParseError err
787 -- XXX we just discard any leftover input at the end
788 stepOuter _ (ParseChunksInitLeftOver _) = return D.Stop
790 -- Buffer is empty, process elements from the stream
791 stepOuter gst (ParseChunksStream st backBuf pst) = do
792 r <- step (adaptState gst) st
793 case r of
794 D.Yield x s -> do
795 pRes <- pstep pst x
796 case pRes of
797 PR.Partial 0 pst1 ->
798 return $ D.Skip $ ParseChunksStream s [] pst1
799 PR.Partial n pst1 -> do
800 assert
801 (n <= sum (map Array.length (x:backBuf)))
802 (return ())
803 let src0 = takeArrayListRev n (x:backBuf)
804 src = Prelude.reverse src0
805 return $ D.Skip $ ParseChunksBuf src s [] pst1
806 PR.Continue 0 pst1 ->
807 return $ D.Skip $ ParseChunksStream s (x:backBuf) pst1
808 PR.Continue n pst1 -> do
809 assert
810 (n <= sum (map Array.length (x:backBuf)))
811 (return ())
812 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
813 src = Prelude.reverse src0
814 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
815 PR.Done 0 b -> do
816 return $ D.Skip $
817 ParseChunksYield b (ParseChunksInit [] s)
818 PR.Done n b -> do
819 assert
820 (n <= sum (map Array.length (x:backBuf)))
821 (return ())
822 let src0 = takeArrayListRev n (x:backBuf)
823 let src = Prelude.reverse src0
824 return $ D.Skip $
825 ParseChunksYield b (ParseChunksInit src s)
826 PR.Error err -> throwM $ ParseError err
827 D.Skip s -> return $ D.Skip $ ParseChunksStream s backBuf pst
828 D.Stop -> do
829 b <- extract pst
830 let src = Prelude.reverse backBuf
831 return $ D.Skip $
832 ParseChunksYield b (ParseChunksInitLeftOver src)
834 -- go back to stream processing mode
835 stepOuter _ (ParseChunksBuf [] s buf pst) =
836 return $ D.Skip $ ParseChunksStream s buf pst
838 -- buffered processing loop
839 stepOuter _ (ParseChunksBuf (x:xs) s backBuf pst) = do
840 pRes <- pstep pst x
841 case pRes of
842 PR.Partial 0 pst1 ->
843 return $ D.Skip $ ParseChunksBuf xs s [] pst1
844 PR.Partial n pst1 -> do
845 assert (n <= sum (map Array.length (x:backBuf))) (return ())
846 let src0 = takeArrayListRev n (x:backBuf)
847 src = Prelude.reverse src0 ++ xs
848 return $ D.Skip $ ParseChunksBuf src s [] pst1
849 PR.Continue 0 pst1 ->
850 return $ D.Skip $ ParseChunksBuf xs s (x:backBuf) pst1
851 PR.Continue n pst1 -> do
852 assert (n <= sum (map Array.length (x:backBuf))) (return ())
853 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
854 src = Prelude.reverse src0 ++ xs
855 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
856 PR.Done 0 b ->
857 return $ D.Skip $ ParseChunksYield b (ParseChunksInit xs s)
858 PR.Done n b -> do
859 assert (n <= sum (map Array.length (x:backBuf))) (return ())
860 let src0 = takeArrayListRev n (x:backBuf)
861 src = Prelude.reverse src0 ++ xs
862 return $ D.Skip $ ParseChunksYield b (ParseChunksInit src s)
863 PR.Error err -> throwM $ ParseError err
865 stepOuter _ (ParseChunksYield a next) = return $ D.Yield a next
867-- | Apply an array stream 'Fold' repeatedly on an array stream and emit the
868-- fold outputs in the output stream.
869--
870-- See "Streamly.Prelude.foldMany" for more details.
871--
872-- /Pre-release/
873{-# INLINE foldArrMany #-}
874foldArrMany
875 :: (IsStream t, MonadThrow m, Storable a)
876 => ASF.Fold m a b
877 -> t m (Array a)
878 -> t m b
879foldArrMany p m = D.fromStreamD $ foldArrManyD p (D.toStreamD m)