@@ -23,6 +23,10 @@ module Streamly.Internal.Data.Stream.Bottom
2323 , fold
2424 , foldContinue
2525 , foldBreak
26+ , foldBreak2
27+ , foldEither
28+ , foldEither2
29+ , foldConcat
2630
2731 -- * Scans
2832 , smapM
@@ -61,10 +65,13 @@ where
6165#include "inline.hs"
6266
6367import Control.Monad.IO.Class (MonadIO(..))
68+import GHC.Types (SPEC(..))
6469import Streamly.Internal.Data.Fold.Type (Fold (..))
6570import Streamly.Internal.Data.Time.Units (AbsTime, RelTime64, addToAbsTime64)
6671import Streamly.Internal.Data.Unboxed (Unbox)
72+import Streamly.Internal.Data.Producer.Type (Producer(..))
6773import Streamly.Internal.System.IO (defaultChunkSize)
74+import Streamly.Internal.Data.SVar.Type (defState)
6875
6976import qualified Streamly.Internal.Data.Array.Unboxed.Type as A
7077import qualified Streamly.Internal.Data.Fold as Fold
@@ -214,21 +221,127 @@ fold fl strm = D.fold fl $ D.fromStreamK $ toStreamK strm
214221--
215222{-# INLINE foldBreak #-}
216223foldBreak :: Monad m => Fold m a b -> Stream m a -> m (b, Stream m a)
217−{-
218−-- XXX This shows quadratic performance when used recursively perhaps because
219−-- of StreamK to StreamD conversions not getting eliminated sue to recursion.
220−foldBreak fl (Stream strm) = fmap f $ D.foldBreak fl $ D.fromStreamK strm
224+foldBreak fl strm = fmap f $ K.foldBreak fl (toStreamK strm)
221225
222226 where
223227
224− f (b, str) = (b, Stream (D.toStreamK str))
225−-}
226−foldBreak fl strm = fmap f $ K.foldBreak fl (toStreamK strm)
228+ f (b, str) = (b, fromStreamK str)
229+
230+-- XXX The quadratic slowdown in recursive use is because recursive function
231+-- cannot be inlined and StreamD/StreamK conversions pile up and cannot be
232+-- eliminated by rewrite rules.
233+
234+-- | Like 'foldBreak' but fuses.
235+--
236+-- /Note:/ Unlike 'foldBreak', recursive application on the resulting stream
237+-- would lead to quadratic slowdown. If you need recursion with fusion (within
238+-- one iteration of recursion) use StreamD.foldBreak directly.
239+--
240+-- /Internal/
241+{-# INLINE foldBreak2 #-}
242+foldBreak2 :: Monad m => Fold m a b -> Stream m a -> m (b, Stream m a)
243+foldBreak2 fl strm = fmap f $ D.foldBreak fl $ toStreamD strm
244+
245+ where
246+
247+ f (b, str) = (b, fromStreamD str)
248+
249+-- | Fold resulting in either breaking the stream or continuation of the fold
250+-- Instead of supplying the input stream in one go we can run the fold multiple
251+-- times each time supplying the next segment of the input stream. If the fold
252+-- has not yet finished it returns a fold that can be run again otherwise it
253+-- returns the fold result and the residual stream.
254+--
255+-- /Internal/
256+{-# INLINE foldEither #-}
257+foldEither :: Monad m =>
258+ Fold m a b -> Stream m a -> m (Either (Fold m a b) (b, Stream m a))
259+foldEither fl strm = fmap (fmap f) $ K.foldEither fl $ toStreamK strm
227260
228261 where
229262
230263 f (b, str) = (b, fromStreamK str)
231264
265+-- | Like 'foldEither' but fuses. However, recursive application on resulting
266+-- stream would lead to quadratic slowdown.
267+--
268+-- /Internal/
269+{-# INLINE foldEither2 #-}
270+foldEither2 :: Monad m =>
271+ Fold m a b -> Stream m a -> m (Either (Fold m a b) (b, Stream m a))
272+foldEither2 fl strm = fmap (fmap f) $ D.foldEither fl $ toStreamD strm
273+
274+ where
275+
276+ f (b, str) = (b, fromStreamD str)
277+
278+-- XXX Array folds can be implemented using this.
279+-- foldContainers? Specialized to foldArrays.
280+
281+-- | Generate streams from individual elements of a stream and fold the
282+-- concatenation of those streams using the supplied fold. Return the result of
283+-- the fold and residual stream.
284+--
285+-- For example, this can be used to efficiently fold an Array Word8 stream
286+-- using Word8 folds.
287+--
288+-- The outer stream forces CPS to allow scalable appends and the inner stream
289+-- forces direct style for stream fusion.
290+--
291+-- /Internal/
292+{-# INLINE foldConcat #-}
293+foldConcat :: Monad m =>
294+ Producer m a b -> Fold m b c -> Stream m a -> m (c, Stream m a)
295+foldConcat
296+ (Producer pstep pinject pextract)
297+ (Fold fstep begin done)
298+ stream = do
299+
300+ res <- begin
301+ case res of
302+ Fold.Partial fs -> go fs streamK
303+ Fold.Done fb -> return (fb, fromStreamK streamK)
304+
305+ where
306+
307+ streamK = toStreamK stream
308+
309+ go !acc m1 = do
310+ let stop = do
311+ r <- done acc
312+ return (r, fromStreamK K.nil)
313+ single a = do
314+ st <- pinject a
315+ res <- go1 SPEC acc st
316+ case res of
317+ Left fs -> do
318+ r <- done fs
319+ return (r, fromStreamK K.nil)
320+ Right (b, s) -> do
321+ x <- pextract s
322+ return (b, fromStreamK (K.fromPure x))
323+ yieldk a r = do
324+ st <- pinject a
325+ res <- go1 SPEC acc st
326+ case res of
327+ Left fs -> go fs r
328+ Right (b, s) -> do
329+ x <- pextract s
330+ return (b, fromStreamK (x `K.cons` r))
331+ in K.foldStream defState yieldk single stop m1
332+
333+ {-# INLINE go1 #-}
334+ go1 !_ !fs st = do
335+ r <- pstep st
336+ case r of
337+ D.Yield x s -> do
338+ res <- fstep fs x
339+ case res of
340+ Fold.Done b -> return $ Right (b, s)
341+ Fold.Partial fs1 -> go1 SPEC fs1 s
342+ D.Skip s -> go1 SPEC fs s
343+ D.Stop -> return $ Left fs
344+
232345------------------------------------------------------------------------------
233346-- Transformation
234347------------------------------------------------------------------------------