Fix a perf issue in iterate/iterateM

It was taking O(n) space, now it takes O(1)

Changelog.md

@@ -96,6 +96,8 @@ New APIs:
 * `bracket`, `handle`, and `finally` now work on streams that aren't
   fully drained. Also, the resource acquisition and release is atomic with
   respect to async exceptions.
+* Fix a performance issue in `Streamly.Prelude.iterate/iterateM` that caused it
+  to consume O(n) space.
 
 ### Deprecations
 

benchmark/Streamly/Benchmark/Data/Stream/StreamK.hs

@@ -24,7 +24,9 @@ import Control.Monad (when)
 import Data.Maybe (isJust)
 import System.Random (randomRIO)
 import Prelude hiding
-    (tail, mapM_, foldl, last, map, mapM, concatMap, zipWith, init)
+    ( tail, mapM_, foldl, last, map, mapM, concatMap, zipWith, init, iterate
+    , repeat, replicate
+    )
 
 import qualified Prelude as P
 import qualified Data.List as List
@@ -44,12 +46,6 @@ import Test.Inspection
 
 Benchmarks that need to be added
 
--- repeat
--- repeatM
--- replicate
--- replicateM
--- iterate
--- iterateM
 -- fromList
 
 -- bindWith
@@ -138,6 +134,30 @@ unfoldrM streamLen n = S.unfoldrM step n
         then return Nothing
         else return (Just (cnt, cnt + 1))
 
+{-# INLINE repeat #-}
+repeat :: Int -> Int -> Stream m Int
+repeat streamLen = S.take streamLen . S.repeat
+
+{-# INLINE repeatM #-}
+repeatM :: S.MonadAsync m => Int -> Int -> Stream m Int
+repeatM streamLen = S.take streamLen . S.repeatM . return
+
+{-# INLINE replicate #-}
+replicate :: Int -> Int -> Stream m Int
+replicate = S.replicate
+
+{-# INLINE replicateM #-}
+replicateM :: S.MonadAsync m => Int -> Int -> Stream m Int
+replicateM streamLen = S.replicateM streamLen . return
+
+{-# INLINE iterate #-}
+iterate :: Int -> Int -> Stream m Int
+iterate streamLen = S.take streamLen . S.iterate (+1)
+
+{-# INLINE iterateM #-}
+iterateM :: S.MonadAsync m => Int -> Int -> Stream m Int
+iterateM streamLen = S.take streamLen . S.iterateM (return . (+1)) . return
+
 {-# INLINE fromFoldable #-}
 fromFoldable :: Int -> Int -> Stream m Int
 fromFoldable streamLen n = S.fromFoldable [n..n+streamLen]
@@ -570,8 +590,14 @@ moduleName = "Data.Stream.StreamK"
 o_1_space_generation :: Int -> Benchmark
 o_1_space_generation streamLen =
     bgroup "generation"
-        [ benchFold "unfoldr"       drain (unfoldr streamLen)
-        , benchFold "unfoldrM"      drain (unfoldrM streamLen)
+        [ benchFold "unfoldr" drain (unfoldr streamLen)
+        , benchFold "unfoldrM" drain (unfoldrM streamLen)
+        , benchFold "repeat" drain (repeat streamLen)
+        , benchFold "repeatM" drain (repeatM streamLen)
+        , benchFold "replicate" drain (replicate streamLen)
+        , benchFold "replicateM" drain (replicateM streamLen)
+        , benchFold "iterate" drain (iterate streamLen)
+        , benchFold "iterateM" drain (iterateM streamLen)
 
         , benchFold "fromFoldable"  drain (fromFoldable streamLen)
         , benchFold "fromFoldableM" drain (fromFoldableM streamLen)

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

@@ -384,7 +384,7 @@ generate n gen = generateM n (return . gen)
 
 {-# INLINE_NORMAL iterateM #-}
 iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
-iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))
+iterateM step = Stream (\_ st -> st >>= \(!x) -> return $ Yield x (step x))
 
 {-# INLINE_NORMAL iterate #-}
 iterate :: Monad m => (a -> a) -> a -> Stream m a

src/Streamly/Internal/Data/Stream/StreamK.hs

@@ -316,16 +316,16 @@ fromIndices gen = go 0
 
 {-# INLINE iterate #-}
 iterate :: IsStream t => (a -> a) -> a -> t m a
-iterate step = fromStream . go
+iterate step = go
     where
-        go s = cons s (go (step s))
+        go !s = cons s (go (step s))
 
 {-# INLINE iterateM #-}
 iterateM :: (IsStream t, MonadAsync m) => (a -> m a) -> m a -> t m a
 iterateM step = go
     where
     go s = mkStream $ \st stp sng yld -> do
-        next <- s
+        !next <- s
         foldStreamShared st stp sng yld (return next |: go (step next))
 
 -------------------------------------------------------------------------------