on master help sign in
Add fromParser and fold_ to array stream folds
102d2fbHarendra Kumar
src/Streamly/Internal/Data/Array/Stream/Fold/Foreign.hs+50
@@ -32,6 +32,7 @@ module Streamly.Internal.Data.Array.Stream.Fold.Foreign
32
32
33 -- * Construction
33 -- * Construction
34 , fromFold
34 , fromFold
35+ , fromParser
35 , fromArrayFold
36 , fromArrayFold
36
37
37 -- * Mapping
38 -- * Mapping
@@ -127,6 +128,55 @@ fromFold (Fold.Fold fstep finitial fextract) =
127 Fold.Partial fs1 ->
128 Fold.Partial fs1 ->
128 goArray SPEC next fs1
129 goArray SPEC next fs1
129
130
131+-- | Convert an element 'Parser' into an array stream fold. If the parser fails
132+-- the fold would throw an exception.
133+--
134+-- /Pre-release/
135+{-# INLINE fromParser #-}
136+fromParser :: forall m a b. (MonadIO m, Storable a) =>
137+ ParserD.Parser m a b -> Fold m a b
138+fromParser (ParserD.Parser step1 initial1 extract1) =
139+ Fold (ParserD.Parser step initial extract1)
140+
141+ where
142+
143+ initial = do
144+ res <- initial1
145+ return
146+ $ case res of
147+ IPartial s1 -> IPartial s1
148+ IDone b -> IDone b
149+ IError err -> IError err
150+
151+ step s (Array fp@(ForeignPtr start _) end) = do
152+ if Ptr start >= end
153+ then return $ Continue 0 s
154+ else goArray SPEC (Ptr start) s
155+
156+ where
157+
158+ goArray !_ !cur !fs = do
159+ x <- liftIO $ peek cur
160+ liftIO $ touchForeignPtr fp
161+ res <- step1 fs x
162+ let elemSize = sizeOf (undefined :: a)
163+ next = cur `plusPtr` elemSize
164+ arrRem = (end `minusPtr` next) `div` elemSize
165+ case res of
166+ ParserD.Done n b -> do
167+ return $ Done (arrRem + n) b
168+ ParserD.Partial n fs1 -> do
169+ let next1 = next `plusPtr` negate (n * elemSize)
170+ if next1 >= Ptr start && cur < end
171+ then goArray SPEC next1 fs1
172+ else return $ Partial (arrRem + n) fs1
173+ ParserD.Continue n fs1 -> do
174+ let next1 = next `plusPtr` negate (n * elemSize)
175+ if next1 >= Ptr start && cur < end
176+ then goArray SPEC next1 fs1
177+ else return $ Continue (arrRem + n) fs1
178+ Error err -> return $ Error err
179+
130-- | Adapt an array stream fold.
180-- | Adapt an array stream fold.
131--
181--
132-- /Pre-release/
182-- /Pre-release/
src/Streamly/Internal/Data/Array/Stream/Foreign.hs+11
@@ -23,6 +23,7 @@ module Streamly.Internal.Data.Array.Stream.Foreign
23
23
24 -- * Elimination
24 -- * Elimination
25 , fold
25 , fold
26+ , fold_
26 -- , parse
27 -- , parse
27 , parseD
28 , parseD
28 , foldMany
29 , foldMany
@@ -43,6 +44,7 @@ where
43
44
44#include "inline.hs"
45#include "inline.hs"
45
46
47+import Data.Bifunctor (second)
46import Control.Exception (assert)
48import Control.Exception (assert)
47import Control.Monad.Catch (MonadThrow, throwM)
49import Control.Monad.Catch (MonadThrow, throwM)
48import Control.Monad.IO.Class (MonadIO(..))
50import Control.Monad.IO.Class (MonadIO(..))
@@ -549,6 +551,15 @@ fold :: (MonadIO m, MonadThrow m, Storable a) =>
549 ASF.Fold m a b -> SerialT m (A.Array a) -> m b
551 ASF.Fold m a b -> SerialT m (A.Array a) -> m b
550fold (ASF.Fold p) s = fst <$> parseD p (D.toStreamD s)
552fold (ASF.Fold p) s = fst <$> parseD p (D.toStreamD s)
551
553
554+-- | Like 'fold' but also returns the remaining stream.
555+--
556+-- /Pre-release/
557+--
558+{-# INLINE fold_ #-}
559+fold_ :: (MonadIO m, MonadThrow m, Storable a) =>
560+ ASF.Fold m a b -> SerialT m (A.Array a) -> m (b, SerialT m (A.Array a))
561+fold_ (ASF.Fold p) s = second D.fromStreamD <$> parseD p (D.toStreamD s)
562+
552{-# ANN type ParseChunksState Fuse #-}
563{-# ANN type ParseChunksState Fuse #-}
553data ParseChunksState x inpBuf st pst =
564data ParseChunksState x inpBuf st pst =
554 ParseChunksInit inpBuf st
565 ParseChunksInit inpBuf st