@@ -24,6 +24,7 @@ module Streamly.Internal.Data.Array.Stream.Foreign
2424 -- * Elimination
2525 -- ** Element Folds
2626 , foldBreak
27+ , foldBreakD
2728 , parseBreak
2829 , parseBreakD
2930
@@ -73,7 +74,7 @@ import Streamly.Internal.Data.Array.Foreign.Type (Array(..))
7374import Streamly.Internal.Data.Array.Stream.Fold.Foreign (ArrayFold(..))
7475import Streamly.Internal.Data.Fold.Type (Fold(..))
7576import Streamly.Internal.Data.Parser (ParseError(..))
76−import Streamly.Internal.Data.Stream.Serial (SerialT)
77+import Streamly.Internal.Data.Stream.Serial (SerialT(..))
7778import Streamly.Internal.Data.Stream.IsStream.Type
7879 (IsStream, fromStreamD, toStreamD)
7980import Streamly.Internal.Data.SVar (adaptState, defState)
@@ -92,6 +93,7 @@ import qualified Streamly.Internal.Data.Parser as PR
9293import qualified Streamly.Internal.Data.Parser.ParserD as PRD
9394import qualified Streamly.Internal.Data.Stream.IsStream as S
9495import qualified Streamly.Internal.Data.Stream.StreamD as D
96+import qualified Streamly.Internal.Data.Stream.StreamK as K
9597
9698-- XXX Since these are immutable arrays MonadIO constraint can be removed from
9799-- most places.
@@ -358,6 +360,40 @@ foldBreakD (Fold fstep initial extract) stream@(D.Stream step state) = do
358360 return $! (b, D.cons arr (D.Stream step st))
359361 FL.Partial fs1 -> goArray SPEC st fp next fs1
360362
363+{-# INLINE_NORMAL foldBreakK #-}
364+foldBreakK :: 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))
366+foldBreakK (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)
371+
372+ where
373+
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
382+
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
396+
361397-- | Fold an array stream using the supplied 'Fold'. Returns the fold result
362398-- and the unconsumed stream.
363399--
@@ -371,10 +407,11 @@ foldBreak ::
371407 => FL.Fold m a b
372408 -> SerialT m (A.Array a)
373409 -> m (b, SerialT m (A.Array a))
374−foldBreak f s = fmap fromStreamD <$> foldBreakD f (toStreamD s)
375−-- If fold performs better than runArrayFoldBreak we can rewrite runArrayFoldBreak to
376−-- fold.
377−-- fold f = runArrayFoldBreak (ArrayFold.fromFold f)
410+-- foldBreak f s = fmap fromStreamD <$> foldBreakD f (toStreamD s)
411+foldBreak 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)
378415
379416-------------------------------------------------------------------------------
380417-- Fold to a single Array