@@ -126,6 +126,7 @@ module Streamly.Internal.Data.Fold
126126 --, lsequence
127127 , lmapM
128128 , scan
129+ , scanMany
129130 , postscan
130131 , indexed
131132
@@ -292,10 +293,9 @@ import qualified Data.IntSet as IntSet
292293import qualified Data.Map.Strict as Map
293294import qualified Data.Set as Set
294295import qualified Streamly.Internal.Data.IsMap as IsMap
296+import qualified Prelude
295297import qualified Streamly.Internal.Data.Pipe.Type as Pipe
296298import qualified Streamly.Internal.Data.Stream.StreamD.Type as StreamD
297−-- import qualified Streamly.Internal.Data.Stream.IsStream.Enumeration as Stream
298−import qualified Prelude
299299
300300import Prelude hiding
301301 ( filter, foldl1, drop, dropWhile, take, takeWhile, zipWith
@@ -417,17 +417,13 @@ transform (Pipe pstep1 pstep2 pinitial) (Fold fstep finitial fextract) =
417417
418418 extract (Tuple' _ fs) = fextract fs
419419
420−-- | Scan the input of a 'Fold' to change it in a stateful manner using another
421−-- 'Fold'.
422−-- /Pre-release/
423−{-# INLINE scan #-}
424−scan :: Monad m => Fold m a b -> Fold m b c -> Fold m a c
425−scan (Fold stepL initialL extractL) (Fold stepR initialR extractR) =
420+{-# INLINE scanWith #-}
421+scanWith :: Monad m => Bool -> Fold m a b -> Fold m b c -> Fold m a c
422+scanWith isMany (Fold stepL initialL extractL) (Fold stepR initialR extractR) =
426423 Fold step initial extract
427424
428425 where
429426
430− -- XXX It can be moved out and used for both scan and postscan
431427 {-# INLINE runStep #-}
432428 runStep actionL sR = do
433429 rL <- actionL
@@ -435,7 +431,10 @@ scan (Fold stepL initialL extractL) (Fold stepR initialR extractR) =
435431 Done bL -> do
436432 rR <- stepR sR bL
437433 case rR of
438− Partial sR1 -> Done <$> extractR sR1
434+ Partial sR1 ->
435+ if isMany
436+ then runStep initialL sR1
437+ else Done <$> extractR sR1
439438 Done bR -> return $ Done bR
440439 Partial sL -> do
441440 !b <- extractL sL
@@ -455,6 +454,24 @@ scan (Fold stepL initialL extractL) (Fold stepR initialR extractR) =
455454
456455 extract = extractR . snd
457456
457+-- | Scan the input of a 'Fold' to change it in a stateful manner using another
458+-- 'Fold'. The scan stops as soon as the fold terminates.
459+--
460+-- /Pre-release/
461+{-# INLINE scan #-}
462+scan :: Monad m => Fold m a b -> Fold m b c -> Fold m a c
463+scan = scanWith False
464+
465+-- XXX This does not fuse beacuse of the recursive step. Need to investigate.
466+--
467+-- | Scan the input of a 'Fold' to change it in a stateful manner using another
468+-- 'Fold'. The scan restarts with a fresh state if the fold terminates.
469+--
470+-- /Pre-release/
471+{-# INLINE scanMany #-}
472+scanMany :: Monad m => Fold m a b -> Fold m b c -> Fold m a c
473+scanMany = scanWith True
474+
458475-- | Postscan the input of a 'Fold' to change it in a stateful manner using
459476-- another 'Fold'.
460477-- /Pre-release/