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::();