From bab90f36224d995c62d4ba7d75c63f87465bab64 Mon Sep 17 00:00:00 2001 From: ddidderr Date: Mon, 20 Jul 2026 19:29:01 +0200 Subject: [PATCH] refactor(compress): move end-stream result policy to Rust Project the worker-count choice in ZSTD_endStream into Rust. Multithreaded streams retain their minimal remaining-output estimate, while single-threaded streams keep the precise end-stream arithmetic. C retains error forwarding, private context reads, and the public wrapper. Test Plan: - `ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo clippy --manifest-path rust/Cargo.toml --all-targets -- -D warnings` -- passed - `ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo clippy --manifest-path rust/cli/Cargo.toml --all-targets -- -D warnings` -- passed - `ulimit -v 41943040; cargo +nightly fmt --manifest-path rust/Cargo.toml --all -- --check` -- passed - `ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo test --manifest-path rust/cli/Cargo.toml --all-targets` -- passed, 190 tests - `ulimit -v 41943040; CARGO_BUILD_JOBS=1 make -j1` -- passed - `ulimit -v 41943040; CARGO_BUILD_JOBS=1 ./tests/rustLibSmoke` -- passed - `ulimit -v 41943040; CARGO_BUILD_JOBS=1 make -j1 -C tests test` -- passed, including large streaming, native, fuzzer, and zstream phases --- lib/compress/zstd_compress.c | 15 ++++++----- rust/src/zstd_compress.rs | 52 ++++++++++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+), 6 deletions(-) diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c index a858fb47a..444b5529d 100644 --- a/lib/compress/zstd_compress.c +++ b/lib/compress/zstd_compress.c @@ -1978,6 +1978,9 @@ ZSTD_inBuffer ZSTD_rust_inBufferForEndFlush(int inBufferMode, size_t expectedPos); size_t ZSTD_rust_endStreamRemaining(size_t remainingToFlush, int frameEnded, int checksumFlag); +size_t ZSTD_rust_endStreamRemainingPolicy(size_t remainingToFlush, + int nbWorkers, + int frameEnded, int checksumFlag); size_t ZSTD_rust_checkBufferStability( int inBufferMode, int outBufferMode, const void* expectedInSrc, size_t expectedInPos, @@ -8676,12 +8679,12 @@ size_t ZSTD_endStream(ZSTD_CStream* zcs, ZSTD_outBuffer* output) ZSTD_inBuffer input = inBuffer_forEndFlush(zcs); size_t const remainingToFlush = ZSTD_compressStream2(zcs, output, &input, ZSTD_e_end); 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 toFlush = ZSTD_rust_endStreamRemaining( - remainingToFlush, zcs->frameEnded, - zcs->appliedParams.fParams.checksumFlag); - DEBUGLOG(4, "ZSTD_endStream : remaining to flush : %u", (unsigned)toFlush); + { size_t const toFlush = ZSTD_rust_endStreamRemainingPolicy( + remainingToFlush, zcs->appliedParams.nbWorkers, + zcs->frameEnded, zcs->appliedParams.fParams.checksumFlag); + if (zcs->appliedParams.nbWorkers <= 0) { + 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 ff951ab48..42f7c56af 100644 --- a/rust/src/zstd_compress.rs +++ b/rust/src/zstd_compress.rs @@ -6495,6 +6495,33 @@ pub extern "C" fn ZSTD_rust_endStreamRemaining( end_stream_remaining(remaining_to_flush, frame_ended, checksum_flag) } +#[inline] +fn end_stream_remaining_policy( + remaining_to_flush: usize, + nb_workers: c_int, + frame_ended: c_int, + checksum_flag: c_int, +) -> usize { + if nb_workers > 0 { + return remaining_to_flush; + } + + ZSTD_rust_endStreamRemaining(remaining_to_flush, frame_ended, checksum_flag) +} + +/// Select the end-stream remaining-output result without crossing C context +/// state. Multithreaded streams keep the minimal estimate; single-threaded +/// streams use the more precise scalar arithmetic above. +#[no_mangle] +pub extern "C" fn ZSTD_rust_endStreamRemainingPolicy( + remaining_to_flush: usize, + nb_workers: c_int, + frame_ended: c_int, + checksum_flag: c_int, +) -> usize { + end_stream_remaining_policy(remaining_to_flush, nb_workers, frame_ended, checksum_flag) +} + #[inline] fn check_buffer_stability( in_buffer_mode: c_int, @@ -15253,6 +15280,31 @@ mod tests { assert_eq!(ZSTD_rust_endStreamRemaining(usize::MAX, 0, 1), 6); } + #[test] + fn end_stream_remaining_policy_returns_minimal_estimate_for_workers() { + for frame_ended in [0, 1] { + for checksum_flag in [0, 1] { + assert_eq!( + end_stream_remaining_policy(17, 1, frame_ended, checksum_flag), + 17 + ); + assert_eq!( + ZSTD_rust_endStreamRemainingPolicy(17, 1, frame_ended, checksum_flag), + 17 + ); + } + } + } + + #[test] + fn end_stream_remaining_policy_estimates_single_thread_combinations() { + assert_eq!(end_stream_remaining_policy(17, 0, 1, 1), 17); + assert_eq!(end_stream_remaining_policy(17, 0, 1, 0), 17); + assert_eq!(end_stream_remaining_policy(17, 0, 0, 0), 20); + assert_eq!(end_stream_remaining_policy(17, 0, 0, 1), 24); + assert_eq!(ZSTD_rust_endStreamRemainingPolicy(17, 0, 0, 1), 24); + } + #[test] fn check_buffer_stability_accepts_matching_stable_input() { let expected_src = b"input".as_ptr().cast::();