From fcec7d3cc6942a057acd59a680d684eabdc3c16a Mon Sep 17 00:00:00 2001
From: ddidderr <ddidderr@paul.network>
Date: Sat, 18 Jul 2026 12:54:13 +0200
Subject: [PATCH] feat(compress): move end-stream estimate to Rust

Keep ZSTD_endStream's streaming call, error propagation, multithreaded
minimal estimate, logging, and context state in C. Move only the final
single-thread estimate behind a scalar Rust ABI helper so the frame-ended
branch and raw checksum flag are explicit while size_t arithmetic wraps as
in C.

Test Plan:
- `cargo clippy --manifest-path rust/Cargo.toml --no-default-features --features compression` -- passed before and after formatting.
- The same clippy command with `--benches` and `--tests` -- passed before and after formatting.
- `cargo +nightly fmt --manifest-path rust/Cargo.toml` -- passed.
- Focused end-stream Rust tests -- 4 passed.
- `make -B -C lib -j2 lib` -- passed.
- `make -C tests test-rust-lib-smoke` -- passed.
- `tests/fuzzer -s4560 -t56 -i57 -v` -- passed.
- `make -C tests -j2 test-zstream` -- passed: 84 named, 6,457, and 8,692 randomized tests. The pre-existing unterminated-string warning remains.
- `git diff --check` and `git diff --cached --check` -- passed.
---
 lib/compress/zstd_compress.c |  8 +++---
 rust/src/zstd_compress.rs    | 49 ++++++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c
index b56ce07d5..048071324 100644
--- a/lib/compress/zstd_compress.c
+++ b/lib/compress/zstd_compress.c
@@ -106,6 +106,8 @@ ZSTD_inBuffer ZSTD_rust_inBufferForEndFlush(int inBufferMode,
                                             const void* expectedSrc,
                                             size_t expectedSize,
                                             size_t expectedPos);
+size_t ZSTD_rust_endStreamRemaining(size_t remainingToFlush,
+                                    int frameEnded, int checksumFlag);
 size_t ZSTD_rust_checkBufferStability(
         int inBufferMode, int outBufferMode,
         const void* expectedInSrc, size_t expectedInPos,
@@ -5655,9 +5657,9 @@ size_t ZSTD_endStream(ZSTD_CStream* zcs, ZSTD_outBuffer* output)
     FORWARD_IF_ERROR(remainingToFlush , "ZSTD_compressStream2(,,ZSTD_e_end) failed");
     if (zcs->appliedParams.nbWorkers > 0) return remainingToFlush;   /* minimal estimation */
     /* single thread mode : attempt to calculate remaining to flush more precisely */
-    {   size_t const lastBlockSize = zcs->frameEnded ? 0 : ZSTD_BLOCKHEADERSIZE;
-        size_t const checksumSize = (size_t)(zcs->frameEnded ? 0 : zcs->appliedParams.fParams.checksumFlag * 4);
-        size_t const toFlush = remainingToFlush + lastBlockSize + checksumSize;
+    {   size_t const toFlush = ZSTD_rust_endStreamRemaining(
+                remainingToFlush, zcs->frameEnded,
+                zcs->appliedParams.fParams.checksumFlag);
         DEBUGLOG(4, "ZSTD_endStream : remaining to flush : %u", (unsigned)toFlush);
         return toFlush;
     }
diff --git a/rust/src/zstd_compress.rs b/rust/src/zstd_compress.rs
index 47a840ca2..a7d675c99 100644
--- a/rust/src/zstd_compress.rs
+++ b/rust/src/zstd_compress.rs
@@ -198,6 +198,31 @@ pub extern "C" fn ZSTD_rust_inBufferForEndFlush(
     }
 }
 
+#[inline]
+fn end_stream_remaining(
+    remaining_to_flush: usize,
+    frame_ended: c_int,
+    checksum_flag: c_int,
+) -> usize {
+    if frame_ended != 0 {
+        return remaining_to_flush;
+    }
+
+    remaining_to_flush
+        .wrapping_add(ZSTD_BLOCK_HEADER_SIZE)
+        .wrapping_add((checksum_flag as usize).wrapping_mul(4))
+}
+
+/// Estimate single-threaded end-stream output without crossing C context state.
+#[no_mangle]
+pub extern "C" fn ZSTD_rust_endStreamRemaining(
+    remaining_to_flush: usize,
+    frame_ended: c_int,
+    checksum_flag: c_int,
+) -> usize {
+    end_stream_remaining(remaining_to_flush, frame_ended, checksum_flag)
+}
+
 #[inline]
 fn check_buffer_stability(
     in_buffer_mode: c_int,
@@ -1471,6 +1496,30 @@ mod tests {
         }
     }
 
+    #[test]
+    fn end_stream_remaining_ignores_estimate_components_after_frame_end() {
+        assert_eq!(end_stream_remaining(17, 1, 1), 17);
+        assert_eq!(ZSTD_rust_endStreamRemaining(17, 1, 1), 17);
+    }
+
+    #[test]
+    fn end_stream_remaining_adds_block_header_without_checksum() {
+        assert_eq!(end_stream_remaining(17, 0, 0), 20);
+        assert_eq!(ZSTD_rust_endStreamRemaining(17, 0, 0), 20);
+    }
+
+    #[test]
+    fn end_stream_remaining_adds_block_header_and_checksum() {
+        assert_eq!(end_stream_remaining(17, 0, 1), 24);
+        assert_eq!(ZSTD_rust_endStreamRemaining(17, 0, 1), 24);
+    }
+
+    #[test]
+    fn end_stream_remaining_wraps_size_t_additions() {
+        assert_eq!(end_stream_remaining(usize::MAX, 0, 1), 6);
+        assert_eq!(ZSTD_rust_endStreamRemaining(usize::MAX, 0, 1), 6);
+    }
+
     #[test]
     fn check_buffer_stability_accepts_matching_stable_input() {
         let expected_src = b"input".as_ptr().cast::<c_void>();