diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c index 93eb854cb..c82705fa3 100644 --- a/lib/compress/zstd_compress.c +++ b/lib/compress/zstd_compress.c @@ -90,6 +90,8 @@ size_t ZSTD_rust_CStreamInSize(void); size_t ZSTD_rust_CStreamOutSize(void); size_t ZSTD_rust_sizeofLocalDict(int dictBufferPresent, size_t dictSize, size_t cdictSize); +size_t ZSTD_rust_sizeofCCtx(size_t objectSize, size_t workspaceSize, + size_t localDictSize, size_t mtctxSize); ZSTD_inBuffer ZSTD_rust_inBufferForEndFlush(int inBufferMode, const void* expectedSrc, size_t expectedSize, @@ -582,11 +584,15 @@ static size_t ZSTD_sizeof_mtctx(const ZSTD_CCtx* cctx) size_t ZSTD_sizeof_CCtx(const ZSTD_CCtx* cctx) { if (cctx==NULL) return 0; /* support sizeof on NULL */ - /* cctx may be in the workspace */ - return (cctx->workspace.workspace == cctx ? 0 : sizeof(*cctx)) - + ZSTD_cwksp_sizeof(&cctx->workspace) - + ZSTD_sizeof_localDict(cctx->localDict) - + ZSTD_sizeof_mtctx(cctx); + { + /* cctx may be in the workspace */ + size_t const objectSize = cctx->workspace.workspace == cctx ? 0 : sizeof(*cctx); + size_t const workspaceSize = ZSTD_cwksp_sizeof(&cctx->workspace); + size_t const localDictSize = ZSTD_sizeof_localDict(cctx->localDict); + size_t const mtctxSize = ZSTD_sizeof_mtctx(cctx); + return ZSTD_rust_sizeofCCtx(objectSize, workspaceSize, localDictSize, + mtctxSize); + } } size_t ZSTD_sizeof_CStream(const ZSTD_CStream* zcs) diff --git a/rust/src/zstd_compress.rs b/rust/src/zstd_compress.rs index bcd200c05..bf179e099 100644 --- a/rust/src/zstd_compress.rs +++ b/rust/src/zstd_compress.rs @@ -430,6 +430,31 @@ pub extern "C" fn ZSTD_rust_sizeofLocalDict( sizeof_local_dict(dict_buffer_present, dict_size, cdict_size) } +#[inline] +fn sizeof_cctx( + object_size: usize, + workspace_size: usize, + local_dict_size: usize, + mtctx_size: usize, +) -> usize { + object_size + .wrapping_add(workspace_size) + .wrapping_add(local_dict_size) + .wrapping_add(mtctx_size) +} + +/// Aggregate C-owned context size components with C `size_t` wrapping +/// semantics. +#[no_mangle] +pub extern "C" fn ZSTD_rust_sizeofCCtx( + object_size: usize, + workspace_size: usize, + local_dict_size: usize, + mtctx_size: usize, +) -> usize { + sizeof_cctx(object_size, workspace_size, local_dict_size, mtctx_size) +} + #[inline] fn reduce_table_internal(table: &mut [u32], reducer_value: u32, preserve_mark: bool) { debug_assert_eq!(table.len() % ZSTD_ROWSIZE, 0); @@ -1409,6 +1434,24 @@ mod tests { assert_eq!(sizeof_local_dict(0, usize::MAX, usize::MAX), usize::MAX); } + #[test] + fn sizeof_cctx_handles_zero_components() { + assert_eq!(sizeof_cctx(0, 0, 0, 0), 0); + assert_eq!(ZSTD_rust_sizeofCCtx(0, 0, 0, 0), 0); + } + + #[test] + fn sizeof_cctx_adds_components_in_order() { + assert_eq!(sizeof_cctx(1, 2, 3, 4), 10); + assert_eq!(ZSTD_rust_sizeofCCtx(1, 2, 3, 4), 10); + } + + #[test] + fn sizeof_cctx_wraps_like_c_size_t_addition() { + assert_eq!(sizeof_cctx(usize::MAX, 1, 2, 3), 5); + assert_eq!(ZSTD_rust_sizeofCCtx(usize::MAX, 1, 2, 3), 5); + } + #[test] fn public_one_shot_abi_is_c_compatible() { let entry: unsafe extern "C" fn(*mut c_void, usize, *const c_void, usize, c_int) -> usize =