Try removing a state but adding a branch in Yield in splitSepBy_

core/src/Streamly/Internal/Data/Stream/Transform.hs

@@ -2154,7 +2154,7 @@ data SplitSepBy s fs b a
     = SplitSepByInit s
     | SplitSepByInitFold0 s
     | SplitSepByInitFold1 s fs
-    | SplitSepByCheck s a fs
+    -- | SplitSepByCheck s a fs
     | SplitSepByNext s fs
     | SplitSepByYield b (SplitSepBy s fs b a)
     | SplitSepByDone
@@ -2233,47 +2233,46 @@ splitSepBy_ predicate (Fold fstep initial _ final) (Stream step1 state1) =
     -- the fold first. Also, in most common cases the fold is not effectful.
     -- On the other hand, in most cases the fold will not terminate without
     -- consuming anything. So both ways are similar.
+
+    {-# INLINE splitSepByCheck #-}
+    splitSepByCheck st x fs = do
+        fres <-
+            if predicate x
+            then do
+                b <- final fs
+                return $ FL.Done b
+            else fstep fs x
+        return
+            $ case fres of
+                  FL.Done b -> Yield b (SplitSepByInitFold0 st)
+                  FL.Partial fs1 -> Skip $ SplitSepByNext st fs1
+
     {-# INLINE_LATE step #-}
     step _ (SplitSepByInit st) = do
         fres <- initial
         return
-            $ Skip
             $ case fres of
-                  FL.Done b -> SplitSepByYield b (SplitSepByInit st)
-                  FL.Partial fs -> SplitSepByInitFold1 st fs
+                  FL.Done b -> Yield b (SplitSepByInit st)
+                  FL.Partial fs -> Skip $ SplitSepByInitFold1 st fs
 
     step _ (SplitSepByInitFold0 st) = do
         fres <- initial
         return
-            $ Skip
             $ case fres of
-                  FL.Done b -> SplitSepByYield b (SplitSepByInitFold0 st)
-                  FL.Partial fs -> SplitSepByNext st fs
+                  FL.Done b -> Yield b (SplitSepByInitFold0 st)
+                  FL.Partial fs -> Skip $ SplitSepByNext st fs
 
     step gst (SplitSepByInitFold1 st fs) = do
         r <- step1 (adaptState gst) st
         case r of
-            Yield x s -> return $ Skip $ SplitSepByCheck s x fs
+            Yield x s -> splitSepByCheck s x fs
             Skip s -> return $ Skip (SplitSepByInitFold1 s fs)
             Stop -> final fs >> return Stop
 
-    step _ (SplitSepByCheck st x fs) = do
-        if predicate x
-        then do
-            b <- final fs
-            return $ Skip $ SplitSepByYield b (SplitSepByInitFold0 st)
-        else do
-            fres <- fstep fs x
-            return
-                $ Skip
-                $ case fres of
-                      FL.Done b -> SplitSepByYield b (SplitSepByInitFold0 st)
-                      FL.Partial fs1 -> SplitSepByNext st fs1
-
     step gst (SplitSepByNext st fs) = do
         r <- step1 (adaptState gst) st
         case r of
-            Yield x s -> return $ Skip $ SplitSepByCheck s x fs
+            Yield x s -> splitSepByCheck s x fs
             Skip s -> return $ Skip (SplitSepByNext s fs)
             Stop -> do
                 b <- final fs