diff --git a/core/src/Streamly/Internal/Data/Unfold.hs b/core/src/Streamly/Internal/Data/Unfold.hs index b025d79d5..3a5e18349 100644 --- a/core/src/Streamly/Internal/Data/Unfold.hs +++ b/core/src/Streamly/Internal/Data/Unfold.hs @@ -201,8 +201,11 @@ module Streamly.Internal.Data.Unfold , map , mapM , mapMWithInput + , postscanlM' , postscan + , scan + , scanMany , foldMany -- pipe @@ -507,6 +510,62 @@ postscan (Fold stepF initial extract) (Unfold stepU injectU) = step Nothing = return Stop +data ScanState s f = ScanInit s | ScanDo s !f | ScanDone + +{-# INLINE_NORMAL scanWith #-} +scanWith :: Monad m => Bool -> Fold m b c -> Unfold m a b -> Unfold m a c +scanWith restart (Fold fstep initial extract) (Unfold stepU injectU) = + Unfold step inject + + where + + inject a = ScanInit <$> injectU a + + {-# INLINE runStep #-} + runStep us action = do + r <- action + case r of + FL.Partial fs -> do + !b <- extract fs + return $ Yield b (ScanDo us fs) + FL.Done b -> + let next = if restart then ScanInit us else ScanDone + in return $ Yield b next + + {-# INLINE_LATE step #-} + step (ScanInit us) = runStep us initial + step (ScanDo us fs) = do + res <- stepU us + case res of + Yield x s -> runStep s (fstep fs x) + Skip s -> return $ Skip $ ScanDo s fs + Stop -> return Stop + step ScanDone = return Stop + +-- | Scan the output of an 'Unfold' to change it in a stateful manner. +-- Once fold is done it will restart from its initial state. +-- +-- >>> u = Unfold.scanMany (Fold.take 2 Fold.sum) Unfold.fromList +-- >>> Unfold.fold Fold.toList u [1,2,3,4,5] +-- [0,1,3,0,3,7,0,5] +-- +-- /Pre-release/ +{-# INLINE_NORMAL scanMany #-} +scanMany :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c +scanMany = scanWith True + +-- | Scan the output of an 'Unfold' to change it in a stateful manner. +-- Once fold is done it will stop. +-- +-- >>> u = Unfold.scan (Fold.take 2 Fold.sum) Unfold.fromList +-- >>> Unfold.fold Fold.toList u [1,2,3,4,5] +-- [0,1,3] +-- +-- /Pre-release/ +{-# INLINE_NORMAL scan #-} +scan :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c +scan = scanWith False + -- | Scan the output of an 'Unfold' to change it in a stateful manner. -- -- /Pre-release/