From 09334d2065787dfb4a41fa3782b3ad4b951ff736 Mon Sep 17 00:00:00 2001 From: ddidderr Date: Sun, 19 Jul 2026 17:21:02 +0200 Subject: [PATCH] feat(compress): move CCtx workspace estimator into Rust The CCtx workspace estimator still assembled its private sizing formula in C, even though the match-state and component-sum leaves already lived in Rust. That left buffering, LDM, static-context, and external-sequence sizing policy outside the rewrite boundary. Move the complete size_t formula into a Rust entry point. C now supplies only private object sizes, sanitizer redzone policy, and LDM scalar fields through a small ABI record. Rust retains the C wrapping/alignment behavior and focused branch coverage while the C adapter preserves the existing public estimator surface. Test Plan: - cargo fmt --manifest-path rust/Cargo.toml -- --check - ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo test --manifest-path rust/Cargo.toml --release estimate_cctx_workspace_size - ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo test --manifest-path rust/Cargo.toml --release (667 passed) - ulimit -v 41943040; CARGO_BUILD_JOBS=1 cargo clippy --manifest-path rust/Cargo.toml --release --all-targets -- -D warnings - ulimit -v 41943040; make -B -C programs -j1 zstd - ulimit -v 41943040; make -C tests -j1 test-zstream ZSTREAM_TESTTIME=-T1s (84 tests and both short fuzzer rounds) --- lib/compress/zstd_compress.c | 85 ++++++------ rust/README.md | 7 +- rust/src/zstd_compress.rs | 258 ++++++++++++++++++++++++++++++++++- 3 files changed, 306 insertions(+), 44 deletions(-) diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c index 452065629..bafe4c47e 100644 --- a/lib/compress/zstd_compress.c +++ b/lib/compress/zstd_compress.c @@ -1124,15 +1124,27 @@ 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); -size_t ZSTD_rust_estimateWorkspaceSize(size_t cctxSpace, - size_t tmpWorkSpace, - size_t blockStateSpace, - size_t ldmSpace, - size_t ldmSeqSpace, - size_t matchStateSize, - size_t tokenSpace, - size_t bufferSpace, - size_t externalSeqSpace); +typedef struct { + size_t cctxSize; + size_t compressedBlockStateSize; + size_t seqDefSize; + size_t rawSeqSize; + size_t externalSequenceSize; + size_t tmpWorkspaceSize; + size_t wildcopyOverlength; + size_t matchTSize; + size_t optimalTSize; + size_t asanRedzoneSize; +} ZSTD_rustCCtxWorkspaceSizing; +size_t ZSTD_rust_estimateCCtxWorkspaceSize( + ZSTD_compressionParameters cParams, + int ldmEnable, U32 ldmHashLog, U32 ldmBucketSizeLog, + U32 ldmMinMatchLength, + int isStatic, int useRowMatchFinder, + size_t buffInSize, size_t buffOutSize, + U64 pledgedSrcSize, int useSequenceProducer, + size_t maxBlockSize, + const ZSTD_rustCCtxWorkspaceSizing* sizing); size_t ZSTD_rust_maxEstimateCCtxSize(size_t estimate0, size_t estimate1, size_t estimate2, size_t estimate3); ZSTD_inBuffer ZSTD_rust_inBufferForEndFlush(int inBufferMode, @@ -3061,35 +3073,32 @@ static size_t ZSTD_estimateCCtxSize_usingCCtxParams_internal( int useSequenceProducer, size_t maxBlockSize) { - size_t const windowSize = (size_t) BOUNDED(1ULL, 1ULL << cParams->windowLog, pledgedSrcSize); - size_t const blockSize = MIN(ZSTD_resolveMaxBlockSize(maxBlockSize), windowSize); - size_t const maxNbSeq = ZSTD_maxNbSeq(blockSize, cParams->minMatch, useSequenceProducer); - size_t const tokenSpace = ZSTD_cwksp_alloc_size(WILDCOPY_OVERLENGTH + blockSize) - + ZSTD_cwksp_aligned64_alloc_size(maxNbSeq * sizeof(SeqDef)) - + 3 * ZSTD_cwksp_alloc_size(maxNbSeq * sizeof(BYTE)); - size_t const tmpWorkSpace = ZSTD_cwksp_alloc_size(TMP_WORKSPACE_SIZE); - size_t const blockStateSpace = 2 * ZSTD_cwksp_alloc_size(sizeof(ZSTD_compressedBlockState_t)); - size_t const matchStateSize = ZSTD_sizeof_matchState(cParams, useRowMatchFinder, /* enableDedicatedDictSearch */ 0, /* forCCtx */ 1); - - size_t const ldmSpace = ZSTD_ldm_getTableSize(*ldmParams); - size_t const maxNbLdmSeq = ZSTD_ldm_getMaxNbSeq(*ldmParams, blockSize); - size_t const ldmSeqSpace = ldmParams->enableLdm == ZSTD_ps_enable ? - ZSTD_cwksp_aligned64_alloc_size(maxNbLdmSeq * sizeof(rawSeq)) : 0; - - - size_t const bufferSpace = ZSTD_cwksp_alloc_size(buffInSize) - + ZSTD_cwksp_alloc_size(buffOutSize); - - size_t const cctxSpace = isStatic ? ZSTD_cwksp_alloc_size(sizeof(ZSTD_CCtx)) : 0; - - size_t const maxNbExternalSeq = ZSTD_sequenceBound(blockSize); - size_t const externalSeqSpace = useSequenceProducer - ? ZSTD_cwksp_aligned64_alloc_size(maxNbExternalSeq * sizeof(ZSTD_Sequence)) - : 0; - - size_t const neededSpace = ZSTD_rust_estimateWorkspaceSize( - cctxSpace, tmpWorkSpace, blockStateSpace, ldmSpace, ldmSeqSpace, - matchStateSize, tokenSpace, bufferSpace, externalSeqSpace); + ZSTD_rustCCtxWorkspaceSizing const sizing = { + sizeof(ZSTD_CCtx), + sizeof(ZSTD_compressedBlockState_t), + sizeof(SeqDef), + sizeof(rawSeq), + sizeof(ZSTD_Sequence), + TMP_WORKSPACE_SIZE, + WILDCOPY_OVERLENGTH, + sizeof(ZSTD_match_t), + sizeof(ZSTD_optimal_t), + ZSTD_RUST_ASAN_REDZONE_SIZE + }; + size_t const neededSpace = ZSTD_rust_estimateCCtxWorkspaceSize( + *cParams, + (int)ldmParams->enableLdm, + ldmParams->hashLog, + ldmParams->bucketSizeLog, + ldmParams->minMatchLength, + isStatic, + (int)useRowMatchFinder, + buffInSize, + buffOutSize, + pledgedSrcSize, + useSequenceProducer, + maxBlockSize, + &sizing); DEBUGLOG(5, "estimate workspace : %u", (U32)neededSpace); return neededSpace; diff --git a/rust/README.md b/rust/README.md index a2039f2dd..2be045db0 100644 --- a/rust/README.md +++ b/rust/README.md @@ -165,9 +165,10 @@ parameter validation and init-then-begin ordering now run in Rust. CDict advanced private workspace construction, private static-CCtx and static-CDict workspace construction and dictionary-content allocation/loading, and advanced-CDict dictionary-content loading remain in C. Rust now owns -advanced-CDict custom-memory validation, workspace-size query/allocation, and -allocation/create/init cleanup ordering; C retains the private workspace -size formula, layout, and allocator callbacks. Reset policy, +advanced-CDict custom-memory validation, workspace-size query/allocation, +allocation/create/init cleanup ordering, and the CCtx workspace-size formula; +C retains private layout-size inputs, workspace layout, and allocator +callbacks. Reset policy, private CCtx/matchfinder/workspace operations, and codec/adaptive-policy callbacks remain in C. CDict initialization ordering and scalar publication, shared compression-begin dictionary selection, CDict reset attach-versus-copy diff --git a/rust/src/zstd_compress.rs b/rust/src/zstd_compress.rs index 4fbcf7dbe..18873d6d9 100644 --- a/rust/src/zstd_compress.rs +++ b/rust/src/zstd_compress.rs @@ -23,10 +23,12 @@ use crate::zstd_compress_frame::{ }; use crate::zstd_compress_literals::min_gain; use crate::zstd_compress_params::{ - ZSTD_compressionParameters, ZSTD_frameParameters, ZSTD_parameters, + ZSTD_compressionParameters, ZSTD_frameParameters, ZSTD_parameters, ZSTD_rustMatchStateSizing, ZSTD_rust_params_adjustCParams, ZSTD_rust_params_checkCParams, ZSTD_rust_params_defaultCLevel, - ZSTD_rust_params_getParamsInternal, ZSTD_rust_params_maxNbSeq, ZSTD_rust_params_selectCParams, - ZSTD_RUST_CPM_NO_ATTACH_DICT, ZSTD_RUST_PS_AUTO, ZSTD_RUST_PS_DISABLE, + ZSTD_rust_params_estimateMatchStateSize, ZSTD_rust_params_getParamsInternal, + ZSTD_rust_params_maxNbSeq, ZSTD_rust_params_resolveMaxBlockSize, + ZSTD_rust_params_selectCParams, ZSTD_RUST_CPM_NO_ATTACH_DICT, ZSTD_RUST_PS_AUTO, + ZSTD_RUST_PS_DISABLE, ZSTD_RUST_PS_ENABLE, }; use crate::zstd_compress_params_api::{ ZSTD_CCtxParams_setParameter, ZSTD_CCtx_params, ZSTD_rust_isUpdateAuthorized, @@ -5325,6 +5327,177 @@ pub extern "C" fn ZSTD_rust_estimateWorkspaceSize( ) } +/// Private C layout sizes needed by the CCtx workspace estimator. +/// +/// The scalar sizes keep `ZSTD_CCtx`, `ldmEntry_t`, and the other private C +/// objects opaque while allowing Rust to own the workspace-sizing formula. +#[repr(C)] +#[derive(Clone, Copy, Debug, Default)] +pub struct ZSTD_rustCCtxWorkspaceSizing { + pub cctxSize: usize, + pub compressedBlockStateSize: usize, + pub seqDefSize: usize, + pub rawSeqSize: usize, + pub externalSequenceSize: usize, + pub tmpWorkspaceSize: usize, + pub wildcopyOverlength: usize, + pub matchTSize: usize, + pub optimalTSize: usize, + pub asanRedzoneSize: usize, +} + +const _: () = { + assert!(size_of::() == size_of::<[usize; 10]>()); +}; + +const ZSTD_HASHLOG3_MAX: u32 = 17; + +#[inline] +fn cctx_cwksp_alloc_size(size: usize, asan_redzone_size: usize) -> usize { + if size == 0 { + 0 + } else { + size.wrapping_add(asan_redzone_size.wrapping_mul(2)) + } +} + +#[inline] +fn cctx_cwksp_align(size: usize, alignment: usize) -> usize { + size.wrapping_add(alignment - 1) & !(alignment - 1) +} + +#[inline] +fn cctx_cwksp_aligned64_alloc_size(size: usize, asan_redzone_size: usize) -> usize { + cctx_cwksp_alloc_size(cctx_cwksp_align(size, 64), asan_redzone_size) +} + +#[inline] +fn cctx_ldm_table_size( + ldm_enable: c_int, + hash_log: u32, + bucket_size_log: u32, + asan_redzone_size: usize, +) -> usize { + if ldm_enable != ZSTD_RUST_PS_ENABLE { + return 0; + } + let hash_size = 1usize << hash_log; + let bucket_log = bucket_size_log.min(hash_log); + let bucket_size = 1usize << (hash_log - bucket_log); + cctx_cwksp_alloc_size(bucket_size, asan_redzone_size).wrapping_add(cctx_cwksp_alloc_size( + hash_size.wrapping_mul(2 * size_of::()), + asan_redzone_size, + )) +} + +/// Estimate the CCtx workspace using the private C layout sizes supplied by +/// the adapter. This owns the orchestration formerly kept in +/// `ZSTD_estimateCCtxSize_usingCCtxParams_internal()`. +#[no_mangle] +pub unsafe extern "C" fn ZSTD_rust_estimateCCtxWorkspaceSize( + cparams: ZSTD_compressionParameters, + ldm_enable: c_int, + ldm_hash_log: u32, + ldm_bucket_size_log: u32, + ldm_min_match_length: u32, + is_static: c_int, + use_row_match_finder: c_int, + buff_in_size: usize, + buff_out_size: usize, + pledged_src_size: u64, + use_sequence_producer: c_int, + max_block_size: usize, + sizing: *const ZSTD_rustCCtxWorkspaceSizing, +) -> usize { + if sizing.is_null() { + return 0; + } + let sizing = unsafe { *sizing }; + let window_limit = 1u64.checked_shl(cparams.windowLog).unwrap_or(u64::MAX); + let window_size = (window_limit.min(pledged_src_size).min(usize::MAX as u64) as usize).max(1); + let block_size = ZSTD_rust_params_resolveMaxBlockSize(max_block_size).min(window_size); + let max_nb_seq = ZSTD_rust_params_maxNbSeq(block_size, cparams.minMatch, use_sequence_producer); + let token_space = cctx_cwksp_alloc_size( + sizing.wildcopyOverlength.wrapping_add(block_size), + sizing.asanRedzoneSize, + ) + .wrapping_add(cctx_cwksp_aligned64_alloc_size( + max_nb_seq.wrapping_mul(sizing.seqDefSize), + sizing.asanRedzoneSize, + )) + .wrapping_add(3usize.wrapping_mul(cctx_cwksp_alloc_size( + max_nb_seq.wrapping_mul(size_of::()), + sizing.asanRedzoneSize, + ))); + let tmp_workspace = cctx_cwksp_alloc_size(sizing.tmpWorkspaceSize, sizing.asanRedzoneSize); + let block_state_space = 2usize.wrapping_mul(cctx_cwksp_alloc_size( + sizing.compressedBlockStateSize, + sizing.asanRedzoneSize, + )); + let match_state_sizing = ZSTD_rustMatchStateSizing { + hashLog3Max: ZSTD_HASHLOG3_MAX, + matchTSize: sizing.matchTSize, + optimalTSize: sizing.optimalTSize, + asanRedzoneSize: sizing.asanRedzoneSize, + }; + let match_state_size = unsafe { + ZSTD_rust_params_estimateMatchStateSize( + cparams, + use_row_match_finder, + 0, + 1, + &match_state_sizing, + ) + }; + let ldm_space = cctx_ldm_table_size( + ldm_enable, + ldm_hash_log, + ldm_bucket_size_log, + sizing.asanRedzoneSize, + ); + let max_nb_ldm_seq = if ldm_enable == ZSTD_RUST_PS_ENABLE && ldm_min_match_length != 0 { + block_size / ldm_min_match_length as usize + } else { + 0 + }; + let ldm_seq_space = if ldm_enable == ZSTD_RUST_PS_ENABLE { + cctx_cwksp_aligned64_alloc_size( + max_nb_ldm_seq.wrapping_mul(sizing.rawSeqSize), + sizing.asanRedzoneSize, + ) + } else { + 0 + }; + let buffer_space = cctx_cwksp_alloc_size(buff_in_size, sizing.asanRedzoneSize) + .wrapping_add(cctx_cwksp_alloc_size(buff_out_size, sizing.asanRedzoneSize)); + let cctx_space = if is_static != 0 { + cctx_cwksp_alloc_size(sizing.cctxSize, sizing.asanRedzoneSize) + } else { + 0 + }; + let external_seq_space = if use_sequence_producer != 0 { + cctx_cwksp_aligned64_alloc_size( + crate::zstd_compress_api::ZSTD_sequenceBound(block_size) + .wrapping_mul(sizing.externalSequenceSize), + sizing.asanRedzoneSize, + ) + } else { + 0 + }; + + estimate_workspace_size( + cctx_space, + tmp_workspace, + block_state_space, + ldm_space, + ldm_seq_space, + match_state_size, + token_space, + buffer_space, + external_seq_space, + ) +} + #[inline] fn max_estimate_cctx_size( estimate0: usize, @@ -12745,6 +12918,85 @@ mod tests { ); } + fn cctx_workspace_test_sizing() -> ZSTD_rustCCtxWorkspaceSizing { + ZSTD_rustCCtxWorkspaceSizing { + cctxSize: 32, + compressedBlockStateSize: 24, + seqDefSize: 8, + rawSeqSize: 12, + externalSequenceSize: 16, + tmpWorkspaceSize: 64, + wildcopyOverlength: 32, + matchTSize: 4, + optimalTSize: 8, + asanRedzoneSize: 0, + } + } + + fn cctx_workspace_test_cparams() -> ZSTD_compressionParameters { + ZSTD_compressionParameters { + windowLog: 10, + chainLog: 6, + hashLog: 6, + searchLog: 1, + minMatch: 4, + targetLength: 16, + strategy: 3, + } + } + + fn cctx_workspace_test_estimate( + sizing: &ZSTD_rustCCtxWorkspaceSizing, + ldm_enable: c_int, + ldm_hash_log: u32, + ldm_bucket_size_log: u32, + ldm_min_match_length: u32, + is_static: c_int, + buff_in_size: usize, + buff_out_size: usize, + use_sequence_producer: c_int, + ) -> usize { + unsafe { + ZSTD_rust_estimateCCtxWorkspaceSize( + cctx_workspace_test_cparams(), + ldm_enable, + ldm_hash_log, + ldm_bucket_size_log, + ldm_min_match_length, + is_static, + ZSTD_RUST_PS_DISABLE, + buff_in_size, + buff_out_size, + 1000, + use_sequence_producer, + 0, + sizing, + ) + } + } + + #[test] + fn estimate_cctx_workspace_size_keeps_static_and_buffer_components_separate() { + let sizing = cctx_workspace_test_sizing(); + let ordinary = cctx_workspace_test_estimate(&sizing, 0, 0, 0, 0, 0, 0, 0, 0); + let static_context = cctx_workspace_test_estimate(&sizing, 0, 0, 0, 0, 1, 0, 0, 0); + let buffered = cctx_workspace_test_estimate(&sizing, 0, 0, 0, 0, 0, 17, 23, 0); + let external = cctx_workspace_test_estimate(&sizing, 0, 0, 0, 0, 0, 0, 0, 1); + + assert_eq!(static_context - ordinary, sizing.cctxSize); + assert_eq!(buffered - ordinary, 17 + 23); + assert_eq!(external - ordinary, 6265); + } + + #[test] + fn estimate_cctx_workspace_size_adds_ldm_table_and_sequence_space_only_when_enabled() { + let sizing = cctx_workspace_test_sizing(); + let ordinary = cctx_workspace_test_estimate(&sizing, 0, 4, 2, 64, 0, 0, 0, 0); + let ldm = cctx_workspace_test_estimate(&sizing, ZSTD_RUST_PS_ENABLE, 4, 2, 64, 0, 0, 0, 0); + + assert_eq!(ldm - ordinary, 324); + } + #[test] fn max_estimate_cctx_size_handles_zero_values() { assert_eq!(max_estimate_cctx_size(0, 0, 0, 0), 0);