ddidderr/zstd-rs
1
0 Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
daeac5238c7181c312f41e317b24a334742e23ed
zstd-rs/rust/src/zstd_compress_superblock.rs
T
ddidderr daeac5238c feat(rust): port block entropy statistics 
Continue the incremental zstd_compress.c migration with its sequence
statistics and seqStore entropy-compression layer.  The following now
live in rust/src/zstd_compress_stats.rs:

- ZSTD_seqToCodes(), exported under its original name because the C
  dictionary builder (zdict.c) and decodecorpus link against it,
- ZSTD_buildSequencesStatistics() and its dummy variant, whose result
  struct no longer crosses the language boundary,
- ZSTD_entropyCompressSeqStore_internal(), _wExtLitBuffer(), and
  ZSTD_entropyCompressSeqStore(),
- ZSTD_buildBlockEntropyStats() with its literals/sequences helpers,
- ZSTD_copyBlockSequences() and the ZSTD_updateRep() rules it shares
  with the superblock writer.

Boundary: ZSTD_CCtx and ZSTD_CCtx_params stay private to C.  These
paths read exactly two parameter fields, so the C shims keep the
original static/exported function names and forward the strategy and
ZSTD_literalsCompressionIsDisabled() as int scalars alongside the
seqStore and entropy-table leaves.  The leaf layouts (SeqDef,
SeqStore_t, ZSTD_hufCTables_t, ZSTD_fseCTables_t, entropy metadata,
SeqCollector, ZSTD_Sequence) are pinned by compile-time asserts in
zstd_compress.c and by both-pointer-width layout tests in Rust.
ZSTD_buildSeqStore, block dispatch/splitting, and the block-size
estimation helpers remain C for a later slice.

The superblock module previously round-tripped through the C export of
ZSTD_buildBlockEntropyStats and mirrored the entropy leaf structs
privately.  It now calls the crate-internal builder directly, and the
shared struct definitions moved to zstd_compress_stats; consequently
ZSTD_rust_compressSuperBlock() takes the two parameter scalars instead
of an opaque ZSTD_CCtx_params pointer, extracted by its C shim.  Its
layout and repcode tests moved with the definitions.

One C helper family gets no shim: the static
ZSTD_entropyCompressSeqStore_wExtLitBuffer() had a single caller and
was folded into the Rust implementation.

Byte-identity was verified against the pre-change compressor: COPYING,
datagen -g5000000 -s7, and datagen -g300000 -s21 -P90, each at levels
1/3/9/19 and --fast=5, plus a superblock-heavy pass at level 19 with
--target-compressed-block-size=1024; all 18 frames are byte-identical,
covering the repeat-mode state machine, longOffsets, RLE/raw fallback,
and dstSize_tooSmall paths through the block splitter and superblock.

Test plan:
- cd rust && cargo fmt --check && cargo clippy --all-targets -- -D
  warnings && cargo test --all-targets (131 tests: new reference
  vectors for seqToCodes, RLE table headers, empty-seqStore repeat
  copies, literal-stats type selection, and repcode resolution in
  copyBlockSequences)
- cargo build --release --no-default-features --features compression
  and --features decompression
- make -C tests fuzzer && ./tests/fuzzer -i1 --no-big-tests (covers
  ZSTD_generateSequences and ZDICT training over the Rust seqToCodes)
- make -C tests test-rust-lib-smoke
- ./tests/zstreamtest -i1 and ./tests/zstreamtest --newapi -t1 -i1
- ./tests/decodecorpus -t -T1s (decodecorpus links Rust seqToCodes)
- /tmp/ref vs /tmp/got frame diff as described above: byte-identical
2026-07-11 23:11:57 +02:00

1071 lines
33 KiB
Rust
Raw Blame History

