diff --git a/lib/compress/zstdmt_compress.c b/lib/compress/zstdmt_compress.c index c37a2b7de..6cfd39637 100644 --- a/lib/compress/zstdmt_compress.c +++ b/lib/compress/zstdmt_compress.c @@ -276,8 +276,6 @@ unsigned ZSTDMT_rust_computeTargetJobLog(unsigned windowLog, unsigned chainLog, int ZSTDMT_rust_overlapLog(int overlapLog, int strategy); size_t ZSTDMT_rust_computeOverlapSize(unsigned windowLog, unsigned chainLog, int strategy, int overlapLog, int enableLdm); -int ZSTDMT_rust_isOverlapped(const void* bufferStart, size_t bufferCapacity, - const void* rangeStart, size_t rangeSize); int ZSTDMT_rust_doesOverlapWindow(const void* bufferStart, size_t bufferCapacity, const void* nextSrc, const void* base, const void* dictBase, U32 dictLimit, U32 lowLimit); @@ -302,12 +300,31 @@ typedef struct { const void* start; size_t size; } ZSTDMT_RustInputRange; +typedef struct { + const void* roundBufferStart; + size_t roundBufferCapacity; + size_t roundBufferPos; + const void* prefixStart; + size_t prefixSize; + size_t targetSectionSize; + const void* inUseStart; + size_t inUseSize; +} ZSTDMT_RustTryGetInputRangeProjection; +typedef struct { + unsigned ready; + unsigned movePrefix; + void* bufferStart; + size_t bufferCapacity; + size_t roundBufferPos; +} ZSTDMT_RustTryGetInputRangeResult; typedef void (*ZSTDMT_jobProjectionFn)(void* opaque, unsigned jobID, ZSTDMT_RustJobProjection* projection); ZSTDMT_RustInputRange ZSTDMT_rust_getInputDataInUse( unsigned firstJobID, unsigned lastJobID, unsigned jobIDMask, size_t roundBufferCapacity, size_t targetSectionSize, void* opaque, ZSTDMT_jobProjectionFn projectJob); +ZSTDMT_RustTryGetInputRangeResult ZSTDMT_rust_tryGetInputRange( + const ZSTDMT_RustTryGetInputRangeProjection* projection); size_t ZSTDMT_rust_toFlushNow(unsigned doneJobID, unsigned nextJobID, unsigned jobIDMask, void* opaque, ZSTDMT_jobProjectionFn projectJob); @@ -1760,15 +1777,6 @@ static Range ZSTDMT_getInputDataInUse(ZSTDMT_CCtx* mtctx) return (Range){ range.start, range.size }; } -/** - * Returns non-zero iff buffer and range overlap. - */ -static int ZSTDMT_isOverlapped(Buffer buffer, Range range) -{ - return ZSTDMT_rust_isOverlapped(buffer.start, buffer.capacity, - range.start, range.size); -} - static int ZSTDMT_doesOverlapWindow(Buffer buffer, ZSTD_window_t window) { DEBUGLOG(5, "ZSTDMT_doesOverlapWindow"); @@ -1811,15 +1819,29 @@ static void ZSTDMT_waitForLdmComplete(ZSTDMT_CCtx* mtctx, Buffer buffer) static int ZSTDMT_tryGetInputRange(ZSTDMT_CCtx* mtctx) { Range const inUse = ZSTDMT_getInputDataInUse(mtctx); - size_t const spaceLeft = mtctx->roundBuff.capacity - mtctx->roundBuff.pos; - size_t const spaceNeeded = mtctx->targetSectionSize; + ZSTDMT_RustTryGetInputRangeProjection const projection = { + mtctx->roundBuff.buffer, + mtctx->roundBuff.capacity, + mtctx->roundBuff.pos, + mtctx->inBuff.prefix.start, + mtctx->inBuff.prefix.size, + mtctx->targetSectionSize, + inUse.start, + inUse.size + }; + ZSTDMT_RustTryGetInputRangeResult const result = + ZSTDMT_rust_tryGetInputRange(&projection); Buffer buffer; DEBUGLOG(5, "ZSTDMT_tryGetInputRange"); assert(mtctx->inBuff.buffer.start == NULL); - assert(mtctx->roundBuff.capacity >= spaceNeeded); - if (spaceLeft < spaceNeeded) { + if (!result.ready) { + DEBUGLOG(5, "Waiting for buffer..."); + return 0; + } + + if (result.movePrefix) { /* ZSTD_invalidateRepCodes() doesn't work for extDict variants. * Simply copy the prefix to the beginning in that case. */ @@ -1828,23 +1850,14 @@ static int ZSTDMT_tryGetInputRange(ZSTDMT_CCtx* mtctx) buffer.start = start; buffer.capacity = prefixSize; - if (ZSTDMT_isOverlapped(buffer, inUse)) { - DEBUGLOG(5, "Waiting for buffer..."); - return 0; - } ZSTDMT_waitForLdmComplete(mtctx, buffer); ZSTD_memmove(start, mtctx->inBuff.prefix.start, prefixSize); mtctx->inBuff.prefix.start = start; - mtctx->roundBuff.pos = prefixSize; + mtctx->roundBuff.pos = result.roundBufferPos; } - buffer.start = mtctx->roundBuff.buffer + mtctx->roundBuff.pos; - buffer.capacity = spaceNeeded; - if (ZSTDMT_isOverlapped(buffer, inUse)) { - DEBUGLOG(5, "Waiting for buffer..."); - return 0; - } - assert(!ZSTDMT_isOverlapped(buffer, mtctx->inBuff.prefix)); + buffer.start = (BYTE*)result.bufferStart; + buffer.capacity = result.bufferCapacity; ZSTDMT_waitForLdmComplete(mtctx, buffer); diff --git a/rust/src/zstdmt_compress.rs b/rust/src/zstdmt_compress.rs index 0ce5e9d60..abc5f5510 100644 --- a/rust/src/zstdmt_compress.rs +++ b/rust/src/zstdmt_compress.rs @@ -170,6 +170,35 @@ pub struct ZSTDMT_inputRange { pub size: usize, } +/// Scalar input-buffer state used by the MT range-selection policy. C keeps +/// the input and round-buffer structs private; Rust decides only whether the +/// next target section can be reused and where that section starts. +#[repr(C)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct ZSTDMT_tryGetInputRangeProjection { + pub roundBufferStart: *const c_void, + pub roundBufferCapacity: usize, + pub roundBufferPos: usize, + pub prefixStart: *const c_void, + pub prefixSize: usize, + pub targetSectionSize: usize, + pub inUseStart: *const c_void, + pub inUseSize: usize, +} + +/// Result of the MT input-range selection. C applies the returned decision, +/// including synchronization waits, prefix copying, and private field +/// updates. +#[repr(C)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct ZSTDMT_tryGetInputRangeResult { + pub ready: c_uint, + pub movePrefix: c_uint, + pub bufferStart: *mut c_void, + pub bufferCapacity: usize, + pub roundBufferPos: usize, +} + pub type ZSTDMT_jobProjectionFn = unsafe extern "C" fn( opaque: *mut c_void, job_id: c_uint, @@ -962,6 +991,93 @@ pub unsafe extern "C" fn ZSTDMT_rust_toFlushNow( }) } +#[inline] +fn try_get_input_range_with( + projection: ZSTDMT_tryGetInputRangeProjection, +) -> ZSTDMT_tryGetInputRangeResult { + let invalid = || ZSTDMT_tryGetInputRangeResult::default(); + + debug_assert!(projection.roundBufferPos <= projection.roundBufferCapacity); + debug_assert!(projection.targetSectionSize > 0); + if projection.roundBufferPos > projection.roundBufferCapacity + || projection.targetSectionSize == 0 + || projection.targetSectionSize > projection.roundBufferCapacity + { + return invalid(); + } + + let space_left = projection.roundBufferCapacity - projection.roundBufferPos; + let (round_buffer_pos, move_prefix) = if space_left < projection.targetSectionSize { + /* The C implementation first checks the prefix range at the start of + * the round buffer before resetting the position there. */ + if is_overlapped( + projection.roundBufferStart, + projection.prefixSize, + projection.inUseStart, + projection.inUseSize, + ) != 0 + { + return invalid(); + } + if projection.prefixSize > projection.roundBufferCapacity - projection.targetSectionSize { + return invalid(); + } + (projection.prefixSize, 1) + } else { + (projection.roundBufferPos, 0) + }; + + let buffer_start = projection + .roundBufferStart + .cast::() + .wrapping_add(round_buffer_pos) + .cast_mut() + .cast::(); + if move_prefix == 0 { + /* This is an internal invariant in the original C helper: a live + * prefix must never be overwritten by the next source section. */ + debug_assert_eq!( + is_overlapped( + buffer_start, + projection.targetSectionSize, + projection.prefixStart, + projection.prefixSize, + ), + 0 + ); + } + if is_overlapped( + buffer_start, + projection.targetSectionSize, + projection.inUseStart, + projection.inUseSize, + ) != 0 + { + return invalid(); + } + + ZSTDMT_tryGetInputRangeResult { + ready: 1, + movePrefix: move_prefix, + bufferStart: buffer_start, + bufferCapacity: projection.targetSectionSize, + roundBufferPos: round_buffer_pos, + } +} + +/// Select the next reusable round-buffer section for MT input. The C adapter +/// retains all synchronization, prefix copying, and private context updates. +#[cfg(not(test))] +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_tryGetInputRange( + projection: *const ZSTDMT_tryGetInputRangeProjection, +) -> ZSTDMT_tryGetInputRangeResult { + let Some(projection) = (unsafe { projection.as_ref() }).copied() else { + return ZSTDMT_tryGetInputRangeResult::default(); + }; + try_get_input_range_with(projection) +} + #[inline] fn cycle_log(chain_log: c_uint, strategy: c_int) -> c_uint { chain_log.wrapping_sub((strategy >= ZSTD_BTLAZY2) as c_uint) @@ -3087,6 +3203,70 @@ mod tests { assert_eq!(result.size, source.len()); } + #[test] + fn try_get_input_range_selects_next_section_without_wrapping() { + let round_buffer = [0u8; 32]; + let prefix = [1u8; 4]; + let projection = ZSTDMT_tryGetInputRangeProjection { + roundBufferStart: round_buffer.as_ptr().cast(), + roundBufferCapacity: round_buffer.len(), + roundBufferPos: 8, + prefixStart: prefix.as_ptr().cast(), + prefixSize: prefix.len(), + targetSectionSize: 8, + ..ZSTDMT_tryGetInputRangeProjection::default() + }; + + let result = try_get_input_range_with(projection); + + assert_eq!(result.ready, 1); + assert_eq!(result.movePrefix, 0); + assert_eq!(result.roundBufferPos, 8); + assert_eq!( + result.bufferStart, + round_buffer.as_ptr().wrapping_add(8).cast_mut().cast() + ); + assert_eq!(result.bufferCapacity, 8); + } + + #[test] + fn try_get_input_range_wraps_after_prefix_and_rejects_live_ranges() { + let round_buffer = [0u8; 32]; + let prefix = [1u8; 4]; + let wrapped = ZSTDMT_tryGetInputRangeProjection { + roundBufferStart: round_buffer.as_ptr().cast(), + roundBufferCapacity: round_buffer.len(), + roundBufferPos: 28, + prefixStart: prefix.as_ptr().cast(), + prefixSize: prefix.len(), + targetSectionSize: 8, + ..ZSTDMT_tryGetInputRangeProjection::default() + }; + let result = try_get_input_range_with(wrapped); + assert_eq!(result.ready, 1); + assert_eq!(result.movePrefix, 1); + assert_eq!(result.roundBufferPos, prefix.len()); + assert_eq!( + result.bufferStart, + round_buffer.as_ptr().wrapping_add(4).cast_mut().cast() + ); + + let blocked_prefix = ZSTDMT_tryGetInputRangeProjection { + inUseStart: round_buffer.as_ptr().cast(), + inUseSize: prefix.len(), + ..wrapped + }; + assert_eq!(try_get_input_range_with(blocked_prefix).ready, 0); + + let blocked_source = ZSTDMT_tryGetInputRangeProjection { + roundBufferPos: 8, + inUseStart: round_buffer.as_ptr().wrapping_add(8).cast(), + inUseSize: 8, + ..wrapped + }; + assert_eq!(try_get_input_range_with(blocked_source).ready, 0); + } + #[test] fn to_flush_now_projects_oldest_slot_and_normalizes_errors() { let projection = ZSTDMT_jobProjection {