git.iko.soy/streamly
on c366e2d help sign in
streamly/src/Streamly/Internal/Data/Array/Stream/Foreign.hs
998 lines 38 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 Foreign.ForeignPtr (touchForeignPtr)
62import Foreign.Ptr (minusPtr, plusPtr, castPtr)
63import Streamly.Internal.Data.Unboxed (Storable, peek, sizeOf)
64import Fusion.Plugin.Types (Fuse(..))
65import GHC.Exts (SpecConstrAnnotation(..))
66import GHC.ForeignPtr (ForeignPtr(..))
67import GHC.Ptr (Ptr(..))
68import GHC.Types (SPEC(..))
69import Prelude hiding (null, last, (!!), read, concat, unlines)
71import Streamly.Internal.Data.Array.Foreign.Mut.Type
72 (arrayToFptrContents, fptrToArrayContents, touch)
73import Streamly.Internal.Data.Array.Foreign.Type (Array(..))
74import Streamly.Internal.Data.Array.Stream.Fold.Foreign (ArrayFold(..))
75import Streamly.Internal.Data.Fold.Type (Fold(..))
76import Streamly.Internal.Data.Parser (ParseError(..))
77import Streamly.Internal.Data.Stream.Serial (SerialT(..))
78import Streamly.Internal.Data.Stream.IsStream.Type
79 (IsStream, fromStreamD, toStreamD)
80import Streamly.Internal.Data.SVar (adaptState, defState)
81import Streamly.Internal.Data.Array.Foreign.Mut.Type
82 (memcpy, allocBytesToElemCount)
84import qualified Streamly.Internal.Data.Array.Foreign as A
85import qualified Streamly.Internal.Data.Array.Foreign as Array
86-- import qualified Streamly.Internal.Data.Array.Stream.Fold.Foreign
87-- as ArrayFold
88import qualified Streamly.Internal.Data.Array.Foreign.Type as A
89import qualified Streamly.Internal.Data.Array.Foreign.Mut.Type as MA
90import qualified Streamly.Internal.Data.Array.Stream.Mut.Foreign as AS
91import qualified Streamly.Internal.Data.Fold.Type as FL
92import qualified Streamly.Internal.Data.Parser as PR
93import qualified Streamly.Internal.Data.Parser.ParserD as PRD
94import qualified Streamly.Internal.Data.Stream.IsStream as S
95import qualified Streamly.Internal.Data.Stream.StreamD as D
96import qualified Streamly.Internal.Data.Stream.StreamK as K
98-- XXX Since these are immutable arrays MonadIO constraint can be removed from
99-- most places.
101-------------------------------------------------------------------------------
102-- Generation
103-------------------------------------------------------------------------------
105-- | @arraysOf n stream@ groups the elements in the input stream into arrays of
106-- @n@ elements each.
107--
108-- > arraysOf n = Stream.chunksOf n (Array.writeN n)
109--
110-- /Pre-release/
111{-# INLINE arraysOf #-}
112arraysOf :: (IsStream t, MonadIO m, Storable a)
113 => Int -> t m a -> t m (Array a)
114arraysOf n str = fromStreamD $ A.arraysOf n (toStreamD str)
116-------------------------------------------------------------------------------
117-- Append
118-------------------------------------------------------------------------------
120-- XXX efficiently compare two streams of arrays. Two streams can have chunks
121-- of different sizes, we can handle that in the stream comparison abstraction.
122-- This could be useful e.g. to fast compare whether two files differ.
124-- | Convert a stream of arrays into a stream of their elements.
125--
126-- Same as the following:
127--
128-- > concat = Stream.unfoldMany Array.read
129--
130-- @since 0.7.0
131{-# INLINE concat #-}
132concat :: (IsStream t, Monad m, Storable a) => t m (Array a) -> t m a
133-- concat m = fromStreamD $ A.flattenArrays (toStreamD m)
134-- concat m = fromStreamD $ D.concatMap A.toStreamD (toStreamD m)
135concat m = fromStreamD $ D.unfoldMany A.read (toStreamD m)
137-- | Convert a stream of arrays into a stream of their elements reversing the
138-- contents of each array before flattening.
139--
140-- > concatRev = Stream.unfoldMany Array.readRev
141--
142-- @since 0.7.0
143{-# INLINE concatRev #-}
144concatRev :: (IsStream t, Monad m, Storable a) => t m (Array a) -> t m a
145-- concatRev m = fromStreamD $ A.flattenArraysRev (toStreamD m)
146concatRev m = fromStreamD $ D.unfoldMany A.readRev (toStreamD m)
148-------------------------------------------------------------------------------
149-- Intersperse and append
150-------------------------------------------------------------------------------
152-- | Flatten a stream of arrays after inserting the given element between
153-- arrays.
154--
155-- /Pre-release/
156{-# INLINE interpose #-}
157interpose :: (Monad m, IsStream t, Storable a) => a -> t m (Array a) -> t m a
158interpose x = S.interpose x A.read
160{-# INLINE intercalateSuffix #-}
161intercalateSuffix :: (Monad m, IsStream t, Storable a)
162 => Array a -> t m (Array a) -> t m a
163intercalateSuffix = S.intercalateSuffix A.read
165-- | Flatten a stream of arrays appending the given element after each
166-- array.
167--
168-- @since 0.7.0
169{-# INLINE interposeSuffix #-}
170interposeSuffix :: (Monad m, IsStream t, Storable a)
171 => a -> t m (Array a) -> t m a
172-- interposeSuffix x = fromStreamD . A.unlines x . toStreamD
173interposeSuffix x = S.interposeSuffix x A.read
175data FlattenState s a =
176 OuterLoop s
177 | InnerLoop s !MA.ArrayContents !(Ptr a) !(Ptr a)
179-- XXX Remove monadIO constraint
180{-# INLINE_NORMAL unlines #-}
181unlines :: forall m a. (MonadIO m, Storable a)
182 => a -> D.Stream m (Array a) -> D.Stream m a
183unlines sep (D.Stream step state) = D.Stream step' (OuterLoop state)
184 where
185 {-# INLINE_LATE step' #-}
186 step' gst (OuterLoop st) = do
187 r <- step (adaptState gst) st
188 return $ case r of
189 D.Yield Array{..} s ->
190 D.Skip (InnerLoop s arrContents arrStart aEnd)
191 D.Skip s -> D.Skip (OuterLoop s)
192 D.Stop -> D.Stop
194 step' _ (InnerLoop st _ p end) | p == end =
195 return $ D.Yield sep $ OuterLoop st
197 step' _ (InnerLoop st contents p end) = do
198 x <- liftIO $ do
199 r <- peek p
200 touch contents
201 return r
202 return $ D.Yield x (InnerLoop st contents (PTR_NEXT(p,a)) end)
204-------------------------------------------------------------------------------
205-- Compact
206-------------------------------------------------------------------------------
208-- XXX These would not be needed once we implement compactLEFold, see
209-- module Streamly.Internal.Data.Array.Stream.Mut.Foreign
210--
211{-# INLINE_NORMAL packArraysChunksOf #-}
212packArraysChunksOf :: (MonadIO m, Storable a)
213 => Int -> D.Stream m (Array a) -> D.Stream m (Array a)
214packArraysChunksOf n str =
215 D.map A.unsafeFreeze $ AS.packArraysChunksOf n $ D.map A.unsafeThaw str
217-- XXX instead of writing two different versions of this operation, we should
218-- write it as a pipe.
219{-# INLINE_NORMAL lpackArraysChunksOf #-}
220lpackArraysChunksOf :: (MonadIO m, Storable a)
221 => Int -> Fold m (Array a) () -> Fold m (Array a) ()
222lpackArraysChunksOf n fld =
223 FL.lmap A.unsafeThaw $ AS.lpackArraysChunksOf n (FL.lmap A.unsafeFreeze fld)
225-- | Coalesce adjacent arrays in incoming stream to form bigger arrays of a
226-- maximum specified size in bytes.
227--
228-- @since 0.7.0
229{-# INLINE compact #-}
230compact :: (MonadIO m, Storable a)
231 => Int -> SerialT m (Array a) -> SerialT m (Array a)
232compact n xs = fromStreamD $ packArraysChunksOf n (toStreamD xs)
234-------------------------------------------------------------------------------
235-- Split
236-------------------------------------------------------------------------------
238data SplitState s arr
239 = Initial s
240 | Buffering s arr
241 | Splitting s arr
242 | Yielding arr (SplitState s arr)
243 | Finishing
245-- | Split a stream of arrays on a given separator byte, dropping the separator
246-- and coalescing all the arrays between two separators into a single array.
247--
248-- @since 0.7.0
249{-# INLINE_NORMAL _splitOn #-}
250_splitOn
251 :: MonadIO m
252 => Word8
253 -> D.Stream m (Array Word8)
254 -> D.Stream m (Array Word8)
255_splitOn byte (D.Stream step state) = D.Stream step' (Initial state)
257 where
259 {-# INLINE_LATE step' #-}
260 step' gst (Initial st) = do
261 r <- step gst st
262 case r of
263 D.Yield arr s -> do
264 (arr1, marr2) <- A.breakOn byte arr
265 return $ case marr2 of
266 Nothing -> D.Skip (Buffering s arr1)
267 Just arr2 -> D.Skip (Yielding arr1 (Splitting s arr2))
268 D.Skip s -> return $ D.Skip (Initial s)
269 D.Stop -> return D.Stop
271 step' gst (Buffering st buf) = do
272 r <- step gst st
273 case r of
274 D.Yield arr s -> do
275 (arr1, marr2) <- A.breakOn byte arr
276 buf' <- A.splice buf arr1
277 return $ case marr2 of
278 Nothing -> D.Skip (Buffering s buf')
279 Just x -> D.Skip (Yielding buf' (Splitting s x))
280 D.Skip s -> return $ D.Skip (Buffering s buf)
281 D.Stop -> return $
282 if A.byteLength buf == 0
283 then D.Stop
284 else D.Skip (Yielding buf Finishing)
286 step' _ (Splitting st buf) = do
287 (arr1, marr2) <- A.breakOn byte buf
288 return $ case marr2 of
289 Nothing -> D.Skip $ Buffering st arr1
290 Just arr2 -> D.Skip $ Yielding arr1 (Splitting st arr2)
292 step' _ (Yielding arr next) = return $ D.Yield arr next
293 step' _ Finishing = return D.Stop
295-- XXX Remove MonadIO constraint.
296-- | Split a stream of arrays on a given separator byte, dropping the separator
297-- and coalescing all the arrays between two separators into a single array.
298--
299-- @since 0.7.0
300{-# INLINE splitOn #-}
301splitOn
302 :: (IsStream t, MonadIO m)
303 => Word8
304 -> t m (Array Word8)
305 -> t m (Array Word8)
306splitOn byte s =
307 fromStreamD $ D.splitInnerBy (A.breakOn byte) A.splice $ toStreamD s
309{-# INLINE splitOnSuffix #-}
310splitOnSuffix
311 :: (IsStream t, MonadIO m)
312 => Word8
313 -> t m (Array Word8)
314 -> t m (Array Word8)
315-- splitOn byte s = fromStreamD $ A.splitOn byte $ toStreamD s
316splitOnSuffix byte s =
317 fromStreamD $ D.splitInnerBySuffix (A.breakOn byte) A.splice $ toStreamD s
319-------------------------------------------------------------------------------
320-- Elimination - Running folds
321-------------------------------------------------------------------------------
323-- XXX This should be written using CPS (as foldK) if we want it to scale wrt
324-- to the number of times it can be called on the same stream.
325--
326{-# INLINE_NORMAL foldBreakD #-}
327foldBreakD :: forall m a b. (MonadIO m, Storable a) =>
328 Fold m a b -> D.Stream m (Array a) -> m (b, D.Stream m (Array a))
329foldBreakD (Fold fstep initial extract) stream@(D.Stream step state) = do
330 res <- initial
331 case res of
332 FL.Partial fs -> go SPEC state fs
333 FL.Done fb -> return $! (fb, stream)
335 where
337 {-# INLINE go #-}
338 go !_ st !fs = do
339 r <- step defState st
340 case r of
341 D.Yield (Array contents start (Ptr end)) s ->
342 let fp = ForeignPtr end (arrayToFptrContents contents)
343 in goArray SPEC s fp start fs
344 D.Skip s -> go SPEC s fs
345 D.Stop -> do
346 b <- extract fs
347 return (b, D.nil)
349 goArray !_ s fp@(ForeignPtr end _) !cur !fs
350 | cur == Ptr end = do
351 liftIO $ touchForeignPtr fp
352 go SPEC s fs
353 goArray !_ st fp@(ForeignPtr end contents) !cur !fs = do
354 x <- liftIO $ peek cur
355 res <- fstep fs x
356 let next = PTR_NEXT(cur,a)
357 case res of
358 FL.Done b -> do
359 let arr = Array (fptrToArrayContents contents) next (Ptr end)
360 return $! (b, D.cons arr (D.Stream step st))
361 FL.Partial fs1 -> goArray SPEC st fp next fs1
363{-# INLINE_NORMAL foldBreakK #-}
364foldBreakK :: forall m a b. (MonadIO m, Storable a) =>
365 Fold m a b -> K.Stream m (Array a) -> m (b, K.Stream m (Array a))
366foldBreakK (Fold fstep initial extract) stream = do
367 res <- initial
368 case res of
369 FL.Partial fs -> go fs stream
370 FL.Done fb -> return (fb, stream)
372 where
374 {-# INLINE go #-}
375 go !fs st = do
376 let stop = (, K.nil) <$> extract fs
377 single a = yieldk a K.nil
378 yieldk (Array contents start (Ptr end)) r =
379 let fp = ForeignPtr end (arrayToFptrContents contents)
380 in goArray fs r fp start
381 in K.foldStream defState yieldk single stop st
383 goArray !fs st fp@(ForeignPtr end _) !cur
384 | cur == Ptr end = do
385 liftIO $ touchForeignPtr fp
386 go fs st
387 goArray !fs st fp@(ForeignPtr end contents) !cur = do
388 x <- liftIO $ peek cur
389 res <- fstep fs x
390 let next = PTR_NEXT(cur,a)
391 case res of
392 FL.Done b -> do
393 let arr = Array (fptrToArrayContents contents) next (Ptr end)
394 return $! (b, K.cons arr st)
395 FL.Partial fs1 -> goArray fs1 st fp next
397-- | Fold an array stream using the supplied 'Fold'. Returns the fold result
398-- and the unconsumed stream.
399--
400-- > foldBreak f = runArrayFoldBreak (ArrayFold.fromFold f)
401--
402-- /Internal/
403--
404{-# INLINE_NORMAL foldBreak #-}
405foldBreak ::
406 (MonadIO m, Storable a)
407 => FL.Fold m a b
408 -> SerialT m (A.Array a)
409 -> m (b, SerialT m (A.Array a))
410-- foldBreak f s = fmap fromStreamD <$> foldBreakD f (toStreamD s)
411foldBreak f (SerialT s) = fmap (fmap SerialT) $ foldBreakK f s
412-- If foldBreak performs better than runArrayFoldBreak we can use a rewrite
413-- rule to rewrite runArrayFoldBreak to fold.
414-- foldBreak f = runArrayFoldBreak (ArrayFold.fromFold f)
416-------------------------------------------------------------------------------
417-- Fold to a single Array
418-------------------------------------------------------------------------------
420-- When we have to take an array partially, take the last part of the array.
421{-# INLINE takeArrayListRev #-}
422takeArrayListRev :: forall a. Storable a => Int -> [Array a] -> [Array a]
423takeArrayListRev = go
425 where
427 go _ [] = []
428 go n _ | n <= 0 = []
429 go n (x:xs) =
430 let len = Array.length x
431 in if n > len
432 then x : go (n - len) xs
433 else if n == len
434 then [x]
435 else let !(Array contents _ end) = x
436 !start = end `plusPtr` negate (n * SIZE_OF(a))
437 in [Array contents start end]
439-- When we have to take an array partially, take the last part of the array in
440-- the first split.
441{-# INLINE splitAtArrayListRev #-}
442splitAtArrayListRev :: forall a. Storable a =>
443 Int -> [Array a] -> ([Array a],[Array a])
444splitAtArrayListRev n ls
445 | n <= 0 = ([], ls)
446 | otherwise = go n ls
447 where
448 go :: Int -> [Array a] -> ([Array a], [Array a])
449 go _ [] = ([], [])
450 go m (x:xs) =
451 let len = Array.length x
452 (xs', xs'') = go (m - len) xs
453 in if m > len
454 then (x:xs', xs'')
455 else if m == len
456 then ([x],xs)
457 else let !(Array contents start end) = x
458 end1 = end `plusPtr` negate (m * SIZE_OF(a))
459 arr2 = Array contents start end1
460 arr1 = Array contents end1 end
461 in ([arr1], arr2:xs)
463-------------------------------------------------------------------------------
464-- Fold to a single Array
465-------------------------------------------------------------------------------
467-- XXX Both of these implementations of splicing seem to perform equally well.
468-- We need to perform benchmarks over a range of sizes though.
470-- CAUTION! length must more than equal to lengths of all the arrays in the
471-- stream.
472{-# INLINE spliceArraysLenUnsafe #-}
473spliceArraysLenUnsafe :: (MonadIO m, Storable a)
474 => Int -> SerialT m (MA.Array a) -> m (MA.Array a)
475spliceArraysLenUnsafe len buffered = do
476 arr <- liftIO $ MA.newArray len
477 end <- S.foldlM' writeArr (return $ MA.aEnd arr) buffered
478 return $ arr {MA.aEnd = end}
480 where
482 writeArr dst (MA.Array ac src ae _) =
483 liftIO $ do
484 let count = ae `minusPtr` src
485 memcpy (castPtr dst) (castPtr src) count
486 touch ac
487 return $ dst `plusPtr` count
489{-# INLINE _spliceArrays #-}
490_spliceArrays :: (MonadIO m, Storable a)
491 => SerialT m (Array a) -> m (Array a)
492_spliceArrays s = do
493 buffered <- S.foldr S.cons S.nil s
494 len <- S.sum (S.map Array.length buffered)
495 arr <- liftIO $ MA.newArray len
496 end <- S.foldlM' writeArr (return $ MA.aEnd arr) s
497 return $ A.unsafeFreeze $ arr {MA.aEnd = end}
499 where
501 writeArr dst (Array ac src ae) =
502 liftIO $ do
503 let count = ae `minusPtr` src
504 memcpy (castPtr dst) (castPtr src) count
505 touch ac
506 return $ dst `plusPtr` count
508{-# INLINE _spliceArraysBuffered #-}
509_spliceArraysBuffered :: (MonadIO m, Storable a)
510 => SerialT m (Array a) -> m (Array a)
511_spliceArraysBuffered s = do
512 buffered <- S.foldr S.cons S.nil s
513 len <- S.sum (S.map Array.length buffered)
514 A.unsafeFreeze <$> spliceArraysLenUnsafe len (S.map A.unsafeThaw s)
516{-# INLINE spliceArraysRealloced #-}
517spliceArraysRealloced :: forall m a. (MonadIO m, Storable a)
518 => SerialT m (Array a) -> m (Array a)
519spliceArraysRealloced s = do
520 let n = allocBytesToElemCount (undefined :: a) (4 * 1024)
521 idst = liftIO $ MA.newArray n
523 arr <- S.foldlM' MA.spliceExp idst (S.map A.unsafeThaw s)
524 liftIO $ A.unsafeFreeze <$> MA.rightSize arr
526-- XXX This should just be "fold A.write"
527--
528-- | Given a stream of arrays, splice them all together to generate a single
529-- array. The stream must be /finite/.
530--
531-- @since 0.7.0
532{-# INLINE toArray #-}
533toArray :: (MonadIO m, Storable a) => SerialT m (Array a) -> m (Array a)
534toArray = spliceArraysRealloced
535-- spliceArrays = _spliceArraysBuffered
537-- exponentially increasing sizes of the chunks upto the max limit.
538-- XXX this will be easier to implement with parsers/terminating folds
539-- With this we should be able to reduce the number of chunks/allocations.
540-- The reallocation/copy based toArray can also be implemented using this.
541--
542{-
543{-# INLINE toArraysInRange #-}
544toArraysInRange :: (IsStream t, MonadIO m, Storable a)
545 => Int -> Int -> Fold m (Array a) b -> Fold m a b
546toArraysInRange low high (Fold step initial extract) =
547-}
549{-
550-- | Fold the input to a pure buffered stream (List) of arrays.
551{-# INLINE _toArraysOf #-}
552_toArraysOf :: (MonadIO m, Storable a)
553 => Int -> Fold m a (SerialT Identity (Array a))
554_toArraysOf n = FL.chunksOf n (A.writeNF n) FL.toStream
555-}
557-------------------------------------------------------------------------------
558-- Elimination - running element parsers
559-------------------------------------------------------------------------------
561-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
562-- to make a newtype.
563{-# ANN type List NoSpecConstr #-}
564newtype List a = List {getList :: [a]}
566-- This can be generalized to any type provided it can be unfolded to a stream
567-- and it can be combined using a semigroup operation.
568--
569-- XXX This should be written using CPS (as parseK) if we want it to scale wrt
570-- to the number of times it can be called on the same stream.
571{-# INLINE_NORMAL parseBreakD #-}
572parseBreakD ::
573 forall m a b. (MonadIO m, MonadThrow m, Storable a)
574 => PRD.Parser m a b
575 -> D.Stream m (Array.Array a)
576 -> m (b, D.Stream m (Array.Array a))
577parseBreakD
578 (PRD.Parser pstep initial extract) stream@(D.Stream step state) = do
580 res <- initial
581 case res of
582 PRD.IPartial s -> go SPEC state (List []) s
583 PRD.IDone b -> return (b, stream)
584 PRD.IError err -> throwM $ ParseError err
586 where
588 -- "backBuf" contains last few items in the stream that we may have to
589 -- backtrack to.
590 --
591 -- XXX currently we are using a dumb list based approach for backtracking
592 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
593 -- That will allow us more efficient random back and forth movement.
594 go !_ st backBuf !pst = do
595 r <- step defState st
596 case r of
597 D.Yield (Array contents start (Ptr end)) s ->
598 gobuf SPEC s backBuf
599 (ForeignPtr end (arrayToFptrContents contents)) start pst
600 D.Skip s -> go SPEC s backBuf pst
601 D.Stop -> do
602 b <- extract pst
603 return (b, D.nil)
605 -- Use strictness on "cur" to keep it unboxed
606 gobuf !_ s backBuf fp@(ForeignPtr end _) !cur !pst
607 | cur == Ptr end = do
608 liftIO $ touchForeignPtr fp
609 go SPEC s backBuf pst
610 gobuf !_ s backBuf fp@(ForeignPtr end contents) !cur !pst = do
611 x <- liftIO $ peek cur
612 pRes <- pstep pst x
613 let next = PTR_NEXT(cur,a)
614 case pRes of
615 PR.Partial 0 pst1 ->
616 gobuf SPEC s (List []) fp next pst1
617 PR.Partial n pst1 -> do
618 assert (n <= Prelude.length (x:getList backBuf)) (return ())
619 let src0 = Prelude.take n (x:getList backBuf)
620 arr0 = A.fromListN n (Prelude.reverse src0)
621 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
622 src = arr0 <> arr1
623 let !(Array cont1 start (Ptr end1)) = src
624 fp1 = ForeignPtr end1 (arrayToFptrContents cont1)
625 gobuf SPEC s (List []) fp1 start pst1
626 PR.Continue 0 pst1 ->
627 gobuf SPEC s (List (x:getList backBuf)) fp next pst1
628 PR.Continue n pst1 -> do
629 assert (n <= Prelude.length (x:getList backBuf)) (return ())
630 let (src0, buf1) = splitAt n (x:getList backBuf)
631 arr0 = A.fromListN n (Prelude.reverse src0)
632 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
633 src = arr0 <> arr1
634 let !(Array cont1 start (Ptr end1)) = src
635 fp1 = ForeignPtr end1 (arrayToFptrContents cont1)
636 gobuf SPEC s (List buf1) fp1 start pst1
637 PR.Done 0 b -> do
638 let arr = Array (fptrToArrayContents contents) next (Ptr end)
639 return (b, D.cons arr (D.Stream step s))
640 PR.Done n b -> do
641 assert (n <= Prelude.length (x:getList backBuf)) (return ())
642 let src0 = Prelude.take n (x:getList backBuf)
643 -- XXX create the array in reverse instead
644 arr0 = A.fromListN n (Prelude.reverse src0)
645 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
646 -- XXX Use StreamK to avoid adding arbitrary layers of
647 -- constructors every time.
648 str = D.cons arr0 (D.cons arr1 (D.Stream step s))
649 return (b, str)
650 PR.Error err -> throwM $ ParseError err
652{-# INLINE_NORMAL parseBreakK #-}
653parseBreakK ::
654 forall m a b. (MonadIO m, MonadThrow m, Storable a)
655 => PRD.Parser m a b
656 -> K.Stream m (Array.Array a)
657 -> m (b, K.Stream m (Array.Array a))
658parseBreakK (PRD.Parser pstep initial extract) stream = do
659 res <- initial
660 case res of
661 PRD.IPartial s -> go s stream []
662 PRD.IDone b -> return (b, stream)
663 PRD.IError err -> throwM $ ParseError err
665 where
667 -- "backBuf" contains last few items in the stream that we may have to
668 -- backtrack to.
669 --
670 -- XXX currently we are using a dumb list based approach for backtracking
671 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
672 -- That will allow us more efficient random back and forth movement.
673 go !pst st backBuf = do
674 let stop = (, K.nil) <$> extract pst
675 single a = yieldk a K.nil
676 yieldk (Array contents start (Ptr end)) r =
677 let fp = ForeignPtr end (arrayToFptrContents contents)
678 in goArray pst backBuf r fp start
679 in K.foldStream defState yieldk single stop st
681 -- Use strictness on "cur" to keep it unboxed
682 goArray !pst backBuf st fp@(ForeignPtr end _) !cur
683 | cur == Ptr end = do
684 liftIO $ touchForeignPtr fp
685 go pst st backBuf
686 goArray !pst backBuf st fp@(ForeignPtr end contents) !cur = do
687 x <- liftIO $ peek cur
688 pRes <- pstep pst x
689 let next = PTR_NEXT(cur,a)
690 case pRes of
691 PR.Partial 0 s ->
692 goArray s [] st fp next
693 PR.Partial n s -> do
694 assert (n <= Prelude.length (x:backBuf)) (return ())
695 let src0 = Prelude.take n (x:backBuf)
696 arr0 = A.fromListN n (Prelude.reverse src0)
697 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
698 src = arr0 <> arr1
699 let !(Array cont1 start (Ptr end1)) = src
700 fp1 = ForeignPtr end1 (arrayToFptrContents cont1)
701 goArray s [] st fp1 start
702 PR.Continue 0 s ->
703 goArray s (x:backBuf) st fp next
704 PR.Continue n s -> do
705 assert (n <= Prelude.length (x:backBuf)) (return ())
706 let (src0, buf1) = splitAt n (x:backBuf)
707 arr0 = A.fromListN n (Prelude.reverse src0)
708 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
709 src = arr0 <> arr1
710 let !(Array cont1 start (Ptr end1)) = src
711 fp1 = ForeignPtr end1 (arrayToFptrContents cont1)
712 goArray s buf1 st fp1 start
713 PR.Done 0 b -> do
714 let arr = Array (fptrToArrayContents contents) next (Ptr end)
715 return (b, K.cons arr st)
716 PR.Done n b -> do
717 assert (n <= Prelude.length (x:backBuf)) (return ())
718 let src0 = Prelude.take n (x:backBuf)
719 -- XXX Use fromListRevN once implemented
720 -- arr0 = A.fromListRevN n src0
721 arr0 = A.fromListN n (Prelude.reverse src0)
722 arr1 = Array (fptrToArrayContents contents) next (Ptr end)
723 str = K.cons arr0 (K.cons arr1 st)
724 return (b, str)
725 PR.Error err -> throwM $ ParseError err
727-- | Parse an array stream using the supplied 'Parser'. Returns the parse
728-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
729--
730-- /Internal/
731--
732{-# INLINE_NORMAL parseBreak #-}
733parseBreak ::
734 (MonadIO m, MonadThrow m, Storable a)
735 => PR.Parser m a b
736 -> SerialT m (A.Array a)
737 -> m (b, SerialT m (A.Array a))
738{-
739parseBreak p s =
740 fmap fromStreamD <$> parseBreakD (PRD.fromParserK p) (toStreamD s)
741-}
742parseBreak p (SerialT s) =
743 fmap (fmap SerialT) $ parseBreakK (PRD.fromParserK p) s
745-------------------------------------------------------------------------------
746-- Elimination - Running Array Folds and parsers
747-------------------------------------------------------------------------------
749-- | Note that this is not the same as using a @Parser m (Array a) b@ with the
750-- regular "Streamly.Internal.Data.IsStream.parse" function. The regular parse
751-- would consume the input arrays as single unit. This parser parses in the way
752-- as described in the ArrayFold module. The input arrays are treated as @n@
753-- element units and can be consumed partially. The remaining elements are
754-- inserted in the source stream as an array.
755--
756{-# INLINE_NORMAL runArrayParserDBreak #-}
757runArrayParserDBreak ::
758 forall m a b. (MonadIO m, MonadThrow m, Storable a)
759 => PRD.Parser m (Array a) b
760 -> D.Stream m (Array.Array a)
761 -> m (b, D.Stream m (Array.Array a))
762runArrayParserDBreak
763 (PRD.Parser pstep initial extract)
764 stream@(D.Stream step state) = do
766 res <- initial
767 case res of
768 PRD.IPartial s -> go SPEC state (List []) s
769 PRD.IDone b -> return (b, stream)
770 PRD.IError err -> throwM $ ParseError err
772 where
774 -- "backBuf" contains last few items in the stream that we may have to
775 -- backtrack to.
776 --
777 -- XXX currently we are using a dumb list based approach for backtracking
778 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
779 -- That will allow us more efficient random back and forth movement.
780 go !_ st backBuf !pst = do
781 r <- step defState st
782 case r of
783 D.Yield x s -> gobuf SPEC [x] s backBuf pst
784 D.Skip s -> go SPEC s backBuf pst
785 D.Stop -> do
786 b <- extract pst
787 return (b, D.nil)
789 gobuf !_ [] s backBuf !pst = go SPEC s backBuf pst
790 gobuf !_ (x:xs) s backBuf !pst = do
791 pRes <- pstep pst x
792 case pRes of
793 PR.Partial 0 pst1 ->
794 gobuf SPEC xs s (List []) pst1
795 PR.Partial n pst1 -> do
796 assert
797 (n <= sum (map Array.length (x:getList backBuf)))
798 (return ())
799 let src0 = takeArrayListRev n (x:getList backBuf)
800 src = Prelude.reverse src0 ++ xs
801 gobuf SPEC src s (List []) pst1
802 PR.Continue 0 pst1 ->
803 gobuf SPEC xs s (List (x:getList backBuf)) pst1
804 PR.Continue n pst1 -> do
805 assert
806 (n <= sum (map Array.length (x:getList backBuf)))
807 (return ())
808 let (src0, buf1) = splitAtArrayListRev n (x:getList backBuf)
809 src = Prelude.reverse src0 ++ xs
810 gobuf SPEC src s (List buf1) pst1
811 PR.Done 0 b ->
812 return (b, D.Stream step s)
813 PR.Done n b -> do
814 assert
815 (n <= sum (map Array.length (x:getList backBuf)))
816 (return ())
817 let src0 = takeArrayListRev n (x:getList backBuf)
818 src = Prelude.reverse src0 ++ xs
819 return (b, D.append (D.fromList src) (D.Stream step s))
820 PR.Error err -> throwM $ ParseError err
822{-
823-- | Parse an array stream using the supplied 'Parser'. Returns the parse
824-- result and the unconsumed stream. Throws 'ParseError' if the parse fails.
825--
826-- /Internal/
827--
828{-# INLINE parseArr #-}
829parseArr ::
830 (MonadIO m, MonadThrow m, Storable a)
831 => ASF.Parser m a b
832 -> SerialT m (A.Array a)
833 -> m (b, SerialT m (A.Array a))
834parseArr p s = fmap fromStreamD <$> parseBreakD p (toStreamD s)
835-}
837-- | Fold an array stream using the supplied array stream 'Fold'.
838--
839-- /Pre-release/
840--
841{-# INLINE runArrayFold #-}
842runArrayFold :: (MonadIO m, MonadThrow m, Storable a) =>
843 ArrayFold m a b -> SerialT m (A.Array a) -> m b
844runArrayFold (ArrayFold p) s = fst <$> runArrayParserDBreak p (toStreamD s)
846-- | Like 'fold' but also returns the remaining stream.
847--
848-- /Pre-release/
849--
850{-# INLINE runArrayFoldBreak #-}
851runArrayFoldBreak :: (MonadIO m, MonadThrow m, Storable a) =>
852 ArrayFold m a b -> SerialT m (A.Array a) -> m (b, SerialT m (A.Array a))
853runArrayFoldBreak (ArrayFold p) s =
854 second fromStreamD <$> runArrayParserDBreak p (toStreamD s)
856{-# ANN type ParseChunksState Fuse #-}
857data ParseChunksState x inpBuf st pst =
858 ParseChunksInit inpBuf st
859 | ParseChunksInitLeftOver inpBuf
860 | ParseChunksStream st inpBuf !pst
861 | ParseChunksBuf inpBuf st inpBuf !pst
862 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
864{-# INLINE_NORMAL runArrayFoldManyD #-}
865runArrayFoldManyD
866 :: (MonadThrow m, Storable a)
867 => ArrayFold m a b
868 -> D.Stream m (Array a)
869 -> D.Stream m b
870runArrayFoldManyD
871 (ArrayFold (PRD.Parser pstep initial extract)) (D.Stream step state) =
873 D.Stream stepOuter (ParseChunksInit [] state)
875 where
877 {-# INLINE_LATE stepOuter #-}
878 -- Buffer is empty, get the first element from the stream, initialize the
879 -- fold and then go to stream processing loop.
880 stepOuter gst (ParseChunksInit [] st) = do
881 r <- step (adaptState gst) st
882 case r of
883 D.Yield x s -> do
884 res <- initial
885 case res of
886 PRD.IPartial ps ->
887 return $ D.Skip $ ParseChunksBuf [x] s [] ps
888 PRD.IDone pb ->
889 let next = ParseChunksInit [x] s
890 in return $ D.Skip $ ParseChunksYield pb next
891 PRD.IError err -> throwM $ ParseError err
892 D.Skip s -> return $ D.Skip $ ParseChunksInit [] s
893 D.Stop -> return D.Stop
895 -- Buffer is not empty, go to buffered processing loop
896 stepOuter _ (ParseChunksInit src st) = do
897 res <- initial
898 case res of
899 PRD.IPartial ps ->
900 return $ D.Skip $ ParseChunksBuf src st [] ps
901 PRD.IDone pb ->
902 let next = ParseChunksInit src st
903 in return $ D.Skip $ ParseChunksYield pb next
904 PRD.IError err -> throwM $ ParseError err
906 -- XXX we just discard any leftover input at the end
907 stepOuter _ (ParseChunksInitLeftOver _) = return D.Stop
909 -- Buffer is empty, process elements from the stream
910 stepOuter gst (ParseChunksStream st backBuf pst) = do
911 r <- step (adaptState gst) st
912 case r of
913 D.Yield x s -> do
914 pRes <- pstep pst x
915 case pRes of
916 PR.Partial 0 pst1 ->
917 return $ D.Skip $ ParseChunksStream s [] pst1
918 PR.Partial n pst1 -> do
919 assert
920 (n <= sum (map Array.length (x:backBuf)))
921 (return ())
922 let src0 = takeArrayListRev n (x:backBuf)
923 src = Prelude.reverse src0
924 return $ D.Skip $ ParseChunksBuf src s [] pst1
925 PR.Continue 0 pst1 ->
926 return $ D.Skip $ ParseChunksStream s (x:backBuf) pst1
927 PR.Continue n pst1 -> do
928 assert
929 (n <= sum (map Array.length (x:backBuf)))
930 (return ())
931 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
932 src = Prelude.reverse src0
933 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
934 PR.Done 0 b -> do
935 return $ D.Skip $
936 ParseChunksYield b (ParseChunksInit [] s)
937 PR.Done n b -> do
938 assert
939 (n <= sum (map Array.length (x:backBuf)))
940 (return ())
941 let src0 = takeArrayListRev n (x:backBuf)
942 let src = Prelude.reverse src0
943 return $ D.Skip $
944 ParseChunksYield b (ParseChunksInit src s)
945 PR.Error err -> throwM $ ParseError err
946 D.Skip s -> return $ D.Skip $ ParseChunksStream s backBuf pst
947 D.Stop -> do
948 b <- extract pst
949 let src = Prelude.reverse backBuf
950 return $ D.Skip $
951 ParseChunksYield b (ParseChunksInitLeftOver src)
953 -- go back to stream processing mode
954 stepOuter _ (ParseChunksBuf [] s buf pst) =
955 return $ D.Skip $ ParseChunksStream s buf pst
957 -- buffered processing loop
958 stepOuter _ (ParseChunksBuf (x:xs) s backBuf pst) = do
959 pRes <- pstep pst x
960 case pRes of
961 PR.Partial 0 pst1 ->
962 return $ D.Skip $ ParseChunksBuf xs s [] pst1
963 PR.Partial n pst1 -> do
964 assert (n <= sum (map Array.length (x:backBuf))) (return ())
965 let src0 = takeArrayListRev n (x:backBuf)
966 src = Prelude.reverse src0 ++ xs
967 return $ D.Skip $ ParseChunksBuf src s [] pst1
968 PR.Continue 0 pst1 ->
969 return $ D.Skip $ ParseChunksBuf xs s (x:backBuf) pst1
970 PR.Continue n pst1 -> do
971 assert (n <= sum (map Array.length (x:backBuf))) (return ())
972 let (src0, buf1) = splitAtArrayListRev n (x:backBuf)
973 src = Prelude.reverse src0 ++ xs
974 return $ D.Skip $ ParseChunksBuf src s buf1 pst1
975 PR.Done 0 b ->
976 return $ D.Skip $ ParseChunksYield b (ParseChunksInit xs s)
977 PR.Done n b -> do
978 assert (n <= sum (map Array.length (x:backBuf))) (return ())
979 let src0 = takeArrayListRev n (x:backBuf)
980 src = Prelude.reverse src0 ++ xs
981 return $ D.Skip $ ParseChunksYield b (ParseChunksInit src s)
982 PR.Error err -> throwM $ ParseError err
984 stepOuter _ (ParseChunksYield a next) = return $ D.Yield a next
986-- | Apply an 'ArrayFold' repeatedly on an array stream and emit the
987-- fold outputs in the output stream.
988--
989-- See "Streamly.Prelude.foldMany" for more details.
990--
991-- /Pre-release/
992{-# INLINE runArrayFoldMany #-}
993runArrayFoldMany
994 :: (IsStream t, MonadThrow m, Storable a)
995 => ArrayFold m a b
996 -> t m (Array a)
997 -> t m b
998runArrayFoldMany p m = fromStreamD $ runArrayFoldManyD p (toStreamD m)