Make the seed in foldIterateM monadic

benchmark/Streamly/Benchmark/Prelude/Serial/Transformation2.hs

@@ -145,7 +145,8 @@ foldIterateM :: Monad m => SerialT m Int -> m ()
 foldIterateM =
     S.drain
         . S.map getSum
-        . Internal.foldIterateM (return . FL.take 2 . FL.sconcat) (Sum 0)
+        . Internal.foldIterateM
+            (return . FL.take 2 . FL.sconcat) (return (Sum 0))
         . S.map Sum
 
 o_1_space_grouping :: Int -> [Benchmark]

src/Streamly/Internal/Data/Stream/IsStream/Reduce.hs

@@ -316,7 +316,7 @@ foldSequence _f _m = undefined
 -- >>> import Data.Monoid (Sum(..))
 -- >>> f x = return (Fold.take 2 (Fold.sconcat x))
 -- >>> s = Stream.map Sum $ Stream.fromList [1..10]
--- >>> Stream.toList $ Stream.map getSum $ Stream.foldIterateM f 0 s
+-- >>> Stream.toList $ Stream.map getSum $ Stream.foldIterateM f (pure 0) s
 -- [3,10,21,36,55,55]
 --
 -- @
@@ -328,7 +328,7 @@ foldSequence _f _m = undefined
 --
 {-# INLINE foldIterateM #-}
 foldIterateM ::
-       (IsStream t, Monad m) => (b -> m (Fold m a b)) -> b -> t m a -> t m b
+       (IsStream t, Monad m) => (b -> m (Fold m a b)) -> m b -> t m a -> t m b
 foldIterateM f i m = fromStreamD $ D.foldIterateM f i (toStreamD m)
 
 ------------------------------------------------------------------------------

src/Streamly/Internal/Data/Stream/StreamD/Nesting.hs

@@ -992,7 +992,7 @@ data FIterState s f m a b
 
 {-# INLINE_NORMAL foldIterateM #-}
 foldIterateM ::
-       Monad m => (b -> m (FL.Fold m a b)) -> b -> Stream m a -> Stream m b
+       Monad m => (b -> m (FL.Fold m a b)) -> m b -> Stream m a -> Stream m b
 foldIterateM func seed0 (Stream step state) =
     Stream stepOuter (FIterInit state seed0)
 
@@ -1005,11 +1005,11 @@ foldIterateM func seed0 (Stream step state) =
             $ Skip
             $ case res of
                   FL.Partial fs -> FIterStream st fstep fs extract
-                  FL.Done fb -> FIterYield fb $ FIterInit st fb
+                  FL.Done fb -> FIterYield fb $ FIterInit st (return fb)
 
     {-# INLINE_LATE stepOuter #-}
     stepOuter _ (FIterInit st seed) = do
-        (FL.Fold fstep initial extract) <- func seed
+        (FL.Fold fstep initial extract) <- seed >>= func
         iterStep initial st fstep extract
     stepOuter gst (FIterStream st fstep fs extract) = do
         r <- step (adaptState gst) st

test/Streamly/Test/Prelude/Serial.hs

@@ -496,7 +496,7 @@ foldIterateM =
                 $ S.map getSum
                 $ IS.foldIterateM
                       (return . FL.take 1 . FL.sconcat)
-                      (Sum 0)
+                      (return (Sum 0))
                 $ S.map Sum strm
         case ms2 of
             Nothing -> assert $ s1 == 0