git.iko.soy/streamly
on master help sign in
Add iterateM and fromIndicesM to Unfolds
f3658d1adithyaov
src/Streamly/Internal/Data/Unfold.hs+30−2
compare with master
unifiedsplit
@@ -76,10 +76,18 @@ module Streamly.Internal.Data.Unfold
76 , singleton
76 , singleton
77 , identity
77 , identity
78 , const
78 , const
79− , replicateM
79+
80− , repeatM
81 , fromList
80 , fromList
82 , fromListM
81 , fromListM
82+
83+ -- ** Specialized Generation
84+ -- | Generate a monadic stream from a seed.
85+ , replicateM
86+ , repeatM
87+ , iterateM
88+ , fromIndicesM
89+
90+ -- ** Enumerations
83 , enumerateFromStepIntegral
91 , enumerateFromStepIntegral
84 , enumerateFromToIntegral
92 , enumerateFromToIntegral
85 , enumerateFromIntegral
93 , enumerateFromIntegral
@@ -453,6 +461,7 @@ fromListM = Unfold step inject
453-- Specialized Generation
461-- Specialized Generation
454------------------------------------------------------------------------------
462------------------------------------------------------------------------------
455
463
464+-- XXX replicateM? Maybe just replicate?
456-- | Generates a stream replicating the seed @n@ times.
465-- | Generates a stream replicating the seed @n@ times.
457--
466--
458{-# INLINE replicateM #-}
467{-# INLINE replicateM #-}
@@ -466,6 +475,7 @@ replicateM n = Unfold step inject
466 then Stop
475 then Stop
467 else Yield x (x, i - 1)
476 else Yield x (x, i - 1)
468
477
478+-- XXX repeatM? Maybe just repeat?
469-- | Generates an infinite stream repeating the seed.
479-- | Generates an infinite stream repeating the seed.
470--
480--
471{-# INLINE repeatM #-}
481{-# INLINE repeatM #-}
@@ -475,6 +485,24 @@ repeatM = Unfold step return
475 {-# INLINE_LATE step #-}
485 {-# INLINE_LATE step #-}
476 step x = return $ Yield x x
486 step x = return $ Yield x x
477
487
488+{-# INLINE iterateM #-}
489+iterateM :: Monad m => (a -> m a) -> Unfold m a a
490+iterateM f = Unfold step f
491+ where
492+ {-# INLINE_LATE step #-}
493+ step x = do
494+ fx <- f x
495+ return $ Yield fx fx
496+
497+{-# INLINE_NORMAL fromIndicesM #-}
498+fromIndicesM :: Monad m => (Int -> m a) -> Unfold m Int a
499+fromIndicesM gen = Unfold step return
500+ where
501+ {-# INLINE_LATE step #-}
502+ step i = do
503+ x <- gen i
504+ return $ Yield x (i + 1)
505+
478-------------------------------------------------------------------------------
506-------------------------------------------------------------------------------
479-- Filtering
507-- Filtering
480-------------------------------------------------------------------------------
508-------------------------------------------------------------------------------