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

feat(compress): move block size estimation to Rust

Browse Source 
Block splitting still estimated literal and sequence section sizes through
three C helpers. Move the pure estimator into zstd_compress_stats.rs so its
histogram, HUF, and FSE cost calculations use the Rust implementations already
behind the compression ABI. Preserve the original literal headers, multi-stream
Huffman surcharge, sequence additional-bit costs, entropy-header accounting,
and fallback estimates; leave C responsible for split recursion and table
ownership.

Test Plan:
- `cargo test --manifest-path rust/Cargo.toml --no-default-features --features compression` -- 229 passed
- `cargo clippy --manifest-path rust/Cargo.toml --no-default-features --features compression` plus benches/tests -- passed before and after formatting
- `cargo +nightly fmt --manifest-path rust/Cargo.toml --all` -- passed
- `make -B -C lib -j2 lib` -- passed
- `make -C tests -j2 test-zstream` -- passed, including 84 deterministic and 15,375 randomized cases
This commit is contained in:
ddidderr
2026-07-18 07:08:27 +02:00
parent d334aaafa3
commit 511882e213
2 changed files with 341 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/src/zstd_compress_stats.rs
+325 -7
View File
@@ -9,15 +9,15 @@
//! `ZSTD_copyBlockSequences()`. The `ZSTD_CCtx` and `ZSTD_CCtx_params`
//! layouts stay private to C: the C shims extract the sequence store, the
//! entropy-table leaves, and the two parameter scalars these paths read
//! (the compression strategy and the literals-compression switch). Block
//! dispatch, block splitting, and the block-size estimation heuristics also
//! remain in C for a later slice.
//! (the compression strategy and the literals-compression switch). Block
//! dispatch and block-splitting recursion remain in C; the pure block-size
//! estimator is owned here.
use crate::bits::ZSTD_highbit32;
use crate::common::{
DEFAULT_MAX_OFF, LL_DEFAULT_NORM, LL_DEFAULT_NORM_LOG, LL_FSE_LOG, LONGNBSEQ, MAX_LL, MAX_ML,
MAX_OFF, MAX_SEQ, MINMATCH, ML_DEFAULT_NORM, ML_DEFAULT_NORM_LOG, ML_FSE_LOG, OFF_FSE_LOG,
OF_DEFAULT_NORM, OF_DEFAULT_NORM_LOG, REP_START_VALUE, ZSTD_MAX_FSE_HEADERS_SIZE,
DEFAULT_MAX_OFF, LL_BITS, LL_DEFAULT_NORM, LL_DEFAULT_NORM_LOG, LL_FSE_LOG, LONGNBSEQ, MAX_LL,
MAX_ML, MAX_OFF, MAX_SEQ, MINMATCH, ML_BITS, ML_DEFAULT_NORM, ML_DEFAULT_NORM_LOG, ML_FSE_LOG,
OFF_FSE_LOG, OF_DEFAULT_NORM, OF_DEFAULT_NORM_LOG, REP_START_VALUE, ZSTD_MAX_FSE_HEADERS_SIZE,
ZSTD_MAX_HUF_HEADER_SIZE, ZSTD_REP_NUM,
};
use crate::errors::{ERR_isError, ZstdErrorCode, ERROR};
@@ -31,7 +31,8 @@ use crate::zstd_compress_literals::{
copy_huf_tables, min_gain, ZSTD_compressLiterals, ZSTD_hufCTables_t,
};
use crate::zstd_compress_sequences::{
SeqDef, ZSTD_buildCTable, ZSTD_encodeSequences, ZSTD_selectEncodingType,
SeqDef, ZSTD_buildCTable, ZSTD_crossEntropyCost, ZSTD_encodeSequences, ZSTD_fseBitCost,
ZSTD_selectEncodingType,
};
use std::ffi::c_void;
use std::mem::{size_of, size_of_val};
@@ -246,6 +247,232 @@ pub unsafe extern "C" fn ZSTD_rust_get1BlockSummary(
invalid()
}
/// Estimates the literal-section size used by the block splitter.
unsafe fn estimate_block_size_literal(
literals: *const u8,
lit_size: usize,
huf: *const ZSTD_hufCTables_t,
huf_metadata: *const ZSTD_hufCTablesMetadata_t,
workspace: *mut c_void,
wksp_size: usize,
write_entropy: bool,
) -> usize {
let metadata = unsafe { &*huf_metadata };
let literal_header_size =
3 + usize::from(lit_size >= 1024) + usize::from(lit_size >= 16 * 1024);
let single_stream = lit_size < 256;
match metadata.hType {
SET_BASIC => lit_size,
SET_RLE => 1,
SET_COMPRESSED | SET_REPEAT => {
let mut max_symbol_value = HUF_SYMBOLVALUE_MAX;
let largest = unsafe {
HIST_count_wksp(
workspace.cast::<u32>(),
&mut max_symbol_value,
literals.cast::<c_void>(),
lit_size,
workspace,
wksp_size,
)
};
if ERR_isError(largest) {
return lit_size;
}
let mut estimate = unsafe {
HUF_estimateCompressedSize(
(*huf).CTable.as_ptr(),
workspace.cast::<u32>(),
max_symbol_value,
)
};
if write_entropy {
estimate = estimate.wrapping_add(metadata.hufDesSize);
}
if !single_stream {
estimate = estimate.wrapping_add(6);
}
estimate.wrapping_add(literal_header_size)
}
_ => {
debug_assert!(false, "invalid literal encoding type");
0
}
}
}
/// Estimates one FSE symbol stream used by the block splitter.
#[allow(clippy::too_many_arguments)]
unsafe fn estimate_block_size_symbol_type(
encoding_type: c_int,
code_table: *const u8,
nb_seq: usize,
max_code: u32,
fse_ctable: *const u32,
additional_bits: *const u8,
default_norm: *const i16,
default_norm_log: u32,
default_max: u32,
workspace: *mut c_void,
wksp_size: usize,
) -> usize {
let mut max = max_code;
unsafe {
let _ = HIST_countFast_wksp(
workspace.cast::<u32>(),
&mut max,
code_table.cast::<c_void>(),
nb_seq,
workspace,
wksp_size,
);
}
let mut estimate_in_bits = match encoding_type {
SET_BASIC => {
debug_assert!(max <= default_max);
unsafe { ZSTD_crossEntropyCost(default_norm, default_norm_log, workspace.cast(), max) }
}
SET_RLE => 0,
SET_COMPRESSED | SET_REPEAT => unsafe {
ZSTD_fseBitCost(fse_ctable, workspace.cast(), max)
},
_ => 0,
};
if ERR_isError(estimate_in_bits) {
return nb_seq.wrapping_mul(10);
}
for index in 0..nb_seq {
let code = unsafe { *code_table.add(index) };
let additional = if additional_bits.is_null() {
code as usize
} else {
unsafe { *additional_bits.add(code as usize) as usize }
};
estimate_in_bits = estimate_in_bits.wrapping_add(additional);
}
estimate_in_bits >> 3
}
/// Estimates the sequence-section size used by the block splitter.
#[allow(clippy::too_many_arguments)]
unsafe fn estimate_block_size_sequences(
of_code_table: *const u8,
ll_code_table: *const u8,
ml_code_table: *const u8,
nb_seq: usize,
fse_tables: *const ZSTD_fseCTables_t,
fse_metadata: *const ZSTD_fseCTablesMetadata_t,
workspace: *mut c_void,
wksp_size: usize,
write_entropy: bool,
) -> usize {
let metadata = unsafe { &*fse_metadata };
let tables = unsafe { &*fse_tables };
let header_size =
1 + 1 + usize::from(nb_seq >= 128) + usize::from(nb_seq >= LONGNBSEQ as usize);
let mut estimate = unsafe {
estimate_block_size_symbol_type(
metadata.ofType,
of_code_table,
nb_seq,
MAX_OFF as u32,
tables.offcodeCTable.as_ptr(),
ptr::null(),
OF_DEFAULT_NORM.as_ptr(),
OF_DEFAULT_NORM_LOG,
DEFAULT_MAX_OFF as u32,
workspace,
wksp_size,
)
};
estimate = estimate.wrapping_add(unsafe {
estimate_block_size_symbol_type(
metadata.llType,
ll_code_table,
nb_seq,
MAX_LL as u32,
tables.litlengthCTable.as_ptr(),
LL_BITS.as_ptr(),
LL_DEFAULT_NORM.as_ptr(),
LL_DEFAULT_NORM_LOG,
MAX_LL as u32,
workspace,
wksp_size,
)
});
estimate = estimate.wrapping_add(unsafe {
estimate_block_size_symbol_type(
metadata.mlType,
ml_code_table,
nb_seq,
MAX_ML as u32,
tables.matchlengthCTable.as_ptr(),
ML_BITS.as_ptr(),
ML_DEFAULT_NORM.as_ptr(),
ML_DEFAULT_NORM_LOG,
MAX_ML as u32,
workspace,
wksp_size,
)
});
if write_entropy {
estimate = estimate.wrapping_add(metadata.fseTablesSize);
}
estimate.wrapping_add(header_size)
}
/// Estimates a block's compressed size for block-split decisions. This is
/// the Rust leaf for C's `ZSTD_estimateBlockSize()` and keeps the workspace and
/// entropy-table representations on the existing C ABI.
#[allow(clippy::too_many_arguments)]
#[no_mangle]
pub unsafe extern "C" fn ZSTD_rust_estimateBlockSize(
literals: *const u8,
lit_size: usize,
of_code_table: *const u8,
ll_code_table: *const u8,
ml_code_table: *const u8,
nb_seq: usize,
entropy: *const ZSTD_entropyCTables_t,
entropy_metadata: *const ZSTD_entropyCTablesMetadata_t,
workspace: *mut c_void,
wksp_size: usize,
write_lit_entropy: c_int,
write_seq_entropy: c_int,
) -> usize {
let entropy = unsafe { &*entropy };
let entropy_metadata = unsafe { &*entropy_metadata };
let literals_size = unsafe {
estimate_block_size_literal(
literals,
lit_size,
ptr::addr_of!(entropy.huf),
ptr::addr_of!(entropy_metadata.hufMetadata),
workspace,
wksp_size,
write_lit_entropy != 0,
)
};
let sequence_size = unsafe {
estimate_block_size_sequences(
of_code_table,
ll_code_table,
ml_code_table,
nb_seq,
ptr::addr_of!(entropy.fse),
ptr::addr_of!(entropy_metadata.fseMetadata),
workspace,
wksp_size,
write_seq_entropy != 0,
)
};
sequence_size.wrapping_add(literals_size).wrapping_add(3)
}
/// Converts a raw sequence offset to the stored offBase representation.
///
/// This is the Rust leaf for C's `ZSTD_finalizeOffBase()`. The repcode
@@ -1994,6 +2221,97 @@ mod tests {
);
}
#[test]
fn estimate_block_size_counts_basic_literals_and_rle_sequences() {
let entropy = ZSTD_entropyCTables_t {
huf: empty_huf_tables(HUF_REPEAT_NONE),
fse: empty_fse_tables(FSE_REPEAT_NONE),
};
let mut metadata = unsafe {
std::mem::MaybeUninit::<ZSTD_entropyCTablesMetadata_t>::zeroed().assume_init()
};
metadata.hufMetadata.hType = SET_BASIC;
metadata.fseMetadata.ofType = SET_RLE;
metadata.fseMetadata.llType = SET_RLE;
metadata.fseMetadata.mlType = SET_RLE;
metadata.fseMetadata.fseTablesSize = 7;
let literals = *b"abcd";
let codes = [0u8];
let mut workspace = vec![0u64; ENTROPY_WORKSPACE_SIZE / size_of::<u64>()];
let without_entropy = unsafe {
ZSTD_rust_estimateBlockSize(
literals.as_ptr(),
literals.len(),
codes.as_ptr(),
codes.as_ptr(),
codes.as_ptr(),
1,
&entropy,
&metadata,
workspace.as_mut_ptr().cast(),
ENTROPY_WORKSPACE_SIZE,
0,
0,
)
};
assert_eq!(without_entropy, 9);
let with_entropy = unsafe {
ZSTD_rust_estimateBlockSize(
literals.as_ptr(),
literals.len(),
codes.as_ptr(),
codes.as_ptr(),
codes.as_ptr(),
1,
&entropy,
&metadata,
workspace.as_mut_ptr().cast(),
ENTROPY_WORKSPACE_SIZE,
0,
1,
)
};
assert_eq!(with_entropy, 16);
}
#[test]
fn estimate_block_size_counts_rle_literals() {
let entropy = ZSTD_entropyCTables_t {
huf: empty_huf_tables(HUF_REPEAT_NONE),
fse: empty_fse_tables(FSE_REPEAT_NONE),
};
let mut metadata = unsafe {
std::mem::MaybeUninit::<ZSTD_entropyCTablesMetadata_t>::zeroed().assume_init()
};
metadata.hufMetadata.hType = SET_RLE;
metadata.fseMetadata.ofType = SET_RLE;
metadata.fseMetadata.llType = SET_RLE;
metadata.fseMetadata.mlType = SET_RLE;
let literals = [b'x'; 1024];
let codes = [0u8];
let mut workspace = vec![0u64; ENTROPY_WORKSPACE_SIZE / size_of::<u64>()];
let estimate = unsafe {
ZSTD_rust_estimateBlockSize(
literals.as_ptr(),
literals.len(),
codes.as_ptr(),
codes.as_ptr(),
codes.as_ptr(),
1,
&entropy,
&metadata,
workspace.as_mut_ptr().cast(),
ENTROPY_WORKSPACE_SIZE,
0,
0,
)
};
assert_eq!(estimate, 6);
}
#[test]
fn compressed_block_state_reset_restores_repcodes_and_repeat_modes() {
let mut block_state =