511fe85 · streamly · git.iko.soy

git.iko.soy/streamly

Use StreamD for iterate and iterateM

commit511fe85cc619352db227dfbb3ce1946e2c31a2cd authoradithyaov <adi.obilisetty@gmail.com> dateThu 5 Dec 2019, 14:23 (7 years ago) parent0d49be3

5 files changed+46−11

compare with master

unifiedsplit

benchmark/BaseStreams.hs+2 benchmark/StreamDOps.hs+4 src/Streamly/Internal/Data/Stream/StreamD.hs+13−1 src/Streamly/Internal/Data/Stream/StreamK.hs+17−1 src/Streamly/Internal/Prelude.hs+10−9

benchmark/BaseStreams.hs+2

@@ -194,6 +194,8 @@ main =

194194 , benchD "dropOne" D.iterateDropOne

195195 , benchD "dropWhileFalse(1/10)" D.iterateDropWhileFalse

196196 , benchD "dropWhileTrue" D.iterateDropWhileTrue

197+ , benchD "iterateM" D.iterateM

198+

197199 ]

198200 ]

199201 , bgroup "list"

benchmark/StreamDOps.hs+4

@@ -251,6 +251,10 @@ iterateTakeAll = iterateSource (S.take maxValue) maxIters

251251iterateDropOne = iterateSource (S.drop 1) maxIters

252252iterateDropWhileTrue = iterateSource (S.dropWhile (<= maxValue)) maxIters

253253

254+{-# INLINE iterateM #-}

255+iterateM :: Monad m => Int -> Stream m Int

256+iterateM i = S.take maxIters (S.iterateM (\x -> return (x + 1)) (return i))

257+

254258-------------------------------------------------------------------------------

255259-- Zipping and concat

256260-------------------------------------------------------------------------------

src/Streamly/Internal/Data/Stream/StreamD.hs+13−1

@@ -78,6 +78,8 @@ module Streamly.Internal.Data.Stream.StreamD

7878 , fromIndicesM

7979 , generate

8080 , generateM

81+ , iterate

82+ , iterateM

8183

8284 -- ** Enumerations

8385 , enumerateFromStepIntegral

@@ -324,7 +326,7 @@ import Prelude

324326 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,

325327 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,

326328 (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat,

327− reverse)

329+ reverse, iterate)

328330

329331import qualified Control.Monad.Catch as MC

330332import qualified Control.Monad.Reader as Reader

@@ -448,12 +450,22 @@ unfold (Unfold ustep inject) seed = Stream step Nothing

448450-- Specialized Generation

449451------------------------------------------------------------------------------

450452

453+{-# INLINE_NORMAL repeatM #-}

451454repeatM :: Monad m => m a -> Stream m a

452455repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()

453456

457+{-# INLINE_NORMAL repeat #-}

454458repeat :: Monad m => a -> Stream m a

455459repeat x = Stream (\_ _ -> return $ Yield x ()) ()

456460

461+{-# INLINE_NORMAL iterateM #-}

462+iterateM :: Monad m => (a -> m a) -> m a -> Stream m a

463+iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))

464+

465+{-# INLINE_NORMAL iterate #-}

466+iterate :: Monad m => (a -> a) -> a -> Stream m a

467+iterate step st = iterateM (return . step) (return st)

468+

457469{-# INLINE_NORMAL replicateM #-}

458470replicateM :: forall m a. Monad m => Int -> m a -> Stream m a

459471replicateM n p = Stream step n

src/Streamly/Internal/Data/Stream/StreamK.hs+17−1

@@ -67,6 +67,8 @@ module Streamly.Internal.Data.Stream.StreamK

6767 , replicateM

6868 , fromIndices

6969 , fromIndicesM

70+ , iterate

71+ , iterateM

7072

7173 -- ** Conversions

7274 , yield

@@ -196,7 +198,7 @@ import Prelude

196198 hiding (foldl, foldr, last, map, mapM, mapM_, repeat, sequence,

197199 take, filter, all, any, takeWhile, drop, dropWhile, minimum,

198200 maximum, elem, notElem, null, head, tail, init, zipWith, lookup,

199− foldr1, (!!), replicate, reverse, concatMap)

201+ foldr1, (!!), replicate, reverse, concatMap, iterate)

200202import qualified Prelude

201203

202204import Streamly.Internal.Data.SVar

@@ -311,6 +313,20 @@ fromIndices gen = go 0

311313 where

312314 go n = (gen n) `cons` go (n + 1)

313315

316+{-# INLINE iterate #-}

317+iterate :: IsStream t => (a -> a) -> a -> t m a

318+iterate step = fromStream . go

319+ where

320+ go s = cons s (go (step s))

321+

322+{-# INLINE iterateM #-}

323+iterateM :: (IsStream t, MonadAsync m) => (a -> m a) -> m a -> t m a

324+iterateM step = go

325+ where

326+ go s = mkStream $ \st stp sng yld -> do

327+ next <- s

328+ foldStreamShared st stp sng yld (return next |: go (step next))

329+

314330-------------------------------------------------------------------------------

315331-- Conversions

316332-------------------------------------------------------------------------------

src/Streamly/Internal/Prelude.hs+10−9

@@ -770,10 +770,9 @@ repeatMSerial = fromStreamS . S.repeatM

770770-- @

771771--

772772-- @since 0.1.2

773−iterate :: IsStream t => (a -> a) -> a -> t m a

774−iterate step = K.fromStream . go

775− where

776− go s = K.cons s (go (step s))

773+{-# INLINE_NORMAL iterate #-}

774+iterate :: (IsStream t, Monad m) => (a -> a) -> a -> t m a

775+iterate step = fromStreamS . S.iterate step

777776

778777-- |

779778-- @

@@ -802,12 +801,14 @@ iterate step = K.fromStream . go

802801-- /Since: 0.7.0 (signature change)/

803802--

804803-- /Since: 0.1.2/

804+{-# INLINE_EARLY iterateM #-}

805805iterateM :: (IsStream t, MonadAsync m) => (a -> m a) -> m a -> t m a

806−iterateM step = go

807− where

808− go s = K.mkStream $ \st stp sng yld -> do

809− next <- s

810− K.foldStreamShared st stp sng yld (return next |: go (step next))

806+iterateM = K.iterateM

807+

808+{-# RULES "iterateM serial" iterateM = iterateMSerial #-}

809+{-# INLINE iterateMSerial #-}

810+iterateMSerial :: MonadAsync m => (a -> m a) -> m a -> SerialT m a

811+iterateMSerial step = fromStreamS . S.iterateM step

811812

812813------------------------------------------------------------------------------

813814-- Conversions