Move parseBreak from StreamK module to ParserDrivers

harendra-kumar · harendra-kumar · commit 14e0f7100abf · 2025-06-24T00:35:58.000+05:30
diff --git a/core/src/Streamly/Internal/Data/ParserDrivers.h b/core/src/Streamly/Internal/Data/ParserDrivers.h
@@ -1,5 +1,6 @@
 #ifndef PARSER_WITH_POS
 #define PARSE_BREAK parseBreak
+#define PARSE_BREAK_STREAMK parseBreakStreamK
 #define PARSE_BREAK_CHUNKS parseBreakChunks
 #define PARSE_BREAK_CHUNKS_GENERIC parseBreakChunksGeneric
 #define PARSE_MANY parseMany
@@ -9,6 +10,8 @@
 #else
 #undef PARSE_BREAK
 #define PARSE_BREAK parseBreakPos
+#undef PARSE_BREAK_STREAMK
+#define PARSE_BREAK_STREAMK parseBreakStreamKPos
 #undef PARSE_BREAK_CHUNKS
 #define PARSE_BREAK_CHUNKS parseBreakChunksPos
 #undef PARSE_BREAK_CHUNKS_GENERIC
@@ -757,6 +760,126 @@ PARSE_BREAK (PRD.Parser pstep initial extract) stream@(Stream step state) = do
                 let src  = Prelude.reverse $ getList buf
                 return (Left (ParseError DEFAULT(i) err), fromList src)
 
