Add foldMany to unfolds

benchmark/Streamly/Benchmark/Data/Unfold.hs

@@ -21,6 +21,8 @@ module Main (main) where
 
 import Control.DeepSeq (NFData(..))
 import Control.Exception (SomeException, ErrorCall, try)
+import Data.Char (ord)
+import Data.Word (Word8)
 import Streamly.Internal.Data.Unfold (Unfold)
 import System.IO (Handle, hClose)
 import System.Random (randomRIO)
@@ -693,8 +695,17 @@ o_1_space_zip size =
           ]
     ]
 
-o_1_space_nested :: Int -> [Benchmark]
-o_1_space_nested size =
+lf :: Word8
+lf = fromIntegral (ord '\n')
+
+-- | Split on line feed.
+foldManySepBy :: Handle -> IO Int
+foldManySepBy =
+    let u = UF.foldMany (FL.takeEndBy_ (== lf) FL.drain) FH.read
+     in UF.fold FL.length u
+
+o_1_space_nested :: BenchEnv -> Int -> [Benchmark]
+o_1_space_nested env size =
     [ bgroup
           "nested"
           [ benchIO "(<*>) (sqrt n x sqrt n)" $ toNullAp size
@@ -711,6 +722,8 @@ o_1_space_nested size =
           , benchIO "filterSome" $ filterSome size
 
           , benchIO "concat" $ concat size
+          , mkBench "foldMany (Fold.takeEndBy_ (== lf) Fold.drain)" env
+            $ \inh _ -> foldManySepBy inh
           ]
     ]
 
@@ -819,7 +832,7 @@ main = do
                   , o_1_space_transformation size
                   , o_1_space_filtering size
                   , o_1_space_zip size
-                  , o_1_space_nested size
+                  , o_1_space_nested env size
                   , o_1_space_copy_read_exceptions env
                   ]
         , bgroup (o_n_space_prefix moduleName)

core/src/Streamly/Internal/Data/Unfold.hs

@@ -362,12 +362,61 @@ fold (Fold fstep initial extract) (Unfold ustep inject) a = do
             Skip s -> go SPEC fs s
             Stop -> extract fs
 
+-- {-# ANN type FoldMany Fuse #-}
+data FoldMany s fs b a
+    = FoldManyStart s
+    | FoldManyFirst fs s
+    | FoldManyLoop s fs
+    | FoldManyYield b (FoldMany s fs b a)
+    | FoldManyDone
+
 -- | Apply a fold multiple times on the output of an unfold.
 --
--- /Unimplemented/
-foldMany :: -- Monad m =>
-    Fold m b c -> Unfold m a b -> Unfold m a c
-foldMany = undefined
+-- /Pre-release/
+{-# INLINE_NORMAL foldMany #-}
+foldMany :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c
+foldMany (Fold fstep initial extract) (Unfold ustep inject1) =
+    Unfold step inject
+
+    where
+
+    inject x = do
+        r <- inject1 x
+        return (FoldManyStart r)
+
+    {-# INLINE consume #-}
+    consume x s fs = do
+        res <- fstep fs x
+        return
+            $ Skip
+            $ case res of
+                  FL.Done b -> FoldManyYield b (FoldManyStart s)
+                  FL.Partial ps -> FoldManyLoop s ps
+
+    {-# INLINE_LATE step #-}
+    step (FoldManyStart st) = do
+        r <- initial
+        return
+            $ Skip
+            $ case r of
+                  FL.Done b -> FoldManyYield b (FoldManyStart st)
+                  FL.Partial fs -> FoldManyFirst fs st
+    step (FoldManyFirst fs st) = do
+        r <- ustep st
+        case r of
+            Yield x s -> consume x s fs
+            Skip s -> return $ Skip (FoldManyFirst fs s)
+            Stop -> return Stop
+    step (FoldManyLoop st fs) = do
+        r <- ustep st
+        case r of
+            Yield x s -> consume x s fs
+            Skip s -> return $ Skip (FoldManyLoop s fs)
+            Stop -> do
+                b <- extract fs
+                return $ Skip (FoldManyYield b FoldManyDone)
+    step (FoldManyYield b next) = return $ Yield b next
+    step FoldManyDone = return Stop
 
 -- | Apply a monadic function to each element of the stream and replace it
 -- with the output of the resulting action.