diff --git a/lib/compress/zstdmt_compress.c b/lib/compress/zstdmt_compress.c index a338910e9..ed63e2163 100644 --- a/lib/compress/zstdmt_compress.c +++ b/lib/compress/zstdmt_compress.c @@ -151,6 +151,10 @@ void ZSTDMT_rust_findSynchronizationPoint(const void* inputSrc, size_t inputSize size_t* toLoad, int* flush); size_t ZSTDMT_rust_nextInputSizeHint(size_t targetSectionSize, size_t inBuffFilled); +size_t ZSTDMT_rust_sizeofCCtx(size_t mtctxSize, size_t factorySize, + size_t bufferPoolSize, size_t jobsSize, + size_t cctxPoolSize, size_t seqPoolSize, + size_t cdictSize, size_t roundBuffSize); typedef struct ZSTDMT_bufferPool_s { ZSTDMT_RustBufferPool* rustPool; @@ -1029,15 +1033,26 @@ size_t ZSTDMT_freeCCtx(ZSTDMT_CCtx* mtctx) size_t ZSTDMT_sizeof_CCtx(ZSTDMT_CCtx* mtctx) { + size_t mtctxSize; + size_t factorySize; + size_t bufferPoolSize; + size_t jobsSize; + size_t cctxPoolSize; + size_t seqPoolSize; + size_t cdictSize; + size_t roundBuffSize; if (mtctx == NULL) return 0; /* supports sizeof NULL */ - return sizeof(*mtctx) - + POOL_sizeof(mtctx->factory) - + ZSTDMT_sizeof_bufferPool(mtctx->bufPool) - + (mtctx->jobIDMask+1) * sizeof(ZSTDMT_jobDescription) - + ZSTDMT_sizeof_CCtxPool(mtctx->cctxPool) - + ZSTDMT_sizeof_seqPool(mtctx->seqPool) - + ZSTD_sizeof_CDict(mtctx->cdictLocal) - + mtctx->roundBuff.capacity; + mtctxSize = sizeof(*mtctx); + factorySize = POOL_sizeof(mtctx->factory); + bufferPoolSize = ZSTDMT_sizeof_bufferPool(mtctx->bufPool); + jobsSize = (mtctx->jobIDMask+1) * sizeof(ZSTDMT_jobDescription); + cctxPoolSize = ZSTDMT_sizeof_CCtxPool(mtctx->cctxPool); + seqPoolSize = ZSTDMT_sizeof_seqPool(mtctx->seqPool); + cdictSize = ZSTD_sizeof_CDict(mtctx->cdictLocal); + roundBuffSize = mtctx->roundBuff.capacity; + return ZSTDMT_rust_sizeofCCtx(mtctxSize, factorySize, bufferPoolSize, + jobsSize, cctxPoolSize, seqPoolSize, + cdictSize, roundBuffSize); } diff --git a/rust/src/zstd_compress.rs b/rust/src/zstd_compress.rs index a6c5aeed3..bcd200c05 100644 --- a/rust/src/zstd_compress.rs +++ b/rust/src/zstd_compress.rs @@ -322,6 +322,52 @@ pub extern "C" fn ZSTDMT_rust_nextInputSizeHint( mt_next_input_size_hint(target_section_size, in_buff_filled) } +#[inline] +fn mt_sizeof_cctx( + mtctx_size: usize, + factory_size: usize, + buffer_pool_size: usize, + jobs_size: usize, + cctx_pool_size: usize, + seq_pool_size: usize, + cdict_size: usize, + round_buff_size: usize, +) -> usize { + mtctx_size + .wrapping_add(factory_size) + .wrapping_add(buffer_pool_size) + .wrapping_add(jobs_size) + .wrapping_add(cctx_pool_size) + .wrapping_add(seq_pool_size) + .wrapping_add(cdict_size) + .wrapping_add(round_buff_size) +} + +/// Aggregate C-owned multithreaded context size components with C `size_t` +/// wrapping semantics. +#[no_mangle] +pub extern "C" fn ZSTDMT_rust_sizeofCCtx( + mtctx_size: usize, + factory_size: usize, + buffer_pool_size: usize, + jobs_size: usize, + cctx_pool_size: usize, + seq_pool_size: usize, + cdict_size: usize, + round_buff_size: usize, +) -> usize { + mt_sizeof_cctx( + mtctx_size, + factory_size, + buffer_pool_size, + jobs_size, + cctx_pool_size, + seq_pool_size, + cdict_size, + round_buff_size, + ) +} + #[inline] fn bitmix(mut val: u64, len: u64) -> u64 { val ^= val.rotate_right(49) ^ val.rotate_right(24); @@ -1237,6 +1283,24 @@ mod tests { assert_eq!(ZSTDMT_rust_nextInputSizeHint(3, 4), usize::MAX); } + #[test] + fn mt_sizeof_cctx_handles_zero_components() { + assert_eq!(mt_sizeof_cctx(0, 0, 0, 0, 0, 0, 0, 0), 0); + assert_eq!(ZSTDMT_rust_sizeofCCtx(0, 0, 0, 0, 0, 0, 0, 0), 0); + } + + #[test] + fn mt_sizeof_cctx_adds_all_components_in_order() { + assert_eq!(mt_sizeof_cctx(1, 2, 3, 4, 5, 6, 7, 8), 36); + assert_eq!(ZSTDMT_rust_sizeofCCtx(1, 2, 3, 4, 5, 6, 7, 8), 36); + } + + #[test] + fn mt_sizeof_cctx_wraps_like_c_size_t_addition() { + assert_eq!(mt_sizeof_cctx(usize::MAX, 1, 2, 3, 4, 5, 6, 7), 27); + assert_eq!(ZSTDMT_rust_sizeofCCtx(usize::MAX, 1, 2, 3, 4, 5, 6, 7), 27); + } + #[test] fn in_buffer_for_end_flush_returns_stable_expected_buffer() { let expected_src = b"input".as_ptr().cast::();