diff --git a/lib/compress/zstdmt_compress.c b/lib/compress/zstdmt_compress.c index 5fea6ff95..71cf84fa2 100644 --- a/lib/compress/zstdmt_compress.c +++ b/lib/compress/zstdmt_compress.c @@ -121,6 +121,13 @@ typedef struct { size_t lastBlockSize; } ZSTDMT_chunkProcessResult; +typedef struct { + int status; + size_t toFlush; + size_t outputPos; + size_t dstFlushed; +} ZSTDMT_flushPublicationResult; + typedef void (*ZSTDMT_chunkProgressFn)(void* opaque, size_t cSize, size_t consumed); ZSTDMT_chunkProcessResult ZSTDMT_rust_compressJobChunks( @@ -128,6 +135,11 @@ ZSTDMT_chunkProcessResult ZSTDMT_rust_compressJobChunks( void* dst, size_t dstCapacity, size_t chunkSize, unsigned lastJob, void* progressContext, ZSTDMT_chunkProgressFn progressCallback); +ZSTDMT_flushPublicationResult ZSTDMT_rust_publishJobOutput( + void* outputDst, size_t outputSize, size_t outputPos, + const void* jobDst, size_t jobCapacity, + size_t cSize, size_t dstFlushed); + typedef struct { rawSeq* seq; size_t pos; @@ -1495,19 +1507,22 @@ static size_t ZSTDMT_flushProduced(ZSTDMT_CCtx* mtctx, ZSTD_outBuffer* output, u } if (cSize > 0) { /* compression is ongoing or completed */ - size_t const toFlush = MIN(cSize - mtctx->jobs[wJobID].dstFlushed, output->size - output->pos); + ZSTDMT_flushPublicationResult const publication = ZSTDMT_rust_publishJobOutput( + output->dst, output->size, output->pos, + mtctx->jobs[wJobID].dstBuff.start, mtctx->jobs[wJobID].dstBuff.capacity, + cSize, mtctx->jobs[wJobID].dstFlushed); + if (publication.status != 0) { + assert(publication.status == 0); + return ERROR(GENERIC); + } DEBUGLOG(5, "ZSTDMT_flushProduced: Flushing %u bytes from job %u (completion:%u/%u, generated:%u)", - (U32)toFlush, mtctx->doneJobID, (U32)srcConsumed, (U32)srcSize, (U32)cSize); + (U32)publication.toFlush, mtctx->doneJobID, + (U32)srcConsumed, (U32)srcSize, (U32)cSize); assert(mtctx->doneJobID < mtctx->nextJobID); assert(cSize >= mtctx->jobs[wJobID].dstFlushed); assert(mtctx->jobs[wJobID].dstBuff.start != NULL); - if (toFlush > 0) { - ZSTD_memcpy((char*)output->dst + output->pos, - (const char*)mtctx->jobs[wJobID].dstBuff.start + mtctx->jobs[wJobID].dstFlushed, - toFlush); - } - output->pos += toFlush; - mtctx->jobs[wJobID].dstFlushed += toFlush; /* can write : this value is only used by mtctx */ + output->pos = publication.outputPos; + mtctx->jobs[wJobID].dstFlushed = publication.dstFlushed; /* can write : this value is only used by mtctx */ if ( (srcConsumed == srcSize) /* job is completed */ && (mtctx->jobs[wJobID].dstFlushed == cSize) ) { /* output buffer fully flushed => free this job position */ diff --git a/rust/src/zstdmt_compress.rs b/rust/src/zstdmt_compress.rs index b8db02830..ea1ea65c6 100644 --- a/rust/src/zstdmt_compress.rs +++ b/rust/src/zstdmt_compress.rs @@ -56,6 +56,73 @@ pub struct ZSTDMT_chunkProcessResult { pub lastBlockSize: usize, } +#[repr(C)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct ZSTDMT_flushPublicationResult { + pub status: c_int, + pub toFlush: usize, + pub outputPos: usize, + pub dstFlushed: usize, +} + +const FLUSH_PUBLICATION_OK: c_int = 0; +const FLUSH_PUBLICATION_INVALID_BOUNDS: c_int = 1; + +#[inline] +fn invalid_flush_publication( + output_pos: usize, + dst_flushed: usize, +) -> ZSTDMT_flushPublicationResult { + ZSTDMT_flushPublicationResult { + status: FLUSH_PUBLICATION_INVALID_BOUNDS, + toFlush: 0, + outputPos: output_pos, + dstFlushed: dst_flushed, + } +} + +#[inline] +unsafe fn publish_job_output_with( + output_dst: *mut c_void, + output_size: usize, + output_pos: usize, + job_dst: *const c_void, + job_capacity: usize, + c_size: usize, + dst_flushed: usize, +) -> ZSTDMT_flushPublicationResult { + let Some(output_available) = output_size.checked_sub(output_pos) else { + return invalid_flush_publication(output_pos, dst_flushed); + }; + let Some(job_available) = c_size.checked_sub(dst_flushed) else { + return invalid_flush_publication(output_pos, dst_flushed); + }; + if c_size > job_capacity { + return invalid_flush_publication(output_pos, dst_flushed); + } + + let to_flush = output_available.min(job_available); + if to_flush > 0 { + if output_dst.is_null() || job_dst.is_null() { + return invalid_flush_publication(output_pos, dst_flushed); + } + unsafe { + ptr::copy_nonoverlapping( + job_dst.cast::().add(dst_flushed), + output_dst.cast::().add(output_pos), + to_flush, + ); + } + } + + ZSTDMT_flushPublicationResult { + status: FLUSH_PUBLICATION_OK, + toFlush: to_flush, + outputPos: output_pos + to_flush, + dstFlushed: dst_flushed + to_flush, + } +} + #[cfg(not(test))] unsafe extern "C" { fn ZSTD_compressContinue_public( @@ -185,6 +252,30 @@ pub unsafe extern "C" fn ZSTDMT_rust_compressJobChunks( } } +#[cfg(not(test))] +#[no_mangle] +pub unsafe extern "C" fn ZSTDMT_rust_publishJobOutput( + output_dst: *mut c_void, + output_size: usize, + output_pos: usize, + job_dst: *const c_void, + job_capacity: usize, + c_size: usize, + dst_flushed: usize, +) -> ZSTDMT_flushPublicationResult { + unsafe { + publish_job_output_with( + output_dst, + output_size, + output_pos, + job_dst, + job_capacity, + c_size, + dst_flushed, + ) + } +} + #[inline] fn cycle_log(chain_log: c_uint, strategy: c_int) -> c_uint { chain_log.wrapping_sub((strategy >= ZSTD_BTLAZY2) as c_uint) @@ -1442,6 +1533,143 @@ mod tests { } } + fn run_output_publication( + output: &mut [u8], + output_size: usize, + output_pos: usize, + job: &[u8], + job_capacity: usize, + c_size: usize, + dst_flushed: usize, + ) -> ZSTDMT_flushPublicationResult { + unsafe { + publish_job_output_with( + output.as_mut_ptr().cast(), + output_size, + output_pos, + job.as_ptr().cast(), + job_capacity, + c_size, + dst_flushed, + ) + } + } + + #[test] + fn output_publication_handles_no_output_space() { + let mut output = [0xa5; 8]; + let original = output; + let job = [1, 2, 3, 4]; + let output_size = output.len(); + + let result = run_output_publication( + &mut output, + output_size, + output_size, + &job, + job.len(), + job.len(), + 0, + ); + + assert_eq!(result.status, FLUSH_PUBLICATION_OK); + assert_eq!(result.toFlush, 0); + assert_eq!(result.outputPos, output.len()); + assert_eq!(result.dstFlushed, 0); + assert_eq!(output, original); + } + + #[test] + fn output_publication_handles_partial_flush_with_offsets() { + let mut output = [0xa5; 10]; + let job = [0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17]; + + let result = run_output_publication(&mut output, 9, 2, &job, job.len(), 6, 1); + + assert_eq!(result.status, FLUSH_PUBLICATION_OK); + assert_eq!(result.toFlush, 5); + assert_eq!(result.outputPos, 7); + assert_eq!(result.dstFlushed, 6); + assert_eq!(&output[..2], &[0xa5, 0xa5]); + assert_eq!(&output[2..7], &job[1..6]); + assert_eq!(&output[7..], &[0xa5, 0xa5, 0xa5]); + } + + #[test] + fn output_publication_handles_complete_flush() { + let mut output = [0xa5; 8]; + let job = [0x20, 0x21, 0x22, 0x23, 0x24, 0x25]; + let output_size = output.len(); + + let result = run_output_publication(&mut output, output_size, 1, &job, job.len(), 6, 2); + + assert_eq!(result.status, FLUSH_PUBLICATION_OK); + assert_eq!(result.toFlush, 4); + assert_eq!(result.outputPos, 5); + assert_eq!(result.dstFlushed, 6); + assert_eq!(&output[1..5], &job[2..6]); + assert_eq!(&output[..1], &[0xa5]); + assert_eq!(&output[5..], &[0xa5, 0xa5, 0xa5]); + } + + #[test] + fn output_publication_handles_repeated_partial_flushes() { + let mut output = [0xa5; 8]; + let job = [0x30, 0x31, 0x32, 0x33, 0x34, 0x35]; + + let first = run_output_publication(&mut output, 3, 0, &job, job.len(), 6, 0); + let second = run_output_publication( + &mut output, + 6, + first.outputPos, + &job, + job.len(), + 6, + first.dstFlushed, + ); + + assert_eq!(first.status, FLUSH_PUBLICATION_OK); + assert_eq!(first.toFlush, 3); + assert_eq!(first.outputPos, 3); + assert_eq!(first.dstFlushed, 3); + assert_eq!(second.status, FLUSH_PUBLICATION_OK); + assert_eq!(second.toFlush, 3); + assert_eq!(second.outputPos, 6); + assert_eq!(second.dstFlushed, 6); + assert_eq!(&output[..6], &job); + assert_eq!(&output[6..], &[0xa5, 0xa5]); + } + + #[test] + fn output_publication_does_not_copy_trailing_sentinels() { + let mut output = [0xa5; 12]; + let job = [0x40, 0x41, 0x42, 0x43, 0x44, 0x45, 0xde, 0xad]; + let output_size = output.len(); + + let result = run_output_publication(&mut output, output_size, 4, &job, job.len(), 5, 2); + + assert_eq!(result.status, FLUSH_PUBLICATION_OK); + assert_eq!(result.toFlush, 3); + assert_eq!(&output[..4], &[0xa5, 0xa5, 0xa5, 0xa5]); + assert_eq!(&output[4..7], &job[2..5]); + assert_eq!(&output[7..], &[0xa5, 0xa5, 0xa5, 0xa5, 0xa5]); + } + + #[test] + fn output_publication_rejects_invalid_bounds_without_copy() { + let mut output = [0xa5; 8]; + let original = output; + let job = [0x50, 0x51, 0x52, 0x53]; + + let result = run_output_publication(&mut output, 4, 5, &job, job.len(), 4, 0); + + assert_eq!(result.status, FLUSH_PUBLICATION_INVALID_BOUNDS); + assert_eq!(result.toFlush, 0); + assert_eq!(result.outputPos, 5); + assert_eq!(result.dstFlushed, 0); + assert_eq!(output, original); + } + #[test] fn chunk_loop_handles_empty_jobs_without_compression() { let mut compressor = MockChunkCompressor {