git.iko.soy/streamly
on master help sign in
Fix a perf issue in iterate/iterateM

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

commit28e17ee82e6d0ce3a8c84cf77fe87f3fb95b42e0 authorHarendra Kumar <harendra@composewell.com> dateWed 2 Jun 2021, 22:18 (5 years ago) parent53aae83
4 files changed+41−13
compare with master
unifiedsplit
Changelog.md+2 benchmark/Streamly/Benchmark/Data/Stream/StreamK.hs+35−9 src/Streamly/Internal/Data/Stream/StreamD/Generate.hs+1−1 src/Streamly/Internal/Data/Stream/StreamK.hs+3−3
Changelog.md+2
@@ -96,6 +96,8 @@ New APIs:
9696* `bracket`, `handle`, and `finally` now work on streams that aren't
9797 fully drained. Also, the resource acquisition and release is atomic with
9898 respect to async exceptions.
99+* Fix a performance issue in `Streamly.Prelude.iterate/iterateM` that caused it
100+ to consume O(n) space.
99101
100102### Deprecations
101103
benchmark/Streamly/Benchmark/Data/Stream/StreamK.hs+35−9
@@ -24,7 +24,9 @@ import Control.Monad (when)
2424import Data.Maybe (isJust)
2525import System.Random (randomRIO)
2626import Prelude hiding
27− (tail, mapM_, foldl, last, map, mapM, concatMap, zipWith, init)
27+ ( tail, mapM_, foldl, last, map, mapM, concatMap, zipWith, init, iterate
28+ , repeat, replicate
29+ )
2830
2931import qualified Prelude as P
3032import qualified Data.List as List
@@ -44,12 +46,6 @@ import Test.Inspection
4446
4547Benchmarks that need to be added
4648
47−-- repeat
48−-- repeatM
49−-- replicate
50−-- replicateM
51−-- iterate
52−-- iterateM
5349-- fromList
5450
5551-- bindWith
@@ -138,6 +134,30 @@ unfoldrM streamLen n = S.unfoldrM step n
138134 then return Nothing
139135 else return (Just (cnt, cnt + 1))
140136
137+{-# INLINE repeat #-}
138+repeat :: Int -> Int -> Stream m Int
139+repeat streamLen = S.take streamLen . S.repeat
140+
141+{-# INLINE repeatM #-}
142+repeatM :: S.MonadAsync m => Int -> Int -> Stream m Int
143+repeatM streamLen = S.take streamLen . S.repeatM . return
144+
145+{-# INLINE replicate #-}
146+replicate :: Int -> Int -> Stream m Int
147+replicate = S.replicate
148+
149+{-# INLINE replicateM #-}
150+replicateM :: S.MonadAsync m => Int -> Int -> Stream m Int
151+replicateM streamLen = S.replicateM streamLen . return
152+
153+{-# INLINE iterate #-}
154+iterate :: Int -> Int -> Stream m Int
155+iterate streamLen = S.take streamLen . S.iterate (+1)
156+
157+{-# INLINE iterateM #-}
158+iterateM :: S.MonadAsync m => Int -> Int -> Stream m Int
159+iterateM streamLen = S.take streamLen . S.iterateM (return . (+1)) . return
160+
141161{-# INLINE fromFoldable #-}
142162fromFoldable :: Int -> Int -> Stream m Int
143163fromFoldable streamLen n = S.fromFoldable [n..n+streamLen]
@@ -570,8 +590,14 @@ moduleName = "Data.Stream.StreamK"
570590o_1_space_generation :: Int -> Benchmark
571591o_1_space_generation streamLen =
572592 bgroup "generation"
573− [ benchFold "unfoldr" drain (unfoldr streamLen)
574− , benchFold "unfoldrM" drain (unfoldrM streamLen)
593+ [ benchFold "unfoldr" drain (unfoldr streamLen)
594+ , benchFold "unfoldrM" drain (unfoldrM streamLen)
595+ , benchFold "repeat" drain (repeat streamLen)
596+ , benchFold "repeatM" drain (repeatM streamLen)
597+ , benchFold "replicate" drain (replicate streamLen)
598+ , benchFold "replicateM" drain (replicateM streamLen)
599+ , benchFold "iterate" drain (iterate streamLen)
600+ , benchFold "iterateM" drain (iterateM streamLen)
575601
576602 , benchFold "fromFoldable" drain (fromFoldable streamLen)
577603 , benchFold "fromFoldableM" drain (fromFoldableM streamLen)
src/Streamly/Internal/Data/Stream/StreamD/Generate.hs+1−1
@@ -384,7 +384,7 @@ generate n gen = generateM n (return . gen)
384384
385385{-# INLINE_NORMAL iterateM #-}
386386iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
387−iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))
387+iterateM step = Stream (\_ st -> st >>= \(!x) -> return $ Yield x (step x))
388388
389389{-# INLINE_NORMAL iterate #-}
390390iterate :: Monad m => (a -> a) -> a -> Stream m a
src/Streamly/Internal/Data/Stream/StreamK.hs+3−3
@@ -316,16 +316,16 @@ fromIndices gen = go 0
316316
317317{-# INLINE iterate #-}
318318iterate :: IsStream t => (a -> a) -> a -> t m a
319−iterate step = fromStream . go
319+iterate step = go
320320 where
321− go s = cons s (go (step s))
321+ go !s = cons s (go (step s))
322322
323323{-# INLINE iterateM #-}
324324iterateM :: (IsStream t, MonadAsync m) => (a -> m a) -> m a -> t m a
325325iterateM step = go
326326 where
327327 go s = mkStream $ \st stp sng yld -> do
328− next <- s
328+ !next <- s
329329 foldStreamShared st stp sng yld (return next |: go (step next))
330330
331331-------------------------------------------------------------------------------