+{-# INLINE_NORMAL PARSE_BREAK_STREAMK #-}
+PARSE_BREAK_STREAMK
+    :: forall m a b. Monad m
+    => ParserK.ParserK a m b
+    -> StreamK m a
+    -> m (Either ParseError b, StreamK m a)
+PARSE_BREAK_STREAMK parser input = do
+    let parserk = ParserK.runParser parser ParserK.parserDone 0 0
+     in go OPTIONAL(0) [] parserk input
+
+    where
+
+    {-# INLINE backtrck #-}
+    -- backtrck :: Int -> [a] -> StreamK m a -> (StreamK m a, [a])
+    backtrck n xs stream =
+        let (pre, post) = Stream.splitAt "Data.StreamK.parseBreak" n xs
+         in (StreamK.append (StreamK.fromList (Prelude.reverse pre)) stream, post)
+
+    {-# INLINE goStop #-}
+    {-
+    goStop
+        :: OPTIONAL(Int ->)
+           [a]
+        -> (ParserK.Input a -> m (ParserK.Step a m b))
+        -> m (Either ParseError b, StreamK m a)
+    -}
+    goStop OPTIONAL(pos) backBuf parserk = do
+        pRes <- parserk ParserK.None
+        case pRes of
+            -- If we stop in an alternative, it will try calling the next
+            -- parser, the next parser may call initial returning Partial and
+            -- then immediately we have to call extract on it.
+            ParserK.Partial 0 cont1 ->
+                 go OPTIONAL(pos) [] cont1 StreamK.nil
+            ParserK.Partial n cont1 -> do
+                let n1 = negate n
+                assertM(n1 >= 0 && n1 <= length backBuf)
+                let (s1, backBuf1) = backtrck n1 backBuf StreamK.nil
+                 in go OPTIONAL(pos + n) backBuf1 cont1 s1
+            ParserK.Continue 0 cont1 ->
+                go OPTIONAL(pos) backBuf cont1 StreamK.nil
+            ParserK.Continue n cont1 -> do
+                let n1 = negate n
+                assertM(n1 >= 0 && n1 <= length backBuf)
+                let (s1, backBuf1) = backtrck n1 backBuf StreamK.nil
+                 in go OPTIONAL(pos + n) backBuf1 cont1 s1
+            ParserK.Done 0 b ->
+                return (Right b, StreamK.nil)
+            ParserK.Done n b -> do
+                let n1 = negate n
+                assertM(n1 >= 0 && n1 <= length backBuf)
+                let (s1, _) = backtrck n1 backBuf StreamK.nil
+                 in return (Right b, s1)
+            ParserK.Error n err ->
+                let strm = StreamK.fromList (Prelude.reverse backBuf)
+                 in return (Left (ParseError (DEFAULT(pos) + n) err), strm)
+
+    {-
+    yieldk
+        :: OPTIONAL(Int ->)
+           [a]
+        -> (ParserK.Input a -> m (ParserK.Step a m b))
+        -> a
+        -> StreamK m a
+        -> m (Either ParseError b, StreamK m a)
+    -}
+    yieldk OPTIONAL(pos) backBuf parserk element stream = do
+        pRes <- parserk (ParserK.Chunk element)
+        -- NOTE: factoring out "StreamK.cons element stream" in a let statement here
+        -- cause big alloc regression.
+        case pRes of
+            ParserK.Partial 1 cont1 -> go OPTIONAL(pos + 1) [] cont1 stream
+            ParserK.Partial 0 cont1 -> go OPTIONAL(pos) [] cont1 (StreamK.cons element stream)
+            ParserK.Partial n cont1 -> do -- n < 0 case
+                let n1 = negate n
+                    bufLen = length backBuf
+                    s = StreamK.cons element stream
+                assertM(n1 >= 0 && n1 <= bufLen)
+                let (s1, _) = backtrck n1 backBuf s
+                go OPTIONAL(pos + n) [] cont1 s1
+            ParserK.Continue 1 cont1 -> go OPTIONAL(pos + 1) (element:backBuf) cont1 stream
+            ParserK.Continue 0 cont1 ->
+                go OPTIONAL(pos) backBuf cont1 (StreamK.cons element stream)
+            ParserK.Continue n cont1 -> do
+                let n1 = negate n
+                    bufLen = length backBuf
+                    s = StreamK.cons element stream
+                assertM(n1 >= 0 && n1 <= bufLen)
+                let (s1, backBuf1) = backtrck n1 backBuf s
+                go OPTIONAL(pos + n) backBuf1 cont1 s1
+            ParserK.Done 1 b -> pure (Right b, stream)
+            ParserK.Done 0 b -> pure (Right b, StreamK.cons element stream)
+            ParserK.Done n b -> do
+                let n1 = negate n
+                    bufLen = length backBuf
+                    s = StreamK.cons element stream
+                assertM(n1 >= 0 && n1 <= bufLen)
+                let (s1, _) = backtrck n1 backBuf s
+                pure (Right b, s1)
+            ParserK.Error n err ->
+                let strm =
+                        StreamK.append
+                            (StreamK.fromList (Prelude.reverse backBuf))
+                            (StreamK.cons element stream)
+                 in return (Left (ParseError (DEFAULT(pos) + n + 1) err), strm)
+
+    {-
+    go
+        :: OPTIONAL(Int ->)
+           [a]
+        -> (ParserK.Input a -> m (ParserK.Step a m b))
+        -> StreamK m a
+        -> m (Either ParseError b, StreamK m a)
+    -}
+    go OPTIONAL(pos) backBuf parserk stream = do
+        let stop = goStop OPTIONAL(pos) backBuf parserk
+            single a = yieldk OPTIONAL(pos) backBuf parserk a StreamK.nil
+         in StreamK.foldStream
+                defState (yieldk OPTIONAL(pos) backBuf parserk) single stop stream
+
 {-# INLINE_NORMAL PARSE_BREAK_CHUNKS #-}
 PARSE_BREAK_CHUNKS
     :: (Monad m, Unbox a)
diff --git a/core/src/Streamly/Internal/Data/ParserDrivers.hs b/core/src/Streamly/Internal/Data/ParserDrivers.hs
@@ -12,6 +12,8 @@ module Streamly.Internal.Data.ParserDrivers
     -- * Running a Parser
       parseBreak
     , parseBreakPos
+    , parseBreakStreamK
+    , parseBreakStreamKPos
     , parseBreakChunks
     , parseBreakChunksPos
     , parseBreakChunksGeneric
diff --git a/core/src/Streamly/Internal/Data/Stream/Nesting.hs b/core/src/Streamly/Internal/Data/Stream/Nesting.hs
@@ -1601,7 +1601,7 @@ parseIterate
     -> Stream m (Either ParseError b)
 parseIterate = Drivers.parseIterate
 
--- | Like 'parseMany' but includes stream position information in the error
+-- | Like 'parseIterate' but includes stream position information in the error
 -- messages.
 --
 {-# INLINE parseIteratePos #-}
diff --git a/core/src/Streamly/Internal/Data/StreamK.hs b/core/src/Streamly/Internal/Data/StreamK.hs
@@ -43,6 +43,7 @@ module Streamly.Internal.Data.StreamK
     , parseDBreak
     , parseD
     , parseBreak
+    , parseBreakPos
     , parse
 
     -- ** Specialized Folds
@@ -153,6 +154,7 @@ import qualified Streamly.Internal.Data.Array as Array
 import qualified Streamly.Internal.Data.Array.Generic as GenArr
 import qualified Streamly.Internal.Data.Fold.Type as FL
 import qualified Streamly.Internal.Data.Parser as Parser
+import qualified Streamly.Internal.Data.ParserDrivers as Drivers
 import qualified Streamly.Internal.Data.Parser.Type as PR
 import qualified Streamly.Internal.Data.ParserK.Type as ParserK
 import qualified Streamly.Internal.Data.Stream as Stream
@@ -1295,131 +1297,26 @@ parseChunks = Array.parse
 -- ParserK Singular
 -------------------------------------------------------------------------------
 
-#ifndef PARSER_WITH_POS
-#define PARSE_BREAK parseBreak
-#define OPTIONAL(x)
-#define DEFAULT(x) 0
-#else
-#define PARSE_BREAK parseBreakPos
-#define OPTIONAL(x) (x)
-#define DEFAULT(x) (x)
-#endif
-
 -- | Similar to 'parseBreak' but works on singular elements.
 --
-{-# INLINE_NORMAL PARSE_BREAK #-}
-PARSE_BREAK
+{-# INLINE parseBreak #-}
+parseBreak
     :: forall m a b. Monad m
     => ParserK.ParserK a m b
     -> StreamK m a
     -> m (Either ParseError b, StreamK m a)
-PARSE_BREAK parser input = do
-    let parserk = ParserK.runParser parser ParserK.parserDone 0 0
-     in go OPTIONAL(0) [] parserk input
-
-    where
+parseBreak = Drivers.parseBreakStreamK
 
-    {-# INLINE backtrack #-}
-    -- backtrack :: Int -> [a] -> StreamK m a -> (StreamK m a, [a])
-    backtrack n xs stream =
-        let (pre, post) = Stream.splitAt "Data.StreamK.parseBreak" n xs
-         in (append (fromList (Prelude.reverse pre)) stream, post)
-
-    {-# INLINE goStop #-}
-    goStop
-        :: OPTIONAL(Int ->)
-           [a]
-        -> (ParserK.Input a -> m (ParserK.Step a m b))
-        -> m (Either ParseError b, StreamK m a)
-    goStop OPTIONAL(pos) backBuf parserk = do
-        pRes <- parserk ParserK.None
-        case pRes of
-            -- If we stop in an alternative, it will try calling the next
-            -- parser, the next parser may call initial returning Partial and
-            -- then immediately we have to call extract on it.
-            ParserK.Partial 0 cont1 ->
-                 go OPTIONAL(pos) [] cont1 nil
-            ParserK.Partial n cont1 -> do
-                let n1 = negate n
-                assertM(n1 >= 0 && n1 <= length backBuf)
-                let (s1, backBuf1) = backtrack n1 backBuf nil
-                 in go OPTIONAL(pos + n) backBuf1 cont1 s1
-            ParserK.Continue 0 cont1 ->
-                go OPTIONAL(pos) backBuf cont1 nil
-            ParserK.Continue n cont1 -> do
-                let n1 = negate n
-                assertM(n1 >= 0 && n1 <= length backBuf)
-                let (s1, backBuf1) = backtrack n1 backBuf nil
-                 in go OPTIONAL(pos + n) backBuf1 cont1 s1
-            ParserK.Done 0 b ->
-                return (Right b, nil)
-            ParserK.Done n b -> do
-                let n1 = negate n
-                assertM(n1 >= 0 && n1 <= length backBuf)
-                let (s1, _) = backtrack n1 backBuf nil
-                 in return (Right b, s1)
-            ParserK.Error n err ->
-                let strm = fromList (Prelude.reverse backBuf)
-                 in return (Left (ParseError (DEFAULT(pos) + n) err), strm)
-
-    yieldk
-        :: OPTIONAL(Int ->)
-           [a]
-        -> (ParserK.Input a -> m (ParserK.Step a m b))
-        -> a
-        -> StreamK m a
-        -> m (Either ParseError b, StreamK m a)
-    yieldk OPTIONAL(pos) backBuf parserk element stream = do
-        pRes <- parserk (ParserK.Chunk element)
-        -- NOTE: factoring out "cons element stream" in a let statement here
-        -- cause big alloc regression.
-        case pRes of
-            ParserK.Partial 1 cont1 -> go OPTIONAL(pos + 1) [] cont1 stream
-            ParserK.Partial 0 cont1 -> go OPTIONAL(pos) [] cont1 (cons element stream)
-            ParserK.Partial n cont1 -> do -- n < 0 case
-                let n1 = negate n
-                    bufLen = length backBuf
-                    s = cons element stream
-                assertM(n1 >= 0 && n1 <= bufLen)
-                let (s1, _) = backtrack n1 backBuf s
-                go OPTIONAL(pos + n) [] cont1 s1
-            ParserK.Continue 1 cont1 -> go OPTIONAL(pos + 1) (element:backBuf) cont1 stream
-            ParserK.Continue 0 cont1 ->
-                go OPTIONAL(pos) backBuf cont1 (cons element stream)
-            ParserK.Continue n cont1 -> do
-                let n1 = negate n
-                    bufLen = length backBuf
-                    s = cons element stream
-                assertM(n1 >= 0 && n1 <= bufLen)
-                let (s1, backBuf1) = backtrack n1 backBuf s
-                go OPTIONAL(pos + n) backBuf1 cont1 s1
-            ParserK.Done 1 b -> pure (Right b, stream)
-            ParserK.Done 0 b -> pure (Right b, cons element stream)
-            ParserK.Done n b -> do
-                let n1 = negate n
-                    bufLen = length backBuf
-                    s = cons element stream
-                assertM(n1 >= 0 && n1 <= bufLen)
-                let (s1, _) = backtrack n1 backBuf s
-                pure (Right b, s1)
-            ParserK.Error n err ->
-                let strm =
-                        append
-                            (fromList (Prelude.reverse backBuf))
-                            (cons element stream)
-                 in return (Left (ParseError (DEFAULT(pos) + n + 1) err), strm)
-
-    go
-        :: OPTIONAL(Int ->)
-           [a]
-        -> (ParserK.Input a -> m (ParserK.Step a m b))
-        -> StreamK m a
-        -> m (Either ParseError b, StreamK m a)
-    go OPTIONAL(pos) backBuf parserk stream = do
-        let stop = goStop OPTIONAL(pos) backBuf parserk
-            single a = yieldk OPTIONAL(pos) backBuf parserk a nil
-         in foldStream
-                defState (yieldk OPTIONAL(pos) backBuf parserk) single stop stream
+-- | Like 'parseBreak' but includes stream position information in the error
+-- messages.
+--
+{-# INLINE parseBreakPos #-}
+parseBreakPos
+    :: forall m a b. Monad m
+    => ParserK.ParserK a m b
+    -> StreamK m a
+    -> m (Either ParseError b, StreamK m a)
+parseBreakPos = Drivers.parseBreakStreamKPos
 
 -- | Run a 'ParserK' over a 'StreamK'. Please use 'parseChunks' where possible,
 -- for better performance.
