on master help sign in
Simplify fold scan implementation
commita8a9ae872abac1d783f4b0dfacbf46f6c24760cd authorHarendra Kumar <harendra@composewell.com> dateSat 6 Nov 2021, 15:41 (5 years ago) parentc864b08
1 file changed+24−33
src/Streamly/Internal/Data/Fold.hs+24−33
src/Streamly/Internal/Data/Fold.hs+24−33
@@ -382,7 +382,6 @@ transform (Pipe pstep1 pstep2 pinitial) (Fold fstep finitial fextract) =
382382
383383-- | Scan the input of a 'Fold' to change it in a stateful manner using another
384384-- 'Fold'.
385−-- XXX if any Fold is done assuming finish the resulting Fold?
386385-- /Pre-release/
387386{-# INLINE scan #-}
388387scan :: Monad m => Fold m a b -> Fold m b c -> Fold m a c
@@ -391,39 +390,31 @@ scan (Fold stepL initialL extractL) (Fold stepR initialR extractR) =
391390
392391 where
393392
393+ runStep actionL sR = do
394+ rL <- actionL
395+ case rL of
396+ Done bL -> do
397+ rR <- stepR sR bL
398+ case rR of
399+ Partial sR1 -> Done <$> extractR sR1
400+ Done bR -> return $ Done bR
401+ Partial sL -> do
402+ !b <- extractL sL
403+ rR <- stepR sR b
404+ return
405+ $ case rR of
406+ Partial sR1 -> Partial (sL, sR1)
407+ Done bR -> Done bR
408+
394409 initial = do
395− r1 <- initialL
396− r2 <- initialR
397− case r1 of
398− Done b1 -> return $ case r2 of
399− Done b2 -> Done b2
400− Partial fs2 -> Partial (Nothing, Just fs2, b1)
401− Partial fs -> return $ case r2 of
402− Done b2 -> Done b2
403− Partial fs2 -> Partial (Just fs, Just fs2, undefined)
404−
405− step (Just s1, s2, _) x = do
406− r1 <- stepL s1 x
407− case r1 of
408− Done b -> do
409− r2 <- stepR (fromJust s2) b
410− case r2 of
411− Partial fs2 -> Done <$> extractR fs2
412− Done b2 -> return $ Done b2
413− Partial fs -> do
414− !b <- extractL fs
415− r2 <- stepR (fromJust s2) b
416− return $ case r2 of
417− Partial fs2 -> Partial (Just fs, Just fs2, undefined)
418− Done b2 -> Done b2
419−
420− step (Nothing, s2, b1) _ = do
421− r2 <- stepR (fromJust s2) b1
422− case r2 of
423− Partial fs2 -> Done <$> extractR fs2
424− Done b2 -> return $ Done b2
425−
426− extract (_, s2, _) = extractR (fromJust s2)
410+ r <- initialR
411+ case r of
412+ Partial sR -> runStep initialL sR
413+ Done b -> return $ Done b
414+
415+ step (sL, sR) x = runStep (stepL sL x) sR
416+
417+ extract = extractR . snd
427418
428419------------------------------------------------------------------------------
429420-- Left folds