@@ -39,6 +39,7 @@ module Streamly.Internal.Data.Stream.StreamD.Transform
3939 -- * Scanning By 'Fold'
4040 , postscanOnce -- XXX rename to postscan
4141 , scanOnce -- XXX rename to scan
42+ , scanMany
4243
4344 -- * Scanning
4445 -- | Left scans. Stateful, mostly one-to-one maps.
@@ -444,36 +445,45 @@ postscanOnce (FL.Fold fstep initial extract) (Stream sstep state) =
444445 Stop -> return Stop
445446 step _ ScanDone = return Stop
446447
447−{-# INLINE scanOnce #-}
448−scanOnce :: Monad m
449− => FL.Fold m a b -> Stream m a -> Stream m b
450−scanOnce (FL.Fold fstep initial extract) (Stream sstep state) =
448+{-# INLINE scanWith #-}
449+scanWith :: Monad m
450+ => Bool -> Fold m a b -> Stream m a -> Stream m b
451+scanWith restart (Fold fstep initial extract) (Stream sstep state) =
451452 Stream step (ScanInit state)
452453
453454 where
454455
455− {-# INLINE_LATE step #-}
456− step _ (ScanInit st) = do
457− res <- initial
456+ {-# INLINE runStep #-}
457+ runStep st action = do
458+ res <- action
458459 case res of
459460 FL.Partial fs -> do
460461 !b <- extract fs
461462 return $ Yield b $ ScanDo st fs
462− FL.Done b -> return $ Yield b ScanDone
463+ FL.Done b ->
464+ let next = if restart then ScanInit st else ScanDone
465+ in return $ Yield b next
466+
467+ {-# INLINE_LATE step #-}
468+ step _ (ScanInit st) = runStep st initial
463469 step gst (ScanDo st fs) = do
464470 res <- sstep (adaptState gst) st
465471 case res of
466− Yield x s -> do
467− r <- fstep fs x
468− case r of
469− FL.Partial fs1 -> do
470− !b <- extract fs1
471− return $ Yield b $ ScanDo s fs1
472− FL.Done b -> return $ Yield b ScanDone
472+ Yield x s -> runStep s (fstep fs x)
473473 Skip s -> return $ Skip $ ScanDo s fs
474474 Stop -> return Stop
475475 step _ ScanDone = return Stop
476476
477+{-# INLINE scanOnce #-}
478+scanOnce :: Monad m
479+ => FL.Fold m a b -> Stream m a -> Stream m b
480+scanOnce = scanWith False
481+
482+{-# INLINE scanMany #-}
483+scanMany :: Monad m
484+ => FL.Fold m a b -> Stream m a -> Stream m b
485+scanMany = scanWith True
486+
477487------------------------------------------------------------------------------
478488-- Scanning - Prescans
479489------------------------------------------------------------------------------