From 327272f8b3e474b6140b7201585ffe849a6d7b75 Mon Sep 17 00:00:00 2001 From: ddidderr Date: Sat, 18 Jul 2026 05:38:54 +0200 Subject: [PATCH] feat(compress): move optimal block policy to Rust Move the frame chunk pre-split policy behind a scalar Rust ABI while retaining C ownership of ZSTD_CCtx and workspace extraction. Test Plan: cargo test --manifest-path rust/Cargo.toml --no-default-features --features compression; cargo clippy --manifest-path rust/Cargo.toml; cargo clippy --manifest-path rust/Cargo.toml --benches; cargo clippy --manifest-path rust/Cargo.toml --tests; make -B -C lib -j2 lib; make -B -C tests -j2 test-zstream (84 deterministic, 6697 and 9217 fuzzer cases). --- lib/compress/zstd_compress.c | 35 ++-------- rust/src/zstd_compress_frame.rs | 114 +++++++++++++++++++++++++++++++- 2 files changed, 120 insertions(+), 29 deletions(-) diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c index ae9898f64..bf12e3614 100644 --- a/lib/compress/zstd_compress.c +++ b/lib/compress/zstd_compress.c @@ -268,6 +268,10 @@ size_t ZSTD_rust_determineBlockSize(int mode, size_t blockSize, size_t remaining size_t ZSTD_rust_validateSequence(U32 offBase, U32 matchLength, U32 minMatch, size_t posInSrc, U32 windowLog, size_t dictSize, int useSequenceProducer); +size_t ZSTD_rust_optimalBlockSize(const void* src, size_t srcSize, + size_t blockSizeMax, int splitLevel, + int strategy, S64 savings, + void* workspace, size_t workspaceSize); typedef char ZSTD_rust_stats_seqdef_layout[(sizeof(SeqDef) == 8) ? 1 : -1]; typedef char ZSTD_rust_stats_seqstore_long_length_pos[ @@ -3239,34 +3243,9 @@ static void ZSTD_overflowCorrectIfNeeded(ZSTD_MatchState_t* ms, static size_t ZSTD_optimalBlockSize(ZSTD_CCtx* cctx, const void* src, size_t srcSize, size_t blockSizeMax, int splitLevel, ZSTD_strategy strat, S64 savings) { - /* split level based on compression strategy, from `fast` to `btultra2` */ - static const int splitLevels[] = { 0, 0, 1, 2, 2, 3, 3, 4, 4, 4 }; - /* note: conservatively only split full blocks (128 KB) currently. - * While it's possible to go lower, let's keep it simple for a first implementation. - * Besides, benefits of splitting are reduced when blocks are already small. - */ - if (srcSize < 128 KB || blockSizeMax < 128 KB) - return MIN(srcSize, blockSizeMax); - /* do not split incompressible data though: - * require verified savings to allow pre-splitting. - * Note: as a consequence, the first full block is not split. - */ - if (savings < 3) { - DEBUGLOG(6, "don't attempt splitting: savings (%i) too low", (int)savings); - return 128 KB; - } - /* apply @splitLevel, or use default value (which depends on @strat). - * note that splitting heuristic is still conditioned by @savings >= 3, - * so the first block will not reach this code path */ - if (splitLevel == 1) return 128 KB; - if (splitLevel == 0) { - assert(ZSTD_fast <= strat && strat <= ZSTD_btultra2); - splitLevel = splitLevels[strat]; - } else { - assert(2 <= splitLevel && splitLevel <= 6); - splitLevel -= 2; - } - return ZSTD_splitBlock(src, blockSizeMax, splitLevel, cctx->tmpWorkspace, cctx->tmpWkspSize); + return ZSTD_rust_optimalBlockSize(src, srcSize, blockSizeMax, splitLevel, + (int)strat, savings, + cctx->tmpWorkspace, cctx->tmpWkspSize); } /*! ZSTD_compress_frameChunk() : diff --git a/rust/src/zstd_compress_frame.rs b/rust/src/zstd_compress_frame.rs index d7dae8a32..b3e6d9c0c 100644 --- a/rust/src/zstd_compress_frame.rs +++ b/rust/src/zstd_compress_frame.rs @@ -9,8 +9,9 @@ //! crosses the language boundary. use crate::errors::{ZstdErrorCode, ERROR}; +use crate::zstd_presplit::ZSTD_splitBlock; use std::ffi::c_void; -use std::os::raw::{c_int, c_uint}; +use std::os::raw::{c_int, c_longlong, c_uint}; use std::ptr; const ZSTD_MAGICNUMBER: u32 = 0xFD2F_B528; @@ -21,6 +22,10 @@ const ZSTD_BLOCKHEADERSIZE: usize = 3; const ZSTD_WINDOWLOG_ABSOLUTEMIN: u32 = 10; const ZSTD_CONTENTSIZE_UNKNOWN: u64 = u64::MAX; const ZSTD_F_ZSTD1: c_int = 0; +const ZSTD_BLOCK_SIZE: usize = 128 << 10; +const ZSTD_FAST: c_int = 1; +const ZSTD_BTULTRA2: c_int = 9; +const SPLIT_LEVELS: [c_int; 10] = [0, 0, 1, 2, 2, 3, 3, 4, 4, 4]; #[inline] unsafe fn write_le16(dst: *mut u8, value: u16) { @@ -209,6 +214,42 @@ pub unsafe extern "C" fn ZSTD_writeLastEmptyBlock(dst: *mut c_void, dst_capacity ZSTD_BLOCKHEADERSIZE } +/// Rust implementation of the private `ZSTD_optimalBlockSize()` policy. +/// +/// The compressor context stays in C. This leaf receives only the source +/// block, scalar policy inputs, and the pre-split workspace projected by the C +/// caller, then delegates the actual split heuristic to `ZSTD_splitBlock`. +#[no_mangle] +pub unsafe extern "C" fn ZSTD_rust_optimalBlockSize( + src: *const c_void, + src_size: usize, + block_size_max: usize, + split_level: c_int, + strategy: c_int, + savings: c_longlong, + workspace: *mut c_void, + workspace_size: usize, +) -> usize { + if src_size < ZSTD_BLOCK_SIZE || block_size_max < ZSTD_BLOCK_SIZE { + return src_size.min(block_size_max); + } + if savings < 3 { + return ZSTD_BLOCK_SIZE; + } + + let split_level = if split_level == 1 { + return ZSTD_BLOCK_SIZE; + } else if split_level == 0 { + debug_assert!((ZSTD_FAST..=ZSTD_BTULTRA2).contains(&strategy)); + SPLIT_LEVELS[strategy as usize] + } else { + debug_assert!((2..=6).contains(&split_level)); + split_level - 2 + }; + + unsafe { ZSTD_splitBlock(src, block_size_max, split_level, workspace, workspace_size) } +} + #[cfg(test)] mod tests { use super::*; @@ -418,4 +459,75 @@ mod tests { ZstdErrorCode::DstSizeTooSmall as i32 ); } + + #[test] + fn optimal_block_size_keeps_small_blocks_intact() { + let result = unsafe { + ZSTD_rust_optimalBlockSize( + ptr::null(), + ZSTD_BLOCK_SIZE - 1, + ZSTD_BLOCK_SIZE, + 0, + ZSTD_FAST, + 3, + ptr::null_mut(), + 0, + ) + }; + assert_eq!(result, ZSTD_BLOCK_SIZE - 1); + } + + #[test] + fn optimal_block_size_requires_savings_before_splitting() { + let result = unsafe { + ZSTD_rust_optimalBlockSize( + ptr::null(), + ZSTD_BLOCK_SIZE, + ZSTD_BLOCK_SIZE, + 0, + ZSTD_FAST, + 2, + ptr::null_mut(), + 0, + ) + }; + assert_eq!(result, ZSTD_BLOCK_SIZE); + } + + #[test] + fn optimal_block_size_honors_explicit_no_split_level() { + let result = unsafe { + ZSTD_rust_optimalBlockSize( + ptr::null(), + ZSTD_BLOCK_SIZE, + ZSTD_BLOCK_SIZE, + 1, + ZSTD_BTULTRA2, + 3, + ptr::null_mut(), + 0, + ) + }; + assert_eq!(result, ZSTD_BLOCK_SIZE); + } + + #[test] + fn optimal_block_size_delegates_valid_split_requests() { + let source = vec![0u8; ZSTD_BLOCK_SIZE]; + let mut workspace = vec![0usize; 8_208usize.div_ceil(std::mem::size_of::())]; + let result = unsafe { + ZSTD_rust_optimalBlockSize( + source.as_ptr().cast(), + source.len(), + source.len(), + 2, + ZSTD_FAST, + 3, + workspace.as_mut_ptr().cast(), + workspace.len() * std::mem::size_of::(), + ) + }; + assert!(result > 0); + assert!(result <= ZSTD_BLOCK_SIZE); + } }