mirror of
https://github.com/ilyakooo0/streamly.git
synced 2024-09-11 08:25:40 +03:00
Fix a perf issue in iterate/iterateM
It was taking O(n) space, now it takes O(1)
This commit is contained in:
parent
53aae83e3a
commit
28e17ee82e
@ -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
|
||||
|
||||
|
@ -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)
|
||||
|
@ -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
|
||||
|
@ -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))
|
||||
|
||||
-------------------------------------------------------------------------------
|
||||
|
Loading…
Reference in New Issue
Block a user