diff --git a/lib/compress/zstdmt_compress.c b/lib/compress/zstdmt_compress.c index c3d912073..e8b42224c 100644 --- a/lib/compress/zstdmt_compress.c +++ b/lib/compress/zstdmt_compress.c @@ -106,6 +106,7 @@ ZSTDMT_RustBufferPool* ZSTDMT_rust_buffer_pool_create(unsigned maxNbBuffers, ZSTD_customMem cMem); void ZSTDMT_rust_buffer_pool_free(ZSTDMT_RustBufferPool* pool); size_t ZSTDMT_rust_buffer_pool_sizeof(const ZSTDMT_RustBufferPool* pool); +size_t ZSTDMT_rust_sizeofBufferPool(size_t wrapperSize, size_t rustPoolSize); void ZSTDMT_rust_buffer_pool_set_size(ZSTDMT_RustBufferPool* pool, size_t bSize); ZSTDMT_RustBufferPool* ZSTDMT_rust_buffer_pool_expand(ZSTDMT_RustBufferPool* pool, unsigned maxNbBuffers); @@ -186,7 +187,8 @@ static ZSTDMT_bufferPool* ZSTDMT_createBufferPool(unsigned maxNbBuffers, ZSTD_cu static size_t ZSTDMT_sizeof_bufferPool(ZSTDMT_bufferPool* bufPool) { if (bufPool == NULL) return 0; - return sizeof(*bufPool) + ZSTDMT_rust_buffer_pool_sizeof(bufPool->rustPool); + return ZSTDMT_rust_sizeofBufferPool( + sizeof(*bufPool), ZSTDMT_rust_buffer_pool_sizeof(bufPool->rustPool)); } /* ZSTDMT_setBufferSize() : diff --git a/rust/src/zstdmt_compress.rs b/rust/src/zstdmt_compress.rs index 056a35e93..d52a8d2df 100644 --- a/rust/src/zstdmt_compress.rs +++ b/rust/src/zstdmt_compress.rs @@ -686,6 +686,16 @@ pub unsafe extern "C" fn ZSTDMT_rust_buffer_pool_sizeof(pool: *const RustBufferP .saturating_add(total_buffer_size) } +#[inline] +fn sizeof_buffer_pool(wrapper_size: usize, rust_pool_size: usize) -> usize { + wrapper_size.wrapping_add(rust_pool_size) +} + +#[no_mangle] +pub extern "C" fn ZSTDMT_rust_sizeofBufferPool(wrapperSize: usize, rustPoolSize: usize) -> usize { + sizeof_buffer_pool(wrapperSize, rustPoolSize) +} + #[no_mangle] pub unsafe extern "C" fn ZSTDMT_rust_buffer_pool_set_size( pool: *mut RustBufferPool, @@ -1608,4 +1618,22 @@ mod tests { assert_eq!(sizeof_cctx_pool(usize::MAX, 1), 0); assert_eq!(sizeof_cctx_pool(usize::MAX - 7, 11), 3); } + + #[test] + fn buffer_pool_size_addition_handles_zero_components() { + assert_eq!(sizeof_buffer_pool(0, 0), 0); + assert_eq!(sizeof_buffer_pool(37, 0), 37); + assert_eq!(sizeof_buffer_pool(0, 53), 53); + } + + #[test] + fn buffer_pool_size_addition_preserves_ordinary_sum() { + assert_eq!(sizeof_buffer_pool(128, 4096), 4224); + } + + #[test] + fn buffer_pool_size_addition_wraps_like_c_size_t() { + assert_eq!(sizeof_buffer_pool(usize::MAX, 1), 0); + assert_eq!(sizeof_buffer_pool(usize::MAX - 7, 11), 3); + } }