From 5946303b27b0709da955ecf0a988c1dd1c64babe Mon Sep 17 00:00:00 2001
From: ZhuchkaTriplesix <mrlololoshka94@gmail.com>
Date: Tue, 29 Sep 2026 10:09:13 +0300
Subject: [PATCH] perf(transport): eliminate full-string heap reallocation in
 write_ndjson framing. Closes #50
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

write_ndjson sanitized raw '\n'/'\r' bytes by collecting the entire
payload into a new String via .chars().map(...).collect(), so every
multi-megabyte tabular response containing a stray newline (e.g. in an
error message or embedded query text) got a full extra heap copy just
to swap a handful of bytes for spaces.

'\n' and '\r' are single-byte ASCII code points, so this rewrites the
sanitization as a byte-slice scan: write the clean slice up to each
newline straight to the writer, then a single space byte, and advance
past it — no intermediate String is ever allocated. A payload with no
newlines (the common case) still goes out in one write_all call, same
as before.
---
 src/transport/framing.rs | 53 +++++++++++++++++++++++++++++++++-------
 1 file changed, 44 insertions(+), 9 deletions(-)

diff --git a/src/transport/framing.rs b/src/transport/framing.rs
index 92d588f..e852fdb 100644
--- a/src/transport/framing.rs
+++ b/src/transport/framing.rs
@@ -12,16 +12,21 @@ pub fn write_ndjson_stdout(payload: &str) -> io::Result<()> {
 /// Generic NDJSON writer that works with any `std::io::Write` sink (useful for testing).
 pub fn write_ndjson<W: Write>(writer: &mut W, payload: &str) -> io::Result<()> {
     // If the JSON payload contains raw '\n' or '\r' bytes (not escaped inside strings),
-    // replace them with spaces or strip to guarantee exact NDJSON framing.
-    if payload.contains('\n') || payload.contains('\r') {
-        let sanitized: String = payload
-            .chars()
-            .map(|c| if c == '\n' || c == '\r' { ' ' } else { c })
-            .collect();
-        writer.write_all(sanitized.as_bytes())?;
-    } else {
-        writer.write_all(payload.as_bytes())?;
+    // replace them with spaces to guarantee exact NDJSON framing. Both are single-byte
+    // ASCII code points, so this writes existing byte slices straight to `writer`
+    // between them instead of collecting a sanitized copy of the whole payload onto
+    // the heap first (issue #50) — a no-op payload (the common case) is written in
+    // one `write_all` call, same as before.
+    let bytes = payload.as_bytes();
+    let mut start = 0;
+    for (i, &b) in bytes.iter().enumerate() {
+        if b == b'\n' || b == b'\r' {
+            writer.write_all(&bytes[start..i])?;
+            writer.write_all(b" ")?;
+            start = i + 1;
+        }
     }
+    writer.write_all(&bytes[start..])?;
     writer.write_all(b"\n")?;
     writer.flush()?;
     Ok(())
@@ -56,6 +61,36 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_write_ndjson_sanitizes_edge_positions_and_runs() {
+        // Regression for issue #50: the byte-slice rewrite must handle a
+        // newline as the very first/last byte and consecutive newlines
+        // (an empty slice between them) without panicking or dropping bytes.
+        let mut buffer = Vec::new();
+        let dirty = "\nleading\r\rmiddle\n\ntrailing\n";
+        write_ndjson(&mut buffer, dirty).unwrap();
+        let output = String::from_utf8(buffer).unwrap();
+        assert_eq!(output, " leading  middle  trailing \n");
+    }
+
+    #[test]
+    fn test_write_ndjson_sanitizes_large_payload() {
+        // A payload well past any small-buffer fast path, to exercise the
+        // byte-slice rewrite over a realistic multi-megabyte tabular result.
+        let mut payload = "x".repeat(2 * 1024 * 1024);
+        payload.push('\n');
+        payload.push_str(&"y".repeat(1024));
+
+        let mut buffer = Vec::new();
+        write_ndjson(&mut buffer, &payload).unwrap();
+        let output = String::from_utf8(buffer).unwrap();
+
+        assert_eq!(output.len(), payload.len() + 1);
+        assert_eq!(&output[..2 * 1024 * 1024], "x".repeat(2 * 1024 * 1024));
+        assert_eq!(output.as_bytes()[2 * 1024 * 1024], b' ');
+        assert!(output.ends_with(&format!("{}\n", "y".repeat(1024))));
+    }
+
     #[test]
     fn test_write_ndjson_error_payload() {
         let mut buffer = Vec::new();