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
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)]
);
}
}