#![allow(non_camel_case_types)]
#![allow(non_snake_case)]
//! Target-size super-block compression.
//!
//! This is the Rust implementation of `zstd_compress_superblock.c`. Its C
//! entry-point shim only extracts fields from the opaque `ZSTD_CCtx`; all
//! sequence partitioning, literal and sequence section writing, entropy
//! fallback, and repcode repair remain here. The C-shaped leaf layouts are
//! shared with `zstd_compress_stats`, which also owns the block entropy
//! statistics this module consumes.
use crate::common::{
LL_BITS, LL_DEFAULT_NORM, LL_DEFAULT_NORM_LOG, MAX_LL, MAX_ML, MAX_OFF, MINMATCH, ML_BITS,
ML_DEFAULT_NORM, ML_DEFAULT_NORM_LOG, OF_DEFAULT_NORM, OF_DEFAULT_NORM_LOG,
};
use crate::errors::{ERR_isError, ZstdErrorCode, ERROR};
use crate::hist::{HIST_countFast_wksp, HIST_count_wksp};
use crate::huf_compress::{
HUF_compress1X_usingCTable, HUF_compress4X_usingCTable, HUF_estimateCompressedSize,
};
use crate::mem::{MEM_32bits, MEM_writeLE16, MEM_writeLE24, MEM_writeLE32};
use crate::zstd_compress_literals::{
ZSTD_compressRleLiteralsBlock, ZSTD_hufCTables_t, ZSTD_noCompressLiterals,
};
use crate::zstd_compress_sequences::{
SeqDef, ZSTD_crossEntropyCost, ZSTD_encodeSequences, ZSTD_fseBitCost,
};
use crate::zstd_compress_stats::{
build_block_entropy_stats, update_rep, SeqStore_t, ZSTD_compressedBlockState_t,
ZSTD_entropyCTablesMetadata_t, ZSTD_entropyCTables_t, ZSTD_fseCTablesMetadata_t,
ZSTD_fseCTables_t, ZSTD_hufCTablesMetadata_t,
};
use std::ffi::c_void;
use std::mem::MaybeUninit;
use std::os::raw::c_int;
use std::ptr;
const SET_BASIC: c_int = 0;
const SET_RLE: c_int = 1;
const SET_COMPRESSED: c_int = 2;
const SET_REPEAT: c_int = 3;
const HUF_FLAGS_BMI2: c_int = 1;
const ZSTD_BLOCK_HEADER_SIZE: usize = 3;
const ZSTD_TARGET_CBLOCK_SIZE_MIN: usize = 1340;
const STREAM_ACCUMULATOR_MIN_32: u32 = 25;
const STREAM_ACCUMULATOR_MIN_64: u32 = 57;
const LONG_NB_SEQ: usize = 0x7f00;
const DEFAULT_MAX_OFF: u32 = 28;
const BYTE_SCALE: usize = 256;
#[inline]
fn remaining_capacity(capacity: usize, written: usize) -> Result<usize, usize> {
capacity
.checked_sub(written)
.ok_or_else(|| ERROR(ZstdErrorCode::DstSizeTooSmall))
}
#[inline]
unsafe fn copy_bytes(dst: *mut u8, src: *const u8, size: usize) {
if size != 0 {
unsafe { ptr::copy_nonoverlapping(src, dst, size) };
}
}
#[inline]
unsafe fn sequence_length(seq_store: *const SeqStore_t, sequence: *const SeqDef) -> (usize, usize) {
let sequence_value = unsafe { sequence.read() };
let mut lit_length = sequence_value.litLength as usize;
let mut match_length = sequence_value.mlBase as usize + MINMATCH;
let sequence_index = unsafe { sequence.offset_from((*seq_store).sequencesStart) as usize };
if unsafe { (*seq_store).longLengthPos } == sequence_index as u32 {
match unsafe { (*seq_store).longLengthType } {
1 => lit_length += 1 << 16,
2 => match_length += 1 << 16,
_ => {}
}
}
(lit_length, match_length)
}
#[inline]
unsafe fn no_compress_block(
dst: *mut u8,
dst_capacity: usize,
src: *const u8,
src_size: usize,
last_block: u32,
) -> usize {
let needed = match src_size.checked_add(ZSTD_BLOCK_HEADER_SIZE) {
Some(value) => value,
None => return ERROR(ZstdErrorCode::DstSizeTooSmall),
};
if needed > dst_capacity {
return ERROR(ZstdErrorCode::DstSizeTooSmall);
}
let block_header = last_block.wrapping_add((src_size as u32).wrapping_shl(3));
unsafe {
MEM_writeLE24(dst.cast::<c_void>(), block_header);
copy_bytes(dst.add(ZSTD_BLOCK_HEADER_SIZE), src, src_size);
}
needed
}
#[allow(clippy::too_many_arguments)]
unsafe fn compress_subblock_literal(
huf_table: *const usize,
huf_metadata: *const ZSTD_hufCTablesMetadata_t,
literals: *const u8,
lit_size: usize,
dst: *mut u8,
dst_size: usize,
bmi2: c_int,
write_entropy: bool,
entropy_written: &mut bool,
) -> usize {
let header_slack = if write_entropy { 200 } else { 0 };
let literal_header_size = 3
+ usize::from(lit_size >= 1024 - header_slack)
+ usize::from(lit_size >= 16 * 1024 - header_slack);
let metadata = unsafe { &*huf_metadata };
*entropy_written = false;
if lit_size == 0 || metadata.hType == SET_BASIC {
return unsafe {
ZSTD_noCompressLiterals(
dst.cast::<c_void>(),
dst_size,
literals.cast::<c_void>(),
lit_size,
)
};
}
if metadata.hType == SET_RLE {
return unsafe {
ZSTD_compressRleLiteralsBlock(
dst.cast::<c_void>(),
dst_size,
literals.cast::<c_void>(),
lit_size,
)
};
}
debug_assert!(metadata.hType == SET_COMPRESSED || metadata.hType == SET_REPEAT);
if literal_header_size > dst_size {
return ERROR(ZstdErrorCode::DstSizeTooSmall);
}
let h_type = if write_entropy {
metadata.hType
} else {
SET_REPEAT
};
let mut written = literal_header_size;
let mut compressed_literals_size = 0usize;
if write_entropy && metadata.hType == SET_COMPRESSED {
let remaining = match remaining_capacity(dst_size, written) {
Ok(value) => value,
Err(error) => return error,
};
if metadata.hufDesSize > remaining {
return ERROR(ZstdErrorCode::DstSizeTooSmall);
}
unsafe {
copy_bytes(
dst.add(written),
metadata.hufDesBuffer.as_ptr(),
metadata.hufDesSize,
);
}
written += metadata.hufDesSize;
compressed_literals_size += metadata.hufDesSize;
}
let flags = if bmi2 != 0 { HUF_FLAGS_BMI2 } else { 0 };
let remaining = match remaining_capacity(dst_size, written) {
Ok(value) => value,
Err(error) => return error,
};
let compressed_size = if literal_header_size == 3 {
unsafe {
HUF_compress1X_usingCTable(
dst.add(written).cast::<c_void>(),
remaining,
literals.cast::<c_void>(),
lit_size,
huf_table,
flags,
)
}
} else {
unsafe {
HUF_compress4X_usingCTable(
dst.add(written).cast::<c_void>(),
remaining,
literals.cast::<c_void>(),
lit_size,
huf_table,
flags,
)
}
};
compressed_literals_size = compressed_literals_size.wrapping_add(compressed_size);
if compressed_size == 0 || ERR_isError(compressed_size) {
return 0;
}
written += compressed_size;
if !write_entropy && compressed_literals_size >= lit_size {
return unsafe {
ZSTD_noCompressLiterals(
dst.cast::<c_void>(),
dst_size,
literals.cast::<c_void>(),
lit_size,
)
};
}
let encoded_header_size = 3
+ usize::from(compressed_literals_size >= 1024)
+ usize::from(compressed_literals_size >= 16 * 1024);
if literal_header_size < encoded_header_size {
return unsafe {
ZSTD_noCompressLiterals(
dst.cast::<c_void>(),
dst_size,
literals.cast::<c_void>(),
lit_size,
)
};
}
unsafe {
match literal_header_size {
3 => {
let header = (h_type as u32)
.wrapping_add(((literal_header_size != 3) as u32) << 2)
.wrapping_add((lit_size as u32).wrapping_shl(4))
.wrapping_add((compressed_literals_size as u32).wrapping_shl(14));
MEM_writeLE24(dst.cast::<c_void>(), header);
}
4 => {
let header = (h_type as u32)
.wrapping_add(2 << 2)
.wrapping_add((lit_size as u32).wrapping_shl(4))
.wrapping_add((compressed_literals_size as u32).wrapping_shl(18));
MEM_writeLE32(dst.cast::<c_void>(), header);
}
5 => {
let header = (h_type as u32)
.wrapping_add(3 << 2)
.wrapping_add((lit_size as u32).wrapping_shl(4))
.wrapping_add((compressed_literals_size as u32).wrapping_shl(22));
MEM_writeLE32(dst.cast::<c_void>(), header);
*dst.add(4) = (compressed_literals_size >> 10) as u8;
}
_ => unreachable!("literal header size is three through five bytes"),
}
}
*entropy_written = true;
written
}
#[allow(clippy::too_many_arguments)]
unsafe fn compress_subblock_sequences(
fse_tables: *const ZSTD_fseCTables_t,
fse_metadata: *const ZSTD_fseCTablesMetadata_t,
sequences: *const SeqDef,
nb_seq: usize,
ll_code: *const u8,
ml_code: *const u8,
of_code: *const u8,
window_log: u32,
dst: *mut u8,
dst_capacity: usize,
bmi2: c_int,
write_entropy: bool,
entropy_written: &mut bool,
) -> usize {
*entropy_written = false;
if dst_capacity < 4 {
return ERROR(ZstdErrorCode::DstSizeTooSmall);
}
let mut written = unsafe {
if nb_seq < 128 {
*dst = nb_seq as u8;
1
} else if nb_seq < LONG_NB_SEQ {
*dst = ((nb_seq >> 8) + 0x80) as u8;
*dst.add(1) = nb_seq as u8;
2
} else {
*dst = 0xff;
MEM_writeLE16(dst.add(1).cast::<c_void>(), (nb_seq - LONG_NB_SEQ) as u16);
3
}
};
if nb_seq == 0 {
return written;
}
let sequence_header = written;
written += 1;
let metadata = unsafe { &*fse_metadata };
if write_entropy {
let remaining = match remaining_capacity(dst_capacity, written) {
Ok(value) => value,
Err(error) => return error,
};
if metadata.fseTablesSize > remaining {
return ERROR(ZstdErrorCode::DstSizeTooSmall);
}
unsafe {
*dst.add(sequence_header) = ((metadata.llType as u8) << 6)
.wrapping_add((metadata.ofType as u8) << 4)
.wrapping_add((metadata.mlType as u8) << 2);
copy_bytes(
dst.add(written),
metadata.fseTablesBuffer.as_ptr(),
metadata.fseTablesSize,
);
}
written += metadata.fseTablesSize;
} else {
unsafe { *dst.add(sequence_header) = 0b1111_1100 };
}
let accumulator_min = if MEM_32bits() {
STREAM_ACCUMULATOR_MIN_32
} else {
STREAM_ACCUMULATOR_MIN_64
};
let bitstream_size = unsafe {
ZSTD_encodeSequences(
dst.add(written).cast::<c_void>(),
dst_capacity - written,
(*fse_tables).matchlengthCTable.as_ptr(),
ml_code,
(*fse_tables).offcodeCTable.as_ptr(),
of_code,
(*fse_tables).litlengthCTable.as_ptr(),
ll_code,
sequences,
nb_seq,
c_int::from(window_log > accumulator_min),
bmi2,
)
};
if ERR_isError(bitstream_size) {
return bitstream_size;
}
written += bitstream_size;
if write_entropy
&& metadata.lastCountSize != 0
&& metadata.lastCountSize.saturating_add(bitstream_size) < 4
{
return 0;
}
if written - sequence_header < 4 {
return 0;
}
*entropy_written = true;
written
}
#[allow(clippy::too_many_arguments)]
unsafe fn compress_subblock(
entropy: *const ZSTD_entropyCTables_t,
entropy_metadata: *const ZSTD_entropyCTablesMetadata_t,
sequences: *const SeqDef,
nb_seq: usize,
literals: *const u8,
lit_size: usize,
ll_code: *const u8,
ml_code: *const u8,
of_code: *const u8,
window_log: u32,
dst: *mut u8,
dst_capacity: usize,
bmi2: c_int,
write_lit_entropy: bool,
write_seq_entropy: bool,
lit_entropy_written: &mut bool,
seq_entropy_written: &mut bool,
last_block: u32,
) -> usize {
if dst_capacity < ZSTD_BLOCK_HEADER_SIZE {
return ERROR(ZstdErrorCode::DstSizeTooSmall);
}
let mut written = ZSTD_BLOCK_HEADER_SIZE;
let literal_size = unsafe {
compress_subblock_literal(
(*entropy).huf.CTable.as_ptr(),
ptr::addr_of!((*entropy_metadata).hufMetadata),
literals,
lit_size,
dst.add(written),
dst_capacity - written,
bmi2,
write_lit_entropy,
lit_entropy_written,
)
};
if ERR_isError(literal_size) {
return literal_size;
}
if literal_size == 0 {
return 0;
}
written += literal_size;
let sequence_size = unsafe {
compress_subblock_sequences(
ptr::addr_of!((*entropy).fse),
ptr::addr_of!((*entropy_metadata).fseMetadata),
sequences,
nb_seq,
ll_code,
ml_code,
of_code,
window_log,
dst.add(written),
dst_capacity - written,
bmi2,
write_seq_entropy,
seq_entropy_written,
)
};
if ERR_isError(sequence_size) {
return sequence_size;
}
if sequence_size == 0 {
return 0;
}
written += sequence_size;
let compressed_size = written - ZSTD_BLOCK_HEADER_SIZE;
let block_header = last_block
.wrapping_add(2 << 1)
.wrapping_add((compressed_size as u32).wrapping_shl(3));
unsafe { MEM_writeLE24(dst.cast::<c_void>(), block_header) };
written
}
unsafe fn estimate_subblock_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 };
if metadata.hType == SET_BASIC {
return lit_size;
}
if metadata.hType == SET_RLE {
return 1;
}
let mut max_symbol_value = 255u32;
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);
}
estimate.wrapping_add(3)
}
#[allow(clippy::too_many_arguments)]
unsafe fn estimate_subblock_symbol_type(
encoding_type: c_int,
code_table: *const u8,
max_code: u32,
nb_seq: usize,
fse_ctable: *const u32,
additional_bits: Option<&[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 size_in_bits = match encoding_type {
SET_BASIC if max <= default_max => unsafe {
ZSTD_crossEntropyCost(default_norm, default_norm_log, workspace.cast::<u32>(), max)
},
SET_BASIC => ERROR(ZstdErrorCode::Generic),
SET_RLE => 0,
SET_COMPRESSED | SET_REPEAT => unsafe {
ZSTD_fseBitCost(fse_ctable, workspace.cast::<u32>(), max)
},
_ => ERROR(ZstdErrorCode::Generic),
};
if ERR_isError(size_in_bits) {
return nb_seq.wrapping_mul(10);
}
for index in 0..nb_seq {
let code = unsafe { *code_table.add(index) };
size_in_bits = size_in_bits.wrapping_add(match additional_bits {
Some(bits) => bits[code as usize] as usize,
None => code as usize,
});
}
size_in_bits / 8
}
#[allow(clippy::too_many_arguments)]
unsafe fn estimate_subblock_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 {
if nb_seq == 0 {
return 3;
}
let metadata = unsafe { &*fse_metadata };
let tables = unsafe { &*fse_tables };
let mut estimate = unsafe {
estimate_subblock_symbol_type(
metadata.ofType,
of_code_table,
MAX_OFF as u32,
nb_seq,
tables.offcodeCTable.as_ptr(),
None,
OF_DEFAULT_NORM.as_ptr(),
OF_DEFAULT_NORM_LOG,
DEFAULT_MAX_OFF,
workspace,
wksp_size,
)
};
estimate = estimate.wrapping_add(unsafe {
estimate_subblock_symbol_type(
metadata.llType,
ll_code_table,
MAX_LL as u32,
nb_seq,
tables.litlengthCTable.as_ptr(),
Some(&LL_BITS),
LL_DEFAULT_NORM.as_ptr(),
LL_DEFAULT_NORM_LOG,
MAX_LL as u32,
workspace,
wksp_size,
)
});
estimate = estimate.wrapping_add(unsafe {
estimate_subblock_symbol_type(
metadata.mlType,
ml_code_table,
MAX_ML as u32,
nb_seq,
tables.matchlengthCTable.as_ptr(),
Some(&ML_BITS),
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(3)
}
struct EstimatedBlockSize {
est_lit_size: usize,
est_block_size: usize,
}
#[allow(clippy::too_many_arguments)]
unsafe fn estimate_subblock_size(
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: bool,
write_seq_entropy: bool,
) -> EstimatedBlockSize {
let est_lit_size = unsafe {
estimate_subblock_literal(
literals,
lit_size,
ptr::addr_of!((*entropy).huf),
ptr::addr_of!((*entropy_metadata).hufMetadata),
workspace,
wksp_size,
write_lit_entropy,
)
};
let est_sequences = unsafe {
estimate_subblock_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,
)
};
EstimatedBlockSize {
est_lit_size,
est_block_size: est_sequences
.wrapping_add(est_lit_size)
.wrapping_add(ZSTD_BLOCK_HEADER_SIZE),
}
}
#[inline]
unsafe fn need_sequence_entropy_tables(metadata: *const ZSTD_fseCTablesMetadata_t) -> bool {
let metadata = unsafe { &*metadata };
[metadata.llType, metadata.mlType, metadata.ofType]
.iter()
.any(|&encoding_type| encoding_type == SET_COMPRESSED || encoding_type == SET_RLE)
}
unsafe fn count_literals(
seq_store: *const SeqStore_t,
sequences: *const SeqDef,
count: usize,
) -> usize {
let mut total = 0usize;
for index in 0..count {
total = total.wrapping_add(unsafe { sequence_length(seq_store, sequences.add(index)).0 });
}
total
}
unsafe fn sequence_decompressed_size(
seq_store: *const SeqStore_t,
sequences: *const SeqDef,
count: usize,
lit_size: usize,
) -> usize {
let mut match_length_sum = 0usize;
for index in 0..count {
match_length_sum = match_length_sum
.wrapping_add(unsafe { sequence_length(seq_store, sequences.add(index)).1 });
}
match_length_sum.wrapping_add(lit_size)
}
unsafe fn size_block_sequences(
sequences: *const SeqDef,
nb_seq: usize,
target_budget: usize,
avg_lit_cost: usize,
avg_seq_cost: usize,
first_subblock: bool,
) -> usize {
debug_assert!(nb_seq > 0);
let mut budget = if first_subblock {
120usize.wrapping_mul(BYTE_SCALE)
} else {
0
};
let first = unsafe { sequences.read() };
budget = budget
.wrapping_add((first.litLength as usize).wrapping_mul(avg_lit_cost))
.wrapping_add(avg_seq_cost);
if budget > target_budget {
return 1;
}
let mut in_size = first.litLength as usize + first.mlBase as usize + MINMATCH;
for index in 1..nb_seq {
let sequence = unsafe { sequences.add(index).read() };
budget = budget
.wrapping_add((sequence.litLength as usize).wrapping_mul(avg_lit_cost))
.wrapping_add(avg_seq_cost);
in_size = in_size
.wrapping_add(sequence.litLength as usize)
.wrapping_add(sequence.mlBase as usize + MINMATCH);
if budget > target_budget && budget < in_size.wrapping_mul(BYTE_SCALE) {
return index;
}
}
nb_seq
}
#[allow(clippy::too_many_arguments, clippy::manual_checked_ops)]
unsafe fn compress_subblock_multi(
seq_store: *const SeqStore_t,
prev_cblock: *const ZSTD_compressedBlockState_t,
next_cblock: *mut ZSTD_compressedBlockState_t,
entropy_metadata: *const ZSTD_entropyCTablesMetadata_t,
window_log: u32,
target_cblock_size: usize,
dst: *mut u8,
dst_capacity: usize,
src: *const u8,
src_size: usize,
bmi2: c_int,
last_block: u32,
workspace: *mut c_void,
wksp_size: usize,
) -> usize {
let sequence_start = unsafe { (*seq_store).sequencesStart };
let sequence_end = unsafe { (*seq_store).sequences };
let nb_sequences = unsafe { sequence_end.offset_from(sequence_start) as usize };
let literals_start = unsafe { (*seq_store).litStart };
let literals_end = unsafe { (*seq_store).lit };
let nb_literals = unsafe { literals_end.offset_from(literals_start) as usize };
let mut sequence_index = 0usize;
let mut literal_offset = 0usize;
let mut source_offset = 0usize;
let mut output_offset = 0usize;
let mut write_lit_entropy = unsafe { (*entropy_metadata).hufMetadata.hType == SET_COMPRESSED };
let mut write_seq_entropy = true;
let target_cblock_size = target_cblock_size.max(ZSTD_TARGET_CBLOCK_SIZE_MIN);
if nb_sequences > 0 {
let estimate = unsafe {
estimate_subblock_size(
literals_start,
nb_literals,
(*seq_store).ofCode,
(*seq_store).llCode,
(*seq_store).mlCode,
nb_sequences,
ptr::addr_of!((*next_cblock).entropy),
entropy_metadata,
workspace,
wksp_size,
write_lit_entropy,
write_seq_entropy,
)
};
if estimate.est_block_size > src_size {
return 0;
}
let avg_lit_cost = if nb_literals != 0 {
estimate.est_lit_size.wrapping_mul(BYTE_SCALE) / nb_literals
} else {
BYTE_SCALE
};
let avg_seq_cost = estimate
.est_block_size
.wrapping_sub(estimate.est_lit_size)
.wrapping_mul(BYTE_SCALE)
/ nb_sequences;
let nb_subblocks = ((estimate.est_block_size.wrapping_add(target_cblock_size / 2))
/ target_cblock_size)
.max(1);
let avg_block_budget = estimate.est_block_size.wrapping_mul(BYTE_SCALE) / nb_subblocks;
let block_budget_supp = 0usize;
for subblock_index in 0..nb_subblocks.saturating_sub(1) {
let remaining_sequences = nb_sequences - sequence_index;
let sequence_count = unsafe {
size_block_sequences(
sequence_start.add(sequence_index),
remaining_sequences,
avg_block_budget.wrapping_add(block_budget_supp),
avg_lit_cost,
avg_seq_cost,
subblock_index == 0,
)
};
if sequence_index + sequence_count == nb_sequences {
break;
}
let literal_size = unsafe {
count_literals(
seq_store,
sequence_start.add(sequence_index),
sequence_count,
)
};
let decompressed_size = unsafe {
sequence_decompressed_size(
seq_store,
sequence_start.add(sequence_index),
sequence_count,
literal_size,
)
};
let mut lit_entropy_written = false;
let mut seq_entropy_written = false;
let compressed_size = unsafe {
compress_subblock(
ptr::addr_of!((*next_cblock).entropy),
entropy_metadata,
sequence_start.add(sequence_index),
sequence_count,
literals_start.add(literal_offset),
literal_size,
(*seq_store).llCode.add(sequence_index),
(*seq_store).mlCode.add(sequence_index),
(*seq_store).ofCode.add(sequence_index),
window_log,
dst.add(output_offset),
dst_capacity - output_offset,
bmi2,
write_lit_entropy,
write_seq_entropy,
&mut lit_entropy_written,
&mut seq_entropy_written,
0,
)
};
if ERR_isError(compressed_size) {
return compressed_size;
}
if compressed_size != 0 && compressed_size < decompressed_size {
source_offset = source_offset.wrapping_add(decompressed_size);
literal_offset = literal_offset.wrapping_add(literal_size);
output_offset = output_offset.wrapping_add(compressed_size);
sequence_index += sequence_count;
if lit_entropy_written {
write_lit_entropy = false;
}
if seq_entropy_written {
write_seq_entropy = false;
}
}
}
}
let literal_size = nb_literals - literal_offset;
let sequence_count = nb_sequences - sequence_index;
let decompressed_size = unsafe {
sequence_decompressed_size(
seq_store,
sequence_start.add(sequence_index),
sequence_count,
literal_size,
)
};
let mut lit_entropy_written = false;
let mut seq_entropy_written = false;
let compressed_size = unsafe {
compress_subblock(
ptr::addr_of!((*next_cblock).entropy),
entropy_metadata,
sequence_start.add(sequence_index),
sequence_count,
literals_start.add(literal_offset),
literal_size,
(*seq_store).llCode.add(sequence_index),
(*seq_store).mlCode.add(sequence_index),
(*seq_store).ofCode.add(sequence_index),
window_log,
dst.add(output_offset),
dst_capacity - output_offset,
bmi2,
write_lit_entropy,
write_seq_entropy,
&mut lit_entropy_written,
&mut seq_entropy_written,
last_block,
)
};
if ERR_isError(compressed_size) {
return compressed_size;
}
if compressed_size != 0 && compressed_size < decompressed_size {
source_offset = source_offset.wrapping_add(decompressed_size);
output_offset = output_offset.wrapping_add(compressed_size);
sequence_index += sequence_count;
if lit_entropy_written {
write_lit_entropy = false;
}
if seq_entropy_written {
write_seq_entropy = false;
}
}
if write_lit_entropy {
unsafe {
ptr::copy_nonoverlapping(
ptr::addr_of!((*prev_cblock).entropy.huf),
ptr::addr_of_mut!((*next_cblock).entropy.huf),
1,
);
}
}
if write_seq_entropy
&& unsafe { need_sequence_entropy_tables(ptr::addr_of!((*entropy_metadata).fseMetadata)) }
{
return 0;
}
if source_offset < src_size {
let raw_size = src_size - source_offset;
let remaining = match remaining_capacity(dst_capacity, output_offset) {
Ok(value) => value,
Err(error) => return error,
};
let raw_compressed_size = unsafe {
no_compress_block(
dst.add(output_offset),
remaining,
src.add(source_offset),
raw_size,
last_block,
)
};
if ERR_isError(raw_compressed_size) {
return raw_compressed_size;
}
output_offset += raw_compressed_size;
if sequence_index < nb_sequences {
let mut reps = unsafe { (*prev_cblock).rep };
for index in 0..sequence_index {
let sequence = unsafe { sequence_start.add(index).read() };
let lit_length_is_zero =
unsafe { sequence_length(seq_store, sequence_start.add(index)).0 == 0 };
update_rep(&mut reps, sequence.offBase, lit_length_is_zero);
}
unsafe { (*next_cblock).rep = reps };
}
}
output_offset
}
/// C ABI implementation called by the declaration-only C superblock shim.
///
/// `ZSTD_CCtx_params` deliberately remains opaque: the C shim extracts the
/// two scalar fields this path reads (the strategy and the
/// literals-compression switch) and the crate-internal entropy statistics
/// builder receives them directly.
#[allow(clippy::too_many_arguments)]
#[no_mangle]
pub unsafe extern "C" fn ZSTD_rust_compressSuperBlock(
seq_store: *const c_void,
prev_cblock: *const c_void,
next_cblock: *mut c_void,
strategy: c_int,
disable_literal_compression: c_int,
workspace: *mut c_void,
wksp_size: usize,
bmi2: c_int,
window_log: u32,
target_cblock_size: usize,
dst: *mut c_void,
dst_capacity: usize,
src: *const c_void,
src_size: usize,
last_block: u32,
) -> usize {
let seq_store = seq_store.cast::<SeqStore_t>();
let prev_cblock = prev_cblock.cast::<ZSTD_compressedBlockState_t>();
let next_cblock = next_cblock.cast::<ZSTD_compressedBlockState_t>();
// The statistics builder only writes the used prefixes of the two
// metadata byte buffers. Start with initialized storage so treating the
// completed struct as a Rust value never exposes uninitialized elements.
let mut entropy_metadata = MaybeUninit::<ZSTD_entropyCTablesMetadata_t>::zeroed();
let entropy_result = unsafe {
build_block_entropy_stats(
seq_store,
ptr::addr_of!((*prev_cblock).entropy),
ptr::addr_of_mut!((*next_cblock).entropy),
strategy,
disable_literal_compression,
entropy_metadata.as_mut_ptr(),
workspace,
wksp_size,
)
};
if ERR_isError(entropy_result) {
return entropy_result;
}
let entropy_metadata = unsafe { entropy_metadata.assume_init() };
unsafe {
compress_subblock_multi(
seq_store,
prev_cblock,
next_cblock,
&entropy_metadata,
window_log,
target_cblock_size,
dst.cast::<u8>(),
dst_capacity,
src.cast::<u8>(),
src_size,
bmi2,
last_block,
workspace,
wksp_size,
)
}
}
#[cfg(test)]
mod tests {
// The shared C leaf layouts and the repcode update rules are asserted in
// `zstd_compress_stats`, which owns those definitions.
use super::*;
#[test]
fn subblock_sizing_respects_the_budget() {
let sequences = [SeqDef {
offBase: 4,
litLength: 10,
mlBase: 10,
}; 8];
// A generous budget keeps every sequence in one sub-block.
let all = unsafe {
size_block_sequences(
sequences.as_ptr(),
sequences.len(),
usize::MAX,
256,
256,
true,
)
};
assert_eq!(all, sequences.len());
// The first sequence always survives, even over budget.
let first =
unsafe { size_block_sequences(sequences.as_ptr(), sequences.len(), 0, 256, 256, true) };
assert_eq!(first, 1);
}
}
Reference in new issue View Git Blame Copy Permalink