on master help sign in
Add scan and scanMany for unfolds
commitaaa775be9f6ab442c26291fbfcc983df7a00d009 authorHarendra Kumar <harendra@composewell.com> dateFri 29 Apr 2022, 14:16 (5 years ago) parentaa2b166
1 file changed+59
core/src/Streamly/Internal/Data/Unfold.hs+59
core/src/Streamly/Internal/Data/Unfold.hs+59
@@ -201,8 +201,11 @@ module Streamly.Internal.Data.Unfold
201201 , map
202202 , mapM
203203 , mapMWithInput
204+
204205 , postscanlM'
205206 , postscan
207+ , scan
208+ , scanMany
206209 , foldMany
207210 -- pipe
208211
@@ -507,6 +510,62 @@ postscan (Fold stepF initial extract) (Unfold stepU injectU) =
507510
508511 step Nothing = return Stop
509512
513+data ScanState s f = ScanInit s | ScanDo s !f | ScanDone
514+
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) =
518+ Unfold step inject
519+
520+ where
521+
522+ inject a = ScanInit <$> injectU a
523+
524+ {-# INLINE runStep #-}
525+ runStep us action = do
526+ r <- action
527+ case r of
528+ FL.Partial fs -> do
529+ !b <- extract fs
530+ return $ Yield b (ScanDo us fs)
531+ FL.Done b ->
532+ let next = if restart then ScanInit us else ScanDone
533+ in return $ Yield b next
534+
535+ {-# INLINE_LATE step #-}
536+ step (ScanInit us) = runStep us initial
537+ step (ScanDo us fs) = do
538+ res <- stepU us
539+ case res of
540+ Yield x s -> runStep s (fstep fs x)
541+ Skip s -> return $ Skip $ ScanDo s fs
542+ Stop -> return Stop
543+ step ScanDone = return Stop
544+
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.
547+--
548+-- >>> u = Unfold.scanMany (Fold.take 2 Fold.sum) Unfold.fromList
549+-- >>> Unfold.fold Fold.toList u [1,2,3,4,5]
550+-- [0,1,3,0,3,7,0,5]
551+--
552+-- /Pre-release/
553+{-# INLINE_NORMAL scanMany #-}
554+scanMany :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c
555+scanMany = scanWith True
556+
557+-- | Scan the output of an 'Unfold' to change it in a stateful manner.
558+-- Once fold is done it will stop.
559+--
560+-- >>> u = Unfold.scan (Fold.take 2 Fold.sum) Unfold.fromList
561+-- >>> Unfold.fold Fold.toList u [1,2,3,4,5]
562+-- [0,1,3]
563+--
564+-- /Pre-release/
565+{-# INLINE_NORMAL scan #-}
566+scan :: Monad m => Fold m b c -> Unfold m a b -> Unfold m a c
567+scan = scanWith False
568+
510569-- | Scan the output of an 'Unfold' to change it in a stateful manner.
511570--
512571-- /Pre-release/