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)
2026-07-19 17:21:02 +02:00
parent e6101bb1af
commit 09334d2065
3 changed files with 306 additions and 44 deletions
@@ -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::<ZSTD_rustCCtxWorkspaceSizing>() == 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::<u32>()),
+        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::<u8>()),
+        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);