diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c index 0becd0266..81baf010f 100644 --- a/lib/compress/zstd_compress.c +++ b/lib/compress/zstd_compress.c @@ -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. diff --git a/rust/src/zstd_compress_stats.rs b/rust/src/zstd_compress_stats.rs index d235d255d..8f75761cb 100644 --- a/rust/src/zstd_compress_stats.rs +++ b/rust/src/zstd_compress_stats.rs @@ -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::(), + &mut max_symbol_value, + literals.cast::(), + lit_size, + workspace, + wksp_size, + ) + }; + if ERR_isError(largest) { + return lit_size; + } + + let mut estimate = unsafe { + HUF_estimateCompressedSize( + (*huf).CTable.as_ptr(), + workspace.cast::(), + 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::(), + &mut max, + code_table.cast::(), + 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::::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::()]; + + 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::::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::()]; + + 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 =