on master help sign in
Inline the consume state in foldMany
commitb8ede3e5e39c1938d8e51c1354edc1f874c7d16a authorHarendra Kumar <harendra@composewell.com> dateFri 5 Feb 2021, 18:33 (6 years ago) parentc0ec60a
1 file changed+17−18
src/Streamly/Internal/Data/Stream/StreamD/Type.hs+17−18
src/Streamly/Internal/Data/Stream/StreamD/Type.hs+17−18
@@ -789,10 +789,9 @@ instance Monad m => Monad (Stream m) where
789789{-# ANN type FoldMany Fuse #-}
790790data FoldMany s fs b a
791791 = FoldManyStart s
792− | FoldManyConsume s fs a
793− | FoldManyBuffer s fs
792+ | FoldManyLoop s fs
794793 | FoldManyYield b (FoldMany s fs b a)
795− | FoldManyFinish
794+ | FoldManyDone
796795
797796-- | This is the stream equivalent of "Data.Fold.Internal.many". The fold
798797-- may consume 0 or more elements. It means:
@@ -810,6 +809,15 @@ foldMany (Fold fstep initial extract) (Stream step state) =
810809
811810 where
812811
812+ {-# INLINE consume #-}
813+ consume x s fs = do
814+ res <- fstep fs x
815+ return
816+ $ Skip
817+ $ case res of
818+ FL.Done b -> FoldManyYield b (FoldManyStart s)
819+ FL.Partial ps -> FoldManyLoop s ps
820+
813821 {-# INLINE_LATE step' #-}
814822 step' _ (FoldManyStart st) = do
815823 r <- initial
@@ -817,26 +825,17 @@ foldMany (Fold fstep initial extract) (Stream step state) =
817825 $ Skip
818826 $ case r of
819827 FL.Done b -> FoldManyYield b (FoldManyStart st)
820− FL.Partial fs -> FoldManyBuffer st fs
821− -- This state is not strictly required but it helps the compiler fuse the
822− -- code.
823− step' _ (FoldManyConsume st fs x) = do
824− r <- fstep fs x
825− return
826− $ Skip
827− $ case r of
828− FL.Done b -> FoldManyYield b (FoldManyStart st)
829− FL.Partial fs1 -> FoldManyBuffer st fs1
830− step' gst (FoldManyBuffer st fs) = do
828+ FL.Partial fs -> FoldManyLoop st fs
829+ step' gst (FoldManyLoop st fs) = do
831830 r <- step (adaptState gst) st
832831 case r of
833− Yield x s -> return $ Skip $ FoldManyConsume s fs x
834− Skip s -> return $ Skip (FoldManyBuffer s fs)
832+ Yield x s -> consume x s fs
833+ Skip s -> return $ Skip (FoldManyLoop s fs)
835834 Stop -> do
836835 b <- extract fs
837− return $ Skip (FoldManyYield b FoldManyFinish)
836+ return $ Skip (FoldManyYield b FoldManyDone)
838837 step' _ (FoldManyYield b next) = return $ Yield b next
839− step' _ FoldManyFinish = return Stop
838+ step' _ FoldManyDone = return Stop
840839
841840{-# ANN type FoldMany1 Fuse #-}
842841data FoldMany1 s fs b a