diff --git a/lib/compress/zstd_compress.c b/lib/compress/zstd_compress.c index 4c8fbebc4..278b1925c 100644 --- a/lib/compress/zstd_compress.c +++ b/lib/compress/zstd_compress.c @@ -288,6 +288,17 @@ size_t ZSTD_rust_countSeqStoreMatchBytes(const SeqStore_t* seqStore); void ZSTD_rust_deriveSeqStoreChunk(SeqStore_t* resultSeqStore, const SeqStore_t* originalSeqStore, size_t startIdx, size_t endIdx); +size_t ZSTD_rust_deriveBlockSplits( + U32* partitions, U32 nbSeq, + const SeqStore_t* originalSeqStore, + SeqStore_t* fullSeqStoreChunk, + SeqStore_t* firstHalfSeqStore, + SeqStore_t* secondHalfSeqStore, + const ZSTD_entropyCTables_t* prevEntropy, + ZSTD_entropyCTables_t* nextEntropy, + int strategy, int disableLiteralCompression, + ZSTD_entropyCTablesMetadata_t* entropyMetadata, + void* workspace, size_t workspaceSize); void ZSTD_rust_seqStore_resolveOffCodes(U32 dRep[ZSTD_REP_NUM], U32 cRep[ZSTD_REP_NUM], const SeqStore_t* seqStore, @@ -2651,7 +2662,7 @@ ZSTD_estimateBlockSize(const BYTE* literals, size_t litSize, * * @return: estimated compressed size of the seqStore, or a zstd error. */ -static size_t +UNUSED_ATTR static size_t ZSTD_buildEntropyStatisticsAndEstimateSubBlockSize(SeqStore_t* seqStore, ZSTD_CCtx* zc) { ZSTD_entropyCTablesMetadata_t* const entropyMetadata = &zc->blockSplitCtx.entropyMetadata; @@ -2791,66 +2802,6 @@ ZSTD_compressSeqStore_singleBlock(ZSTD_CCtx* zc, return cSize; } -/* Struct to keep track of where we are in our recursive calls. */ -typedef struct { - U32* splitLocations; /* Array of split indices */ - size_t idx; /* The current index within splitLocations being worked on */ -} seqStoreSplits; - -#define MIN_SEQUENCES_BLOCK_SPLITTING 300 - -/* Helper function to perform the recursive search for block splits. - * Estimates the cost of seqStore prior to split, and estimates the cost of splitting the sequences in half. - * If advantageous to split, then we recurse down the two sub-blocks. - * If not, or if an error occurred in estimation, then we do not recurse. - * - * Note: The recursion depth is capped by a heuristic minimum number of sequences, - * defined by MIN_SEQUENCES_BLOCK_SPLITTING. - * In theory, this means the absolute largest recursion depth is 10 == log2(maxNbSeqInBlock/MIN_SEQUENCES_BLOCK_SPLITTING). - * In practice, recursion depth usually doesn't go beyond 4. - * - * Furthermore, the number of splits is capped by ZSTD_MAX_NB_BLOCK_SPLITS. - * At ZSTD_MAX_NB_BLOCK_SPLITS == 196 with the current existing blockSize - * maximum of 128 KB, this value is actually impossible to reach. - */ -static void -ZSTD_deriveBlockSplitsHelper(seqStoreSplits* splits, size_t startIdx, size_t endIdx, - ZSTD_CCtx* zc, const SeqStore_t* origSeqStore) -{ - SeqStore_t* const fullSeqStoreChunk = &zc->blockSplitCtx.fullSeqStoreChunk; - SeqStore_t* const firstHalfSeqStore = &zc->blockSplitCtx.firstHalfSeqStore; - SeqStore_t* const secondHalfSeqStore = &zc->blockSplitCtx.secondHalfSeqStore; - size_t estimatedOriginalSize; - size_t estimatedFirstHalfSize; - size_t estimatedSecondHalfSize; - size_t midIdx = (startIdx + endIdx)/2; - - DEBUGLOG(5, "ZSTD_deriveBlockSplitsHelper: startIdx=%zu endIdx=%zu", startIdx, endIdx); - assert(endIdx >= startIdx); - if (endIdx - startIdx < MIN_SEQUENCES_BLOCK_SPLITTING || splits->idx >= ZSTD_MAX_NB_BLOCK_SPLITS) { - DEBUGLOG(6, "ZSTD_deriveBlockSplitsHelper: Too few sequences (%zu)", endIdx - startIdx); - return; - } - ZSTD_deriveSeqStoreChunk(fullSeqStoreChunk, origSeqStore, startIdx, endIdx); - ZSTD_deriveSeqStoreChunk(firstHalfSeqStore, origSeqStore, startIdx, midIdx); - ZSTD_deriveSeqStoreChunk(secondHalfSeqStore, origSeqStore, midIdx, endIdx); - estimatedOriginalSize = ZSTD_buildEntropyStatisticsAndEstimateSubBlockSize(fullSeqStoreChunk, zc); - estimatedFirstHalfSize = ZSTD_buildEntropyStatisticsAndEstimateSubBlockSize(firstHalfSeqStore, zc); - estimatedSecondHalfSize = ZSTD_buildEntropyStatisticsAndEstimateSubBlockSize(secondHalfSeqStore, zc); - DEBUGLOG(5, "Estimated original block size: %zu -- First half split: %zu -- Second half split: %zu", - estimatedOriginalSize, estimatedFirstHalfSize, estimatedSecondHalfSize); - if (ZSTD_isError(estimatedOriginalSize) || ZSTD_isError(estimatedFirstHalfSize) || ZSTD_isError(estimatedSecondHalfSize)) { - return; - } - if (estimatedFirstHalfSize + estimatedSecondHalfSize < estimatedOriginalSize) { - DEBUGLOG(5, "split decided at seqNb:%zu", midIdx); - ZSTD_deriveBlockSplitsHelper(splits, startIdx, midIdx, zc, origSeqStore); - splits->splitLocations[splits->idx] = (U32)midIdx; - splits->idx++; - ZSTD_deriveBlockSplitsHelper(splits, midIdx, endIdx, zc, origSeqStore); - } -} - /* Base recursive function. * Populates a table with intra-block partition indices that can improve compression ratio. * @@ -2858,18 +2809,19 @@ ZSTD_deriveBlockSplitsHelper(seqStoreSplits* splits, size_t startIdx, size_t end */ static size_t ZSTD_deriveBlockSplits(ZSTD_CCtx* zc, U32 partitions[], U32 nbSeq) { - seqStoreSplits splits; - splits.splitLocations = partitions; - splits.idx = 0; - if (nbSeq <= 4) { - DEBUGLOG(5, "ZSTD_deriveBlockSplits: Too few sequences to split (%u <= 4)", nbSeq); - /* Refuse to try and split anything with less than 4 sequences */ - return 0; - } - ZSTD_deriveBlockSplitsHelper(&splits, 0, nbSeq, zc, &zc->seqStore); - splits.splitLocations[splits.idx] = nbSeq; - DEBUGLOG(5, "ZSTD_deriveBlockSplits: final nb partitions: %zu", splits.idx+1); - return splits.idx; + return ZSTD_rust_deriveBlockSplits( + partitions, nbSeq, + &zc->seqStore, + &zc->blockSplitCtx.fullSeqStoreChunk, + &zc->blockSplitCtx.firstHalfSeqStore, + &zc->blockSplitCtx.secondHalfSeqStore, + &zc->blockState.prevCBlock->entropy, + &zc->blockState.nextCBlock->entropy, + (int)zc->appliedParams.cParams.strategy, + ZSTD_literalsCompressionIsDisabled(&zc->appliedParams), + &zc->blockSplitCtx.entropyMetadata, + zc->tmpWorkspace, + zc->tmpWkspSize); } /* ZSTD_compressBlock_splitBlock(): diff --git a/rust/src/lib.rs b/rust/src/lib.rs index e899341ff..b72620007 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -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; diff --git a/rust/src/zstd_compress_block_split.rs b/rust/src/zstd_compress_block_split.rs new file mode 100644 index 000000000..0a7ed73ee --- /dev/null +++ b/rust/src/zstd_compress_block_split.rs @@ -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( + 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( + 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( + 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)] + ); + } +}