git.iko.soy/streamly
on master help sign in
Implement array stream foldBreak using CPS
d6b509eHarendra Kumar
src/Streamly/Internal/Data/Array/Stream/Foreign.hs+42−5
compare with master
unifiedsplit
@@ -24,6 +24,7 @@ module Streamly.Internal.Data.Array.Stream.Foreign
24 -- * Elimination
24 -- * Elimination
25 -- ** Element Folds
25 -- ** Element Folds
26 , foldBreak
26 , foldBreak
27+ , foldBreakD
27 , parseBreak
28 , parseBreak
28 , parseBreakD
29 , parseBreakD
29
30
@@ -73,7 +74,7 @@ import Streamly.Internal.Data.Array.Foreign.Type (Array(..))
73import Streamly.Internal.Data.Array.Stream.Fold.Foreign (ArrayFold(..))
74import Streamly.Internal.Data.Array.Stream.Fold.Foreign (ArrayFold(..))
74import Streamly.Internal.Data.Fold.Type (Fold(..))
75import Streamly.Internal.Data.Fold.Type (Fold(..))
75import Streamly.Internal.Data.Parser (ParseError(..))
76import Streamly.Internal.Data.Parser (ParseError(..))
76−import Streamly.Internal.Data.Stream.Serial (SerialT)
77+import Streamly.Internal.Data.Stream.Serial (SerialT(..))
77import Streamly.Internal.Data.Stream.IsStream.Type
78import Streamly.Internal.Data.Stream.IsStream.Type
78 (IsStream, fromStreamD, toStreamD)
79 (IsStream, fromStreamD, toStreamD)
79import Streamly.Internal.Data.SVar (adaptState, defState)
80import Streamly.Internal.Data.SVar (adaptState, defState)
@@ -92,6 +93,7 @@ import qualified Streamly.Internal.Data.Parser as PR
92import qualified Streamly.Internal.Data.Parser.ParserD as PRD
93import qualified Streamly.Internal.Data.Parser.ParserD as PRD
93import qualified Streamly.Internal.Data.Stream.IsStream as S
94import qualified Streamly.Internal.Data.Stream.IsStream as S
94import qualified Streamly.Internal.Data.Stream.StreamD as D
95import qualified Streamly.Internal.Data.Stream.StreamD as D
96+import qualified Streamly.Internal.Data.Stream.StreamK as K
95
97
96-- XXX Since these are immutable arrays MonadIO constraint can be removed from
98-- XXX Since these are immutable arrays MonadIO constraint can be removed from
97-- most places.
99-- most places.
@@ -358,6 +360,40 @@ foldBreakD (Fold fstep initial extract) stream@(D.Stream step state) = do
358 return $! (b, D.cons arr (D.Stream step st))
360 return $! (b, D.cons arr (D.Stream step st))
359 FL.Partial fs1 -> goArray SPEC st fp next fs1
361 FL.Partial fs1 -> goArray SPEC st fp next fs1
360
362
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+
361-- | Fold an array stream using the supplied 'Fold'. Returns the fold result
397-- | Fold an array stream using the supplied 'Fold'. Returns the fold result
362-- and the unconsumed stream.
398-- and the unconsumed stream.
363--
399--
@@ -371,10 +407,11 @@ foldBreak ::
371 => FL.Fold m a b
407 => FL.Fold m a b
372 -> SerialT m (A.Array a)
408 -> SerialT m (A.Array a)
373 -> m (b, SerialT m (A.Array a))
409 -> m (b, SerialT m (A.Array a))
374−foldBreak f s = fmap fromStreamD <$> foldBreakD f (toStreamD s)
410+-- foldBreak f s = fmap fromStreamD <$> foldBreakD f (toStreamD s)
411+foldBreak f (SerialT s) = fmap (fmap SerialT) $ foldBreakK f s
375−-- If fold performs better than runArrayFoldBreak we can rewrite runArrayFoldBreak to
412+-- If foldBreak performs better than runArrayFoldBreak we can use a rewrite
376−-- fold.
413+-- rule to rewrite runArrayFoldBreak to fold.
377−-- fold f = runArrayFoldBreak (ArrayFold.fromFold f)
414+-- foldBreak f = runArrayFoldBreak (ArrayFold.fromFold f)
378
415
379-------------------------------------------------------------------------------
416-------------------------------------------------------------------------------
380-- Fold to a single Array
417-- Fold to a single Array