ddidderr/zstd-rs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(compress): move block split search to Rust

Browse Source 
Port the recursive post-block-split partition search behind a narrow Rust ABI leaf while keeping CCtx ownership, entropy-state projection, and block emission in C. Preserve left-first midpoint insertion, minimum-sequence gating, the split cap, and error-as-no-split behavior.

Test Plan:

- cargo test --manifest-path rust/Cargo.toml --no-default-features --features compression

- cargo clippy --manifest-path rust/Cargo.toml --no-default-features --features compression --all-targets

- cargo +nightly fmt --manifest-path rust/Cargo.toml

- make -B -C lib -j2 lib

- tests/fuzzer -s4560 -t47 -i48 -v
This commit is contained in:
ddidderr
2026-07-18 08:08:40 +02:00
parent 7ee445c431
commit b81109095f
3 changed files with 472 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
lib/compress/zstd_compress.c
+25 -73
View File
@@ -288,6 +288,17 @@ size_t ZSTD_rust_countSeqStoreMatchBytes(const SeqStore_t* seqStore);
void ZSTD_rust_deriveSeqStoreChunk(SeqStore_t* resultSeqStore,
const SeqStore_t* originalSeqStore,
size_t startIdx, size_t endIdx);
size_t ZSTD_rust_deriveBlockSplits(
U32* partitions, U32 nbSeq,
const SeqStore_t* originalSeqStore,
SeqStore_t* fullSeqStoreChunk,
SeqStore_t* firstHalfSeqStore,
SeqStore_t* secondHalfSeqStore,
const ZSTD_entropyCTables_t* prevEntropy,
ZSTD_entropyCTables_t* nextEntropy,
int strategy, int disableLiteralCompression,
ZSTD_entropyCTablesMetadata_t* entropyMetadata,
void* workspace, size_t workspaceSize);
void ZSTD_rust_seqStore_resolveOffCodes(U32 dRep[ZSTD_REP_NUM],
U32 cRep[ZSTD_REP_NUM],
const SeqStore_t* seqStore,
@@ -2651,7 +2662,7 @@ ZSTD_estimateBlockSize(const BYTE* literals, size_t litSize,
*
* @return: estimated compressed size of the seqStore, or a zstd error.
*/
static size_t
UNUSED_ATTR static size_t
ZSTD_buildEntropyStatisticsAndEstimateSubBlockSize(SeqStore_t* seqStore, ZSTD_CCtx* zc)
{
ZSTD_entropyCTablesMetadata_t* const entropyMetadata = &zc->blockSplitCtx.entropyMetadata;
@@ -2791,66 +2802,6 @@ ZSTD_compressSeqStore_singleBlock(ZSTD_CCtx* zc,
return cSize;
}
/* Struct to keep track of where we are in our recursive calls. */
typedef struct {
U32* splitLocations; /* Array of split indices */
size_t idx; /* The current index within splitLocations being worked on */
} seqStoreSplits;
#define MIN_SEQUENCES_BLOCK_SPLITTING 300
/* Helper function to perform the recursive search for block splits.
* Estimates the cost of seqStore prior to split, and estimates the cost of splitting the sequences in half.
* If advantageous to split, then we recurse down the two sub-blocks.
* If not, or if an error occurred in estimation, then we do not recurse.
*
* Note: The recursion depth is capped by a heuristic minimum number of sequences,
* defined by MIN_SEQUENCES_BLOCK_SPLITTING.
* In theory, this means the absolute largest recursion depth is 10 == log2(maxNbSeqInBlock/MIN_SEQUENCES_BLOCK_SPLITTING).
* In practice, recursion depth usually doesn't go beyond 4.
*
* Furthermore, the number of splits is capped by ZSTD_MAX_NB_BLOCK_SPLITS.
* At ZSTD_MAX_NB_BLOCK_SPLITS == 196 with the current existing blockSize
* maximum of 128 KB, this value is actually impossible to reach.
*/
static void
ZSTD_deriveBlockSplitsHelper(seqStoreSplits* splits, size_t startIdx, size_t endIdx,
ZSTD_CCtx* zc, const SeqStore_t* origSeqStore)
{
SeqStore_t* const fullSeqStoreChunk = &zc->blockSplitCtx.fullSeqStoreChunk;
SeqStore_t* const firstHalfSeqStore = &zc->blockSplitCtx.firstHalfSeqStore;
SeqStore_t* const secondHalfSeqStore = &zc->blockSplitCtx.secondHalfSeqStore;
size_t estimatedOriginalSize;
size_t estimatedFirstHalfSize;
size_t estimatedSecondHalfSize;
size_t midIdx = (startIdx + endIdx)/2;
DEBUGLOG(5, "ZSTD_deriveBlockSplitsHelper: startIdx=%zu endIdx=%zu", startIdx, endIdx);
assert(endIdx >= startIdx);
if (endIdx - startIdx < MIN_SEQUENCES_BLOCK_SPLITTING || splits->idx >= ZSTD_MAX_NB_BLOCK_SPLITS) {
DEBUGLOG(6, "ZSTD_deriveBlockSplitsHelper: Too few sequences (%zu)", endIdx - startIdx);
return;
}
ZSTD_deriveSeqStoreChunk(fullSeqStoreChunk, origSeqStore, startIdx, endIdx);
ZSTD_deriveSeqStoreChunk(firstHalfSeqStore, origSeqStore, startIdx, midIdx);
ZSTD_deriveSeqStoreChunk(secondHalfSeqStore, origSeqStore, midIdx, endIdx);
estimatedOriginalSize = ZSTD_buildEntropyStatisticsAndEstimateSubBlockSize(fullSeqStoreChunk, zc);
estimatedFirstHalfSize = ZSTD_buildEntropyStatisticsAndEstimateSubBlockSize(firstHalfSeqStore, zc);
estimatedSecondHalfSize = ZSTD_buildEntropyStatisticsAndEstimateSubBlockSize(secondHalfSeqStore, zc);
DEBUGLOG(5, "Estimated original block size: %zu -- First half split: %zu -- Second half split: %zu",
estimatedOriginalSize, estimatedFirstHalfSize, estimatedSecondHalfSize);
if (ZSTD_isError(estimatedOriginalSize) || ZSTD_isError(estimatedFirstHalfSize) || ZSTD_isError(estimatedSecondHalfSize)) {
return;
}
if (estimatedFirstHalfSize + estimatedSecondHalfSize < estimatedOriginalSize) {
DEBUGLOG(5, "split decided at seqNb:%zu", midIdx);
ZSTD_deriveBlockSplitsHelper(splits, startIdx, midIdx, zc, origSeqStore);
splits->splitLocations[splits->idx] = (U32)midIdx;
splits->idx++;
ZSTD_deriveBlockSplitsHelper(splits, midIdx, endIdx, zc, origSeqStore);
}
}
/* Base recursive function.
* Populates a table with intra-block partition indices that can improve compression ratio.
*
@@ -2858,18 +2809,19 @@ ZSTD_deriveBlockSplitsHelper(seqStoreSplits* splits, size_t startIdx, size_t end
*/
static size_t ZSTD_deriveBlockSplits(ZSTD_CCtx* zc, U32 partitions[], U32 nbSeq)
{
seqStoreSplits splits;
splits.splitLocations = partitions;
splits.idx = 0;
if (nbSeq <= 4) {
DEBUGLOG(5, "ZSTD_deriveBlockSplits: Too few sequences to split (%u <= 4)", nbSeq);
/* Refuse to try and split anything with less than 4 sequences */
return 0;
}
ZSTD_deriveBlockSplitsHelper(&splits, 0, nbSeq, zc, &zc->seqStore);
splits.splitLocations[splits.idx] = nbSeq;
DEBUGLOG(5, "ZSTD_deriveBlockSplits: final nb partitions: %zu", splits.idx+1);
return splits.idx;
return ZSTD_rust_deriveBlockSplits(
partitions, nbSeq,
&zc->seqStore,
&zc->blockSplitCtx.fullSeqStoreChunk,
&zc->blockSplitCtx.firstHalfSeqStore,
&zc->blockSplitCtx.secondHalfSeqStore,
&zc->blockState.prevCBlock->entropy,
&zc->blockState.nextCBlock->entropy,
(int)zc->appliedParams.cParams.strategy,
ZSTD_literalsCompressionIsDisabled(&zc->appliedParams),
&zc->blockSplitCtx.entropyMetadata,
zc->tmpWorkspace,
zc->tmpWkspSize);
}
/* ZSTD_compressBlock_splitBlock():
rust/src/lib.rs
+2
View File
@@ -37,6 +37,8 @@ pub mod zstd_compress;
#[cfg(feature = "compression")]
pub mod zstd_compress_api;
#[cfg(feature = "compression")]
pub mod zstd_compress_block_split;
#[cfg(feature = "compression")]
pub mod zstd_compress_dictionary;
#[cfg(feature = "compression")]
pub mod zstd_compress_frame;
rust/src/zstd_compress_block_split.rs
+445
View File
@@ -0,0 +1,445 @@
//! Post-block-split partition search.
//!
//! The C compressor owns `ZSTD_CCtx`, block-state ownership, entropy-table
//! selection, and block emission. This module owns only the recursive search
//! over shallow `SeqStore_t` views. The estimator is supplied with projected
//! entropy state and parameter scalars rather than the private C context.
use crate::errors::ERR_isError;
use crate::zstd_compress_stats::{
SeqStore_t, ZSTD_entropyCTablesMetadata_t, ZSTD_entropyCTables_t,
ZSTD_rust_buildBlockEntropyStats, ZSTD_rust_deriveSeqStoreChunk, ZSTD_rust_estimateBlockSize,
};
use std::ffi::c_void;
use std::os::raw::c_int;
const MIN_SEQUENCES_BLOCK_SPLITTING: usize = 300;
const MAX_NB_BLOCK_SPLITS: usize = 196;
const SET_COMPRESSED: c_int = 2;
/// Recursively searches one sequence range for useful midpoint partitions.
///
/// The callback receives each derived view in the same order as the C
/// implementation: full range, left half, right half. It returns either an
/// estimated size or a zstd error code. Errors stop the current search path
/// and therefore cannot create a split.
#[allow(clippy::too_many_arguments)]
unsafe fn derive_block_splits_helper<F>(
partitions: *mut u32,
split_index: &mut usize,
start_index: usize,
end_index: usize,
original_seq_store: *const SeqStore_t,
full_seq_store_chunk: *mut SeqStore_t,
first_half_seq_store: *mut SeqStore_t,
second_half_seq_store: *mut SeqStore_t,
estimate: &mut F,
) where
F: FnMut(*mut SeqStore_t) -> usize,
{
debug_assert!(end_index >= start_index);
let sequence_count = end_index - start_index;
let midpoint = start_index.wrapping_add(end_index) / 2;
if sequence_count < MIN_SEQUENCES_BLOCK_SPLITTING || *split_index >= MAX_NB_BLOCK_SPLITS {
return;
}
unsafe {
ZSTD_rust_deriveSeqStoreChunk(
full_seq_store_chunk,
original_seq_store,
start_index,
end_index,
);
ZSTD_rust_deriveSeqStoreChunk(
first_half_seq_store,
original_seq_store,
start_index,
midpoint,
);
ZSTD_rust_deriveSeqStoreChunk(
second_half_seq_store,
original_seq_store,
midpoint,
end_index,
);
}
let estimated_original_size = estimate(full_seq_store_chunk);
let estimated_first_half_size = estimate(first_half_seq_store);
let estimated_second_half_size = estimate(second_half_seq_store);
if ERR_isError(estimated_original_size)
|| ERR_isError(estimated_first_half_size)
|| ERR_isError(estimated_second_half_size)
{
return;
}
if estimated_first_half_size.wrapping_add(estimated_second_half_size) < estimated_original_size
{
unsafe {
derive_block_splits_helper(
partitions,
split_index,
start_index,
midpoint,
original_seq_store,
full_seq_store_chunk,
first_half_seq_store,
second_half_seq_store,
estimate,
);
}
/* Keep the configured split limit a true upper bound even if a left
* subtree reaches it before its parent gets to midpoint insertion. */
if *split_index >= MAX_NB_BLOCK_SPLITS {
return;
}
unsafe {
*partitions.add(*split_index) = midpoint as u32;
}
*split_index += 1;
unsafe {
derive_block_splits_helper(
partitions,
split_index,
midpoint,
end_index,
original_seq_store,
full_seq_store_chunk,
first_half_seq_store,
second_half_seq_store,
estimate,
);
}
}
}
/// Runs the partition search with a caller-provided estimator.
unsafe fn derive_block_splits_with_estimator<F>(
partitions: *mut u32,
nb_sequences: u32,
original_seq_store: *const SeqStore_t,
full_seq_store_chunk: *mut SeqStore_t,
first_half_seq_store: *mut SeqStore_t,
second_half_seq_store: *mut SeqStore_t,
mut estimate: F,
) -> usize
where
F: FnMut(*mut SeqStore_t) -> usize,
{
if nb_sequences <= 4 {
return 0;
}
let mut split_index = 0usize;
unsafe {
derive_block_splits_helper(
partitions,
&mut split_index,
0,
nb_sequences as usize,
original_seq_store,
full_seq_store_chunk,
first_half_seq_store,
second_half_seq_store,
&mut estimate,
);
*partitions.add(split_index) = nb_sequences;
}
split_index
}
/// Rust leaf for C's `ZSTD_deriveBlockSplits()`.
///
/// The entropy and parameter arguments are the state projected from the C
/// context by the narrow shim. Each estimate rebuilds `next_entropy` and the
/// shared metadata exactly as the original C helper did before comparing the
/// three candidate sizes.
#[allow(clippy::too_many_arguments)]
#[no_mangle]
pub unsafe extern "C" fn ZSTD_rust_deriveBlockSplits(
partitions: *mut u32,
nb_sequences: u32,
original_seq_store: *const SeqStore_t,
full_seq_store_chunk: *mut SeqStore_t,
first_half_seq_store: *mut SeqStore_t,
second_half_seq_store: *mut SeqStore_t,
prev_entropy: *const ZSTD_entropyCTables_t,
next_entropy: *mut ZSTD_entropyCTables_t,
strategy: c_int,
disable_literal_compression: c_int,
entropy_metadata: *mut ZSTD_entropyCTablesMetadata_t,
workspace: *mut c_void,
workspace_size: usize,
) -> usize {
let mut estimate = |seq_store: *mut SeqStore_t| -> usize {
let entropy_stats = unsafe {
ZSTD_rust_buildBlockEntropyStats(
seq_store,
prev_entropy,
next_entropy,
strategy,
disable_literal_compression,
entropy_metadata,
workspace,
workspace_size,
)
};
if ERR_isError(entropy_stats) {
return entropy_stats;
}
let seq_store = unsafe { &*seq_store };
let lit_size = unsafe { seq_store.lit.offset_from(seq_store.litStart) } as usize;
let nb_seq = unsafe { seq_store.sequences.offset_from(seq_store.sequencesStart) } as usize;
let write_lit_entropy =
unsafe { c_int::from((*entropy_metadata).hufMetadata.hType == SET_COMPRESSED) };
unsafe {
ZSTD_rust_estimateBlockSize(
seq_store.litStart,
lit_size,
seq_store.ofCode,
seq_store.llCode,
seq_store.mlCode,
nb_seq,
next_entropy,
entropy_metadata,
workspace,
workspace_size,
write_lit_entropy,
1,
)
}
};
unsafe {
derive_block_splits_with_estimator(
partitions,
nb_sequences,
original_seq_store,
full_seq_store_chunk,
first_half_seq_store,
second_half_seq_store,
&mut estimate,
)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::errors::{ZstdErrorCode, ERROR};
use crate::zstd_compress_sequences::SeqDef;
fn empty_seq_store() -> SeqStore_t {
SeqStore_t {
sequencesStart: std::ptr::null_mut(),
sequences: std::ptr::null_mut(),
litStart: std::ptr::null_mut(),
lit: std::ptr::null_mut(),
llCode: std::ptr::null_mut(),
mlCode: std::ptr::null_mut(),
ofCode: std::ptr::null_mut(),
maxNbSeq: 0,
maxNbLit: 0,
longLengthType: 0,
longLengthPos: 0,
}
}
fn make_seq_store(
sequences: &mut [SeqDef],
codes: &mut [u8],
literals: &mut [u8],
long_length_type: c_int,
long_length_pos: u32,
) -> SeqStore_t {
for sequence in sequences.iter_mut() {
*sequence = SeqDef {
offBase: 4,
litLength: 0,
mlBase: 0,
};
}
let sequences_start = sequences.as_mut_ptr();
let literals_start = literals.as_mut_ptr();
let codes_start = codes.as_mut_ptr();
SeqStore_t {
sequencesStart: sequences_start,
sequences: unsafe { sequences_start.add(sequences.len()) },
litStart: literals_start,
lit: literals_start,
llCode: codes_start,
mlCode: codes_start,
ofCode: codes_start,
maxNbSeq: sequences.len(),
maxNbLit: literals.len(),
longLengthType: long_length_type,
longLengthPos: long_length_pos,
}
}
unsafe fn sequence_range(
seq_store: *mut SeqStore_t,
original_sequences: *mut SeqDef,
) -> (usize, usize) {
let seq_store = unsafe { &*seq_store };
let start = unsafe { seq_store.sequencesStart.offset_from(original_sequences) } as usize;
let end = unsafe { seq_store.sequences.offset_from(original_sequences) } as usize;
(start, end)
}
unsafe fn run_search<F>(
sequences: &mut [SeqDef],
long_length_type: c_int,
long_length_pos: u32,
partitions: &mut [u32],
estimate: F,
) -> usize
where
F: FnMut(*mut SeqStore_t) -> usize,
{
let mut codes = vec![0u8; sequences.len()];
let mut literals = vec![0u8; sequences.len() + 1];
let original = make_seq_store(
sequences,
&mut codes,
&mut literals,
long_length_type,
long_length_pos,
);
let mut full = empty_seq_store();
let mut first = empty_seq_store();
let mut second = empty_seq_store();
derive_block_splits_with_estimator(
partitions.as_mut_ptr(),
sequences.len() as u32,
&original,
&mut full,
&mut first,
&mut second,
estimate,
)
}
#[test]
fn recursive_search_is_left_first_with_midpoint_insertion() {
let mut sequences = vec![SeqDef::default(); 1200];
let original_sequences = sequences.as_mut_ptr();
let mut partitions = vec![u32::MAX; 8];
let mut calls = Vec::new();
let mut call_index = 0usize;
let splits = unsafe {
run_search(&mut sequences, 0, 0, &mut partitions, |chunk| {
let range = sequence_range(chunk, original_sequences);
calls.push(range);
let estimate = if call_index.is_multiple_of(3) {
1000
} else {
0
};
call_index += 1;
estimate
})
};
assert_eq!(splits, 7);
assert_eq!(&partitions[..splits], &[150, 300, 450, 600, 750, 900, 1050]);
assert_eq!(
&calls[..6],
&[
(0, 1200),
(0, 600),
(600, 1200),
(0, 600),
(0, 300),
(300, 600)
]
);
assert_eq!(partitions[splits], 1200);
}
#[test]
fn minimum_sequence_count_skips_estimation() {
let mut sequences = vec![SeqDef::default(); 299];
let mut partitions = [u32::MAX; 2];
let mut calls = 0usize;
let splits = unsafe {
run_search(&mut sequences, 0, 0, &mut partitions, |_| {
calls += 1;
0
})
};
assert_eq!(splits, 0);
assert_eq!(calls, 0);
assert_eq!(partitions[0], 299);
}
#[test]
fn split_count_is_capped_at_configured_limit() {
let mut sequences = vec![SeqDef::default(); 60_000];
let mut partitions = vec![u32::MAX; MAX_NB_BLOCK_SPLITS + 1];
let mut call_index = 0usize;
let splits = unsafe {
run_search(&mut sequences, 0, 0, &mut partitions, |_| {
let estimate = if call_index.is_multiple_of(3) {
1000
} else {
0
};
call_index += 1;
estimate
})
};
assert_eq!(splits, MAX_NB_BLOCK_SPLITS);
assert!(partitions[..splits]
.windows(2)
.all(|window| window[0] < window[1]));
assert_eq!(partitions[splits], 60_000);
}
#[test]
fn estimator_errors_leave_the_current_range_unsplit() {
let mut sequences = vec![SeqDef::default(); 600];
let mut partitions = [u32::MAX; 2];
let mut calls = 0usize;
let splits = unsafe {
run_search(&mut sequences, 0, 0, &mut partitions, |_| {
calls += 1;
ERROR(ZstdErrorCode::Generic)
})
};
assert_eq!(splits, 0);
assert_eq!(calls, 3);
assert_eq!(partitions[0], 600);
}
#[test]
fn derived_views_keep_long_length_metadata_in_their_ranges() {
let mut sequences = vec![SeqDef::default(); 600];
let original_sequences = sequences.as_mut_ptr();
let mut partitions = [u32::MAX; 2];
let mut metadata = Vec::new();
let splits = unsafe {
run_search(&mut sequences, 1, 450, &mut partitions, |chunk| {
let range = sequence_range(chunk, original_sequences);
let view = &*chunk;
metadata.push((range, view.longLengthType, view.longLengthPos));
0
})
};
assert_eq!(splits, 0);
assert_eq!(
metadata,
vec![((0, 600), 1, 450), ((0, 300), 0, 450), ((300, 600), 1, 150)]
);
}
}