@@ -201,8 +201,11 @@ module Streamly.Internal.Data.Unfold
203 , mapMWithInput
203 , mapMWithInput
204 , postscanlM'
205 , postscanlM'
205 , postscan
206 , postscan
206 , foldMany
209 , foldMany
@@ -507,6 +510,62 @@ postscan (Fold stepF initial extract) (Unfold stepU injectU) =
508 step Nothing = return Stop
511 step Nothing = return Stop
513+data ScanState s f = ScanInit s | ScanDo s !f | ScanDone
515+{-# INLINE_NORMAL scanWith #-}
516+scanWith :: Monad m => Bool -> Fold m b c -> Unfold m a b -> Unfold m a c
517+scanWith restart (Fold fstep initial extract) (Unfold stepU injectU) =
522+ inject a = ScanInit <$> injectU a
524+ {-# INLINE runStep #-}
525+ runStep us action = do
530+ return $ Yield b (ScanDo us fs)
532+ let next = if restart then ScanInit us else ScanDone
533+ in return $ Yield b next
535+ {-# INLINE_LATE step #-}
536+ step (ScanInit us) = runStep us initial
537+ step (ScanDo us fs) = do
540+ Yield x s -> runStep s (fstep fs x)
541+ Skip s -> return $ Skip $ ScanDo s fs
543+ step ScanDone = return Stop
545+-- | Scan the output of an 'Unfold' to change it in a stateful manner.
546+-- Once fold is done it will restart from its initial state.
548+-- >>> u = Unfold.scanMany (Fold.take 2 Fold.sum) Unfold.fromList
549+-- >>> Unfold.fold Fold.toList u [1,2,3,4,5]
553+{-# INLINE_NORMAL scanMany #-}
554+scanMany :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c
555+scanMany = scanWith True
557+-- | Scan the output of an 'Unfold' to change it in a stateful manner.
558+-- Once fold is done it will stop.
560+-- >>> u = Unfold.scan (Fold.take 2 Fold.sum) Unfold.fromList
561+-- >>> Unfold.fold Fold.toList u [1,2,3,4,5]
565+{-# INLINE_NORMAL scan #-}
566+scan :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c
567+scan = scanWith False
510-- | Scan the output of an 'Unfold' to change it in a stateful manner.
569-- | Scan the output of an 'Unfold' to change it in a stateful manner.
512-- /Pre-release/
571-- /Pre-release/