git.iko.soy/streamly
on master help sign in
Fix a perf issue in iterate/iterateM
28e17eeHarendra Kumar
Changelog.md+2
compare with master
unifiedsplit
@@ -96,6 +96,8 @@ New APIs:
96* `bracket`, `handle`, and `finally` now work on streams that aren't
96* `bracket`, `handle`, and `finally` now work on streams that aren't
97 fully drained. Also, the resource acquisition and release is atomic with
97 fully drained. Also, the resource acquisition and release is atomic with
98 respect to async exceptions.
98 respect to async exceptions.
99+* Fix a performance issue in `Streamly.Prelude.iterate/iterateM` that caused it
100+ to consume O(n) space.
99
101
100### Deprecations
102### Deprecations
101
103
benchmark/Streamly/Benchmark/Data/Stream/StreamK.hs+35−9
@@ -24,7 +24,9 @@ import Control.Monad (when)
24import Data.Maybe (isJust)
24import Data.Maybe (isJust)
25import System.Random (randomRIO)
25import System.Random (randomRIO)
26import Prelude hiding
26import 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+ )
28
30
29import qualified Prelude as P
31import qualified Prelude as P
30import qualified Data.List as List
32import qualified Data.List as List
@@ -44,12 +46,6 @@ import Test.Inspection
44
46
45Benchmarks that need to be added
47Benchmarks that need to be added
46
48
47−-- repeat
48−-- repeatM
49−-- replicate
50−-- replicateM
51−-- iterate
52−-- iterateM
53-- fromList
49-- fromList
54
50
55-- bindWith
51-- bindWith
@@ -138,6 +134,30 @@ unfoldrM streamLen n = S.unfoldrM step n
138 then return Nothing
134 then return Nothing
139 else return (Just (cnt, cnt + 1))
135 else return (Just (cnt, cnt + 1))
140
136
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+
141{-# INLINE fromFoldable #-}
161{-# INLINE fromFoldable #-}
142fromFoldable :: Int -> Int -> Stream m Int
162fromFoldable :: Int -> Int -> Stream m Int
143fromFoldable streamLen n = S.fromFoldable [n..n+streamLen]
163fromFoldable streamLen n = S.fromFoldable [n..n+streamLen]
@@ -570,8 +590,14 @@ moduleName = "Data.Stream.StreamK"
570o_1_space_generation :: Int -> Benchmark
590o_1_space_generation :: Int -> Benchmark
571o_1_space_generation streamLen =
591o_1_space_generation streamLen =
572 bgroup "generation"
592 bgroup "generation"
573− [ benchFold "unfoldr" drain (unfoldr streamLen)
593+ [ benchFold "unfoldr" drain (unfoldr streamLen)
574− , benchFold "unfoldrM" drain (unfoldrM 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)
575
601
576 , benchFold "fromFoldable" drain (fromFoldable streamLen)
602 , benchFold "fromFoldable" drain (fromFoldable streamLen)
577 , benchFold "fromFoldableM" drain (fromFoldableM streamLen)
603 , 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)
384
384
385{-# INLINE_NORMAL iterateM #-}
385{-# INLINE_NORMAL iterateM #-}
386iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
386iterateM :: 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))
388
388
389{-# INLINE_NORMAL iterate #-}
389{-# INLINE_NORMAL iterate #-}
390iterate :: Monad m => (a -> a) -> a -> Stream m a
390iterate :: 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
316
316
317{-# INLINE iterate #-}
317{-# INLINE iterate #-}
318iterate :: IsStream t => (a -> a) -> a -> t m a
318iterate :: IsStream t => (a -> a) -> a -> t m a
319−iterate step = fromStream . go
319+iterate step = go
320 where
320 where
321− go s = cons s (go (step s))
321+ go !s = cons s (go (step s))
322
322
323{-# INLINE iterateM #-}
323{-# INLINE iterateM #-}
324iterateM :: (IsStream t, MonadAsync m) => (a -> m a) -> m a -> t m a
324iterateM :: (IsStream t, MonadAsync m) => (a -> m a) -> m a -> t m a
325iterateM step = go
325iterateM step = go
326 where
326 where
327 go s = mkStream $ \st stp sng yld -> do
327 go s = mkStream $ \st stp sng yld -> do
328− next <- s
328+ !next <- s
329 foldStreamShared st stp sng yld (return next |: go (step next))
329 foldStreamShared st stp sng yld (return next |: go (step next))
330
330
331-------------------------------------------------------------------------------
331-------------------------------------------------------------------------------