diff --git a/lib/compress/zstdmt_compress.c b/lib/compress/zstdmt_compress.c index ddcd8c0ef..a338910e9 100644 --- a/lib/compress/zstdmt_compress.c +++ b/lib/compress/zstdmt_compress.c @@ -149,6 +149,8 @@ void ZSTDMT_rust_findSynchronizationPoint(const void* inputSrc, size_t inputSize const void* inBuffStart, size_t inBuffFilled, int rsyncable, U64 primePower, U64 hitMask, size_t* toLoad, int* flush); +size_t ZSTDMT_rust_nextInputSizeHint(size_t targetSectionSize, + size_t inBuffFilled); typedef struct ZSTDMT_bufferPool_s { ZSTDMT_RustBufferPool* rustPool; @@ -1665,9 +1667,8 @@ findSynchronizationPoint(ZSTDMT_CCtx const* mtctx, ZSTD_inBuffer const input) size_t ZSTDMT_nextInputSizeHint(const ZSTDMT_CCtx* mtctx) { - size_t hintInSize = mtctx->targetSectionSize - mtctx->inBuff.filled; - if (hintInSize==0) hintInSize = mtctx->targetSectionSize; - return hintInSize; + return ZSTDMT_rust_nextInputSizeHint(mtctx->targetSectionSize, + mtctx->inBuff.filled); } /** ZSTDMT_compressStream_generic() : diff --git a/rust/src/zstd_compress.rs b/rust/src/zstd_compress.rs index 95b73ca68..a6c5aeed3 100644 --- a/rust/src/zstd_compress.rs +++ b/rust/src/zstd_compress.rs @@ -303,6 +303,25 @@ pub extern "C" fn ZSTD_rust_nextInputSizeHint( ) } +#[inline] +fn mt_next_input_size_hint(target_section_size: usize, in_buff_filled: usize) -> usize { + let hint_in_size = target_section_size.wrapping_sub(in_buff_filled); + if hint_in_size == 0 { + target_section_size + } else { + hint_in_size + } +} + +/// Return the next input size required by the C multithreaded streaming state. +#[no_mangle] +pub extern "C" fn ZSTDMT_rust_nextInputSizeHint( + target_section_size: usize, + in_buff_filled: usize, +) -> usize { + mt_next_input_size_hint(target_section_size, in_buff_filled) +} + #[inline] fn bitmix(mut val: u64, len: u64) -> u64 { val ^= val.rotate_right(49) ^ val.rotate_right(24); @@ -1191,6 +1210,33 @@ mod tests { ); } + #[test] + fn mt_next_input_size_hint_handles_empty_input_buffer() { + assert_eq!(mt_next_input_size_hint(128, 0), 128); + assert_eq!(ZSTDMT_rust_nextInputSizeHint(128, 0), 128); + } + + #[test] + fn mt_next_input_size_hint_returns_remaining_capacity() { + assert_eq!(mt_next_input_size_hint(128, 37), 91); + assert_eq!(ZSTDMT_rust_nextInputSizeHint(128, 37), 91); + } + + #[test] + fn mt_next_input_size_hint_replaces_full_buffer_with_target_size() { + assert_eq!(ZSTDMT_rust_nextInputSizeHint(128, 128), 128); + } + + #[test] + fn mt_next_input_size_hint_preserves_zero_target_behavior() { + assert_eq!(ZSTDMT_rust_nextInputSizeHint(0, 0), 0); + } + + #[test] + fn mt_next_input_size_hint_preserves_wrapped_overfill() { + assert_eq!(ZSTDMT_rust_nextInputSizeHint(3, 4), usize::MAX); + } + #[test] fn in_buffer_for_end_flush_returns_stable_expected_buffer() { let expected_src = b"input".as_ptr().cast::();