on master help sign in
Implement foldBreak using CPS
commit93d8e0139eb7f4afa4e7e5fc6fbd4ddfe26e92a3 authorHarendra Kumar <harendra@composewell.com> dateThu 19 May 2022, 00:01 (4 years ago) parentd68e0e6
2 files changed+39−2
core/src/Streamly/Internal/Data/Stream/StreamK.hs+25 src/Streamly/Internal/Data/Stream/IsStream/Common.hs+14−2
core/src/Streamly/Internal/Data/Stream/StreamK.hs+25
@@ -82,6 +82,7 @@ module Streamly.Internal.Data.Stream.StreamK
8282 , foldlx'
8383 , foldlMx'
8484 , fold
85+ , foldBreak
8586
8687 -- ** Specialized Folds
8788 , drain
@@ -327,6 +328,30 @@ fold (FL.Fold step begin done) m = do
327328 FL.Done b1 -> return b1
328329 in foldStream defState yieldk single stop m1
329330
331+{-# INLINE foldBreak #-}
332+foldBreak :: Monad m => FL.Fold m a b -> Stream m a -> m (b, Stream m a)
333+foldBreak (FL.Fold step begin done) m = do
334+ res <- begin
335+ case res of
336+ FL.Partial fs -> go fs m
337+ FL.Done fb -> return (fb, m)
338+
339+ where
340+
341+ go !acc m1 =
342+ let stop = (, nil) <$> done acc
343+ single a =
344+ step acc a
345+ >>= \case
346+ FL.Partial s -> (, nil) <$> done s
347+ FL.Done b1 -> return (b1, nil)
348+ yieldk a r =
349+ step acc a
350+ >>= \case
351+ FL.Partial s -> go s r
352+ FL.Done b1 -> return (b1, r)
353+ in foldStream defState yieldk single stop m1
354+
330355-- | Like 'foldl'' but with a monadic step function.
331356{-# INLINE foldlM' #-}
332357foldlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> m b
src/Streamly/Internal/Data/Stream/IsStream/Common.hs+14−2
@@ -83,6 +83,7 @@ import qualified Streamly.Internal.Data.Stream.Async as Async
8383import qualified Streamly.Internal.Data.Stream.IsStream.Type as IsStream
8484import qualified Streamly.Internal.Data.Stream.Parallel as Par
8585import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
86+import qualified Streamly.Internal.Data.Stream.StreamK as K
8687import qualified Streamly.Internal.Data.Stream.StreamD as D
8788#ifdef USE_STREAMK_ONLY
8889import qualified Streamly.Internal.Data.Stream.StreamK as S
@@ -309,15 +310,26 @@ foldContinue f s = D.foldContinue f $ IsStream.toStreamD s
309310fold :: Monad m => Fold m a b -> SerialT m a -> m b
310311fold fl (SerialT strm) = D.fold fl $ D.fromStreamK strm
311312
312−-- XXX This shows quadratic performance when used recursively perhaps because
313−-- of StreamK to StreamD conversions not getting eliminated sue to recursion.
313+-- | Like 'fold' but also returns the remaining stream.
314+--
315+-- /Inhibits stream fusion/
316+--
314317{-# INLINE foldBreak #-}
315318foldBreak :: Monad m => Fold m a b -> SerialT m a -> m (b, SerialT m a)
319+{-
320+-- XXX This shows quadratic performance when used recursively perhaps because
321+-- of StreamK to StreamD conversions not getting eliminated sue to recursion.
316322foldBreak fl (SerialT strm) = fmap f $ D.foldBreak fl $ D.fromStreamK strm
317323
318324 where
319325
320326 f (b, str) = (b, SerialT (D.toStreamK str))
327+-}
328+foldBreak fl (SerialT strm) = fmap f $ K.foldBreak fl strm
329+
330+ where
331+
332+ f (b, str) = (b, SerialT str)
321333
322334------------------------------------------------------------------------------
323335-- Transformation