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
lib/compress/zstd_compress.c
+16 -102
View File
@@ -257,6 +257,14 @@ size_t ZSTD_rust_buildBlockEntropyStats(
int strategy, int disableLiteralCompression,
ZSTD_entropyCTablesMetadata_t* entropyMetadata,
void* workspace, size_t wkspSize);
size_t ZSTD_rust_estimateBlockSize(
const BYTE* literals, size_t litSize,
const BYTE* ofCodeTable, const BYTE* llCodeTable,
const BYTE* mlCodeTable, size_t nbSeq,
const ZSTD_entropyCTables_t* entropy,
const ZSTD_entropyCTablesMetadata_t* entropyMetadata,
void* workspace, size_t wkspSize,
int writeLitEntropy, int writeSeqEntropy);
size_t ZSTD_rust_copyBlockSequences(
SeqCollector* seqCollector, const SeqStore_t* seqStore,
const U32 prevRepcodes[ZSTD_REP_NUM]);
@@ -2600,101 +2608,8 @@ size_t ZSTD_buildBlockEntropyStats(
workspace, wkspSize);
}
/* Returns the size estimate for the literals section (header + content) of a block */
static size_t
ZSTD_estimateBlockSize_literal(const BYTE* literals, size_t litSize,
const ZSTD_hufCTables_t* huf,
const ZSTD_hufCTablesMetadata_t* hufMetadata,
void* workspace, size_t wkspSize,
int writeEntropy)
{
unsigned* const countWksp = (unsigned*)workspace;
unsigned maxSymbolValue = HUF_SYMBOLVALUE_MAX;
size_t literalSectionHeaderSize = 3 + (litSize >= 1 KB) + (litSize >= 16 KB);
U32 singleStream = litSize < 256;
if (hufMetadata->hType == set_basic) return litSize;
else if (hufMetadata->hType == set_rle) return 1;
else if (hufMetadata->hType == set_compressed || hufMetadata->hType == set_repeat) {
size_t const largest = HIST_count_wksp (countWksp, &maxSymbolValue, (const BYTE*)literals, litSize, workspace, wkspSize);
if (ZSTD_isError(largest)) return litSize;
{ size_t cLitSizeEstimate = HUF_estimateCompressedSize((const HUF_CElt*)huf->CTable, countWksp, maxSymbolValue);
if (writeEntropy) cLitSizeEstimate += hufMetadata->hufDesSize;
if (!singleStream) cLitSizeEstimate += 6; /* multi-stream huffman uses 6-byte jump table */
return cLitSizeEstimate + literalSectionHeaderSize;
} }
assert(0); /* impossible */
return 0;
}
/* Returns the size estimate for the FSE-compressed symbols (of, ml, ll) of a block */
static size_t
ZSTD_estimateBlockSize_symbolType(SymbolEncodingType_e type,
const BYTE* codeTable, size_t nbSeq, unsigned maxCode,
const FSE_CTable* fseCTable,
const U8* additionalBits,
short const* defaultNorm, U32 defaultNormLog, U32 defaultMax,
void* workspace, size_t wkspSize)
{
unsigned* const countWksp = (unsigned*)workspace;
const BYTE* ctp = codeTable;
const BYTE* const ctStart = ctp;
const BYTE* const ctEnd = ctStart + nbSeq;
size_t cSymbolTypeSizeEstimateInBits = 0;
unsigned max = maxCode;
HIST_countFast_wksp(countWksp, &max, codeTable, nbSeq, workspace, wkspSize); /* can't fail */
if (type == set_basic) {
/* We selected this encoding type, so it must be valid. */
assert(max <= defaultMax);
(void)defaultMax;
cSymbolTypeSizeEstimateInBits = ZSTD_crossEntropyCost(defaultNorm, defaultNormLog, countWksp, max);
} else if (type == set_rle) {
cSymbolTypeSizeEstimateInBits = 0;
} else if (type == set_compressed || type == set_repeat) {
cSymbolTypeSizeEstimateInBits = ZSTD_fseBitCost(fseCTable, countWksp, max);
}
if (ZSTD_isError(cSymbolTypeSizeEstimateInBits)) {
return nbSeq * 10;
}
while (ctp < ctEnd) {
if (additionalBits) cSymbolTypeSizeEstimateInBits += additionalBits[*ctp];
else cSymbolTypeSizeEstimateInBits += *ctp; /* for offset, offset code is also the number of additional bits */
ctp++;
}
return cSymbolTypeSizeEstimateInBits >> 3;
}
/* Returns the size estimate for the sequences section (header + content) of a block */
static size_t
ZSTD_estimateBlockSize_sequences(const BYTE* ofCodeTable,
const BYTE* llCodeTable,
const BYTE* mlCodeTable,
size_t nbSeq,
const ZSTD_fseCTables_t* fseTables,
const ZSTD_fseCTablesMetadata_t* fseMetadata,
void* workspace, size_t wkspSize,
int writeEntropy)
{
size_t sequencesSectionHeaderSize = 1 /* seqHead */ + 1 /* min seqSize size */ + (nbSeq >= 128) + (nbSeq >= LONGNBSEQ);
size_t cSeqSizeEstimate = 0;
cSeqSizeEstimate += ZSTD_estimateBlockSize_symbolType(fseMetadata->ofType, ofCodeTable, nbSeq, MaxOff,
fseTables->offcodeCTable, NULL,
OF_defaultNorm, OF_defaultNormLog, DefaultMaxOff,
workspace, wkspSize);
cSeqSizeEstimate += ZSTD_estimateBlockSize_symbolType(fseMetadata->llType, llCodeTable, nbSeq, MaxLL,
fseTables->litlengthCTable, LL_bits,
LL_defaultNorm, LL_defaultNormLog, MaxLL,
workspace, wkspSize);
cSeqSizeEstimate += ZSTD_estimateBlockSize_symbolType(fseMetadata->mlType, mlCodeTable, nbSeq, MaxML,
fseTables->matchlengthCTable, ML_bits,
ML_defaultNorm, ML_defaultNormLog, MaxML,
workspace, wkspSize);
if (writeEntropy) cSeqSizeEstimate += fseMetadata->fseTablesSize;
return cSeqSizeEstimate + sequencesSectionHeaderSize;
}
/* Returns the size estimate for a given stream of literals, of, ll, ml */
/* The block-size estimator is implemented in Rust; C retains the stateful
* block-splitting recursion and the entropy-table ownership. */
static size_t
ZSTD_estimateBlockSize(const BYTE* literals, size_t litSize,
const BYTE* ofCodeTable,
@@ -2706,13 +2621,12 @@ ZSTD_estimateBlockSize(const BYTE* literals, size_t litSize,
void* workspace, size_t wkspSize,
int writeLitEntropy, int writeSeqEntropy)
{
size_t const literalsSize = ZSTD_estimateBlockSize_literal(literals, litSize,
&entropy->huf, &entropyMetadata->hufMetadata,
workspace, wkspSize, writeLitEntropy);
size_t const seqSize = ZSTD_estimateBlockSize_sequences(ofCodeTable, llCodeTable, mlCodeTable,
nbSeq, &entropy->fse, &entropyMetadata->fseMetadata,
workspace, wkspSize, writeSeqEntropy);
return seqSize + literalsSize + ZSTD_blockHeaderSize;
return ZSTD_rust_estimateBlockSize(
literals, litSize,
ofCodeTable, llCodeTable, mlCodeTable, nbSeq,
entropy, entropyMetadata,
workspace, wkspSize,
writeLitEntropy, writeSeqEntropy);
}
/* Builds entropy statistics and uses them for blocksize estimation.
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 =