From 894c59e2d1ae042314df432e34acf14c5ff1d321 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Tue, 29 Sep 2026 05:25:05 +0300 Subject: [PATCH 01/20] feat(dictionary): raw dictionaries from the FastCOVER search `create_raw_dict_from_dir`, `_source` and `_slice` built their content with a reservoir-sampled local-maximum-coverage trainer. Measured on the repository files against upstream zstd's `--train` dictionary, held-out samples compressed by upstream zstd: it came out 1.38% larger at level 3, 1.80% at level 19 and 8.14% on 1 KiB samples, where the FastCOVER search is within 0.14% of upstream; one training on 1.4 MB took minutes where the search takes under a second. The three functions now return what `optimize_fastcover_dict` picks at its defaults, as `zstd --train` does, without entropy tables: a directory is one sample per file, and a corpus without sizes is cut into at least sixteen samples of at most 128 KiB. A corpus no larger than the dictionary is its own content, and one the trainer refuses gives its last `dict_size` bytes, as before for tiny corpora. The old trainer, its reservoir sampler and frequency estimate are removed. The frame compressor's dict-builder test now compresses a small frame, which a dictionary exists for: a long frame of near-identical lines compresses smaller without one, in upstream zstd too (245 against 261 bytes with this dictionary at level 1). Part of #128 --- zstd/benches/dict_builder_fastcover.rs | 2 +- zstd/src/dictionary/frequency.rs | 55 ----- zstd/src/dictionary/frequency/tests.rs | 21 -- zstd/src/dictionary/lmc.rs | 136 ------------ zstd/src/dictionary/mod.rs | 227 +++++++++----------- zstd/src/dictionary/reservoir.rs | 111 ---------- zstd/src/dictionary/reservoir/tests.rs | 32 --- zstd/src/dictionary/tests.rs | 73 +++++++ zstd/src/encoding/frame_compressor/tests.rs | 15 +- 9 files changed, 177 insertions(+), 495 deletions(-) delete mode 100644 zstd/src/dictionary/frequency.rs delete mode 100644 zstd/src/dictionary/frequency/tests.rs delete mode 100644 zstd/src/dictionary/lmc.rs delete mode 100644 zstd/src/dictionary/reservoir.rs delete mode 100644 zstd/src/dictionary/reservoir/tests.rs diff --git a/zstd/benches/dict_builder_fastcover.rs b/zstd/benches/dict_builder_fastcover.rs index 4e6468a5b..ab2d510b4 100644 --- a/zstd/benches/dict_builder_fastcover.rs +++ b/zstd/benches/dict_builder_fastcover.rs @@ -36,7 +36,7 @@ fn bench_dict_builder(c: &mut Criterion) { ..FastCoverOptions::default() }; - c.bench_function("dict_builder/lmc_raw", |b| { + c.bench_function("dict_builder/raw_content", |b| { b.iter(|| { let mut out = Vec::new(); create_raw_dict_from_source( diff --git a/zstd/src/dictionary/frequency.rs b/zstd/src/dictionary/frequency.rs deleted file mode 100644 index 9bc0b6ba4..000000000 --- a/zstd/src/dictionary/frequency.rs +++ /dev/null @@ -1,55 +0,0 @@ -//! Contains `compute_frequency`, a function -//! that uses a rolling Karp-Rabin hash to -//! efficiently count the number of occurences -//! of a given k-mer within a set. - -/// Computes a best effort guess as to how many times `pattern` occurs within -/// `body`. While not 100% accurate, it will be accurate the vast majority of time -pub fn estimate_frequency(pattern: &[u8], body: &[u8]) -> usize { - if pattern.is_empty() { - return 0; - } - assert!(body.len() >= pattern.len()); - // A prime number for modulo operations to reduce collisions (q) - const PRIME: i64 = 2654435761; - // Number of characters in the input alphabet (d) - const ALPHABET_SIZE: i64 = 256; - // Hash of input pattern (p) - let mut pattern_hash: i64 = 0; - // Hash of the current window of text (t) - let mut window_hash: i64 = 0; - // High-order digit multiplier: h = ALPHABET_SIZE^(pattern.len()-1) mod PRIME - let mut h: i64 = 1; - for _ in 0..pattern.len().saturating_sub(1) { - h = (h * ALPHABET_SIZE) % PRIME; - } - - // Compute initial hash values - for i in 0..pattern.len() { - pattern_hash = (ALPHABET_SIZE * pattern_hash + pattern[i] as i64) % PRIME; - window_hash = (ALPHABET_SIZE * window_hash + body[i] as i64) % PRIME; - } - - let mut num_occurrences = 0; - for i in 0..=body.len() - pattern.len() { - if pattern_hash == window_hash { - num_occurrences += 1; - } - - // Compute hash for next window using rolling hash - if i < body.len() - pattern.len() { - window_hash = (ALPHABET_SIZE * (window_hash - body[i] as i64 * h) - + body[i + pattern.len()] as i64) - % PRIME; - // Ensure non-negative (Euclidean modulo) - if window_hash < 0 { - window_hash += PRIME; - } - } - } - - num_occurrences -} - -#[cfg(test)] -mod tests; diff --git a/zstd/src/dictionary/frequency/tests.rs b/zstd/src/dictionary/frequency/tests.rs deleted file mode 100644 index a919aa047..000000000 --- a/zstd/src/dictionary/frequency/tests.rs +++ /dev/null @@ -1,21 +0,0 @@ -use super::estimate_frequency; -#[test] -fn dead_beef() { - assert_eq!( - estimate_frequency(&[0xde, 0xad], &[0xde, 0xad, 0xbe, 0xef, 0xde, 0xad]), - 2 - ); -} - -#[test] -fn smallest_body() { - assert_eq!(estimate_frequency(&[0x00, 0xff], &[0x00, 0xff]), 1); -} - -#[test] -fn no_match() { - assert_eq!( - estimate_frequency(&[0xff, 0xff], &[0xde, 0xad, 0xbe, 0xef, 0xde, 0xad]), - 0 - ); -} diff --git a/zstd/src/dictionary/lmc.rs b/zstd/src/dictionary/lmc.rs deleted file mode 100644 index e2794ffa5..000000000 --- a/zstd/src/dictionary/lmc.rs +++ /dev/null @@ -1,136 +0,0 @@ -//! An implementation of the local maximum coverage algorithm -//! described in the paper "Effective Construction of Relative Lempel-Ziv Dictionaries", -//! by Liao, Petri, Moffat, and Wirth, published under the University of Melbourne. -//! -//! See: -//! -//! Facebook's implementation was also used as a reference. -//! - -use super::DictParams; -use crate::dictionary::frequency::estimate_frequency; -use core::convert::TryInto; -use std::collections::HashMap; -use std::vec::Vec; - -/// The size of each k-mer -pub(super) const K: usize = 16; - -///As found under "4: Experiments - Varying k-mer Size" in the original paper, -/// "when k = 16, across all our text collections, there is a reasonable spread" -/// -/// Reasonable range: [6, 16] -pub(super) type KMer = [u8; K]; - -pub struct Segment { - /// The actual contents of the segment. - pub raw: Vec, - /// A measure of how "ideal" a given segment would be to include in the dictionary - /// - /// Higher is better, there's no upper limit. This number is determined by - /// estimating the number of occurrences in a given epoch - pub score: usize, -} - -impl Eq for Segment {} - -impl PartialEq for Segment { - fn eq(&self, other: &Self) -> bool { - // We only really care about score in regards to heap order - self.score == other.score - } -} - -impl PartialOrd for Segment { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for Segment { - fn cmp(&self, other: &Self) -> core::cmp::Ordering { - self.score.cmp(&other.score) - } -} - -/// A re-usable allocation containing large allocations -/// that are used multiple times during dictionary construction (once per epoch) -pub struct Context { - /// Keeps track of the number of occurrences of a particular k-mer within an epoch. - /// - /// Reset for each epoch. - pub frequencies: HashMap, -} - -/// Returns the highest scoring segment in an epoch -/// as a slice of that epoch. -pub fn pick_best_segment( - params: &DictParams, - ctx: &mut Context, - epoch: &'_ [u8], - collection_sample: &'_ [u8], -) -> Segment { - let mut segments = epoch.chunks(params.segment_size as usize).peekable(); - let mut best_segment: &[u8] = segments.peek().expect("at least one segment"); - let mut top_segment_score: usize = 0; - // Iterate over segments and score each segment, keeping track of the best segment - for segment in segments { - let segment_score = score_segment(ctx, collection_sample, segment); - if segment_score > top_segment_score { - best_segment = segment; - top_segment_score = segment_score; - } - } - - Segment { - raw: best_segment.into(), - score: top_segment_score, - } -} - -/// Given a segment, compute the score (or usefulness) of that segment against the entire epoch. -/// -/// `score_segment` modifies `ctx.frequencies`. -fn score_segment(ctx: &mut Context, collection_sample: &[u8], segment: &[u8]) -> usize { - if segment.len() < K { - return 0; - } - let mut segment_score = 0; - // Determine the score of each overlapping k-mer - for i in 0..=(segment.len() - K) { - let kmer: &KMer = (&segment[i..i + K]) - .try_into() - .expect("Failed to make kmer"); - // if the kmer is already in the pool, it recieves a score of zero - if ctx.frequencies.contains_key(kmer) { - continue; - } - let kmer_score = estimate_frequency(kmer, collection_sample); - ctx.frequencies.insert(*kmer, kmer_score); - segment_score += kmer_score; - } - - segment_score -} - -/// Computes the number of epochs and the size of each epoch. -/// -/// Returns a (number of epochs, epoch size) tuple. -/// -/// A translation of `COVER_epoch_info_t COVER_computeEpochs()` from facebook/zstd. -pub fn compute_epoch_info( - params: &DictParams, - max_dict_size: usize, - num_kmers: usize, -) -> (usize, usize) { - let min_epoch_size = 10_000; // 10 KiB - let mut num_epochs: usize = usize::max(1, max_dict_size / params.segment_size as usize); - let mut epoch_size: usize = num_kmers / num_epochs; - if epoch_size >= min_epoch_size { - assert!(epoch_size * num_epochs <= num_kmers); - return (num_epochs, epoch_size); - } - epoch_size = usize::min(min_epoch_size, num_kmers); - num_epochs = num_kmers / epoch_size; - (num_epochs, epoch_size) -} diff --git a/zstd/src/dictionary/mod.rs b/zstd/src/dictionary/mod.rs index dcf8b02ac..b7d77dc20 100644 --- a/zstd/src/dictionary/mod.rs +++ b/zstd/src/dictionary/mod.rs @@ -17,16 +17,12 @@ //! The `optimize_*` forms search segment and dmer sizes, scoring each candidate //! by the total size of the scoring samples compressed with it. //! -//! [`create_raw_dict_from_slice`] and its reader forms build raw content from an -//! undivided corpus instead, estimating segment value by k-mer frequency in a -//! reservoir sample (Liao, Petri, Moffat and Wirth, "Effective construction of -//! Relative Lempel-Ziv Dictionaries"). +//! [`create_raw_dict_from_slice`] and its reader and directory forms build raw +//! content, with no entropy tables, from a corpus without sample sizes: the +//! FastCOVER search `zstd --train` runs, over the corpus cut into samples. mod cover; mod fastcover; -mod frequency; mod legacy; -mod lmc; -mod reservoir; mod samples; mod selection; mod suffix_array; @@ -37,18 +33,13 @@ use crate::blocks::sequence_section::{ }; use crate::decoding::dictionary::MAGIC_NUM as DICT_MAGIC_NUM; use crate::decoding::sequence_section_decoder::{LL_MAX_LOG, ML_MAX_LOG, OF_MAX_LOG}; -use crate::dictionary::reservoir::create_sample; use crate::fse::fse_encoder::{self, build_table_from_symbol_counts}; use crate::huff0::HuffmanTable as HuffmanDecoderTable; use crate::huff0::huff0_encoder::{HuffmanEncoder, HuffmanTable as HuffmanEncoderTable}; -use core::cmp::Reverse; pub use legacy::DEFAULT_SELECTIVITY; -use lmc::*; pub use samples::TrainingError; use samples::refuse; use std::{ - boxed::Box, - collections::{BinaryHeap, HashMap}, format, fs::{self, File}, io::{self, Read}, @@ -152,32 +143,17 @@ pub struct FinalizeOptions { pub dict_id: Option, } -/// A set of values that are used during dictionary construction. -/// -/// Changing these values can improve the resulting dictionary size for certain datasets. -// TODO: move `k` here. -pub(super) struct DictParams { - /// Segment size. - /// - /// As found under "4. Experiments - Varying Segment Size" in the original paper, a - /// segment size of 2 kiB was effective. - /// - /// "We explored a range of \[`segment_size`\] values and found the performance of LMC is insensitive - /// to \[`segment_size`\]. We fix \[`segment_size`\] to 2kiB - /// - /// Reasonable range: [16, 2048+] - pub segment_size: u32, -} - -/// Creates a "raw content" dictionary, training off of every file in this directory and all -/// sub-directories. +/// Create a "raw content" dictionary of at most `dict_size` bytes, with no +/// entropy tables, from every file in this directory and its subdirectories, +/// and write it to `output`. /// -/// The resulting dictionary will be approximately `dict_size` or less, and written to `output`. +/// Each file is one sample, and the content is what [`optimize_fastcover_dict`] +/// picks at its defaults, as `zstd --train` does; too few files to search is +/// trained as [`create_raw_dict_from_slice`] trains an undivided corpus. /// /// # Errors -/// This function returns `Ok(())` if the dictionary was created successfully, and an -/// `Err(io::Error)` if an error was encountered reading the input directory or -/// writing dictionary bytes to `output`. +/// Returns an error reading the directory or its files, writing to `output`, +/// or allocating the trainer's tables. /// /// # Examples /// ```no_run @@ -209,22 +185,15 @@ pub fn create_raw_dict_from_dir, W: io::Write>( } recurse_read(dir, &mut file_paths)?; - // Open each file and chain the readers together - let mut total_file_len: u64 = 0; - let mut file_handles: Vec = Vec::new(); + // Every file is one sample. + let mut corpus = Vec::new(); + let mut sizes = Vec::with_capacity(file_paths.len()); for path in file_paths { - let handle = File::open(path)?; - total_file_len += handle.metadata()?.len(); - file_handles.push(handle); + let before = corpus.len(); + File::open(path)?.read_to_end(&mut corpus)?; + sizes.push(corpus.len() - before); } - let empty_reader: Box = Box::new(io::empty()); - let chained_files = file_handles - .iter() - .fold(empty_reader, |acc, reader| Box::new(acc.chain(reader))); - - // Create a dict using the new reader - create_raw_dict_from_source(chained_files, total_file_len as usize, output, dict_size)?; - Ok(()) + output.write_all(&fastcover_raw_content(&corpus, Some(&sizes), dict_size)?) } /// Read from `source` to create a "raw content" dictionary of `dict_size`. @@ -242,10 +211,11 @@ pub fn create_raw_dict_from_dir, W: io::Write>( /// sources too large to fit comfortably in memory. /// /// A corpus already in memory trains without this copy through -/// [`create_raw_dict_from_slice`]. +/// [`create_raw_dict_from_slice`], whose training this is. /// -/// # API note -/// This public API returns `io::Result<()>` and propagates source/output I/O failures. +/// # Errors +/// Returns an error reading `source`, writing to `output`, or allocating the +/// trainer's tables. pub fn create_raw_dict_from_source( mut source: R, source_size: usize, @@ -261,15 +231,18 @@ pub fn create_raw_dict_from_source( create_raw_dict_from_slice(&all, output, dict_size) } -/// Create a "raw content" dictionary of at most `dict_size` bytes from a -/// corpus already in memory, writing it to `output`. +/// Create a "raw content" dictionary of at most `dict_size` bytes, with no +/// entropy tables, from a corpus already in memory, writing it to `output`. /// -/// The same training as [`create_raw_dict_from_source`], reading `source` in -/// place: a caller that holds the samples anyway does not pay for a second -/// copy of them. +/// The corpus has no sample sizes, so it is cut into at least sixteen samples +/// of at most 128 KiB, and the content is what [`optimize_fastcover_dict`] +/// picks over them at its defaults, as `zstd --train` does. A corpus no larger +/// than `dict_size` is its own content; one the trainer refuses (a +/// `dict_size` under [`TRAINER_DICT_SIZE_MIN`], too little to search) gives +/// its last `dict_size` bytes. /// /// # Errors -/// Returns the error `output` reports while the dictionary is written. +/// Returns an error writing to `output` or allocating the trainer's tables. /// /// # Examples /// ``` @@ -287,83 +260,77 @@ pub fn create_raw_dict_from_slice( output: &mut W, dict_size: usize, ) -> io::Result<()> { - if dict_size == 0 || all.is_empty() { - return Ok(()); - } + output.write_all(&fastcover_raw_content(all, None, dict_size)?) +} - if all.len() < K { - let keep = usize::min(all.len(), dict_size); - output.write_all(&all[all.len() - keep..])?; - return Ok(()); +/// Largest sample an undivided corpus is cut into: upstream zstd's largest +/// block, the most of a sample a dictionary's statistics read. +const RAW_SAMPLE_MAX: usize = 128 << 10; +/// Fewest samples an undivided corpus is cut into, so the search has samples +/// to build from and samples to score on. +const RAW_SAMPLES_MIN: usize = 16; + +/// Raw content of at most `dict_size` bytes from `corpus`, as `zstd --train` +/// picks it: FastCOVER searching `k` over the samples `sizes` cuts, or over +/// even cuts of the corpus when there are none or too few to search. +fn fastcover_raw_content( + corpus: &[u8], + sizes: Option<&[usize]>, + dict_size: usize, +) -> io::Result> { + if corpus.len() <= dict_size { + return Ok(corpus.to_vec()); } - - let source_size = all.len(); - vprintln!("create_dict: creating {dict_size} byte dict from {source_size} byte source"); - - let params = DictParams { segment_size: 2048 }; - let num_segments = usize::max(1, source_size / params.segment_size as usize); - // According to 4. Experiments - Varying Reservoir Sampler Thresholds, - // setting reservoir size to collection size / min{collection size / (2 * number of segments), - // 256} was effective - let denom = usize::max(1, source_size / (2 * num_segments)); - let sample_scale = usize::max(1, usize::min(denom, 256)); - let mut sample_size = source_size / sample_scale; - sample_size = usize::max(sample_size, usize::min(source_size, 16)); - vprintln!("create_dict: creating {sample_size} byte sample of collection"); - let mut sample_reader = all; - let collection_sample = create_sample(&mut sample_reader, sample_size); - - // A collection of segments to be used in the final dictionary. - // - // Contains the best segment from every epoch. - // Reverse is used because we want a min heap, where - // the lowest scoring items come first - let mut pool: BinaryHeap> = BinaryHeap::new(); - let (num_epochs, epoch_size_kmers) = compute_epoch_info(¶ms, dict_size, source_size / K); - // Plain `*`/`+` throughout the epoch walk below: epochs partition the - // training source, so `epoch_size_kmers * K`, `epoch_idx * epoch_size`, and - // `start + epoch_size` are all bounded by the source length (<= isize::MAX) - // and cannot overflow usize. - let epoch_size = usize::max(K, epoch_size_kmers * K); - vprintln!("create_dict: computed epoch info, using {num_epochs} epochs of {epoch_size} bytes"); - let mut epoch_counter = 0; - let mut ctx = Context { - frequencies: HashMap::with_capacity(epoch_size / K), - }; - // Score each segment in each planned epoch and select the highest-scoring - // segment for the pool. Keep exactly `num_epochs` windows to avoid - // emitting more segments than the requested dictionary budget allows. - for epoch_idx in 0..num_epochs { - let start = epoch_idx * epoch_size; - if start >= all.len() { - break; + let refused = + |result: &io::Result>| result.as_ref().err().and_then(TrainingError::of).is_some(); + if let Some(sizes) = sizes { + let trained = search_raw_content(corpus, sizes, dict_size); + if !refused(&trained) { + return trained; } - let end = if epoch_idx + 1 == num_epochs { - all.len() - } else { - usize::min(start + epoch_size, all.len()) - }; - let epoch = &all[start..end]; - epoch_counter += 1; - let best_segment = pick_best_segment(¶ms, &mut ctx, epoch, &collection_sample); - vprintln!( - "\tcreate_dict: epoch {epoch_counter}/{num_epochs} has best segment score {}", - best_segment.score - ); - pool.push(Reverse(best_segment)); - // Wipe frequency list for next epoch - ctx.frequencies.clear(); } - vprintln!( - "create_dict: {epoch_counter} epochs written, writing {} segments", - pool.len() - ); - // Write the dictionary with the highest scoring segment last because - // closer items can be represented with a smaller offset - while let Some(segment) = pool.pop() { - output.write_all(&segment.0.raw)?; + let cut = corpus.len().div_ceil(RAW_SAMPLES_MIN).min(RAW_SAMPLE_MAX); + let even: Vec = corpus.chunks(cut).map(<[u8]>::len).collect(); + let trained = search_raw_content(corpus, &even, dict_size); + if refused(&trained) { + return Ok(corpus[corpus.len() - dict_size..].to_vec()); } - Ok(()) + trained +} + +/// The content of the dictionary [`optimize_fastcover_dict`] picks at its +/// defaults, rebuilt from the parameters it chose: a build is deterministic, +/// so this is the content the winner was finalized from. +fn search_raw_content(corpus: &[u8], sizes: &[usize], dict_size: usize) -> io::Result> { + let (_, chosen) = optimize_fastcover_dict( + corpus, + sizes, + dict_size, + &FastCoverOptions::default(), + FinalizeOptions::default(), + )?; + let set = samples::SampleSet::new(corpus, sizes)?; + let split = set.split(chosen.cover.split_point)?; + let ctx = fastcover::FastCoverContext::new( + &set, + split.train, + chosen.cover.d as usize, + chosen.f, + chosen.accel, + )?; + let mut content = Vec::new(); + // The build fills its buffer from the back; what it returns is the tail. + let kept = ctx + .build( + &mut Vec::new(), + &mut fastcover::WindowCounts::default(), + &mut content, + dict_size, + chosen.cover.k as usize, + )? + .len(); + content.drain(..content.len() - kept); + Ok(content) } /// The `i`th of [`MAX_HUFFMAN_STATS_BYTES`] samples spread evenly over `len` diff --git a/zstd/src/dictionary/reservoir.rs b/zstd/src/dictionary/reservoir.rs deleted file mode 100644 index 06a79b0b9..000000000 --- a/zstd/src/dictionary/reservoir.rs +++ /dev/null @@ -1,111 +0,0 @@ -use super::lmc::K; -use alloc::vec::Vec; -use core::f64::consts::E; -use fastrand; -use std::{io, vec}; - -/// Creates a representative sample of `input` of `size` bytes. -pub fn create_sample(input: &mut R, size: usize) -> Vec { - let reservoir = Reservoir::new(size); - reservoir.fill(input) -} - -/// A reservoir is created from an input stream. -/// -/// Once filled, it will contain a best effort sample of a dataset, -/// where each input value has an equivalent probability of being included. -struct Reservoir { - /// Where the sampled data is stored. - /// - /// Once the lake is filled, then this should contain a representative sample - /// of the larger dataset. - lake: Vec, - /// K is the size of each sample. - /// - /// The original Zstd dictionary implementation states that values - /// between 16 and 2048+ are reasonable. - k: u16, -} - -impl Reservoir { - /// Initialize a new empty reservoir, creating an allocation of `size`. - pub fn new(size: usize) -> Self { - assert!(size >= 16, "Reservoirs cannot be below 16 bytes in size"); - let lake: Vec = vec![0; size]; - let k = K as u16; - Self { lake, k } - } - - /// Filling the reservoir is performed using Algorithm L. - /// - /// The return value is the populated reservoir. - pub fn fill(mut self, source: &mut R) -> Vec { - // https://en.wikipedia.org/wiki/Reservoir_sampling#:~:text=end%0A%20%20end%0Aend-,Optimal%3A%20Algorithm,-L%5Bedit - // https://richardstartin.github.io/posts/reservoir-sampling#algorithm-l:~:text=%3B%0A%20%20%20%20%7D%0A%7D-,Algorithm%20L,-Algorithm%20L%20was - - // First fill the reservoir with the start of the input stream - let mut total_bytes_read: usize = 0; - while let Ok(num_bytes) = source.read(self.lake.as_mut_slice()) { - total_bytes_read += num_bytes; - // Stop when we've completely filled the buffer - if total_bytes_read == self.lake.len() { - break; - } - // If we haven't filled the lake all the way, resize it - if num_bytes == 0 { - self.lake.resize(total_bytes_read, 0); - } - } - - let mut threshold = E.powf(fastrand::f64().ln() / f64::from(self.k)); - // An index into the stream of the next sample to take - let mut next = self.lake.len(); - // Because we're sampling k-mers of size K into the lake, - // split the lake into chunks of k size for simplicity - let mut lake_chunks = self - .lake - .chunks_mut(self.k as usize) - .collect::>(); - // Used when discarding chunks - let end_of_lake = lake_chunks.len(); - let mut counter = end_of_lake / self.k as usize; - // Algorithm L is considered better than algorithm R because it - // determines how many inputs can be skipped, rather than - // processing every input. - - // This is done by abusing the statistics in ways - // I do not understand. - - // Items with a weight smaller than the threshold enter the lake, - // replacing the item in the lake with the largest threshold - - let mut dumpster = Vec::with_capacity(self.k as usize); - loop { - // `num_bytes_read` is kept track of to watch for EOD. - let num_bytes_read: u64; - if counter == next { - num_bytes_read = source - .read(lake_chunks[fastrand::usize(0..end_of_lake)]) - .unwrap() as u64; - // Advance at least to the next sample, skipping forward a few samples - next += ((fastrand::f64().ln() / f64::ln(1.0 - threshold)).floor() as usize + 1) - * self.k as usize; - // Update the threshold to reflect changes - threshold *= E.powf(fastrand::f64().ln() / f64::from(end_of_lake as u32)) - } else { - // Drop the next chunk - num_bytes_read = source.read(&mut dumpster).unwrap() as u64; - //source.seek_relative(self.k.into()).unwrap(); - } - if num_bytes_read == 0 { - break; - } - counter += self.k as usize; - } - self.lake.shrink_to_fit(); - self.lake - } -} - -#[cfg(test)] -mod tests; diff --git a/zstd/src/dictionary/reservoir/tests.rs b/zstd/src/dictionary/reservoir/tests.rs deleted file mode 100644 index 788200f6a..000000000 --- a/zstd/src/dictionary/reservoir/tests.rs +++ /dev/null @@ -1,32 +0,0 @@ -use super::Reservoir; -use alloc::vec; - -#[test] -fn initial_fill() { - // Create a reservoir 16 bytes in size and read - // 16 bytes into it - let r = Reservoir::new(16); - let test_data = vec![0_u8; 16]; - let output = r.fill(&mut test_data.as_slice()); - assert_eq!(test_data, output); -} - -#[test] -fn shrinks_for_small_sample() { - // Create a reservoir larger than the sample. - // The output should be smaller. - let r = Reservoir::new(32); - let test_data = vec![0_u8; 28]; - let output = r.fill(&mut test_data.as_slice()); - assert!(output.len() == 28); -} - -#[test] -fn lake_doesnt_grow() { - // Create a sample larger than the reservoir - // The output should be smaller. - let r = Reservoir::new(32); - let test_data = vec![0_u8; 16_000_000]; - let output = r.fill(&mut test_data.as_slice()); - assert!(output.len() == 32); -} diff --git a/zstd/src/dictionary/tests.rs b/zstd/src/dictionary/tests.rs index 8d831b623..d871aea49 100644 --- a/zstd/src/dictionary/tests.rs +++ b/zstd/src/dictionary/tests.rs @@ -678,6 +678,79 @@ fn create_raw_dict_from_source_never_exceeds_requested_size() { ); } +/// Raw content is what `zstd --train`'s search picks over the corpus cut into +/// samples: the dictionary the search finalizes ends with it (the finalizer +/// only cuts the front to make room for its header). +#[test] +fn raw_content_is_the_fastcover_search_winner() { + let corpus = training_data(); + let dict_size = 4096; + let mut raw = Vec::new(); + create_raw_dict_from_slice(&corpus, &mut raw, dict_size).unwrap(); + assert!(!raw.is_empty() && raw.len() <= dict_size); + + let cut = corpus.len().div_ceil(RAW_SAMPLES_MIN).min(RAW_SAMPLE_MAX); + let sizes: Vec = corpus.chunks(cut).map(<[u8]>::len).collect(); + let (dict, _) = optimize_fastcover_dict( + &corpus, + &sizes, + dict_size, + &FastCoverOptions::default(), + FinalizeOptions::default(), + ) + .unwrap(); + let finalized = Dictionary::decode_dict(&dict).unwrap(); + assert!(raw.ends_with(&finalized.dict_content)); +} + +/// A corpus no larger than the dictionary is all of its content. +#[test] +fn a_corpus_that_fits_is_its_own_raw_content() { + let corpus = &training_data()[..3000]; + let mut raw = Vec::new(); + create_raw_dict_from_slice(corpus, &mut raw, 4096).unwrap(); + assert_eq!(raw, corpus); +} + +/// A directory trains one sample per file, and its raw dictionary shortens a +/// small frame of the same kind of data. +#[test] +fn raw_content_from_a_directory_helps_a_small_frame() { + let dir = std::env::temp_dir().join(std::format!("szstd-raw-dir-{}", std::process::id())); + fs::create_dir_all(&dir).unwrap(); + for file in 0..40u32 { + let mut body = std::string::String::new(); + for line in 0..30u32 { + body.push_str(&std::format!( + "tenant=demo table=orders key={} region=eu status=shipped\n", + file * 30 + line + )); + } + fs::write(dir.join(std::format!("s{file:02}")), body).unwrap(); + } + let mut raw = Vec::new(); + let trained = create_raw_dict_from_dir(&dir, &mut raw, 2048); + fs::remove_dir_all(&dir).unwrap(); + trained.unwrap(); + assert!(!raw.is_empty() && raw.len() <= 2048); + + let frame = b"tenant=demo table=orders key=77 region=eu status=shipped\n"; + let compress = |dict: Option<&[u8]>| { + let mut out = Vec::new(); + let mut compressor = FrameCompressor::new(CompressionLevel::Default); + if let Some(dict) = dict { + compressor + .set_dictionary(Dictionary::from_raw_content(7, dict.to_vec()).unwrap()) + .unwrap(); + } + compressor.set_source(&frame[..]); + compressor.set_drain(&mut out); + compressor.compress(); + out.len() + }; + assert!(compress(Some(&raw)) < compress(None)); +} + /// The entropy tables a search builds once finalize every candidate into the /// same bytes `finalize_raw_dict` writes for it, including the content cut /// to fit and an explicit id; and where the samples are too thin to decide the diff --git a/zstd/src/encoding/frame_compressor/tests.rs b/zstd/src/encoding/frame_compressor/tests.rs index 16ede6f83..c402bc59d 100644 --- a/zstd/src/encoding/frame_compressor/tests.rs +++ b/zstd/src/encoding/frame_compressor/tests.rs @@ -428,16 +428,13 @@ fn dictionary_compression_roundtrips_with_dict_builder_dictionary() { let decoder_dict = crate::decoding::Dictionary::from_raw_content(dict_id, raw_dict.clone()).unwrap(); - // Payload the trained dictionary genuinely covers: training lines in a - // permuted order, so whole lines are dictionary matches while the - // payload does not simply repeat its previous line. (A payload of - // unseen `key=` values is NOT such a case: every line then matches the - // previous one through the repeat offset, and upstream too compresses - // it smaller WITHOUT the dictionary, the dictionary id and the - // dictionary-tier parameters costing more than the first line's - // literals.) + // A small frame of training lines in a permuted order: the case a + // dictionary is for, where the frame alone has too little history to + // match against. (A long frame of such lines is NOT such a case: past the + // first few lines every line matches the previous one through the repeat + // offset, and upstream too compresses it smaller WITHOUT the dictionary.) let mut payload = Vec::new(); - for i in 0..96u32 { + for i in 0..12u32 { let idx = (i * 37 + 11) % 256; payload.extend_from_slice( format!("tenant=demo table=orders key={idx} region=eu\n").as_bytes(), From e8b6599c6adb2e45a50cf74497cb7deff25573b9 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Mon, 28 Sep 2026 04:01:55 +0300 Subject: [PATCH 02/20] feat(dict)!: entropy tables measured from the samples The finalizer built its tables from the samples' raw bytes, counting byte values modulo each alphabet as if they were sequence codes. Such tables made a dictionary compress its own samples worse than its bare content did. - finalize_raw_dict follows upstream's ZDICT_analyzeEntropy: the first block of each sample is compressed with the content as a raw dictionary, and the literals, literal/match lengths and offset codes those blocks hold are counted; tables at upstream's logs (literals up to 11 bits, offsets 8, lengths 9), a mostly flat stand-in when literals are flat - blocks written raw teach nothing and are left out, as upstream does - the literals table is scaled to its longest code, as upstream writes it - FinalizeOptions carries the level (upstream's zParams.compressionLevel), used both for the analysis and for scoring candidates; CoverOptions no longer has one - MIN_TRAINED_DICT_SIZE is upstream's ZDICT_DICTSIZE_MIN, 256 - ZDICT_finalizeDictionary passes the sample sizes and honours compressionLevel Measured on 321 repository files, 16 KiB dictionaries, compressed with upstream zstd -3: the dictionary gap to upstream's went from +2.3..3.0% to +0.09..0.6% across --train, fixed COVER and fixed FastCOVER. BREAKING CHANGE: finalize_raw_dict takes the sample sizes; FinalizeOptions gains `level` and CoverOptions loses it; MIN_TRAINED_DICT_SIZE is 256. Part of #128 --- c-api/src/dict.rs | 64 ++-- zstd/src/bin/structured-zstd/main.rs | 38 +-- zstd/src/bin/structured-zstd/tests.rs | 50 +-- zstd/src/dictionary/finalize.rs | 419 +++++++++++++++++++++++++ zstd/src/dictionary/finalize/tests.rs | 150 +++++++++ zstd/src/dictionary/legacy.rs | 2 +- zstd/src/dictionary/mod.rs | 350 ++++----------------- zstd/src/dictionary/samples.rs | 7 +- zstd/src/dictionary/selection.rs | 32 +- zstd/src/dictionary/tests.rs | 133 +++----- zstd/src/encoding/blocks/compressed.rs | 8 +- zstd/src/encoding/blocks/mod.rs | 5 + zstd/src/encoding/frame_compressor.rs | 6 + zstd/src/fse/fse_encoder.rs | 29 +- zstd/src/huff0/huff0_encoder.rs | 9 + 15 files changed, 805 insertions(+), 497 deletions(-) create mode 100644 zstd/src/dictionary/finalize.rs create mode 100644 zstd/src/dictionary/finalize/tests.rs diff --git a/c-api/src/dict.rs b/c-api/src/dict.rs index d5ede57b7..69476ecee 100644 --- a/c-api/src/dict.rs +++ b/c-api/src/dict.rs @@ -163,7 +163,6 @@ impl ZDICT_cover_params_t { self.splitPoint, self.shrinkDict, self.shrinkDictMaxRegression, - self.zParams.compressionLevel, ) } @@ -181,7 +180,6 @@ fn cover_options( split_point: f64, shrink: c_uint, max_regression: c_uint, - level: c_int, ) -> CoverOptions { CoverOptions { k, @@ -189,13 +187,15 @@ fn cover_options( steps, split_point, shrink: (shrink != 0).then_some(max_regression), - level, } } +/// `ZDICT_params_t` as the codec takes it: the id (zero derives one) and the +/// level the dictionary is analysed and scored at (zero is the default). fn finalize_options(params: &ZDICT_params_t) -> FinalizeOptions { FinalizeOptions { dict_id: (params.dictID != 0).then_some(params.dictID), + level: params.compressionLevel, } } @@ -330,7 +330,6 @@ impl ZDICT_fastCover_params_t { self.splitPoint, self.shrinkDict, self.shrinkDictMaxRegression, - self.zParams.compressionLevel, ), f: self.f, accel: self.accel, @@ -436,14 +435,12 @@ pub unsafe extern "C" fn ZDICT_optimizeTrainFromBuffer_fastCover( /// const void* dictContent, size_t dictContentSize, const void* samplesBuffer, /// const size_t* samplesSizes, unsigned nbSamples, ZDICT_params_t parameters)`. /// -/// Wraps raw `dictContent` (plus entropy tables analysed from the samples) into -/// a full zstd dictionary, writing up to `maxDictSize` bytes into -/// `dstDictBuffer`. Returns the dictionary size or an error code. -/// -/// Of `parameters`, only `dictID` is honoured (0 derives a compliant ID). The -/// FastCOVER finalizer builds the entropy tables directly from the samples, so -/// `compressionLevel` does not tune them, and `notificationLevel` (builder -/// verbosity) has no effect here; both are accepted for ABI compatibility. +/// Wraps raw `dictContent` into a full zstd dictionary, writing up to +/// `maxDictSize` bytes into `dstDictBuffer`, with entropy tables measured on +/// the samples compressed with the content at `compressionLevel` (zero is the +/// default level). `dictID` zero derives a compliant id; `notificationLevel` +/// (builder verbosity) has no effect here. Returns the dictionary size or an +/// error code. /// /// # Safety /// All buffers valid for their stated lengths; `samplesSizes` valid for @@ -460,36 +457,29 @@ pub unsafe extern "C" fn ZDICT_finalizeDictionary( nb_samples: c_uint, parameters: ZDICT_params_t, ) -> usize { - let Some(total) = (unsafe { total_sample_len(samples_sizes, nb_samples) }) else { - return encode(ZSTD_ErrorCode::ZSTD_error_dictionaryCreation_failed); - }; // NULL + non-zero length is caller error, not a slice to build. - if (samples_buffer.is_null() && total > 0) - || (dict_content.is_null() && dict_content_size > 0) - || (dst_dict_buffer.is_null() && max_dict_size > 0) - { + if dict_content.is_null() && dict_content_size > 0 { return encode(ZSTD_ErrorCode::ZSTD_error_dictionaryCreation_failed); } let content = unsafe { in_slice(dict_content, dict_content_size) }; - let samples = unsafe { in_slice(samples_buffer, total) }; - // dictID 0 means "derive a compliant id"; any non-zero value is forced. - let finalize = FinalizeOptions { - dict_id: (parameters.dictID != 0).then_some(parameters.dictID), - }; - - let outcome = catch_unwind(AssertUnwindSafe(|| { - codec::dictionary::finalize_raw_dict(content, samples, max_dict_size, finalize) - })); - let dict = match outcome { - Ok(Ok(dict)) => dict, - _ => return encode(ZSTD_ErrorCode::ZSTD_error_dictionaryCreation_failed), - }; - if dict.len() > max_dict_size { - return encode(ZSTD_ErrorCode::ZSTD_error_dstSize_tooSmall); + unsafe { + train_into( + dst_dict_buffer, + max_dict_size, + samples_buffer, + samples_sizes, + nb_samples, + |samples, sizes| { + codec::dictionary::finalize_raw_dict( + content, + samples, + sizes, + max_dict_size, + finalize_options(¶meters), + ) + }, + ) } - let out = unsafe { crate::ffi::out_slice(dst_dict_buffer, max_dict_size) }; - out[..dict.len()].copy_from_slice(&dict); - dict.len() } /// `unsigned ZDICT_getDictID(const void* dictBuffer, size_t dictSize)` — the diff --git a/zstd/src/bin/structured-zstd/main.rs b/zstd/src/bin/structured-zstd/main.rs index 09463af60..d0c3aad76 100644 --- a/zstd/src/bin/structured-zstd/main.rs +++ b/zstd/src/bin/structured-zstd/main.rs @@ -3288,10 +3288,10 @@ fn train_dictionary(opts: &Options) -> Result<()> { // real bound is higher and depends on the entropy tables the corpus // produces, which is why the trainer still checks; this only settles the // part that is knowable without reading anything. - if opts.max_dict < structured_zstd::dictionary::TRAINER_DICT_SIZE_MIN { + if opts.max_dict < structured_zstd::dictionary::MIN_TRAINED_DICT_SIZE { bail!( "--maxdict must be at least {} bytes", - structured_zstd::dictionary::TRAINER_DICT_SIZE_MIN + structured_zstd::dictionary::MIN_TRAINED_DICT_SIZE ); } // Whether the trainer takes the tuning it was given is a question about the @@ -3303,8 +3303,8 @@ fn train_dictionary(opts: &Options) -> Result<()> { Legacy, } let plan = match opts.trainer { - Trainer::FastCover => Plan::FastCover(fastcover_options(&opts.trainer_params, opts.level)?), - Trainer::Cover => Plan::Cover(cover_options(&opts.trainer_params, opts.level)?), + Trainer::FastCover => Plan::FastCover(fastcover_options(&opts.trainer_params)?), + Trainer::Cover => Plan::Cover(cover_options(&opts.trainer_params)?), // The legacy trainer is tuned by selectivity alone; a cover tuning list // given before `--train-legacy` names a trainer that no longer runs, // as it does in the reference. @@ -3371,8 +3371,11 @@ fn train_dictionary(opts: &Options) -> Result<()> { } } + // The dictionary is built for the command's level, as the reference's is + // for its `-#` (zstdcli.c, `dictCLevel`). let finalize = FinalizeOptions { dict_id: opts.dict_id, + level: opts.level, }; // Every trainer counts samples, so they are loaded as the reference's // command loads them: shuffled, capped per file, and cut by `-B`. The same @@ -3722,17 +3725,12 @@ fn place_trained_dictionary(output: &Path, dict: &[u8], samples: &[fs::Metadata] placed } -/// The COVER tuning `--train-cover=...` asked for, scored at `level` as the -/// reference's command scores at its `-#` (zstdcli.c, `dictCLevel`). A listed -/// tuning starts from zero, a bare flag from the command's defaults. `split` -/// is a percentage and `k` at least `d` (the most a search over `d` tries, -/// 8, when `d` is not given); both are settled here, before any sample is -/// read. -fn cover_options( - params: &TrainerParams, - level: i32, -) -> Result { - cover_tuning(params, level, "--train-cover") +/// The COVER tuning `--train-cover=...` asked for. A listed tuning starts from +/// zero, a bare flag from the command's defaults. `split` is a percentage and +/// `k` at least `d` (the most a search over `d` tries, 8, when `d` is not +/// given); both are settled here, before any sample is read. +fn cover_options(params: &TrainerParams) -> Result { + cover_tuning(params, "--train-cover") } /// The split a tuning report shows: the percentage `split=` was given as. @@ -3744,7 +3742,6 @@ fn split_percent(split_point: f64) -> u32 { fn cover_tuning( params: &TrainerParams, - level: i32, flag: &str, ) -> Result { use structured_zstd::dictionary::CoverOptions; @@ -3756,13 +3753,9 @@ fn cover_tuning( steps: 0, split_point: 0.0, shrink: None, - level, } } else { - CoverOptions { - level, - ..CoverOptions::default() - } + CoverOptions::default() }; if let Some(split) = params.split_percent { if split > 100 { @@ -3788,7 +3781,6 @@ fn cover_tuning( /// in `1..=10`, and the rest as [`cover_options`] checks it. fn fastcover_options( params: &TrainerParams, - level: i32, ) -> Result { use structured_zstd::dictionary::FastCoverOptions; @@ -3800,7 +3792,7 @@ fn fastcover_options( bail!("{FLAG} d must be 6 or 8, got {d}"); } let mut options = FastCoverOptions { - cover: cover_tuning(params, level, FLAG)?, + cover: cover_tuning(params, FLAG)?, ..FastCoverOptions::default() }; if params.listed { diff --git a/zstd/src/bin/structured-zstd/tests.rs b/zstd/src/bin/structured-zstd/tests.rs index b47553217..0c64b470c 100644 --- a/zstd/src/bin/structured-zstd/tests.rs +++ b/zstd/src/bin/structured-zstd/tests.rs @@ -270,7 +270,10 @@ fn an_archive_built_from_two_dictionaries_names_neither() { &[4_000; 10], 8 * 1024, &options, - FinalizeOptions { dict_id: Some(id) }, + FinalizeOptions { + dict_id: Some(id), + ..FinalizeOptions::default() + }, ) .expect("training the fixture dictionary must succeed"); let mut out = Vec::new(); @@ -729,7 +732,7 @@ fn the_reported_split_is_the_percentage_given() { split_percent: Some(percent), ..TrainerParams::default() }; - let options = cover_tuning(¶ms, 3, "--train-cover").unwrap(); + let options = cover_tuning(¶ms, "--train-cover").unwrap(); assert_eq!(split_percent(options.split_point), percent); } } @@ -4811,35 +4814,23 @@ fn trainer_parameters_parse_and_build_options() { assert!(parse_trainer_params("k=x", true).is_err(), "not a number"); assert!(parse_trainer_params("zzz=1", true).is_err(), "unknown key"); - let fixed = fastcover_options(¶ms, 5).unwrap(); + let fixed = fastcover_options(¶ms).unwrap(); assert_eq!( (fixed.cover.k, fixed.cover.d, fixed.f, fixed.accel), (200, 8, 20, 2) ); - assert_eq!( - ( - fixed.cover.steps, - fixed.cover.split_point, - fixed.cover.level - ), - (4, 0.75, 5) - ); - let shrunk = cover_options( - &parse_trainer_params("k=64,d=8,shrink=3", false).unwrap(), - 3, - ) - .unwrap(); + assert_eq!((fixed.cover.steps, fixed.cover.split_point), (4, 0.75)); + let shrunk = cover_options(&parse_trainer_params("k=64,d=8,shrink=3", false).unwrap()).unwrap(); assert_eq!((shrunk.k, shrunk.d, shrunk.shrink), (64, 8, Some(3))); - let bad = - |text: &str| fastcover_options(&parse_trainer_params(text, true).unwrap(), 3).is_err(); + let bad = |text: &str| fastcover_options(&parse_trainer_params(text, true).unwrap()).is_err(); assert!(bad("d=7")); assert!(bad("f=32")); assert!(bad("accel=11")); assert!(bad("k=4,d=8")); assert!(bad("k=6"), "a search over d reaches 8"); assert!(bad("split=101")); - assert!(cover_options(&parse_trainer_params("k=4,d=8", false).unwrap(), 3).is_err()); + assert!(cover_options(&parse_trainer_params("k=4,d=8", false).unwrap()).is_err()); let opts = parse(&["--train-fastcover=k=200,d=8", "s1"]).unwrap(); assert_eq!(opts.mode, Mode::Train); @@ -4908,7 +4899,7 @@ fn a_listed_tuning_starts_from_zero() { use structured_zstd::dictionary::{CoverOptions, FastCoverOptions}; let zeros = parse_trainer_params("k=0,d=0,f=0,steps=0,split=0,accel=0", true).unwrap(); - let options = fastcover_options(&zeros, 3).expect("zero knobs are valid"); + let options = fastcover_options(&zeros).expect("zero knobs are valid"); assert_eq!( (options.cover.k, options.cover.d, options.cover.steps), (0, 0, 0) @@ -4917,24 +4908,13 @@ fn a_listed_tuning_starts_from_zero() { (options.f, options.accel, options.cover.split_point), (0, 0, 0.0) ); - let bare = fastcover_options(&TrainerParams::default(), 3).unwrap(); assert_eq!( - bare, - FastCoverOptions { - cover: CoverOptions { - level: 3, - ..FastCoverOptions::default().cover - }, - ..FastCoverOptions::default() - } + fastcover_options(&TrainerParams::default()).unwrap(), + FastCoverOptions::default() ); - let bare_cover = cover_options(&TrainerParams::default(), 3).unwrap(); assert_eq!( - bare_cover, - CoverOptions { - level: 3, - ..CoverOptions::default() - } + cover_options(&TrainerParams::default()).unwrap(), + CoverOptions::default() ); // A zero beside a real value leaves that value in force. let mixed = parse_trainer_params("k=0,d=6", true).unwrap(); diff --git a/zstd/src/dictionary/finalize.rs b/zstd/src/dictionary/finalize.rs new file mode 100644 index 000000000..9eafcd4da --- /dev/null +++ b/zstd/src/dictionary/finalize.rs @@ -0,0 +1,419 @@ +//! Finalizing raw dictionary content into a dictionary (upstream zstd +//! `zdict.c`, `ZDICT_finalizeDictionary` and `ZDICT_analyzeEntropy`). +//! +//! The entropy tables a dictionary carries seed the first block of every frame +//! compressed with it, so they are drawn from what such blocks actually hold: +//! each sample's first block is compressed with the content as a raw +//! dictionary, and the literals and sequence codes it produces are counted. + +use super::FinalizeOptions; +use super::samples::{SampleSet, TrainingError, refuse}; +use crate::bit_io::BitWriter; +use crate::decoding::dictionary::MAGIC_NUM as DICT_MAGIC_NUM; +use crate::encoding::blocks::{ + encode_literal_length, encode_match_len, encode_offset, encode_offset_with_history, +}; +use crate::encoding::workspace::Workspace; +use crate::encoding::{ + CompressionLevel, DictionarySizes, EncoderDictionary, FrameCompressor, HistoryBuf, + MatchGeneratorDriver, Matcher, Sequence, +}; +use crate::fse::fse_encoder::{FSETable, build_table_at_log}; +use crate::huff0::HuffmanTable as HuffmanDecoderTable; +use crate::huff0::huff0_encoder::{HuffmanEncoder, HuffmanTable}; +use std::{format, io, vec, vec::Vec}; + +/// Smallest dictionary finalized (upstream zstd `ZDICT_DICTSIZE_MIN`). +pub(super) const DICT_SIZE_MIN: usize = 256; +/// Repeat offsets a dictionary starts frames with (upstream `repStartValue`). +/// The content has to reach the largest of them. +const START_REPS: [u32; 3] = [1, 4, 8]; +const MIN_CONTENT_SIZE: usize = 8; +/// Longest literals code (upstream `ZDICT_analyzeEntropy`, `huffLog`). +const HUF_MAX_BITS: usize = 11; +/// Table logs of the three sequence streams (upstream `OffFSELog`, +/// `MLFSELog`, `LLFSELog`). +const OF_LOG: u8 = 8; +const ML_LOG: u8 = 9; +const LL_LOG: u8 = 9; +/// Highest offset code a first block can need (upstream `OFFCODE_MAX`). +const OFFCODE_MAX: u32 = 30; +/// Most of a sample the statistics read: its first block. +const BLOCK_SIZE_MAX: usize = 128 << 10; + +/// Finalize `content` into a dictionary of at most `dict_size` bytes, drawing +/// the entropy tables from the first `count` samples. Content that does not +/// fit after the header is cut from the front, keeping the best segments. +pub(super) fn finalize( + content: &[u8], + samples: &SampleSet<'_>, + count: usize, + dict_size: usize, + options: FinalizeOptions, +) -> io::Result> { + if content.is_empty() { + return Err(invalid("raw dictionary content must not be empty")); + } + if dict_size < DICT_SIZE_MIN { + return Err(too_small()); + } + let dict_id = options.dict_id.unwrap_or_else(|| derive_dict_id(content)); + if dict_id == 0 { + return Err(invalid("dictionary id must be non-zero")); + } + let mut out = Vec::with_capacity(dict_size); + out.extend_from_slice(&DICT_MAGIC_NUM); + out.extend_from_slice(&dict_id.to_le_bytes()); + analyze_entropy(&mut out, content, samples, count, options.level)?; + for rep in START_REPS { + out.extend_from_slice(&rep.to_le_bytes()); + } + let room = match dict_size.checked_sub(out.len()) { + Some(room) if room >= MIN_CONTENT_SIZE => room, + _ => return Err(too_small()), + }; + let content = &content[content.len() - content.len().min(room)..]; + // Zeros ahead of short content: the last byte is the best position. + if content.len() < MIN_CONTENT_SIZE { + out.resize(out.len() + MIN_CONTENT_SIZE - content.len(), 0); + } + out.extend_from_slice(content); + Ok(out) +} + +fn too_small() -> io::Error { + refuse( + TrainingError::DictionaryTooSmall, + "dictionary size too small to fit header and offset history", + ) +} + +/// A refusal upstream reports as a failed creation, with no finer cause. +fn invalid(reason: &str) -> io::Error { + io::Error::new(io::ErrorKind::InvalidInput, reason) +} + +fn derive_dict_id(raw_content: &[u8]) -> u32 { + let mut h = 0xcbf2_9ce4_8422_2325u64; + for &b in raw_content { + h ^= u64::from(b); + h = h.wrapping_mul(0x0100_0000_01b3); + } + // Upstream's compliant range: ids below 32768 are reserved. + ((h % ((1u64 << 31) - 32768)) + 32768) as u32 +} + +/// Literal and sequence-code counts, every symbol starting at one so each can +/// be described (upstream zstd `ZDICT_analyzeEntropy`). +struct EntropyCounts { + literals: [usize; 256], + offset_codes: Vec, + match_lengths: [usize; 53], + literal_lengths: [usize; 36], +} + +/// Append the literals table, the three sequence tables and nothing else to +/// `out`, from the statistics of the first `count` samples compressed with +/// `content` as a raw dictionary at `level`. +fn analyze_entropy( + out: &mut Vec, + content: &[u8], + samples: &SampleSet<'_>, + count: usize, + level: i32, +) -> io::Result<()> { + // The largest offset a first block can reach sets the offset alphabet. + let reach = content.len() as u64 + BLOCK_SIZE_MAX as u64; + let offcode_max = reach.ilog2(); + if offcode_max > OFFCODE_MAX { + return Err(invalid( + "dictionary content is too large to describe its offsets", + )); + } + let mut counts = EntropyCounts { + literals: [1; 256], + offset_codes: vec![1; offcode_max as usize + 1], + match_lengths: [1; 53], + literal_lengths: [1; 36], + }; + count_samples(&mut counts, content, samples, count, level)?; + + let mut literals = literals_table(&counts.literals); + if literals.table_log() == 8 { + // Every symbol at eight bits describes nothing and cannot be written; + // a mostly flat distribution that still compresses stands in + // (upstream zstd `ZDICT_flatLit`). + counts.literals = [2; 256]; + counts.literals[0] = 4; + counts.literals[253] = 1; + counts.literals[254] = 1; + literals = literals_table(&counts.literals); + debug_assert_eq!(literals.table_log(), 9); + } + out.extend_from_slice(&huffman_description(&literals)?); + for table in [ + build_table_at_log(&counts.offset_codes, OF_LOG), + build_table_at_log(&counts.match_lengths, ML_LOG), + build_table_at_log(&counts.literal_lengths, LL_LOG), + ] { + out.extend_from_slice(&fse_description(&table)); + } + Ok(()) +} + +/// Compress the first block of each of the first `count` samples with +/// `content` as a raw dictionary and count what the compressed blocks hold. +fn count_samples( + counts: &mut EntropyCounts, + content: &[u8], + samples: &SampleSet<'_>, + count: usize, + level: i32, +) -> io::Result<()> { + let dictionary = crate::decoding::Dictionary::from_raw_content(0, content.to_vec()) + .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?; + let recorder = Recorder { + inner: MatchGeneratorDriver::new(BLOCK_SIZE_MAX, 1), + literals: [0; 256], + sequences: Vec::new(), + }; + let mut compressor: FrameCompressor<&[u8], Vec, Recorder> = + FrameCompressor::new_with_matcher(recorder, CompressionLevel::from_level(level)); + compressor + .set_encoder_dictionary(EncoderDictionary::from_dictionary(dictionary)) + .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?; + let mut frame = Vec::new(); + for index in 0..count { + let sample = samples.sample(index); + let block = &sample[..sample.len().min(BLOCK_SIZE_MAX)]; + if block.is_empty() { + continue; + } + frame.clear(); + compressor.set_source(block); + compressor.set_drain(frame); + compressor.set_source_size_hint(block.len() as u64); + compressor.compress(); + frame = compressor.take_drain().expect("the drain was set above"); + let recorder = compressor.matcher_mut(); + // A block written raw holds no sequences to learn from, and its bytes + // are no literals of any compressed block (upstream skips it too). + if first_block_is_compressed(&frame) { + for (total, &seen) in counts.literals.iter_mut().zip(&recorder.literals) { + *total += seen as usize; + } + let mut reps = START_REPS; + for &(ll, offset, ml) in &recorder.sequences { + let off_base = encode_offset_with_history(offset, ll, &mut reps); + counts.literal_lengths[encode_literal_length(ll).0 as usize] += 1; + counts.match_lengths[encode_match_len(ml).0 as usize] += 1; + if let Some(slot) = counts + .offset_codes + .get_mut(encode_offset(off_base).0 as usize) + { + *slot += 1; + } + } + } + recorder.literals = [0; 256]; + recorder.sequences.clear(); + } + Ok(()) +} + +/// Whether `frame`'s first block is a compressed one. +fn first_block_is_compressed(frame: &[u8]) -> bool { + let Ok((_, header_len)) = crate::decoding::frame::read_frame_header_with_format(frame, false) + else { + return false; + }; + // Block_Type is bits 1-2 of the block header (RFC 8878 3.1.1.2). + frame + .get(usize::from(header_len)) + .is_some_and(|&byte| (byte >> 1) & 3 == 2) +} + +/// The literals code for `counts`, its weights scaled to its longest code +/// rather than to the length limit, as upstream writes it (`HUF_writeCTable` +/// takes the `maxNbBits` its build returned): the same codes, and the table +/// log a decoder derives from the weights is the one the code needs. +fn literals_table(counts: &[usize; 256]) -> HuffmanTable { + let table = HuffmanTable::build_limited(counts, HUF_MAX_BITS); + let longest = (0..=255u8) + .filter_map(|symbol| table.num_bits_for_symbol(symbol)) + .max() + .map_or(HUF_MAX_BITS, usize::from); + if longest < HUF_MAX_BITS { + HuffmanTable::build_limited(counts, longest) + } else { + table + } +} + +/// The literals table's description as a dictionary stores it. +fn huffman_description(table: &HuffmanTable) -> io::Result> { + // The encoder writes the description ahead of the data it codes; one + // symbol of data, then the decoder reports where the description ends. + let mut writer = BitWriter::new(); + let mut encoder = HuffmanEncoder::new(table, &mut writer); + encoder.encode(&[0], true); + let encoded = writer.dump(); + let mut decoder = HuffmanDecoderTable::new(); + let description_len = decoder + .build_decoder(encoded.as_slice()) + .map_err(|e| io::Error::other(format!("failed to decode generated huffman table: {e}")))?; + Ok(encoded[..description_len as usize].to_vec()) +} + +/// A sequence table's normalized counts, as a dictionary stores them. +fn fse_description(table: &FSETable) -> Vec { + let mut writer = BitWriter::new(); + table.write_table(&mut writer); + writer.dump() +} + +/// The production matcher, recording the literals and sequences each frame's +/// blocks are built from. Everything else is forwarded, so frames compress as +/// any compressor's would, dictionary reuse included. +struct Recorder { + inner: MatchGeneratorDriver, + literals: [u32; 256], + /// Literal length, offset and match length of each sequence. + sequences: Vec<(u32, u32, u32)>, +} + +impl Matcher for Recorder { + fn get_last_space(&mut self) -> &[u8] { + self.inner.get_last_space() + } + + fn fill_in_place( + &mut self, + capacity: usize, + fill: &mut dyn FnMut(&mut HistoryBuf) -> (usize, bool), + ) -> (usize, bool) { + self.inner.fill_in_place(capacity, fill) + } + + fn uncommitted_input(&self) -> &[u8] { + self.inner.uncommitted_input() + } + + fn commit_filled(&mut self, len: usize) { + self.inner.commit_filled(len); + } + + fn skip_matching(&mut self) { + self.inner.skip_matching(); + } + + fn skip_matching_with_hint(&mut self, incompressible_hint: Option) { + self.inner.skip_matching_with_hint(incompressible_hint); + } + + fn start_matching(&mut self, mut handle_sequence: impl for<'a> FnMut(Sequence<'a>)) { + let Self { + inner, + literals, + sequences, + } = self; + inner.start_matching(|sequence| { + match &sequence { + Sequence::Literals { literals: bytes } => { + for &byte in *bytes { + literals[usize::from(byte)] += 1; + } + } + Sequence::Triple { + literals: bytes, + offset, + match_len, + } => { + for &byte in *bytes { + literals[usize::from(byte)] += 1; + } + // Lengths and offsets of one block, far below `u32::MAX`. + sequences.push((bytes.len() as u32, *offset as u32, *match_len as u32)); + } + } + handle_sequence(sequence); + }); + } + + fn reset(&mut self, level: CompressionLevel) { + self.inner.reset(level); + } + + fn reset_in_workspace(&mut self, level: CompressionLevel, workspace: &mut Workspace) { + self.inner.reset_in_workspace(level, workspace); + } + + fn leave_workspace(&mut self) { + self.inner.leave_workspace(); + } + + fn set_source_size_hint(&mut self, size: u64) { + self.inner.set_source_size_hint(size); + } + + fn set_dictionary_size_hint(&mut self, sizes: DictionarySizes) { + self.inner.set_dictionary_size_hint(sizes); + } + + fn clear_param_overrides(&mut self) { + self.inner.clear_param_overrides(); + } + + fn prime_with_dictionary(&mut self, dict_content: &[u8], offset_hist: [u32; 3]) { + self.inner.prime_with_dictionary(dict_content, offset_hist); + } + + fn dictionary_is_resident(&self) -> bool { + self.inner.dictionary_is_resident() + } + + fn reapply_resident_dictionary(&mut self, offset_hist: [u32; 3]) { + self.inner.reapply_resident_dictionary(offset_hist); + } + + fn restore_primed_dictionary(&mut self, level: CompressionLevel) -> bool { + self.inner.restore_primed_dictionary(level) + } + + fn capture_primed_dictionary(&mut self, level: CompressionLevel) { + self.inner.capture_primed_dictionary(level); + } + + fn invalidate_primed_dictionary(&mut self) { + self.inner.invalidate_primed_dictionary(); + } + + fn seed_dictionary_entropy( + &mut self, + huff: Option<&HuffmanTable>, + ll: Option<&FSETable>, + ml: Option<&FSETable>, + of: Option<&FSETable>, + ) { + self.inner.seed_dictionary_entropy(huff, ll, ml, of); + } + + fn supports_dictionary_priming(&self) -> bool { + self.inner.supports_dictionary_priming() + } + + fn block_samples_match_dict(&self, block: &[u8]) -> bool { + self.inner.block_samples_match_dict(block) + } + + fn heap_size(&self) -> usize { + self.inner.heap_size() + } + + fn window_size(&self) -> u64 { + self.inner.window_size() + } +} + +#[cfg(test)] +mod tests; diff --git a/zstd/src/dictionary/finalize/tests.rs b/zstd/src/dictionary/finalize/tests.rs new file mode 100644 index 000000000..90ee510b4 --- /dev/null +++ b/zstd/src/dictionary/finalize/tests.rs @@ -0,0 +1,150 @@ +use super::*; +use crate::encoding::compress_slice_to_vec; +use std::string::String; + +/// Log lines with a shared vocabulary, one sample each. +fn log_samples() -> (Vec, Vec) { + let mut data = Vec::new(); + let mut sizes = Vec::new(); + for i in 0..400u32 { + let line = format!( + "ts=2026-09-{:02} level={} svc=orders tenant=t{} msg=\"{}\" latency_ms={}\n", + i % 28 + 1, + ["INFO", "WARN", "ERROR"][i as usize % 3], + i % 17, + [ + "flush memtable", + "rotate segment", + "compact level", + "write block" + ][i as usize % 4], + (i * 37) % 1000 + ); + sizes.push(line.len()); + data.extend_from_slice(line.as_bytes()); + } + (data, sizes) +} + +/// Every sample compressed at the default level with `dictionary`, frames +/// without the dictionary id (raw content has none to record, so the sides +/// would otherwise differ by four bytes a frame). +fn compressed_total(dictionary: &[u8], samples: &SampleSet<'_>) -> usize { + let mut compressor: FrameCompressor = FrameCompressor::new(CompressionLevel::Default); + compressor.set_dictionary_id_flag(false); + compressor + .set_encoder_dictionary( + EncoderDictionary::from_serialized_or_raw_content(dictionary).expect("dictionary"), + ) + .unwrap(); + (0..samples.count()) + .map(|index| { + compressor + .compress_independent_frame(samples.sample(index)) + .len() + }) + .sum() +} + +/// The entropy tables describe what compressing the samples with the content +/// actually produces, so a finalized dictionary compresses them at least as +/// well as its content alone does. Tables guessed from the samples' raw bytes +/// (counting byte values as if they were sequence codes) made it compress them +/// worse than the bare content. +#[test] +fn finalized_tables_compress_the_samples_better_than_the_content_alone() { + let (data, sizes) = log_samples(); + let samples = SampleSet::new(&data, &sizes).unwrap(); + // A plausible content: the most common line shapes, best last. + let content: Vec = (0..40) + .flat_map(|i| samples.sample(i * 7).iter().copied()) + .collect(); + let dict = finalize( + &content, + &samples, + samples.count(), + 8192, + FinalizeOptions::default(), + ) + .unwrap(); + assert!(dict.ends_with(&content)); + let with_tables = compressed_total(&dict, &samples); + let content_only = compressed_total(&content, &samples); + assert!( + with_tables < content_only, + "tables {with_tables} vs content alone {content_only}" + ); +} + +/// Samples that do not compress teach nothing: their blocks are written raw, +/// so the counts stay at their floor of one, and the flat literal distribution +/// that leaves is replaced by one a table description can express. +#[test] +fn samples_that_do_not_compress_leave_a_describable_flat_table() { + let mut state = 0x9E37_79B9_7F4A_7C15u64; + let data: Vec = (0..20_000) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state as u8 + }) + .collect(); + let samples = SampleSet::new(&data, &[2_000; 10]).unwrap(); + let dict = finalize( + &data[..1024], + &samples, + samples.count(), + 4096, + FinalizeOptions::default(), + ) + .unwrap(); + crate::decoding::Dictionary::decode_dict(&dict).expect("the dictionary parses back"); +} + +/// The literals table is scaled to its longest code, not to the length limit: +/// a flat alphabet comes out at eight bits, which is how the finalizer tells +/// it apart, and the stand-in distribution at nine, with a description that +/// can be written (every weight the same would have none). +#[test] +fn literal_tables_take_the_depth_of_their_longest_code() { + assert_eq!(literals_table(&[1; 256]).table_log(), 8); + let mut counts = [2usize; 256]; + counts[0] = 4; + counts[253] = 1; + counts[254] = 1; + let table = literals_table(&counts); + assert_eq!(table.table_log(), 9); + assert!(!huffman_description(&table).unwrap().is_empty()); +} + +/// Only a frame whose first block is compressed counts as one. +#[test] +fn first_block_kind_is_read_from_the_frame() { + let text: String = core::iter::repeat_n("compressible text ", 200).collect(); + assert!(first_block_is_compressed(&compress_slice_to_vec( + text.as_bytes(), + CompressionLevel::Default + ))); + let mut state = 1u32; + let noise: Vec = (0..4096) + .map(|_| { + state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + (state >> 24) as u8 + }) + .collect(); + assert!(!first_block_is_compressed(&compress_slice_to_vec( + &noise, + CompressionLevel::Default + ))); + assert!(!first_block_is_compressed(b"not a frame")); +} + +/// The id is derived from the content when none is given, inside the range +/// upstream reserves for them, and stays the same for the same content. +#[test] +fn derived_ids_are_stable_and_compliant() { + let id = derive_dict_id(b"some content"); + assert_eq!(id, derive_dict_id(b"some content")); + assert!((32_768..1 << 31).contains(&id)); +} diff --git a/zstd/src/dictionary/legacy.rs b/zstd/src/dictionary/legacy.rs index 86249e127..ab13df894 100644 --- a/zstd/src/dictionary/legacy.rs +++ b/zstd/src/dictionary/legacy.rs @@ -39,7 +39,7 @@ pub(crate) const MIN_SAMPLES_SIZE: usize = CONTENT_SIZE_MIN * MIN_RATIO as usize pub const DEFAULT_SELECTIVITY: u32 = 9; /// Smallest dictionary the trainer is asked for (`ZDICT_DICTSIZE_MIN`). -pub(crate) const DICT_SIZE_MIN: usize = super::TRAINER_DICT_SIZE_MIN; +pub(crate) const DICT_SIZE_MIN: usize = super::MIN_TRAINED_DICT_SIZE; /// What was too small for the legacy trainer to produce a dictionary. #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/zstd/src/dictionary/mod.rs b/zstd/src/dictionary/mod.rs index b7d77dc20..43c7ebd85 100644 --- a/zstd/src/dictionary/mod.rs +++ b/zstd/src/dictionary/mod.rs @@ -22,20 +22,12 @@ //! FastCOVER search `zstd --train` runs, over the corpus cut into samples. mod cover; mod fastcover; +mod finalize; mod legacy; mod samples; mod selection; mod suffix_array; -use crate::bit_io::BitWriter; -use crate::blocks::sequence_section::{ - MAX_LITERAL_LENGTH_CODE, MAX_MATCH_LENGTH_CODE, MAX_OFFSET_CODE, -}; -use crate::decoding::dictionary::MAGIC_NUM as DICT_MAGIC_NUM; -use crate::decoding::sequence_section_decoder::{LL_MAX_LOG, ML_MAX_LOG, OF_MAX_LOG}; -use crate::fse::fse_encoder::{self, build_table_from_symbol_counts}; -use crate::huff0::HuffmanTable as HuffmanDecoderTable; -use crate::huff0::huff0_encoder::{HuffmanEncoder, HuffmanTable as HuffmanEncoderTable}; pub use legacy::DEFAULT_SELECTIVITY; pub use samples::TrainingError; use samples::refuse; @@ -44,25 +36,18 @@ use std::{ fs::{self, File}, io::{self, Read}, path::{Path, PathBuf}, - // `vec` import covers the `vec![..]` macro used below: this crate is - // no_std-with-std-feature, so the std prelude isn't pulled in implicitly - // for top-level items in this module. Removing this import fails the - // build with `cannot find macro 'vec' in this scope` — verified. - vec, vec::Vec, }; const MAX_TRAINING_PREALLOC_BYTES: usize = 8 * 1024 * 1024; -const MAX_HUFFMAN_STATS_BYTES: usize = 64 * 1024; -/// Smallest size a trained dictionary can occupy, whatever it was trained on. -/// -/// The magic number, the dictionary ID, the three repeat offsets and the -/// shortest content the writers emit are unconditional; a real dictionary is -/// larger still, since the entropy tables between them are never empty. Use it -/// to reject an impossible `dict_size` before spending the corpus: the training -/// entry points can only discover the true bound once those tables are built. -pub const MIN_TRAINED_DICT_SIZE: usize = DICT_MAGIC_NUM.len() + 4 + 12 + 8; +/// Smallest `dict_size` a dictionary is trained or finalized into (upstream +/// zstd `ZDICT_DICTSIZE_MIN`). +/// +/// Use it to reject an impossible `dict_size` before spending the corpus. A +/// size at or above it can still be too small once the entropy tables the +/// samples produce are built, which only training can tell. +pub const MIN_TRAINED_DICT_SIZE: usize = finalize::DICT_SIZE_MIN; /// Tuning for COVER training, the knobs of the reference's /// `ZDICT_cover_params_t`. @@ -87,8 +72,6 @@ pub struct CoverOptions { /// (upstream zstd `COVER_selectDict`). A dictionary so found is its header /// plus that tail, so none is smaller than the header plus 256 bytes. pub shrink: Option, - /// Compression level candidates are scored at; zero is the default level. - pub level: i32, } impl Default for CoverOptions { @@ -99,7 +82,6 @@ impl Default for CoverOptions { steps: 4, split_point: 1.0, shrink: None, - level: 0, } } } @@ -112,8 +94,8 @@ impl Default for CoverOptions { /// building and the rest scoring. #[derive(Debug, Clone, Copy, PartialEq)] pub struct FastCoverOptions { - /// Segment size, dmer size, search, split, shrink and scoring level, as for - /// COVER. `d` is at least 4 here; the reference takes 6 and 8. + /// Segment size, dmer size, search, split and shrink, as for COVER. `d` is + /// at least 4 here; the reference takes 6 and 8. pub cover: CoverOptions, /// Width of the dmer frequency table in bits, `1..=31`; its memory grows /// as `2^f`. Zero is 20. @@ -136,11 +118,16 @@ impl Default for FastCoverOptions { } } -/// Header options for a finalized dictionary. +/// How a dictionary is finalized and trained against, the reference's +/// `ZDICT_params_t`. #[derive(Debug, Clone, Copy, Default)] pub struct FinalizeOptions { /// The dictionary id; `None` derives one from the content. pub dict_id: Option, + /// The compression level the dictionary is built for: the entropy tables + /// come from samples compressed at it, and the trainers score candidates + /// at it. Zero is the default level. + pub level: i32, } /// Create a "raw content" dictionary of at most `dict_size` bytes, with no @@ -333,252 +320,52 @@ fn search_raw_content(corpus: &[u8], sizes: &[usize], dict_size: usize) -> io::R Ok(content) } -/// The `i`th of [`MAX_HUFFMAN_STATS_BYTES`] samples spread evenly over `len` -/// bytes. -/// -/// Computed in 64 bits: `i * len` reaches 2^48 for an addressable corpus, which -/// a 32-bit `usize` cannot hold — the product overflows for any corpus past -/// 64 KiB there, and the multiply panics rather than sampling. The quotient is -/// always below `len`, so the narrowing back is exact. -fn strided_index(i: usize, len: usize) -> usize { - ((i as u64 * len as u64) / MAX_HUFFMAN_STATS_BYTES as u64) as usize -} - -fn serialize_huffman_table(sample_data: &[u8], raw_content: &[u8]) -> io::Result> { - fn bounded_huffman_stats(data: &[u8]) -> Vec { - if data.len() <= MAX_HUFFMAN_STATS_BYTES { - return data.to_vec(); - } - - let mut stats = Vec::with_capacity(MAX_HUFFMAN_STATS_BYTES); - for i in 0..MAX_HUFFMAN_STATS_BYTES { - stats.push(data[strided_index(i, data.len())]); - } - stats - } - - let source = if sample_data.len() >= 2 { - sample_data - } else { - raw_content - }; - let mut stats = bounded_huffman_stats(source); - if stats.len() < 2 || stats.iter().all(|b| *b == stats[0]) { - // A corpus with no distribution to measure gets a synthetic one. It - // stops at 128 symbols because a perfectly flat alphabet gives every - // symbol the same weight: FSE cannot encode that (an RLE weight - // stream), and the direct nibble form addresses at most 128 symbols, so - // a full 0..=255 alphabet would have no description at all. - stats = (0u8..128).collect(); - } - - let mut table = HuffmanEncoderTable::build_from_data(stats.as_slice()); - if table - .writeable_table_description_size(&mut crate::fse::fse_encoder::FSETable::blank()) - .is_none() - { - // Sampled real data can land on the same shape: a flat alphabet wider - // than 128 symbols. Fall back to the synthetic narrow one, which always - // has a description. - stats = (0u8..128).collect(); - table = HuffmanEncoderTable::build_from_data(stats.as_slice()); - } - let mut writer = BitWriter::new(); - let mut encoder = HuffmanEncoder::new(&table, &mut writer); - encoder.encode(&[stats[0]], true); - let encoded = writer.dump(); - - let mut decoder = HuffmanDecoderTable::new(); - let table_size = decoder - .build_decoder(encoded.as_slice()) - .map_err(|e| io::Error::other(format!("failed to decode generated huffman table: {e}")))?; - Ok(encoded[..table_size as usize].to_vec()) -} - -fn serialize_fse_table(table: &fse_encoder::FSETable) -> Vec { - let mut writer = BitWriter::new(); - table.write_table(&mut writer); - writer.dump() -} - -fn bounded_fse_symbols(data: &[u8], max_symbol: u8) -> Vec { - let modulo = u16::from(max_symbol) + 1; - if data.is_empty() { - return Vec::from([0u8]); - } - if data.len() <= MAX_HUFFMAN_STATS_BYTES { - return data - .iter() - .map(|b| (u16::from(*b) % modulo) as u8) - .collect(); - } - - let mut out = Vec::with_capacity(MAX_HUFFMAN_STATS_BYTES); - for i in 0..MAX_HUFFMAN_STATS_BYTES { - let idx = strided_index(i, data.len()); - out.push((u16::from(data[idx]) % modulo) as u8); - } - out -} - -fn serialize_fse_table_from_corpus( - sample_data: &[u8], - raw_content: &[u8], - max_symbol: u8, - max_log: u8, -) -> io::Result> { - fn counts_total_for_source(source: &[u8], max_symbol: u8, counts: &mut [usize]) -> usize { - counts.fill(0); - for symbol in bounded_fse_symbols(source, max_symbol) { - counts[usize::from(symbol)] += 1; - } - counts.iter().sum::() - } - - let mut counts = vec![0usize; usize::from(max_symbol) + 1]; - let using_sample = !sample_data.is_empty(); - let mut total = counts_total_for_source( - if using_sample { - sample_data - } else { - raw_content - }, - max_symbol, - &mut counts, - ); - if total <= 1 && using_sample && !raw_content.is_empty() { - total = counts_total_for_source(raw_content, max_symbol, &mut counts); - } - if total <= 1 { - return Err(io::Error::new( - io::ErrorKind::InvalidInput, - "insufficient symbol statistics for FSE table", - )); - } - let table = build_table_from_symbol_counts(&counts, max_log, false); - Ok(serialize_fse_table(&table)) -} - -fn derive_dict_id(raw_content: &[u8]) -> u32 { - let mut h = 0xcbf29ce484222325u64; - for &b in raw_content { - h ^= u64::from(b); - h = h.wrapping_mul(0x100000001b3); - } - let compliant = (h % ((1u64 << 31) - 32768)) + 32768; - compliant as u32 -} - -/// Finalize raw dictionary content into a full zstd dictionary binary -/// (`magic + dict_id + entropy tables + offset history + content`). +/// Finalize raw dictionary content into a dictionary of at most `dict_size` +/// bytes: magic, id, entropy tables, repeat offsets and the content (the +/// reference's `ZDICT_finalizeDictionary`). +/// +/// The entropy tables are measured, not guessed: the first block of each +/// sample is compressed with `raw_content` as a raw dictionary at +/// [`FinalizeOptions::level`], and the tables describe the literals and +/// sequences those blocks produced. `samples` is every sample back to back and +/// `sample_sizes` their lengths. Content that does not fit after the header is +/// cut from the front: trainers place their best segments last. +/// +/// # Errors +/// +/// `InvalidInput` when `raw_content` is empty, when `dict_size` is below +/// [`MIN_TRAINED_DICT_SIZE`] or leaves less than eight bytes of content after +/// the header, when the id is zero, or when `sample_sizes` does not add up to +/// `samples.len()`. +/// +/// # Examples +/// +/// ``` +/// use structured_zstd::dictionary::{FinalizeOptions, finalize_raw_dict}; +/// +/// let mut samples = Vec::new(); +/// let mut sizes = Vec::new(); +/// for i in 0..50u32 { +/// let line = format!("tenant=demo table=orders key={i} status=shipped\n"); +/// sizes.push(line.len()); +/// samples.extend_from_slice(line.as_bytes()); +/// } +/// let content = b"tenant=demo table=orders key= status=shipped\n"; +/// let dict = finalize_raw_dict(content, &samples, &sizes, 1024, FinalizeOptions::default()) +/// .unwrap(); +/// assert!(dict.ends_with(content)); +/// ``` pub fn finalize_raw_dict( raw_content: &[u8], - sample_data: &[u8], + samples: &[u8], + sample_sizes: &[usize], dict_size: usize, options: FinalizeOptions, ) -> io::Result> { - if raw_content.is_empty() { - return Err(io::Error::new( - io::ErrorKind::InvalidInput, - "raw dictionary content must not be empty", - )); - } - let mut tables = serialize_huffman_table(sample_data, raw_content)?; - for (max_symbol, max_log) in ENTROPY_STREAMS { - tables.extend_from_slice(&serialize_fse_table_from_corpus( - sample_data, - raw_content, - max_symbol, - max_log, - )?); - } - let mut out = Vec::new(); - assemble_dict(&mut out, raw_content, &tables, dict_size, options)?; - Ok(out) -} - -/// The offset, match-length and literal-length streams, in the order their -/// tables follow the literals table in a dictionary. -const ENTROPY_STREAMS: [(u8, u8); 3] = [ - (MAX_OFFSET_CODE, OF_MAX_LOG), - (MAX_MATCH_LENGTH_CODE, ML_MAX_LOG), - (MAX_LITERAL_LENGTH_CODE, LL_MAX_LOG), -]; - -/// The entropy tables [`finalize_raw_dict`] writes, when `sample_data` alone -/// decides them; `None` when the samples are too thin and the tables would -/// fall back on the content, which then has to be finalized in full. -fn sample_entropy_tables(sample_data: &[u8]) -> Option> { - if sample_data.len() < 2 { - return None; - } - let mut tables = serialize_huffman_table(sample_data, &[]).ok()?; - for (max_symbol, max_log) in ENTROPY_STREAMS { - tables.extend_from_slice( - &serialize_fse_table_from_corpus(sample_data, &[], max_symbol, max_log).ok()?, - ); - } - Some(tables) -} - -/// A dictionary of `raw_content` behind already serialized entropy `tables`, -/// written over `out`, whose allocation a caller finalizing many candidates -/// keeps from one to the next. -fn assemble_dict( - out: &mut Vec, - raw_content: &[u8], - tables: &[u8], - dict_size: usize, - options: FinalizeOptions, -) -> io::Result<()> { - out.clear(); - out.reserve(dict_size.max(256)); - out.extend_from_slice(&DICT_MAGIC_NUM); - let dict_id = options - .dict_id - .unwrap_or_else(|| derive_dict_id(raw_content)); - if dict_id == 0 { - return Err(io::Error::new( - io::ErrorKind::InvalidInput, - "dictionary id must be non-zero", - )); - } - out.extend_from_slice(&dict_id.to_le_bytes()); - out.extend_from_slice(tables); - - // Repeat offsets: keep default bootstrap history. - out.extend_from_slice(&1u32.to_le_bytes()); - out.extend_from_slice(&4u32.to_le_bytes()); - out.extend_from_slice(&8u32.to_le_bytes()); - - let min_content_size = 8usize; - let max_content_budget = dict_size.saturating_sub(out.len()); - if max_content_budget < min_content_size { - return Err(io::Error::new( - io::ErrorKind::InvalidInput, - "dictionary size too small to fit header and offset history", - )); - } - - let content = if raw_content.len() > max_content_budget { - &raw_content[raw_content.len() - max_content_budget..] - } else { - raw_content - }; - if content.len() < min_content_size { - out.resize(out.len() + (min_content_size - content.len()), 0); - } - out.extend_from_slice(content); - Ok(()) + let set = samples::SampleSet::new(samples, sample_sizes)?; + finalize::finalize(raw_content, &set, set.count(), dict_size, options) } -/// Smallest dictionary, in bytes, any trainer here builds (upstream zstd -/// `ZDICT_DICTSIZE_MIN`); a smaller one is refused with -/// [`TrainingError::DictionaryTooSmall`]. Unlike [`MIN_TRAINED_DICT_SIZE`] it -/// bounds a training request, so a caller can refuse one before loading the -/// samples. -pub const TRAINER_DICT_SIZE_MIN: usize = 256; - /// The `k` and `d` values a training run tries, resolved from the options the /// way the reference's optimizers resolve theirs. struct SearchSpace { @@ -700,7 +487,7 @@ impl SearchSpace { /// `InvalidInput` for tuning out of range or for a `k` and `d` that no /// dictionary of `dict_size` bytes holds, where [`TrainingError::of`] reports /// [`TrainingError::Parameter`], and for a `dict_size` under -/// [`TRAINER_DICT_SIZE_MIN`], where it reports +/// [`MIN_TRAINED_DICT_SIZE`], where it reports /// [`TrainingError::DictionaryTooSmall`]. /// /// # Examples @@ -845,10 +632,10 @@ fn check_samples_and_dict_size<'s>( /// Refuse a dictionary smaller than any trainer builds. fn check_dict_size(dict_size: usize) -> io::Result<()> { - if dict_size < TRAINER_DICT_SIZE_MIN { + if dict_size < MIN_TRAINED_DICT_SIZE { return Err(refuse( TrainingError::DictionaryTooSmall, - &format!("a dictionary must be at least {TRAINER_DICT_SIZE_MIN} bytes"), + &format!("a dictionary must be at least {MIN_TRAINED_DICT_SIZE} bytes"), )); } Ok(()) @@ -865,7 +652,7 @@ fn check_dict_size(dict_size: usize) -> io::Result<()> { /// # Errors /// /// `InvalidInput` when `k` or `d` is zero or `d > k`, when `k` exceeds -/// `dict_size`, when `dict_size` is under [`TRAINER_DICT_SIZE_MIN`], when there +/// `dict_size`, when `dict_size` is under [`MIN_TRAINED_DICT_SIZE`], when there /// are fewer than five samples or they do not add up to `samples.len()`; /// [`TrainingError::of`] tells these causes apart. /// @@ -962,14 +749,8 @@ fn run_cover( // A plain run with no shrinking prices nothing: its one dictionary is the // answer, as the reference's plain trainer returns it unscored. let scored = !plain || options.shrink.is_some() || space.split_point < 1.0; - let mut evaluator = selection::Evaluator::new( - &set, - split.train, - split.test.clone(), - dict_size, - options.level, - finalize, - ); + let mut evaluator = + selection::Evaluator::new(&set, split.train, split.test.clone(), dict_size, finalize); let mut best = selection::Best::new(); let mut state = Vec::new(); let mut content_scratch = Vec::new(); @@ -1171,7 +952,6 @@ fn run_fastcover( ctx.finalize_samples(split.train), split.test.clone(), dict_size, - options.cover.level, finalize, ) }); @@ -1256,7 +1036,13 @@ pub fn create_legacy_dict_from_slice( // the corpus to its size limit inside the search alone, for its suffix // sort (zdict.c, `ZDICT_trainBuffer_legacy`), and builds the entropy tables // from all of them (`ZDICT_trainFromBuffer_unsafe_legacy`). - let finalized = finalize_raw_dict(content.as_slice(), samples, dict_size, finalize)?; + let finalized = finalize_raw_dict( + content.as_slice(), + samples, + sample_sizes, + dict_size, + finalize, + )?; output.write_all(finalized.as_slice()) } diff --git a/zstd/src/dictionary/samples.rs b/zstd/src/dictionary/samples.rs index d6bbc6899..8471c55c5 100644 --- a/zstd/src/dictionary/samples.rs +++ b/zstd/src/dictionary/samples.rs @@ -128,9 +128,8 @@ pub(super) struct Split { pub(super) test: Range, } -/// Why a segment trainer refused its input. It rides inside the -/// `InvalidInput` error the trainer returns; [`TrainingError::of`] reads it -/// back. +/// Why a segment trainer or the finalizer refused its input. It rides inside +/// the `InvalidInput` error returned; [`TrainingError::of`] reads it back. #[derive(Clone, Copy, Debug, PartialEq, Eq)] #[non_exhaustive] pub enum TrainingError { @@ -140,7 +139,7 @@ pub enum TrainingError { /// bytes. Samples, /// The dictionary asked for is smaller than - /// [`TRAINER_DICT_SIZE_MIN`](super::TRAINER_DICT_SIZE_MIN). + /// [`MIN_TRAINED_DICT_SIZE`](super::MIN_TRAINED_DICT_SIZE). DictionaryTooSmall, } diff --git a/zstd/src/dictionary/selection.rs b/zstd/src/dictionary/selection.rs index 4c7c691ba..33c4faa72 100644 --- a/zstd/src/dictionary/selection.rs +++ b/zstd/src/dictionary/selection.rs @@ -2,8 +2,9 @@ //! scoring samples with it, and count the bytes (upstream zstd `cover.c`, //! `COVER_selectDict` and `COVER_checkTotalCompressedSize`). +use super::FinalizeOptions; +use super::finalize::finalize; use super::samples::SampleSet; -use super::{FinalizeOptions, assemble_dict, finalize_raw_dict, sample_entropy_tables}; use crate::encoding::{CompressionLevel, EncoderDictionary, FrameCompressor}; use core::ops::Range; use std::{io, vec::Vec}; @@ -31,12 +32,7 @@ pub(super) struct Evaluator<'s> { /// Samples a dictionary is scored on. scoring: Range, capacity: usize, - level: i32, finalize: FinalizeOptions, - /// The entropy tables, built once: they come from the finalize samples, - /// the same for every candidate. `None` where the samples are too thin - /// and each candidate's content decides them. - tables: Option>, compressor: Option, frame: Vec, /// The candidate being priced, and a shrunk one being tried against it: @@ -52,7 +48,6 @@ impl<'s> Evaluator<'s> { finalize_samples: usize, scoring: Range, capacity: usize, - level: i32, finalize: FinalizeOptions, ) -> Self { Self { @@ -60,9 +55,7 @@ impl<'s> Evaluator<'s> { finalize_samples, scoring, capacity, - level, finalize, - tables: sample_entropy_tables(samples.leading(finalize_samples)), compressor: None, frame: Vec::new(), dict: Vec::new(), @@ -77,21 +70,12 @@ impl<'s> Evaluator<'s> { Ok(out) } - /// [`Self::finalize`] over `out`, keeping its allocation. + /// [`Self::finalize`] over `out`. fn finalize_into(&self, content: &[u8], out: &mut Vec) -> io::Result<()> { - if let Some(tables) = &self.tables { - if content.is_empty() { - // The full path owns that refusal and its wording. - *out = finalize_raw_dict(content, &[], self.capacity, self.finalize)?; - return Ok(()); - } - return assemble_dict(out, content, tables, self.capacity, self.finalize); - } - // Samples too thin to decide the tables alone: the content decides - // them, a path a search takes only on corpora of a handful of bytes. - *out = finalize_raw_dict( + *out = finalize( content, - self.samples.leading(self.finalize_samples), + self.samples, + self.finalize_samples, self.capacity, self.finalize, )?; @@ -114,7 +98,9 @@ impl<'s> Evaluator<'s> { // matcher for it, which belongs to the compressor, not here. let prepared = EncoderDictionary::from_bytes(dict) .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?; - let level = self.level; + // Scored at the level the dictionary is finalized for, as upstream uses + // its one `compressionLevel` for both. + let level = self.finalize.level; let compressor = self .compressor .get_or_insert_with(|| FrameCompressor::new(CompressionLevel::from_level(level))); diff --git a/zstd/src/dictionary/tests.rs b/zstd/src/dictionary/tests.rs index d871aea49..d96a99a6b 100644 --- a/zstd/src/dictionary/tests.rs +++ b/zstd/src/dictionary/tests.rs @@ -3,41 +3,7 @@ use crate::decoding::Dictionary; use crate::encoding::{CompressionLevel, EncoderDictionary, FrameCompressor}; use std::io::Cursor; use std::string::ToString; - -/// The Huffman statistics are gathered by striding across the corpus, and the -/// stride is `i * len` — a product that reaches 2^48 for an addressable corpus -/// and so cannot be held by a 32-bit `usize`. Past 64 KiB of samples it -/// overflowed there and the multiply panicked instead of sampling, which no -/// amount of 64-bit testing shows. The index is computed in 64 bits now: the -/// last sample of a corpus far larger than the statistics buffer still lands -/// inside it, and the spread is even. -#[test] -fn statistics_are_strided_without_overflowing_the_index() { - let len = 4 * MAX_HUFFMAN_STATS_BYTES + 12_345; - assert_eq!( - strided_index(0, len), - 0, - "the first sample is the first byte" - ); - assert!( - strided_index(MAX_HUFFMAN_STATS_BYTES - 1, len) < len, - "and the last one is still inside the corpus" - ); - // Evenly spread: each step advances by the corpus over the sample count. - assert_eq!(strided_index(1, len), len / MAX_HUFFMAN_STATS_BYTES); - assert_eq!( - strided_index(MAX_HUFFMAN_STATS_BYTES / 2, len), - len / 2, - "the middle sample is the middle of the corpus" - ); - // The product that overflows a 32-bit `usize`: 2^16 samples over a corpus - // of 2^32 bytes is 2^48, which only 64-bit arithmetic holds. - assert_eq!( - strided_index(MAX_HUFFMAN_STATS_BYTES - 1, u32::MAX as usize), - ((MAX_HUFFMAN_STATS_BYTES as u64 - 1) * u32::MAX as u64 / MAX_HUFFMAN_STATS_BYTES as u64) - as usize - ); -} +use std::vec; fn training_data() -> Vec { training_samples().0 @@ -100,7 +66,10 @@ fn fixed_fastcover(k: u32, d: u32) -> FastCoverOptions { #[test] fn plain_trainers_write_a_parseable_dictionary() { let (data, sizes) = training_samples(); - let finalize = FinalizeOptions { dict_id: Some(77) }; + let finalize = FinalizeOptions { + dict_id: Some(77), + ..FinalizeOptions::default() + }; let cover = train_cover_dict(&data, &sizes, 4096, &fixed_cover(128, 8), finalize).unwrap(); let fast = train_fastcover_dict(&data, &sizes, 4096, &fixed_fastcover(128, 8), finalize).unwrap(); @@ -471,7 +440,7 @@ fn the_preflight_names_the_cause_the_trainer_names() { #[test] fn the_preflight_refuses_a_dictionary_under_the_trainer_minimum() { let (data, sizes) = training_samples(); - let dict_size = TRAINER_DICT_SIZE_MIN - 1; + let dict_size = MIN_TRAINED_DICT_SIZE - 1; let cover = CoverOptions::default(); let fastcover = FastCoverOptions::default(); let trained = optimize_cover_dict(&data, &sizes, dict_size, &cover, FinalizeOptions::default()) @@ -549,11 +518,10 @@ fn a_frequency_table_too_wide_for_the_target_is_an_error() { assert_eq!(err.kind(), io::ErrorKind::OutOfMemory); } -/// A corpus of one repeated byte carries no symbol distribution to describe, so -/// the serializer substitutes a synthetic alphabet. That substitute has to be -/// one the Huffman table description can actually express: a flat alphabet -/// wider than 128 symbols has neither an FSE nor a direct representation, and -/// training on such a corpus must not take the process down with it. +/// A corpus of one repeated byte leaves the literal counts flat: every symbol +/// at eight bits, which no Huffman description can express. The finalizer then +/// describes a mostly flat distribution that can be written, and training on +/// such a corpus must not take the process down with it. #[test] fn training_on_a_single_repeated_byte_does_not_crash() { let sample = vec![7u8; 4096]; @@ -562,7 +530,10 @@ fn training_on_a_single_repeated_byte_does_not_crash() { &[512; 8], 4096, &fixed_fastcover(256, 8), - FinalizeOptions { dict_id: Some(1) }, + FinalizeOptions { + dict_id: Some(1), + ..FinalizeOptions::default() + }, ) .expect("a uniform corpus must train"); Dictionary::decode_dict(&dict).expect("the trained dictionary must parse back"); @@ -751,60 +722,31 @@ fn raw_content_from_a_directory_helps_a_small_frame() { assert!(compress(Some(&raw)) < compress(None)); } -/// The entropy tables a search builds once finalize every candidate into the -/// same bytes `finalize_raw_dict` writes for it, including the content cut -/// to fit and an explicit id; and where the samples are too thin to decide the -/// tables, each candidate's content still does. -#[test] -fn cached_entropy_tables_finalize_as_the_full_path_does() { - let (data, sizes) = training_samples(); - let set = samples::SampleSet::new(&data, &sizes).unwrap(); - for finalize in [ - FinalizeOptions::default(), - FinalizeOptions { - dict_id: Some(0x1234_5678), - }, - ] { - let evaluator = selection::Evaluator::new(&set, sizes.len(), 0..0, 1024, 3, finalize); - for content in [&data[..300], &data[1000..5000], &data[..8]] { - assert_eq!( - evaluator.finalize(content).unwrap(), - finalize_raw_dict(content, &data, 1024, finalize).unwrap() - ); - } - } - let thin = samples::SampleSet::new(&data[..1], &[1]).unwrap(); - let evaluator = selection::Evaluator::new(&thin, 1, 0..0, 1024, 3, FinalizeOptions::default()); - let content = &data[..600]; - assert_eq!( - evaluator.finalize(content).ok(), - finalize_raw_dict(content, &data[..1], 1024, FinalizeOptions::default()).ok() - ); -} - #[test] fn finalize_raw_dict_rejects_empty_raw_content() { - let sample = training_data(); - let err = finalize_raw_dict(&[], sample.as_slice(), 4096, FinalizeOptions::default()) + let (data, sizes) = training_samples(); + let err = finalize_raw_dict(&[], &data, &sizes, 4096, FinalizeOptions::default()) .expect_err("empty raw dictionary must be rejected"); assert_eq!(err.kind(), io::ErrorKind::InvalidInput); } #[test] fn finalize_raw_dict_rejects_too_small_budget() { - let sample = training_data(); + let (data, sizes) = training_samples(); let raw = b"some-raw-bytes"; - let err = finalize_raw_dict(raw, sample.as_slice(), 32, FinalizeOptions::default()) - .expect_err("tiny dict_size must fail"); - assert_eq!(err.kind(), io::ErrorKind::InvalidInput); - assert!(err.to_string().contains("dictionary size too small")); + for dict_size in [32, MIN_TRAINED_DICT_SIZE - 1] { + let err = finalize_raw_dict(raw, &data, &sizes, dict_size, FinalizeOptions::default()) + .expect_err("tiny dict_size must fail"); + assert_eq!(err.kind(), io::ErrorKind::InvalidInput); + assert!(err.to_string().contains("dictionary size too small")); + } } #[test] fn finalize_raw_dict_pads_to_minimum_content_size() { - let sample = training_data(); + let (data, sizes) = training_samples(); let raw = b"x"; - let finalized = finalize_raw_dict(raw, sample.as_slice(), 4096, FinalizeOptions::default()) + let finalized = finalize_raw_dict(raw, &data, &sizes, 4096, FinalizeOptions::default()) .expect("finalize should pad small raw content"); let parsed = Dictionary::decode_dict(finalized.as_slice()).expect("finalized dict parses"); assert!(parsed.dict_content.len() >= 8); @@ -813,15 +755,34 @@ fn finalize_raw_dict_pads_to_minimum_content_size() { #[test] fn finalize_raw_dict_rejects_zero_dict_id() { - let sample = training_data(); + let (data, sizes) = training_samples(); let raw = b"raw-fastcover-bytes"; let err = finalize_raw_dict( raw, - sample.as_slice(), + &data, + &sizes, 4096, - FinalizeOptions { dict_id: Some(0) }, + FinalizeOptions { + dict_id: Some(0), + ..FinalizeOptions::default() + }, ) .expect_err("dict_id=0 must be rejected"); assert_eq!(err.kind(), io::ErrorKind::InvalidInput); assert_eq!(err.to_string(), "dictionary id must be non-zero"); } + +/// The sample sizes must describe the samples exactly. +#[test] +fn finalize_raw_dict_rejects_sizes_that_do_not_add_up() { + let (data, sizes) = training_samples(); + let err = finalize_raw_dict( + b"content", + &data, + &sizes[1..], + 4096, + FinalizeOptions::default(), + ) + .unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::InvalidInput); +} diff --git a/zstd/src/encoding/blocks/compressed.rs b/zstd/src/encoding/blocks/compressed.rs index 8c719cc85..bd7c7f0f7 100644 --- a/zstd/src/encoding/blocks/compressed.rs +++ b/zstd/src/encoding/blocks/compressed.rs @@ -2932,7 +2932,7 @@ const ML_EXTRA_BITS: [u8; 53] = [ /// table index. Every code's baseline is a multiple of its own extra-bit width, /// so masking off those bits is the same subtraction the ranges spelled out. #[inline] -fn encode_literal_length(len: u32) -> (u8, u32, usize) { +pub(crate) fn encode_literal_length(len: u32) -> (u8, u32, usize) { debug_assert!(len < 131_072, "literal length {len} out of encodable range"); let code = if len < 64 { LL_CODE[len as usize] @@ -2949,7 +2949,7 @@ fn encode_literal_length(len: u32) -> (u8, u32, usize) { /// /// Table-driven for the same reason as [`encode_literal_length`]. #[inline] -fn encode_match_len(len: u32) -> (u8, u32, usize) { +pub(crate) fn encode_match_len(len: u32) -> (u8, u32, usize) { debug_assert!( (3..131_075).contains(&len), "match length {len} out of encodable range", @@ -2968,7 +2968,7 @@ fn encode_match_len(len: u32) -> (u8, u32, usize) { /// history per RFC 8878 §3.1.2.5. Updates `offset_hist` in place. /// /// Encoded offset codes: 1/2/3 = repeat offsets, N+3 = new absolute offset N. -pub(in crate::encoding) fn encode_offset_with_history( +pub(crate) fn encode_offset_with_history( actual_offset: u32, lit_len: u32, offset_hist: &mut [u32; 3], @@ -3075,7 +3075,7 @@ fn low_bits(value: u32, bits: usize) -> u32 { value & ((1u32 << bits) - 1) } -fn encode_offset(len: u32) -> (u8, u32, usize) { +pub(crate) fn encode_offset(len: u32) -> (u8, u32, usize) { let log = len.ilog2(); let lower = len & ((1 << log) - 1); (log as u8, lower, log as usize) diff --git a/zstd/src/encoding/blocks/mod.rs b/zstd/src/encoding/blocks/mod.rs index 671ff6fb8..c47e1661e 100644 --- a/zstd/src/encoding/blocks/mod.rs +++ b/zstd/src/encoding/blocks/mod.rs @@ -6,3 +6,8 @@ mod compressed; pub(super) use compressed::*; +// The dictionary finalizer derives sequence codes the way blocks do. +#[cfg(feature = "dict-builder")] +pub(crate) use compressed::{ + encode_literal_length, encode_match_len, encode_offset, encode_offset_with_history, +}; diff --git a/zstd/src/encoding/frame_compressor.rs b/zstd/src/encoding/frame_compressor.rs index 08c43e3ed..8fca059f3 100644 --- a/zstd/src/encoding/frame_compressor.rs +++ b/zstd/src/encoding/frame_compressor.rs @@ -3231,6 +3231,12 @@ impl FrameCompressor { self.compressed_data.take() } + /// The matcher, for a caller that reads what it recorded between frames. + #[cfg(feature = "dict-builder")] + pub(crate) fn matcher_mut(&mut self) -> &mut M { + &mut self.state.matcher + } + /// Before calling [FrameCompressor::compress] you can replace the matcher pub fn replace_matcher(&mut self, mut match_generator: M) -> M { core::mem::swap(&mut match_generator, &mut self.state.matcher); diff --git a/zstd/src/fse/fse_encoder.rs b/zstd/src/fse/fse_encoder.rs index e8cbbc347..008ae4d76 100644 --- a/zstd/src/fse/fse_encoder.rs +++ b/zstd/src/fse/fse_encoder.rs @@ -425,7 +425,7 @@ pub fn build_table_from_data( build_table_from_counts(&counts[..=max_symbol], max_log, avoid_0_numbit) } -#[cfg(any(test, feature = "fuzz-exports", feature = "dict-builder"))] +#[cfg(any(test, feature = "fuzz-exports"))] pub(crate) fn build_table_from_symbol_counts( counts: &[usize], max_log: u8, @@ -499,7 +499,32 @@ pub(crate) fn build_seq_ctable_into( build_table_from_probabilities_into(&probs[..=max_symbol], table_log, out); } -#[cfg(any(test, feature = "fuzz-exports", feature = "dict-builder"))] +/// An FSE table at exactly `table_log` for `counts`, with low-probability +/// symbols kept at `-1` (upstream zstd `FSE_normalizeCount` given the log +/// outright, as the dictionary finalizer calls it, zdict.c +/// `ZDICT_analyzeEntropy`). `counts` must not be too spread for the log: +/// every symbol present needs a slot. +#[cfg(feature = "dict-builder")] +pub(crate) fn build_table_at_log(counts: &[usize], table_log: u8) -> FSETable { + let total = counts.iter().sum::(); + let max_symbol = counts + .iter() + .rposition(|&count| count > 0) + .unwrap_or_default(); + debug_assert!(total > 1 && table_log >= min_table_log(total, max_symbol)); + let mut probs = [0i32; 256]; + normalize_counts( + &mut probs[..=max_symbol], + table_log, + &counts[..=max_symbol], + total, + max_symbol, + true, + ); + build_table_from_probabilities(&probs[..=max_symbol], table_log) +} + +#[cfg(any(test, feature = "fuzz-exports"))] fn build_table_from_counts(counts: &[usize], max_log: u8, avoid_0_numbit: bool) -> FSETable { let mut out = FSETable::blank(); build_table_from_counts_into(counts, max_log, avoid_0_numbit, &mut out); diff --git a/zstd/src/huff0/huff0_encoder.rs b/zstd/src/huff0/huff0_encoder.rs index 9ed0897ca..4e2aec021 100644 --- a/zstd/src/huff0/huff0_encoder.rs +++ b/zstd/src/huff0/huff0_encoder.rs @@ -994,6 +994,15 @@ impl HuffmanTable { Self::build_from_weights_reusing(weights, None) } + /// The optimal code for `counts` with no code longer than `max_bits` + /// (upstream zstd `HUF_buildCTable_wksp`, as the dictionary finalizer + /// calls it). Unlike [`Self::build_from_counts`] it runs no table-log + /// search: the table describes future literals, not these. + #[cfg(feature = "dict-builder")] + pub(crate) fn build_limited(counts: &[usize], max_bits: usize) -> Self { + Self::build_from_weights(&build_limited_weights(counts, max_bits)) + } + /// [`Self::build_from_weights`] filling a discarded table's buffers instead /// of taking new ones. A table is built per block and, where the splitter /// probes, per split candidate; most are measured and thrown away, so From 2ea6a21868abad6e2d4f0ff08b9870bb02d1ce9e Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Mon, 28 Sep 2026 04:21:50 +0300 Subject: [PATCH 03/20] perf(dict): FastCOVER hashes each position once, unchecked - the counting and selecting loops are monomorphised on the number of bytes a dmer hashes, once above the loop, instead of branching on d at every position - the window keeps each entering position's table index in a ring and reads it back as the position leaves; upstream hashes it a second time - reads and table indexing go through pointers: every position below nb_dmers has eight readable bytes and hashes below 2^f, both tables' length; the ring is bounded by the epoch, which a window never exceeds Dictionaries are byte-identical. x86 (runner2), fixed FastCOVER k=256 d=8 on 6420 samples, task-clock over three interleaved rounds: 5.29 s -> 4.98 s. Part of #128 --- zstd/src/dictionary/fastcover.rs | 207 +++++++++++++++++++------------ 1 file changed, 129 insertions(+), 78 deletions(-) diff --git a/zstd/src/dictionary/fastcover.rs b/zstd/src/dictionary/fastcover.rs index c3f3cac1c..bdaac99e9 100644 --- a/zstd/src/dictionary/fastcover.rs +++ b/zstd/src/dictionary/fastcover.rs @@ -29,54 +29,50 @@ fn dmer_read_len(d: usize) -> usize { d.max(8) } -/// The shift that takes a `W`-byte dmer hash to an `f`-bit table index. -fn hash_shift(f: u32) -> u32 { - if W == 4 { 32 - f } else { 64 - f } -} - -/// Upstream zstd `FASTCOVER_hashPtrToIndex`: hash the first `W = min(d, 8)` -/// bytes of the dmer at `pos` into a table index, `shift` from -/// [`hash_shift`]. Caller guarantees `pos + dmer_read_len(d) <= sample.len()`. +/// Upstream zstd `FASTCOVER_hashPtrToIndex`: hash the first `H` bytes of the +/// dmer at `ptr` (`H = min(d, 8)`) into an `f`-bit table index. +/// +/// # Safety +/// +/// `ptr` is readable for eight bytes, which every dmer position is: the +/// positions stop [`dmer_read_len`] bytes before the end. #[inline(always)] -fn hash_dmer_index(sample: &[u8], pos: usize, shift: u32) -> usize { - if W == 4 { - let v = u32::from_le_bytes(sample[pos..pos + 4].try_into().unwrap()); - return (v.wrapping_mul(PRIME_4_BYTES) >> shift) as usize; +unsafe fn hash_at(ptr: *const u8, f: u32) -> usize { + // SAFETY: the caller's eight readable bytes. + let v = u64::from_le(unsafe { ptr.cast::().read_unaligned() }); + match H { + 4 => ((v as u32).wrapping_mul(PRIME_4_BYTES) >> (32 - f)) as usize, + 5 => ((v << 24).wrapping_mul(PRIME_5_BYTES) >> (64 - f)) as usize, + 6 => ((v << 16).wrapping_mul(PRIME_6_BYTES) >> (64 - f)) as usize, + 7 => ((v << 8).wrapping_mul(PRIME_7_BYTES) >> (64 - f)) as usize, + _ => (v.wrapping_mul(PRIME_8_BYTES) >> (64 - f)) as usize, } - let v = u64::from_le_bytes(sample[pos..pos + 8].try_into().unwrap()); - let h = match W { - 5 => (v << 24).wrapping_mul(PRIME_5_BYTES), - 6 => (v << 16).wrapping_mul(PRIME_6_BYTES), - 7 => (v << 8).wrapping_mul(PRIME_7_BYTES), - _ => v.wrapping_mul(PRIME_8_BYTES), - }; - (h >> shift) as usize } -/// Evaluate `$body` with the const `$w` bound to the hash width of `d`-byte -/// dmers, so the loops inside are compiled once per width and carry no -/// width branch. -macro_rules! with_hash_width { - ($d:expr, $w:ident => $body:expr) => { +/// Run `$body` with `$h` bound to the number of bytes a dmer of `$d` hashes, +/// as a constant: the counting and selecting loops are monomorphised on it +/// once, rather than branching on `d` at every position. +macro_rules! with_hashed_bytes { + ($d:expr, $h:ident => $body:expr) => { match $d.min(8) { 4 => { - const $w: usize = 4; + const $h: usize = 4; $body } 5 => { - const $w: usize = 5; + const $h: usize = 5; $body } 6 => { - const $w: usize = 6; + const $h: usize = 6; $body } 7 => { - const $w: usize = 7; + const $h: usize = 7; $body } _ => { - const $w: usize = 8; + const $h: usize = 8; $body } } @@ -124,29 +120,6 @@ fn zeroed_counts(len: usize) -> Result, TableTooLarge> { Ok(unsafe { Vec::from_raw_parts(pointer.cast::(), len, len) }) } -/// Count every `accel`-th dmer that lies wholly inside its sample, the samples -/// bounded by `offsets`. -fn count_dmers( - freqs: &mut [u32], - data: &[u8], - offsets: &[usize], - read_len: usize, - f: u32, - accel: u32, -) { - let shift = hash_shift::(f); - let step = accel as usize; - for bounds in offsets.windows(2) { - let (mut start, end) = (bounds[0], bounds[1]); - while start + read_len <= end { - // Bounded by the number of dmers, which the caller checked fits - // `u32`. - freqs[hash_dmer_index::(data, start, shift)] += 1; - start += step; - } - } -} - /// Share of the training samples the entropy tables are drawn from, and dmers /// skipped between counted ones, per acceleration (upstream zstd /// `FASTCOVER_defaultAccelParameters`). @@ -199,7 +172,7 @@ impl<'s> FastCoverContext<'s> { samples.check_holds_dmer(train, read_len)?; let mut freqs = zeroed_counts::(1usize << f)?; let offsets = &samples.offsets()[..=train]; - with_hash_width!(d, W => count_dmers::(&mut freqs, data, offsets, read_len, f, accel)); + with_hashed_bytes!(d, H => count_dmers::(data, offsets, read_len, f, accel as usize, &mut freqs)); Ok(Self { data, nb_dmers, @@ -246,31 +219,52 @@ impl<'s> FastCoverContext<'s> { // A window holds at most `dmers_in_k + 1` occurrences of one index (one // past the segment before the oldest leaves). Upstream zstd keeps them in // `u16` for any `k`; a longer segment than that counts in `u32`. + let WindowCounts { narrow, wide, ring } = window; let tail = if layout.dmers_in_k < usize::from(u16::MAX) { - let counts = window.narrow.get_or_insert_with_result(self.f)?; - with_hash_width!(self.d, W => self.select_segments::<_, W>(out, freqs, counts, layout)) + let counts = narrow.get_or_insert_with_result(self.f)?; + with_hashed_bytes!(self.d, H => self.select_segments::(out, freqs, counts, ring, layout)) } else { - let counts = window.wide.get_or_insert_with_result(self.f)?; - with_hash_width!(self.d, W => self.select_segments::<_, W>(out, freqs, counts, layout)) + let counts = wide.get_or_insert_with_result(self.f)?; + with_hashed_bytes!(self.d, H => self.select_segments::(out, freqs, counts, ring, layout)) }; Ok(&out[tail..]) } /// Pick a segment per epoch visit until `out` is filled from its end, and - /// return where the content starts in it. - fn select_segments( + /// return where the content starts in it. `H` is the number of bytes a + /// dmer hashes. + fn select_segments( &self, out: &mut [u8], freqs: &mut [u32], segment_freqs: &mut [C], + ring: &mut Vec, layout: EpochLayout, ) -> usize { let EpochLayout { dmers_in_k, epochs } = layout; - let (sample, d) = (self.data, self.d); - let shift = hash_shift::(self.f); - let index = |pos: usize| hash_dmer_index::(sample, pos, shift); + let (sample, f, d) = (self.data, self.f, self.d); let zero = C::from(0); let one = C::from(1); + debug_assert!(freqs.len() == 1 << f && segment_freqs.len() == 1 << f); + debug_assert!(epochs.num * epochs.size <= self.nb_dmers); + // Every position below `nb_dmers` has eight readable bytes and hashes + // below `2^f`, the length of both tables, so the loops below read and + // index without checks. + let base = sample.as_ptr(); + let freq_ptr = freqs.as_mut_ptr(); + let window_ptr = segment_freqs.as_mut_ptr(); + // A window reaching this many positions sheds its oldest. + let window_limit = dmers_in_k + 1; + // The window's dmer indices in arrival order, so each position is + // hashed once as it enters and read back as it leaves. Upstream hashes + // it again on the way out. A window never spans more than its epoch, + // so a segment longer than the longest epoch, the last one, needs no + // more room than that. + let longest_epoch = self.nb_dmers - (epochs.num - 1) * epochs.size; + let ring_len = window_limit.min(longest_epoch); + ring.clear(); + ring.resize(ring_len, 0); + let ring_ptr = ring.as_mut_ptr(); // Fill from the back (upstream zstd layout) so the best segments sit at // the end of the dictionary and get referenced with the smallest offsets. let mut tail = out.len(); @@ -293,18 +287,36 @@ impl<'s> FastCoverContext<'s> { let mut active_begin = epoch_begin; let mut active_end = epoch_begin; let mut active_score = 0u64; + let mut ring_head = 0usize; + let mut ring_tail = 0usize; while active_end < epoch_end { - let idx = index(active_end); - if segment_freqs[idx] == zero { - active_score += u64::from(freqs[idx]); + // SAFETY: `active_end < epoch_end <= nb_dmers`; see above. + // `ring_tail < ring_len`, the ring's length. + let idx = unsafe { hash_at::(base.add(active_end), f) }; + unsafe { *ring_ptr.add(ring_tail) = idx as u32 }; + ring_tail += 1; + if ring_tail == ring_len { + ring_tail = 0; } + // SAFETY: `idx < 2^f`, both tables' length. + let count = unsafe { &mut *window_ptr.add(idx) }; + if *count == zero { + active_score += u64::from(unsafe { *freq_ptr.add(idx) }); + } + *count += one; active_end += 1; - segment_freqs[idx] += one; - if active_end - active_begin == dmers_in_k + 1 { - let del = index(active_begin); - segment_freqs[del] -= one; - if segment_freqs[del] == zero { - active_score -= u64::from(freqs[del]); + if active_end - active_begin == window_limit { + // SAFETY: `ring_head < ring_len`; the index it holds was + // hashed from a position, so it is below `2^f`. + let del = unsafe { *ring_ptr.add(ring_head) } as usize; + ring_head += 1; + if ring_head == ring_len { + ring_head = 0; + } + let count = unsafe { &mut *window_ptr.add(del) }; + *count -= one; + if *count == zero { + active_score -= u64::from(unsafe { *freq_ptr.add(del) }); } active_begin += 1; } @@ -316,24 +328,34 @@ impl<'s> FastCoverContext<'s> { } // Reset the window counts for the next epoch. while active_begin < epoch_end { - let del = index(active_begin); - segment_freqs[del] -= one; + // SAFETY: as for the leaving position above. + let del = unsafe { *ring_ptr.add(ring_head) } as usize; + ring_head += 1; + if ring_head == ring_len { + ring_head = 0; + } + unsafe { *window_ptr.add(del) -= one }; active_begin += 1; } // Drop the segment's leading dmers that earlier segments already // cover: they score nothing and would only spend dictionary bytes. // The tail needs no trim, the best segment is recorded as a dmer // that scores enters it. - while best_begin < best_end && freqs[index(best_begin)] == 0 { + // SAFETY (both reads): `best_begin < best_end <= epoch_end <= nb_dmers`. + while best_begin < best_end + && unsafe { *freq_ptr.add(hash_at::(base.add(best_begin), f)) } == 0 + { best_begin += 1; } debug_assert!( - best_score == 0 || freqs[index(best_end - 1)] != 0, + best_score == 0 + || unsafe { *freq_ptr.add(hash_at::(base.add(best_end - 1), f)) } != 0, "the best segment ends on a dmer that scores" ); // Zero the chosen segment's frequencies: its dmers are covered. for pos in best_begin..best_end { - freqs[index(pos)] = 0; + // SAFETY: `pos < best_end <= epoch_end <= nb_dmers`. + unsafe { *freq_ptr.add(hash_at::(base.add(pos), f)) = 0 }; } if best_score == 0 { @@ -359,6 +381,33 @@ impl<'s> FastCoverContext<'s> { } } +/// Count every `step`-th dmer lying wholly inside its sample (upstream zstd +/// `FASTCOVER_computeFrequency`). `offsets` bounds the samples of `data`. +fn count_dmers( + data: &[u8], + offsets: &[usize], + read_len: usize, + f: u32, + step: usize, + freqs: &mut [u32], +) { + debug_assert_eq!(freqs.len(), 1 << f); + debug_assert_eq!(offsets.last(), Some(&data.len())); + let base = data.as_ptr(); + let counts = freqs.as_mut_ptr(); + for bounds in offsets.windows(2) { + let end = bounds[1]; + let mut start = bounds[0]; + while start + read_len <= end { + // SAFETY: `start + read_len <= end <= data.len()` and `read_len` + // is at least eight; the index is below `2^f`, the table's length. + // A count is bounded by the dmer count, far below `u32::MAX`. + unsafe { *counts.add(hash_at::(base.add(start), f)) += 1 }; + start += step; + } + } +} + /// How the corpus is walked: the dmers a segment spans, and the epochs it is /// split into. #[derive(Clone, Copy)] @@ -373,11 +422,13 @@ impl WindowCount for u16 {} impl WindowCount for u32 {} /// Per-window occurrence counts (upstream zstd `segmentFreqs`), kept across -/// builds: every build leaves them zero, so one table serves every `k`. +/// builds: every build leaves them zero, so one table serves every `k`. With +/// them, the ring of the window's dmer indices. #[derive(Default)] pub(super) struct WindowCounts { narrow: LazyCounts, wide: LazyCounts, + ring: Vec, } struct LazyCounts(Option>); From a3031abeb712b1c5e73031ac27321c3f69a28f4c Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Mon, 28 Sep 2026 04:22:05 +0300 Subject: [PATCH 04/20] perf(dict): analyse every sample at the average sample's parameters The finalizer hinted each sample's own size, so every sample ran its own table sizes and the dictionary was indexed again for each one. Upstream analyses them all with the parameters of the average sample (ZSTD_getParams(level, averageSampleSize, dictSize)); doing the same keeps one set of tables for the pass and the dictionary resident between samples. x86 (runner2), fixed FastCOVER k=256 d=8 on 6420 samples, task-clock over three interleaved rounds: 4.99 s -> 4.38 s; the dictionary's total over the corpus at upstream -3 went from 1851466 to 1850248 (upstream's own: 1849508). Part of #128 --- zstd/src/dictionary/finalize.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/zstd/src/dictionary/finalize.rs b/zstd/src/dictionary/finalize.rs index 9eafcd4da..22c7586fe 100644 --- a/zstd/src/dictionary/finalize.rs +++ b/zstd/src/dictionary/finalize.rs @@ -182,6 +182,11 @@ fn count_samples( compressor .set_encoder_dictionary(EncoderDictionary::from_dictionary(dictionary)) .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?; + // Every sample runs the parameters of the average one, as upstream's + // analysis does (`ZSTD_getParams(level, averageSampleSize, dictSize)`): + // one set of tables for the whole pass, which also keeps the dictionary + // resident from one sample to the next instead of indexing it again. + let average = samples.leading(count).len() / count.max(1); let mut frame = Vec::new(); for index in 0..count { let sample = samples.sample(index); @@ -192,7 +197,7 @@ fn count_samples( frame.clear(); compressor.set_source(block); compressor.set_drain(frame); - compressor.set_source_size_hint(block.len() as u64); + compressor.set_source_size_hint(average as u64); compressor.compress(); frame = compressor.take_drain().expect("the drain was set above"); let recorder = compressor.matcher_mut(); From be2653ee7a36f54b6689ec58491fca1c36075bf9 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Mon, 28 Sep 2026 04:37:00 +0300 Subject: [PATCH 05/20] perf(dict): shrink the search's winner, not every candidate Upstream's selection shrinks each candidate of its parameter search (and then never runs it). The size a dictionary is cut to is a separate choice from the k and d that built it, and cutting every candidate multiplies the whole search by the number of sizes tried: a zero-k, zero-d FastCOVER search with shrink took 56.6 s where the search alone takes a few. The search now keeps the winner's content and cuts only that dictionary. Part of #128 --- zstd/src/dictionary/fastcover.rs | 18 ++-- zstd/src/dictionary/fastcover/tests.rs | 11 +- zstd/src/dictionary/mod.rs | 34 +++--- zstd/src/dictionary/selection.rs | 141 ++++++++++++++----------- zstd/src/dictionary/tests.rs | 32 ++++++ 5 files changed, 138 insertions(+), 98 deletions(-) diff --git a/zstd/src/dictionary/fastcover.rs b/zstd/src/dictionary/fastcover.rs index bdaac99e9..b53ff8d45 100644 --- a/zstd/src/dictionary/fastcover.rs +++ b/zstd/src/dictionary/fastcover.rs @@ -120,12 +120,17 @@ fn zeroed_counts(len: usize) -> Result, TableTooLarge> { Ok(unsafe { Vec::from_raw_parts(pointer.cast::(), len, len) }) } -/// Share of the training samples the entropy tables are drawn from, and dmers -/// skipped between counted ones, per acceleration (upstream zstd -/// `FASTCOVER_defaultAccelParameters`). +/// Share of the training samples the entropy tables are drawn from, per +/// acceleration (upstream zstd `FASTCOVER_defaultAccelParameters`; the dmers +/// skipped between counted ones are `accel - 1`). const ACCEL_FINALIZE_PERCENT: [usize; MAX_ACCEL as usize + 1] = [100, 100, 50, 34, 25, 20, 17, 14, 13, 11, 10]; +/// How many of `train` samples the entropy tables are drawn from at `accel`. +pub(super) fn finalize_samples(train: usize, accel: u32) -> usize { + train * ACCEL_FINALIZE_PERCENT[accel as usize] / 100 +} + /// Every trainable position of the training samples hashed into a frequency /// table. Depends on `d`, `f` and `accel`, so one context serves every `k`. pub(super) struct FastCoverContext<'s> { @@ -134,7 +139,6 @@ pub(super) struct FastCoverContext<'s> { freqs: Vec, d: usize, f: u32, - finalize_percent: usize, } impl<'s> FastCoverContext<'s> { @@ -179,15 +183,9 @@ impl<'s> FastCoverContext<'s> { freqs, d, f, - finalize_percent: ACCEL_FINALIZE_PERCENT[accel as usize], }) } - /// How many of `train` samples the entropy tables are drawn from. - pub(super) fn finalize_samples(&self, train: usize) -> usize { - train * self.finalize_percent / 100 - } - /// Build content of at most `capacity` bytes from segments of `k` bytes /// (upstream zstd `FASTCOVER_buildDictionary`). `freqs` is scratch the /// build refills with the counted frequencies and spends, so one table diff --git a/zstd/src/dictionary/fastcover/tests.rs b/zstd/src/dictionary/fastcover/tests.rs index 26e2239dc..0c55d8683 100644 --- a/zstd/src/dictionary/fastcover/tests.rs +++ b/zstd/src/dictionary/fastcover/tests.rs @@ -39,14 +39,9 @@ fn accel_strides_the_count_and_shrinks_the_finalize_share() { // Five positions per 12-byte sample; every other one is positions 0, 2, 4. assert_eq!(total(&full), 25); assert_eq!(total(&strided), 15); - assert_eq!(full.finalize_samples(100), 100); - assert_eq!(strided.finalize_samples(100), 50); - assert_eq!( - FastCoverContext::new(&set, 5, 8, 20, 10) - .unwrap() - .finalize_samples(100), - 10 - ); + assert_eq!(finalize_samples(100, 1), 100); + assert_eq!(finalize_samples(100, 2), 50); + assert_eq!(finalize_samples(100, 10), 10); } /// Training bytes shorter than one dmer read are refused rather than yielding diff --git a/zstd/src/dictionary/mod.rs b/zstd/src/dictionary/mod.rs index 43c7ebd85..0553b20cd 100644 --- a/zstd/src/dictionary/mod.rs +++ b/zstd/src/dictionary/mod.rs @@ -70,7 +70,9 @@ pub struct CoverOptions { /// dictionary of its own, and keep the first whose scoring samples /// compress to at most this many percent more than with the whole content /// (upstream zstd `COVER_selectDict`). A dictionary so found is its header - /// plus that tail, so none is smaller than the header plus 256 bytes. + /// plus that tail, so none is smaller than the header plus 256 bytes. A + /// search cuts its winner down this way; the parameters it chooses are + /// unchanged. pub shrink: Option, } @@ -781,7 +783,7 @@ fn run_cover( let built = match cover::CoverContext::new(&set, split.train, d) { Ok(ctx) => Some(ctx), Err(err) => { - best.offer(Err(err), chosen); + best.offer(Err(err), &[], chosen); None } }; @@ -794,9 +796,9 @@ fn run_cover( if !scored { return Ok((evaluator.finalize(content)?, chosen)); } - best.offer(evaluator.select(content, options.shrink), chosen); + best.offer(evaluator.score(content), content, chosen); } - best.finish() + best.finish(&mut evaluator, options.shrink) } /// Train a FastCOVER dictionary of at most `dict_size` bytes with the `k` and @@ -909,7 +911,13 @@ fn run_fastcover( let mut best = selection::Best::new(); // `None` beside a `d` is a dmer size these samples cannot count. let mut context: Option<(usize, Option>)> = None; - let mut evaluator: Option> = None; + let mut evaluator = selection::Evaluator::new( + &set, + fastcover::finalize_samples(split.train, accel), + split.test.clone(), + dict_size, + finalize, + ); for (d, k) in space.pairs() { if !segment_fits(k, d, dict_size) { continue; @@ -936,7 +944,7 @@ fn run_fastcover( let built = match fastcover::FastCoverContext::new(&set, split.train, d, f, accel) { Ok(ctx) => Some(ctx), Err(err) => { - best.offer(Err(err), chosen); + best.offer(Err(err), &[], chosen); None } }; @@ -945,23 +953,13 @@ fn run_fastcover( let Some(ctx) = context.as_ref().and_then(|(_, built)| built.as_ref()) else { continue; }; - // The finalize share depends on `accel` alone, so every context agrees. - let evaluator = evaluator.get_or_insert_with(|| { - selection::Evaluator::new( - &set, - ctx.finalize_samples(split.train), - split.test.clone(), - dict_size, - finalize, - ) - }); let content = ctx.build(&mut freqs, &mut window, &mut content_scratch, dict_size, k)?; if !scored { return Ok((evaluator.finalize(content)?, chosen)); } - best.offer(evaluator.select(content, options.cover.shrink), chosen); + best.offer(evaluator.score(content), content, chosen); } - best.finish() + best.finish(&mut evaluator, options.cover.shrink) } /// Train and finalize a dictionary with the reference's original trainer, the diff --git a/zstd/src/dictionary/selection.rs b/zstd/src/dictionary/selection.rs index 33c4faa72..819af5c6f 100644 --- a/zstd/src/dictionary/selection.rs +++ b/zstd/src/dictionary/selection.rs @@ -115,42 +115,48 @@ impl<'s> Evaluator<'s> { Ok(total) } - /// Finalize `content` and price it. With `shrink`, also try the content's - /// last 256, 512, ... bytes and keep the first that costs at most - /// `shrink` percent more than the whole: a smaller dictionary that does - /// nearly as well. - /// - /// Upstream zstd runs the same search but forces it off in the only path - /// that reaches it, so its `shrink` has no effect; here it takes effect. - pub(super) fn select(&mut self, content: &[u8], shrink: Option) -> io::Result> { + /// Finalize `content` and price it, in the evaluator's candidate buffer. + pub(super) fn score(&mut self, content: &[u8]) -> io::Result> { let mut dict = core::mem::take(&mut self.dict); let priced = self.finalize_and_price(content, &mut dict); self.dict = dict; - let total = priced?; - let Some(regression) = shrink else { - return Ok(Priced { - dict: &self.dict, - total, - }); - }; + Ok(Priced { + total: priced?, + dict: &self.dict, + }) + } + + /// Try the last 256, 512, ... bytes of `content` and keep the first that + /// costs at most `regression` percent more than `full`, the whole + /// content's score: a smaller dictionary that does nearly as well. + /// + /// Upstream zstd runs the same search for every candidate of its + /// parameter search, but forces it off in the only path that reaches it, + /// so its `shrink` has no effect. Here it runs once, on the winner: the + /// size a dictionary is cut to is a separate choice from the `k` and `d` + /// that built it, and searching it per candidate multiplies the cost of + /// the whole search by the number of sizes tried. + pub(super) fn shrink( + &mut self, + content: &[u8], + full: Scored, + regression: u32, + ) -> io::Result { let mut size = SHRINK_START; while size < content.len() { let mut shrunk = core::mem::take(&mut self.shrunk); let priced = self.finalize_and_price(&content[content.len() - size..], &mut shrunk); - self.shrunk = shrunk; - let candidate_total = priced?; - if within_regression(candidate_total, total, regression) { - return Ok(Priced { - dict: &self.shrunk, - total: candidate_total, + let total = priced?; + if within_regression(total, full.total, regression) { + return Ok(Scored { + dict: shrunk, + total, }); } + self.shrunk = shrunk; size *= 2; } - Ok(Priced { - dict: &self.dict, - total, - }) + Ok(full) } } @@ -163,16 +169,17 @@ pub(super) fn within_regression(total: usize, full: usize, regression: u32) -> b total as u128 * 100 <= full as u128 * (u128::from(regression) + 100) } -/// A candidate [`Evaluator::select`] priced, still in the evaluator's buffer. +/// A candidate [`Evaluator::score`] priced, still in the evaluator's buffer. pub(super) struct Priced<'e> { pub(super) dict: &'e [u8], pub(super) total: usize, } -/// The cheapest dictionary seen, and the parameters that built it. A tie keeps -/// the earlier one, as upstream's strict comparison does. +/// The cheapest dictionary seen, the content it was finalized from and the +/// parameters that built it. A tie keeps the earlier one, as upstream's strict +/// comparison does. pub(super) struct Best

{ - found: Option<(Scored, P)>, + found: Option<(Scored, Vec, P)>, last_error: Option, } @@ -184,44 +191,54 @@ impl

Best

{ } } - /// Keep `candidate` if it beats the best so far, copying its bytes into the - /// buffer the previous winner held; a candidate that loses is not copied. - pub(super) fn offer(&mut self, candidate: io::Result>, params: P) { + /// Keep `candidate` if it is the cheapest so far. Its dictionary and + /// `content` are copied only then, into the buffers the previous winner + /// held; a candidate that loses is not copied at all. + pub(super) fn offer(&mut self, candidate: io::Result>, content: &[u8], params: P) { match candidate { - Ok(priced) => { - if self - .found - .as_ref() - .is_none_or(|(best, _)| priced.total < best.total) - { - let mut dict = self - .found - .take() - .map_or_else(Vec::new, |(previous, _)| previous.dict); - dict.clear(); - dict.extend_from_slice(priced.dict); - self.found = Some(( - Scored { - dict, - total: priced.total, - }, - params, - )); + Ok(priced) => match &mut self.found { + Some((best, _, _)) if priced.total >= best.total => {} + Some((best, kept, kept_params)) => { + best.total = priced.total; + best.dict.clear(); + best.dict.extend_from_slice(priced.dict); + kept.clear(); + kept.extend_from_slice(content); + *kept_params = params; } - } + None => { + let best = Scored { + dict: priced.dict.to_vec(), + total: priced.total, + }; + self.found = Some((best, content.to_vec(), params)); + } + }, Err(err) => self.last_error = Some(err), } } - /// The winner, or why no candidate could be priced. - pub(super) fn finish(self) -> io::Result<(Vec, P)> { - match (self.found, self.last_error) { - (Some((scored, params)), _) => Ok((scored.dict, params)), - (None, Some(err)) => Err(err), - (None, None) => Err(super::samples::refuse( - super::TrainingError::Parameter, - "no parameter combination is valid for this dictionary size", - )), - } + /// The winner's dictionary, cut down with `shrink` when given, and its + /// parameters; or why no candidate could be priced. + pub(super) fn finish( + self, + evaluator: &mut Evaluator<'_>, + shrink: Option, + ) -> io::Result<(Vec, P)> { + let (full, content, params) = match (self.found, self.last_error) { + (Some(found), _) => found, + (None, Some(err)) => return Err(err), + (None, None) => { + return Err(super::samples::refuse( + super::TrainingError::Parameter, + "no parameter combination is valid for this dictionary size", + )); + } + }; + let chosen = match shrink { + Some(regression) => evaluator.shrink(&content, full, regression)?, + None => full, + }; + Ok((chosen.dict, params)) } } diff --git a/zstd/src/dictionary/tests.rs b/zstd/src/dictionary/tests.rs index d96a99a6b..79c5e4608 100644 --- a/zstd/src/dictionary/tests.rs +++ b/zstd/src/dictionary/tests.rs @@ -263,6 +263,38 @@ fn shrink_keeps_a_smaller_dictionary_within_the_regression() { assert!(price(&strict, &data, &sizes) <= full_price); } +/// In a search, `shrink` cuts the winner down rather than weighing every +/// candidate at every size: the parameters chosen are the ones the search +/// chooses without it, and the dictionary is that winner's, cut or whole. +#[test] +fn shrink_cuts_the_search_winner() { + let (data, sizes) = training_samples(); + let options = CoverOptions { + split_point: 0.75, + ..CoverOptions::default() + }; + let (whole, chosen) = + optimize_cover_dict(&data, &sizes, 8192, &options, FinalizeOptions::default()).unwrap(); + let (cut, chosen_with_shrink) = optimize_cover_dict( + &data, + &sizes, + 8192, + &CoverOptions { + shrink: Some(1000), + ..options + }, + FinalizeOptions::default(), + ) + .unwrap(); + assert_eq!( + (chosen_with_shrink.k, chosen_with_shrink.d), + (chosen.k, chosen.d) + ); + let content = |dict: &[u8]| Dictionary::decode_dict(dict).unwrap().dict_content; + assert!(cut.len() < whole.len()); + assert!(content(&whole).ends_with(&content(&cut))); +} + /// A split below 1 scores on the trailing samples only: a split that leaves /// none of them, or fewer than five to build from, is refused. #[test] From 491903d07c9066fc679ac150ef9be93865846f70 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Mon, 28 Sep 2026 17:15:50 +0300 Subject: [PATCH 06/20] fix(dict): count what the analysed block wrote - Only a sample's first block reaches the entropy tables, as upstream's ZDICT_countEStats compresses MIN(128 KiB, window) bytes as one block. A sample longer than the window used to add every later block, including ones written raw. Once a frame has shown the window, the rest of each sample is not fed at all. Regression test: only_the_first_block_of_a_sample_is_counted. - Offset codes are counted with the repeat policy the frame's blocks used: at the fast strategy an offset equal to rep[1] or rep[2] is written explicitly, and the finalizer counted it as repeat 2 or 3. One function, uses_fast_offset_codes, now decides this for the block encoder, its size estimate and the finalizer. Regression test: offset_codes_follow_the_frames_repeat_policy. - The four table descriptions are written straight into the dictionary: the sequence tables are normalized and described without building encoder tables (write_ncount, shared with FSETable::write_table), and the literals table writes its description alone instead of encoding a symbol and parsing it back. - The finalizer's too-small refusal carries TrainingError::DictionaryTooSmall and is checked first, as upstream does, so ZDICT_finalizeDictionary returns dstSize_tooSmall. --- c-api/src/tests.rs | 15 +++ zstd/src/bit_io/bit_writer.rs | 8 +- zstd/src/dictionary/finalize.rs | 85 ++++++++-------- zstd/src/dictionary/finalize/tests.rs | 96 ++++++++++++++++++- zstd/src/dictionary/mod.rs | 3 +- zstd/src/encoding/blocks/compressed.rs | 17 ++-- zstd/src/encoding/blocks/mod.rs | 1 + zstd/src/encoding/frame_compressor.rs | 7 ++ zstd/src/fse/fse_encoder.rs | 128 ++++++++++++++----------- zstd/src/fse/tests.rs | 30 ++++++ zstd/src/huff0/huff0_encoder.rs | 3 +- 11 files changed, 280 insertions(+), 113 deletions(-) diff --git a/c-api/src/tests.rs b/c-api/src/tests.rs index 4812b8f7c..8a795c2e6 100644 --- a/c-api/src/tests.rs +++ b/c-api/src/tests.rs @@ -534,6 +534,21 @@ fn zdict_trainers_report_upstreams_error_codes() { ) }; assert_eq!(code_from_result(n), ZSTD_error_parameter_outOfBound); + // The finalizer checks the buffer before anything else, empty content + // included. + let n = unsafe { + crate::dict::ZDICT_finalizeDictionary( + dict.as_mut_ptr(), + 255, + samples.as_ptr(), + 0, + samples.as_ptr(), + sizes.as_ptr(), + sizes.len() as u32, + zparams, + ) + }; + assert_eq!(code_from_result(n), ZSTD_error_dstSize_tooSmall); // A zero `k` asks the plain trainers for a search they do not run. let n = unsafe { ZDICT_trainFromBuffer_cover( diff --git a/zstd/src/bit_io/bit_writer.rs b/zstd/src/bit_io/bit_writer.rs index e4ad9ad07..5e4d32445 100644 --- a/zstd/src/bit_io/bit_writer.rs +++ b/zstd/src/bit_io/bit_writer.rs @@ -18,9 +18,9 @@ pub(crate) struct BitWriter>> { impl BitWriter> { /// Initialize a new writer. // Production encode paths write into caller-owned buffers via `from`; - // the owned-`Vec` constructor (and `dump` below) only serve the - // dictionary trainer and test/fuzz round-trips. - #[cfg(any(test, feature = "fuzz-exports", feature = "dict-builder"))] + // the owned-`Vec` constructor (and `dump` below) only serve test/fuzz + // round-trips. + #[cfg(any(test, feature = "fuzz-exports"))] pub fn new() -> Self { Self { output: Vec::new(), @@ -428,7 +428,7 @@ impl>> BitWriter { /// /// This function consumes the writer, so it cannot be used after /// dumping - #[cfg(any(test, feature = "fuzz-exports", feature = "dict-builder"))] + #[cfg(any(test, feature = "fuzz-exports"))] pub fn dump(mut self) -> V { if self.misaligned() != 0 { panic!( diff --git a/zstd/src/dictionary/finalize.rs b/zstd/src/dictionary/finalize.rs index 22c7586fe..fb26bc60c 100644 --- a/zstd/src/dictionary/finalize.rs +++ b/zstd/src/dictionary/finalize.rs @@ -12,16 +12,16 @@ use crate::bit_io::BitWriter; use crate::decoding::dictionary::MAGIC_NUM as DICT_MAGIC_NUM; use crate::encoding::blocks::{ encode_literal_length, encode_match_len, encode_offset, encode_offset_with_history, + encode_offset_with_history_fast, }; use crate::encoding::workspace::Workspace; use crate::encoding::{ CompressionLevel, DictionarySizes, EncoderDictionary, FrameCompressor, HistoryBuf, MatchGeneratorDriver, Matcher, Sequence, }; -use crate::fse::fse_encoder::{FSETable, build_table_at_log}; -use crate::huff0::HuffmanTable as HuffmanDecoderTable; +use crate::fse::fse_encoder::{FSETable, write_ncount_at_log}; use crate::huff0::huff0_encoder::{HuffmanEncoder, HuffmanTable}; -use std::{format, io, vec, vec::Vec}; +use std::{io, vec, vec::Vec}; /// Smallest dictionary finalized (upstream zstd `ZDICT_DICTSIZE_MIN`). pub(super) const DICT_SIZE_MIN: usize = 256; @@ -51,12 +51,13 @@ pub(super) fn finalize( dict_size: usize, options: FinalizeOptions, ) -> io::Result> { - if content.is_empty() { - return Err(invalid("raw dictionary content must not be empty")); - } + // Upstream zstd's order (`ZDICT_finalizeDictionary`): the size first. if dict_size < DICT_SIZE_MIN { return Err(too_small()); } + if content.is_empty() { + return Err(invalid("raw dictionary content must not be empty")); + } let dict_id = options.dict_id.unwrap_or_else(|| derive_dict_id(content)); if dict_id == 0 { return Err(invalid("dictionary id must be non-zero")); @@ -150,14 +151,14 @@ fn analyze_entropy( literals = literals_table(&counts.literals); debug_assert_eq!(literals.table_log(), 9); } - out.extend_from_slice(&huffman_description(&literals)?); - for table in [ - build_table_at_log(&counts.offset_codes, OF_LOG), - build_table_at_log(&counts.match_lengths, ML_LOG), - build_table_at_log(&counts.literal_lengths, LL_LOG), - ] { - out.extend_from_slice(&fse_description(&table)); - } + // Each description goes straight into the dictionary; the sequence tables + // are normalized and described without being built. + let mut writer = BitWriter::from(&mut *out); + HuffmanEncoder::new(&literals, &mut writer).write_table(); + write_ncount_at_log(&counts.offset_codes, OF_LOG, &mut writer); + write_ncount_at_log(&counts.match_lengths, ML_LOG, &mut writer); + write_ncount_at_log(&counts.literal_lengths, LL_LOG, &mut writer); + writer.flush(); Ok(()) } @@ -174,6 +175,7 @@ fn count_samples( .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?; let recorder = Recorder { inner: MatchGeneratorDriver::new(BLOCK_SIZE_MAX, 1), + armed: false, literals: [0; 256], sequences: Vec::new(), }; @@ -187,20 +189,31 @@ fn count_samples( // one set of tables for the whole pass, which also keeps the dictionary // resident from one sample to the next instead of indexing it again. let average = samples.leading(count).len() / count.max(1); + // Only a sample's first block is counted, `MIN(128 KiB, window)` bytes + // (upstream zstd `ZDICT_countEStats`). The window follows the hint alone, + // so once one frame has shown it, the rest of each sample is not fed at + // all. + let mut block_size = BLOCK_SIZE_MAX; let mut frame = Vec::new(); for index in 0..count { let sample = samples.sample(index); - let block = &sample[..sample.len().min(BLOCK_SIZE_MAX)]; + let block = &sample[..sample.len().min(block_size)]; if block.is_empty() { continue; } frame.clear(); + compressor.matcher_mut().armed = true; compressor.set_source(block); compressor.set_drain(frame); compressor.set_source_size_hint(average as u64); compressor.compress(); frame = compressor.take_drain().expect("the drain was set above"); + let fast_offset_codes = compressor.uses_fast_offset_codes(); let recorder = compressor.matcher_mut(); + let window = recorder.inner.window_size() as usize; + if window > 0 { + block_size = block_size.min(window); + } // A block written raw holds no sequences to learn from, and its bytes // are no literals of any compressed block (upstream skips it too). if first_block_is_compressed(&frame) { @@ -209,7 +222,12 @@ fn count_samples( } let mut reps = START_REPS; for &(ll, offset, ml) in &recorder.sequences { - let off_base = encode_offset_with_history(offset, ll, &mut reps); + // The code the block wrote for this offset, under its policy. + let off_base = if fast_offset_codes { + encode_offset_with_history_fast(offset, ll, &mut reps) + } else { + encode_offset_with_history(offset, ll, &mut reps) + }; counts.literal_lengths[encode_literal_length(ll).0 as usize] += 1; counts.match_lengths[encode_match_len(ml).0 as usize] += 1; if let Some(slot) = counts @@ -255,33 +273,13 @@ fn literals_table(counts: &[usize; 256]) -> HuffmanTable { } } -/// The literals table's description as a dictionary stores it. -fn huffman_description(table: &HuffmanTable) -> io::Result> { - // The encoder writes the description ahead of the data it codes; one - // symbol of data, then the decoder reports where the description ends. - let mut writer = BitWriter::new(); - let mut encoder = HuffmanEncoder::new(table, &mut writer); - encoder.encode(&[0], true); - let encoded = writer.dump(); - let mut decoder = HuffmanDecoderTable::new(); - let description_len = decoder - .build_decoder(encoded.as_slice()) - .map_err(|e| io::Error::other(format!("failed to decode generated huffman table: {e}")))?; - Ok(encoded[..description_len as usize].to_vec()) -} - -/// A sequence table's normalized counts, as a dictionary stores them. -fn fse_description(table: &FSETable) -> Vec { - let mut writer = BitWriter::new(); - table.write_table(&mut writer); - writer.dump() -} - -/// The production matcher, recording the literals and sequences each frame's -/// blocks are built from. Everything else is forwarded, so frames compress as -/// any compressor's would, dictionary reuse included. +/// The production matcher, recording the literals and sequences of the first +/// block it matches after being armed. Everything else is forwarded, so frames +/// compress as any compressor's would, dictionary reuse included. struct Recorder { inner: MatchGeneratorDriver, + /// Whether the next matched block is recorded; cleared once it is. + armed: bool, literals: [u32; 256], /// Literal length, offset and match length of each sequence. sequences: Vec<(u32, u32, u32)>, @@ -317,10 +315,15 @@ impl Matcher for Recorder { } fn start_matching(&mut self, mut handle_sequence: impl for<'a> FnMut(Sequence<'a>)) { + if !core::mem::take(&mut self.armed) { + self.inner.start_matching(handle_sequence); + return; + } let Self { inner, literals, sequences, + .. } = self; inner.start_matching(|sequence| { match &sequence { diff --git a/zstd/src/dictionary/finalize/tests.rs b/zstd/src/dictionary/finalize/tests.rs index 90ee510b4..7d1cc4c8e 100644 --- a/zstd/src/dictionary/finalize/tests.rs +++ b/zstd/src/dictionary/finalize/tests.rs @@ -1,6 +1,6 @@ use super::*; use crate::encoding::compress_slice_to_vec; -use std::string::String; +use std::{format, string::String}; /// Log lines with a shared vocabulary, one sample each. fn log_samples() -> (Vec, Vec) { @@ -115,7 +115,99 @@ fn literal_tables_take_the_depth_of_their_longest_code() { counts[254] = 1; let table = literals_table(&counts); assert_eq!(table.table_log(), 9); - assert!(!huffman_description(&table).unwrap().is_empty()); + let mut writer = BitWriter::new(); + HuffmanEncoder::new(&table, &mut writer).write_table(); + let description = writer.dump(); + let parsed = crate::huff0::HuffmanTable::new() + .build_decoder(&description) + .expect("the description parses back"); + assert_eq!(parsed as usize, description.len()); +} + +/// Only a sample's first block is counted (upstream zstd `ZDICT_countEStats` +/// compresses `MIN(128 KiB, window)` bytes as one block). Samples averaging a +/// line keep the window, and so the block, far under 4 KiB; two versions of an +/// 8 KiB sample that differ only past 4 KiB must finalize to the same bytes. +#[test] +fn only_the_first_block_of_a_sample_is_counted() { + let (mut data, mut sizes) = log_samples(); + let content: Vec = data[..4096].to_vec(); + let head = data[..4096].to_vec(); + let mut state = 0x2545_F491_4F6C_DD1Du64; + let noise: Vec = (0..4096) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state as u8 + }) + .collect(); + sizes.push(8192); + let finalized = |tail: &[u8], data: &mut Vec| { + let base = data.len(); + data.extend_from_slice(&head); + data.extend_from_slice(tail); + let samples = SampleSet::new(data, &sizes).unwrap(); + let dict = finalize( + &content, + &samples, + samples.count(), + 8192, + FinalizeOptions::default(), + ) + .unwrap(); + data.truncate(base); + dict + }; + let with_runs = finalized(&[b'z'; 4096], &mut data); + let with_noise = finalized(&noise, &mut data); + assert!( + with_runs == with_noise, + "a sample's second block reached the tables" + ); +} + +/// Offset codes are counted with the repeat policy the frame's blocks used. +/// Two tokens alternating behind short varied gaps make an offset often equal +/// the one before last: the full search writes it as repeat 2 or 3 (offset +/// code 1), which the fast strategy never emits, so at level 1 that code keeps +/// its floor of one while at level 3 it is counted. +#[test] +fn offset_codes_follow_the_frames_repeat_policy() { + let mut state = 0x9E37_79B9_7F4A_7C15u64; + let mut next = || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state + }; + let tokens = [*b"alpha-token-0001", *b"bravo-token-0002"]; + let mut data = Vec::new(); + let mut sizes = Vec::new(); + for _ in 0..64 { + let start = data.len(); + for i in 0..40 { + data.extend_from_slice(&tokens[i % 2]); + for _ in 0..1 + next() % 2 { + data.push(b'!' + (next() % 60) as u8); + } + } + sizes.push(data.len() - start); + } + let samples = SampleSet::new(&data, &sizes).unwrap(); + let content = samples.sample(0).to_vec(); + let offset_code_one = |level| { + let mut counts = EntropyCounts { + literals: [1; 256], + offset_codes: vec![1; OFFCODE_MAX as usize + 1], + match_lengths: [1; 53], + literal_lengths: [1; 36], + }; + count_samples(&mut counts, &content, &samples, samples.count(), level).unwrap(); + counts.offset_codes[1] + }; + assert!(offset_code_one(3) > 1, "the input exercises deeper repeats"); + assert_eq!(offset_code_one(1), 1); } /// Only a frame whose first block is compressed counts as one. diff --git a/zstd/src/dictionary/mod.rs b/zstd/src/dictionary/mod.rs index 0553b20cd..432be6d5f 100644 --- a/zstd/src/dictionary/mod.rs +++ b/zstd/src/dictionary/mod.rs @@ -338,7 +338,8 @@ fn search_raw_content(corpus: &[u8], sizes: &[usize], dict_size: usize) -> io::R /// `InvalidInput` when `raw_content` is empty, when `dict_size` is below /// [`MIN_TRAINED_DICT_SIZE`] or leaves less than eight bytes of content after /// the header, when the id is zero, or when `sample_sizes` does not add up to -/// `samples.len()`. +/// `samples.len()`; a size too small carries +/// [`TrainingError::DictionaryTooSmall`]. /// /// # Examples /// diff --git a/zstd/src/encoding/blocks/compressed.rs b/zstd/src/encoding/blocks/compressed.rs index bd7c7f0f7..b15a85f48 100644 --- a/zstd/src/encoding/blocks/compressed.rs +++ b/zstd/src/encoding/blocks/compressed.rs @@ -773,10 +773,7 @@ fn encode_block_parts( ml: &mut ml_counts, of: &mut of_counts, }; - let (ll_max, ml_max, of_max) = if matches!( - state.strategy_tag, - crate::encoding::strategy::StrategyTag::Fast - ) { + let (ll_max, ml_max, of_max) = if uses_fast_offset_codes(state.strategy_tag) { fill_and_count::(raw_sequences, &mut state.offset_hist, counts, codes) } else { fill_and_count::(raw_sequences, &mut state.offset_hist, counts, codes) @@ -1266,7 +1263,7 @@ fn estimate_sequences_section_bytes( // Upstream zstd: OF code's value equals its additional-bits width. of_bits += of as usize; }; - if matches!(strategy, crate::encoding::strategy::StrategyTag::Fast) { + if uses_fast_offset_codes(strategy) { for seq in sequences { count_offset(encode_offset_with_history_fast( seq.off_base, @@ -3033,6 +3030,14 @@ pub(crate) fn encode_offset_with_history( encoded } +/// Whether blocks at `strategy` code offsets with +/// [`encode_offset_with_history_fast`] rather than the full repeat search of +/// [`encode_offset_with_history`]. +#[inline] +pub(crate) fn uses_fast_offset_codes(strategy: crate::encoding::strategy::StrategyTag) -> bool { + matches!(strategy, crate::encoding::strategy::StrategyTag::Fast) +} + /// Fast-matcher offset→offBase conversion, mirroring upstream zstd's /// `ZSTD_compressBlock_fast`: emit offBase 1 only for the immediate repeat /// offset (`rep[0]` when `lit_len > 0`, `rep[1]` when `lit_len == 0` — the @@ -3045,7 +3050,7 @@ pub(crate) fn encode_offset_with_history( /// per-sequence probe the fast matcher never pays and can shift the FSE symbol /// histogram the wrong way). The repeat-offset history update follows directly /// from the emitted code, identical to the full converter's rules. -pub(in crate::encoding) fn encode_offset_with_history_fast( +pub(crate) fn encode_offset_with_history_fast( actual_offset: u32, lit_len: u32, offset_hist: &mut [u32; 3], diff --git a/zstd/src/encoding/blocks/mod.rs b/zstd/src/encoding/blocks/mod.rs index c47e1661e..149f9c23b 100644 --- a/zstd/src/encoding/blocks/mod.rs +++ b/zstd/src/encoding/blocks/mod.rs @@ -10,4 +10,5 @@ pub(super) use compressed::*; #[cfg(feature = "dict-builder")] pub(crate) use compressed::{ encode_literal_length, encode_match_len, encode_offset, encode_offset_with_history, + encode_offset_with_history_fast, uses_fast_offset_codes, }; diff --git a/zstd/src/encoding/frame_compressor.rs b/zstd/src/encoding/frame_compressor.rs index 8fca059f3..19b5c36a4 100644 --- a/zstd/src/encoding/frame_compressor.rs +++ b/zstd/src/encoding/frame_compressor.rs @@ -3237,6 +3237,13 @@ impl FrameCompressor { &mut self.state.matcher } + /// Whether the last frame's blocks coded offsets with the fast matcher's + /// repeat policy, for a caller counting offset codes as they did. + #[cfg(feature = "dict-builder")] + pub(crate) fn uses_fast_offset_codes(&self) -> bool { + crate::encoding::blocks::uses_fast_offset_codes(self.state.strategy_tag) + } + /// Before calling [FrameCompressor::compress] you can replace the matcher pub fn replace_matcher(&mut self, mut match_generator: M) -> M { core::mem::swap(&mut match_generator, &mut self.state.matcher); diff --git a/zstd/src/fse/fse_encoder.rs b/zstd/src/fse/fse_encoder.rs index 008ae4d76..807e1b043 100644 --- a/zstd/src/fse/fse_encoder.rs +++ b/zstd/src/fse/fse_encoder.rs @@ -292,61 +292,69 @@ impl FSETable { } pub(crate) fn write_table>>(&self, writer: &mut BitWriter) { - assert!( - writer.index().is_multiple_of(8), - "FSE table headers must start on a byte boundary" - ); - #[cfg(debug_assertions)] - let start_idx = writer.index(); - writer.write_bits(self.acc_log() - 5, 4); - let mut probability_counter = 0usize; - let probability_sum = 1 << self.acc_log(); - - let mut prob_idx = 0; - while probability_counter < probability_sum { - let max_remaining_value = probability_sum - probability_counter + 1; - let bits_to_write = max_remaining_value.ilog2() + 1; - let low_threshold = ((1 << bits_to_write) - 1) - (max_remaining_value); - let mask = (1 << (bits_to_write - 1)) - 1; - - let prob = self.states[prob_idx].probability; - prob_idx += 1; - let value = (prob + 1) as u32; - if value < low_threshold as u32 { - writer.write_bits(value, bits_to_write as usize - 1); - } else if value > mask { - writer.write_bits(value + low_threshold as u32, bits_to_write as usize); - } else { - writer.write_bits(value, bits_to_write as usize); - } + let probs: [i32; 256] = core::array::from_fn(|i| self.states[i].probability); + write_ncount(&probs, self.acc_log(), writer); + } +} + +/// Write the NCount description of the normalized distribution `probs` at +/// `acc_log` (RFC 8878 4.1.1), padded to a byte; what [`FSETable::write_table`] +/// writes for a built table, from the normalized counts alone. +fn write_ncount>>(probs: &[i32], acc_log: u8, writer: &mut BitWriter) { + assert!( + writer.index().is_multiple_of(8), + "FSE table headers must start on a byte boundary" + ); + #[cfg(debug_assertions)] + let start_idx = writer.index(); + writer.write_bits(acc_log - 5, 4); + let mut probability_counter = 0usize; + let probability_sum = 1 << acc_log; + + let mut prob_idx = 0; + while probability_counter < probability_sum { + let max_remaining_value = probability_sum - probability_counter + 1; + let bits_to_write = max_remaining_value.ilog2() + 1; + let low_threshold = ((1 << bits_to_write) - 1) - (max_remaining_value); + let mask = (1 << (bits_to_write - 1)) - 1; + + let prob = probs[prob_idx]; + prob_idx += 1; + let value = (prob + 1) as u32; + if value < low_threshold as u32 { + writer.write_bits(value, bits_to_write as usize - 1); + } else if value > mask { + writer.write_bits(value + low_threshold as u32, bits_to_write as usize); + } else { + writer.write_bits(value, bits_to_write as usize); + } - if prob == -1 { - probability_counter += 1; - } else if prob > 0 { - probability_counter += prob as usize; - } else { - let mut zeros = 0u8; - while prob_idx < self.states.len() && self.states[prob_idx].probability == 0 { - zeros += 1; - prob_idx += 1; - if zeros == 3 { - writer.write_bits(3u8, 2); - zeros = 0; - } + if prob == -1 { + probability_counter += 1; + } else if prob > 0 { + probability_counter += prob as usize; + } else { + let mut zeros = 0u8; + while prob_idx < probs.len() && probs[prob_idx] == 0 { + zeros += 1; + prob_idx += 1; + if zeros == 3 { + writer.write_bits(3u8, 2); + zeros = 0; } - writer.write_bits(zeros, 2); } + writer.write_bits(zeros, 2); } - writer.write_bits(0u8, writer.misaligned()); - #[cfg(debug_assertions)] - { - let written_bits = writer.index() - start_idx; - let computed = self.table_header_bits(); - debug_assert_eq!( - written_bits, computed, - "table_header_bits() mismatch: written={written_bits}, computed={computed}" - ); - } + } + writer.write_bits(0u8, writer.misaligned()); + #[cfg(debug_assertions)] + { + let written_bits = writer.index() - start_idx; + let computed = ncount_header_bits(probs, acc_log); + debug_assert_eq!( + written_bits, computed, + "ncount_header_bits() mismatch: written={written_bits}, computed={computed}" + ); } } @@ -499,13 +507,17 @@ pub(crate) fn build_seq_ctable_into( build_table_from_probabilities_into(&probs[..=max_symbol], table_log, out); } -/// An FSE table at exactly `table_log` for `counts`, with low-probability -/// symbols kept at `-1` (upstream zstd `FSE_normalizeCount` given the log -/// outright, as the dictionary finalizer calls it, zdict.c -/// `ZDICT_analyzeEntropy`). `counts` must not be too spread for the log: -/// every symbol present needs a slot. +/// Write the description of `counts` normalized at exactly `table_log`, with +/// low-probability symbols kept at `-1` (upstream zstd `FSE_normalizeCount` +/// given the log outright, then `FSE_writeNCount`, as the dictionary finalizer +/// calls them, zdict.c `ZDICT_analyzeEntropy`). No table is built. `counts` +/// must not be too spread for the log: every symbol present needs a slot. #[cfg(feature = "dict-builder")] -pub(crate) fn build_table_at_log(counts: &[usize], table_log: u8) -> FSETable { +pub(crate) fn write_ncount_at_log>>( + counts: &[usize], + table_log: u8, + writer: &mut BitWriter, +) { let total = counts.iter().sum::(); let max_symbol = counts .iter() @@ -521,7 +533,7 @@ pub(crate) fn build_table_at_log(counts: &[usize], table_log: u8) -> FSETable { max_symbol, true, ); - build_table_from_probabilities(&probs[..=max_symbol], table_log) + write_ncount(&probs[..=max_symbol], table_log, writer); } #[cfg(any(test, feature = "fuzz-exports"))] diff --git a/zstd/src/fse/tests.rs b/zstd/src/fse/tests.rs index d85a4e278..8478b278a 100644 --- a/zstd/src/fse/tests.rs +++ b/zstd/src/fse/tests.rs @@ -54,6 +54,36 @@ fn rebuilding_a_used_table_matches_a_fresh_build() { assert_eq!(reused.table_header_bits(), fresh.table_header_bits()); } +/// A description written from counts at a given log, with low-probability +/// symbols and zero runs in it, is the one the table it describes writes: the +/// decoder reads exactly the bytes written, and the encoder table rebuilt from +/// what it read writes the same bytes back. +#[cfg(feature = "dict-builder")] +#[test] +fn a_description_written_from_counts_is_the_tables_own() { + let mut counts = [0usize; 31]; + counts[0] = 900; + counts[1] = 1; + counts[4] = 3; + counts[9] = 1; + counts[30] = 40; + let mut buf = alloc::vec::Vec::new(); + let mut writer = crate::bit_io::BitWriter::from(&mut buf); + fse_encoder::write_ncount_at_log(&counts, 8, &mut writer); + writer.flush(); + let mut dec_table = FSETable::new(255); + let read = dec_table.read_table_probabilities(&buf, 8).unwrap(); + assert_eq!(read, buf.len()); + let table = dec_table + .to_encoder_table() + .expect("log 8 has an encoder table"); + let mut again = alloc::vec::Vec::new(); + let mut writer = crate::bit_io::BitWriter::from(&mut again); + table.write_table(&mut writer); + writer.flush(); + assert_eq!(again, buf); +} + #[test] fn decoder_entry_layout_is_four_bytes_for_huffman_weights() { assert_eq!(core::mem::size_of::(), 4); diff --git a/zstd/src/huff0/huff0_encoder.rs b/zstd/src/huff0/huff0_encoder.rs index 4e2aec021..d7e0009a3 100644 --- a/zstd/src/huff0/huff0_encoder.rs +++ b/zstd/src/huff0/huff0_encoder.rs @@ -348,7 +348,8 @@ impl>> HuffmanEncoder<'_, '_, V> { self.weights_into(&mut buf).to_vec() } - fn write_table(&mut self) { + /// Write the table's description alone, ending on a byte boundary. + pub(crate) fn write_table(&mut self) { #[cfg(feature = "std")] { // Cached path: the size query that precedes every emit has already From 73e326b4988360fe063d9830912d0522c40f8b3c Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Mon, 28 Sep 2026 17:27:25 +0300 Subject: [PATCH 07/20] fix(dict): size the analysed block as upstream does The entropy analysis runs upstream's parameters for the average sample and the content, ZSTD_getParams(level, averageSampleSize, dictSize), set explicitly, and hints the frames as a source of both together, the size upstream fits the window to. A frame with a dictionary fits its window to the source alone, so hinting the average sample cut the analysed block to the sample's window: on the repository files the fixed COVER dictionary lost 999 bytes against the one before. Matcher gains apply_parameters, a default no-op the built-in matcher implements, so FrameCompressor::set_parameters works with any matcher. Regression test: the_first_block_spans_the_window_of_sample_and_content. --- zstd/src/dictionary/finalize.rs | 53 ++++++++---- zstd/src/dictionary/finalize/tests.rs | 50 ++++++++++- zstd/src/encoding/frame_compressor.rs | 106 +++++++++++------------ zstd/src/encoding/match_generator/mod.rs | 4 + zstd/src/encoding/mod.rs | 6 ++ 5 files changed, 147 insertions(+), 72 deletions(-) diff --git a/zstd/src/dictionary/finalize.rs b/zstd/src/dictionary/finalize.rs index fb26bc60c..26e3a5131 100644 --- a/zstd/src/dictionary/finalize.rs +++ b/zstd/src/dictionary/finalize.rs @@ -16,8 +16,8 @@ use crate::encoding::blocks::{ }; use crate::encoding::workspace::Workspace; use crate::encoding::{ - CompressionLevel, DictionarySizes, EncoderDictionary, FrameCompressor, HistoryBuf, - MatchGeneratorDriver, Matcher, Sequence, + CompressionLevel, CompressionParameters, DictionarySizes, EncoderDictionary, FrameCompressor, + HistoryBuf, LevelParameters, MatchGeneratorDriver, Matcher, Sequence, }; use crate::fse::fse_encoder::{FSETable, write_ncount_at_log}; use crate::huff0::huff0_encoder::{HuffmanEncoder, HuffmanTable}; @@ -184,16 +184,37 @@ fn count_samples( compressor .set_encoder_dictionary(EncoderDictionary::from_dictionary(dictionary)) .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?; - // Every sample runs the parameters of the average one, as upstream's - // analysis does (`ZSTD_getParams(level, averageSampleSize, dictSize)`): - // one set of tables for the whole pass, which also keeps the dictionary - // resident from one sample to the next instead of indexing it again. + // Every sample runs the parameters of the average one with this content, + // as upstream's analysis does (`ZSTD_getParams(level, averageSampleSize, + // dictSize)`): one set of tables for the whole pass, which also keeps the + // dictionary resident from one sample to the next instead of indexing it + // again. let average = samples.leading(count).len() / count.max(1); - // Only a sample's first block is counted, `MIN(128 KiB, window)` bytes - // (upstream zstd `ZDICT_countEStats`). The window follows the hint alone, - // so once one frame has shown it, the rest of each sample is not fed at - // all. - let mut block_size = BLOCK_SIZE_MAX; + let params = LevelParameters::for_level( + level, + Some(average as u64).filter(|&size| size != 0), + content.len(), + ); + compressor.set_parameters( + &CompressionParameters::builder(CompressionLevel::from_level(level)) + .window_log(params.window_log) + .hash_log(params.hash_log) + .chain_log(params.chain_log) + .search_log(params.search_log) + .min_match(params.min_match) + .target_length(params.target_length) + .strategy(params.strategy) + .build() + .expect("the level table's parameters are in range"), + ); + // Only a sample's first block is counted, `MIN(128 KiB, 1 << windowLog)` + // bytes compressed as one block (upstream zstd `ZDICT_countEStats`). The + // rest of a sample is not fed at all. + let block_size = BLOCK_SIZE_MAX.min(1 << params.window_log); + // A frame's window is fitted to its source, and upstream's analysis sizes + // it for the sample and the content together (`ZSTD_adjustCParams`, + // `tSize = srcSize + dictSize`), so that is the size every frame hints. + let sized_as = (average + content.len()) as u64; let mut frame = Vec::new(); for index in 0..count { let sample = samples.sample(index); @@ -205,15 +226,11 @@ fn count_samples( compressor.matcher_mut().armed = true; compressor.set_source(block); compressor.set_drain(frame); - compressor.set_source_size_hint(average as u64); + compressor.set_source_size_hint(sized_as); compressor.compress(); frame = compressor.take_drain().expect("the drain was set above"); let fast_offset_codes = compressor.uses_fast_offset_codes(); let recorder = compressor.matcher_mut(); - let window = recorder.inner.window_size() as usize; - if window > 0 { - block_size = block_size.min(window); - } // A block written raw holds no sequences to learn from, and its bytes // are no literals of any compressed block (upstream skips it too). if first_block_is_compressed(&frame) { @@ -372,6 +389,10 @@ impl Matcher for Recorder { self.inner.clear_param_overrides(); } + fn apply_parameters(&mut self, params: &CompressionParameters) { + self.inner.apply_parameters(params); + } + fn prime_with_dictionary(&mut self, dict_content: &[u8], offset_hist: [u32; 3]) { self.inner.prime_with_dictionary(dict_content, offset_hist); } diff --git a/zstd/src/dictionary/finalize/tests.rs b/zstd/src/dictionary/finalize/tests.rs index 7d1cc4c8e..5518b2d5e 100644 --- a/zstd/src/dictionary/finalize/tests.rs +++ b/zstd/src/dictionary/finalize/tests.rs @@ -125,13 +125,14 @@ fn literal_tables_take_the_depth_of_their_longest_code() { } /// Only a sample's first block is counted (upstream zstd `ZDICT_countEStats` -/// compresses `MIN(128 KiB, window)` bytes as one block). Samples averaging a -/// line keep the window, and so the block, far under 4 KiB; two versions of an +/// compresses `MIN(128 KiB, 1 << windowLog)` bytes as one block, the window of +/// the average sample plus the content). Samples averaging a line with half a +/// KiB of content keep that window at its 1 KiB floor; two versions of an /// 8 KiB sample that differ only past 4 KiB must finalize to the same bytes. #[test] fn only_the_first_block_of_a_sample_is_counted() { let (mut data, mut sizes) = log_samples(); - let content: Vec = data[..4096].to_vec(); + let content: Vec = data[..512].to_vec(); let head = data[..4096].to_vec(); let mut state = 0x2545_F491_4F6C_DD1Du64; let noise: Vec = (0..4096) @@ -167,6 +168,49 @@ fn only_the_first_block_of_a_sample_is_counted() { ); } +/// The first block spans the window of the average sample plus the content, +/// not of the sample alone: with 4 KiB of content that window is 8 KiB, so the +/// second half of an 8 KiB sample is counted and two versions differing there +/// finalize differently. +#[test] +fn the_first_block_spans_the_window_of_sample_and_content() { + let (mut data, mut sizes) = log_samples(); + let content: Vec = data[..4096].to_vec(); + let head = data[..4096].to_vec(); + let mut state = 0x2545_F491_4F6C_DD1Du64; + let noise: Vec = (0..4096) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state as u8 + }) + .collect(); + sizes.push(8192); + let finalized = |tail: &[u8], data: &mut Vec| { + let base = data.len(); + data.extend_from_slice(&head); + data.extend_from_slice(tail); + let samples = SampleSet::new(data, &sizes).unwrap(); + let dict = finalize( + &content, + &samples, + samples.count(), + 8192, + FinalizeOptions::default(), + ) + .unwrap(); + data.truncate(base); + dict + }; + let with_runs = finalized(&[b'z'; 4096], &mut data); + let with_noise = finalized(&noise, &mut data); + assert!( + with_runs != with_noise, + "the second half of the sample's first block was not counted" + ); +} + /// Offset codes are counted with the repeat policy the frame's blocks used. /// Two tokens alternating behind short varied gaps make an offset often equal /// the one before last: the full search writes it as repeat 2 or 3 (offset diff --git a/zstd/src/encoding/frame_compressor.rs b/zstd/src/encoding/frame_compressor.rs index 19b5c36a4..4b70fb32c 100644 --- a/zstd/src/encoding/frame_compressor.rs +++ b/zstd/src/encoding/frame_compressor.rs @@ -1802,59 +1802,6 @@ impl FrameCompressor { } } - /// Configure fine-grained compression parameters. - /// - /// Resets the base [`CompressionLevel`] to the parameters' level and installs the per-knob overrides - /// (window/hash/chain/search logs, strategy, LDM) applied at the next - /// frame. Pass `None`-equivalent (a builder that overrides nothing) - /// to fall back to plain level-based compression. - /// - /// ```rust - /// use structured_zstd::encoding::{ - /// CompressionLevel, CompressionParameters, FrameCompressor, Strategy, - /// }; - /// let params = CompressionParameters::builder(CompressionLevel::Level(19)) - /// .strategy(Strategy::Btultra2) - /// .enable_long_distance_matching(true) - /// .build() - /// .unwrap(); - /// let mut compressor: FrameCompressor = FrameCompressor::new(CompressionLevel::Default); - /// compressor.set_parameters(¶ms); - /// let compressed = compressor.compress_independent_frame(b"some data to compress"); - /// assert!(!compressed.is_empty()); - /// ``` - pub fn set_parameters(&mut self, params: &crate::encoding::CompressionParameters) { - self.compression_level = params.level(); - let overrides = params.overrides(); - self.tuning = FrameTuning::from_overrides(&overrides); - // Keep `state.strategy_tag` consistent immediately so the borrowed - // one-shot eligibility gate (`borrowed_eligible`) and literal gates - // are correct even before the next `compress()` re-sync. Resolve it - // size-adaptively (same `resolve_level_params` path `prepare_frame` - // uses) so a hint already set here yields the same strategy the matcher - // will run, not the bare level-only mapping. - // The dictionary counts only when the frame will prime it (same gate - // as `prepare_frame`'s `use_dictionary_state`): uncompressed mode - // ignores an attached dictionary and has no CDict tier to resolve. - let with_dictionary = !matches!(self.compression_level, CompressionLevel::Uncompressed) - && self.state.matcher.supports_dictionary_priming(); - let params = self.resolve_frame_params(self.source_size_hint, with_dictionary); - self.sync_effective_strategy(¶ms); - self.state.huf_optimal_search = - huf_search_enabled(self.state.strategy_tag, self.source_size_hint); - self.state.literal_compression_disabled = literal_compression_disabled( - self.state.strategy_tag, - self.compression_level, - gate_target_length( - self.compression_level, - &self.tuning, - self.dictionary.as_ref().filter(|_| with_dictionary), - ), - self.tuning.literal_compression, - ); - self.state.matcher.set_param_overrides(Some(overrides)); - } - /// Whether the borrowed (no per-block history copy) one-shot loop is /// valid for an `input_len`-byte slice under the resolved `prep`. /// @@ -2378,6 +2325,59 @@ impl FrameCompressor { ); } + /// Configure fine-grained compression parameters. + /// + /// Resets the base [`CompressionLevel`] to the parameters' level and installs the per-knob overrides + /// (window/hash/chain/search logs, strategy, LDM) applied at the next + /// frame. Pass `None`-equivalent (a builder that overrides nothing) + /// to fall back to plain level-based compression. A custom matcher + /// receives them through [`Matcher::apply_parameters`]. + /// + /// ```rust + /// use structured_zstd::encoding::{ + /// CompressionLevel, CompressionParameters, FrameCompressor, Strategy, + /// }; + /// let params = CompressionParameters::builder(CompressionLevel::Level(19)) + /// .strategy(Strategy::Btultra2) + /// .enable_long_distance_matching(true) + /// .build() + /// .unwrap(); + /// let mut compressor: FrameCompressor = FrameCompressor::new(CompressionLevel::Default); + /// compressor.set_parameters(¶ms); + /// let compressed = compressor.compress_independent_frame(b"some data to compress"); + /// assert!(!compressed.is_empty()); + /// ``` + pub fn set_parameters(&mut self, params: &crate::encoding::CompressionParameters) { + self.compression_level = params.level(); + self.tuning = FrameTuning::from_overrides(¶ms.overrides()); + // Keep `state.strategy_tag` consistent immediately so the borrowed + // one-shot eligibility gate (`borrowed_eligible`) and literal gates + // are correct even before the next `compress()` re-sync. Resolve it + // size-adaptively (same `resolve_level_params` path `prepare_frame` + // uses) so a hint already set here yields the same strategy the matcher + // will run, not the bare level-only mapping. + // The dictionary counts only when the frame will prime it (same gate + // as `prepare_frame`'s `use_dictionary_state`): uncompressed mode + // ignores an attached dictionary and has no CDict tier to resolve. + let with_dictionary = !matches!(self.compression_level, CompressionLevel::Uncompressed) + && self.state.matcher.supports_dictionary_priming(); + let level_params = self.resolve_frame_params(self.source_size_hint, with_dictionary); + self.sync_effective_strategy(&level_params); + self.state.huf_optimal_search = + huf_search_enabled(self.state.strategy_tag, self.source_size_hint); + self.state.literal_compression_disabled = literal_compression_disabled( + self.state.strategy_tag, + self.compression_level, + gate_target_length( + self.compression_level, + &self.tuning, + self.dictionary.as_ref().filter(|_| with_dictionary), + ), + self.tuning.literal_compression, + ); + self.state.matcher.apply_parameters(params); + } + /// Provide a hint about the total uncompressed size for the next frame. /// /// When set, the encoder selects smaller hash tables and windows for diff --git a/zstd/src/encoding/match_generator/mod.rs b/zstd/src/encoding/match_generator/mod.rs index a0bfe5831..dc82ddc21 100644 --- a/zstd/src/encoding/match_generator/mod.rs +++ b/zstd/src/encoding/match_generator/mod.rs @@ -1110,6 +1110,10 @@ impl Matcher for MatchGeneratorDriver { self.param_overrides = None; } + fn apply_parameters(&mut self, params: &super::CompressionParameters) { + self.param_overrides = Some(params.overrides()); + } + fn reset(&mut self, level: CompressionLevel) { // On its own the driver lays its tables out in a workspace of its own, // which is moved out for the call and back: moving it leaves the diff --git a/zstd/src/encoding/mod.rs b/zstd/src/encoding/mod.rs index add752c44..aedfdb9bd 100644 --- a/zstd/src/encoding/mod.rs +++ b/zstd/src/encoding/mod.rs @@ -562,6 +562,12 @@ pub trait Matcher { /// so switching back to a bare level after a customized frame does not /// keep the old overrides sticky. Default no-op for custom matchers. fn clear_param_overrides(&mut self) {} + /// Install the fine-grained overrides of `params` (window, table logs, + /// strategy, LDM) for the next frame. Called by + /// [`FrameCompressor::set_parameters`](crate::encoding::FrameCompressor::set_parameters); + /// consumed at the next [`reset`](Self::reset). Default no-op for custom + /// matchers, which size themselves. + fn apply_parameters(&mut self, _params: &CompressionParameters) {} /// Prime matcher state with dictionary history before compressing the next frame. /// Default implementation is a no-op for custom matchers that do not support this. fn prime_with_dictionary(&mut self, _dict_content: &[u8], _offset_hist: [u32; 3]) {} From bf33753d56540cd44394cc4f5b9668f0f9e37368 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Mon, 28 Sep 2026 17:40:35 +0300 Subject: [PATCH 08/20] fix(dict): count every compressed block, skip raw ones The entropy analysis counts every block of a sample again, and now leaves out each block written raw or as one repeated byte: its bytes are no literals of any compressed block and it holds no sequences. The recorder keeps one entry per block, matched or skipped, and the frame's block headers say which were compressed. Counting the first block alone, as upstream's ZDICT_countEStats does, measured worse on the repository files: at our window the samples compressed 0.013% larger in total across fixed COVER, fixed FastCOVER and the default search, at upstream's window (average sample plus content, set with its parameters) 0.05% larger. This replaces both, including the Matcher::apply_parameters hook the second one needed. Regression test: a_block_written_raw_is_not_counted. --- zstd/src/dictionary/finalize.rs | 161 +++++++++++++---------- zstd/src/dictionary/finalize/tests.rs | 105 +++++---------- zstd/src/encoding/frame_compressor.rs | 106 +++++++-------- zstd/src/encoding/match_generator/mod.rs | 4 - zstd/src/encoding/mod.rs | 6 - 5 files changed, 175 insertions(+), 207 deletions(-) diff --git a/zstd/src/dictionary/finalize.rs b/zstd/src/dictionary/finalize.rs index 26e3a5131..b7e4dc1e7 100644 --- a/zstd/src/dictionary/finalize.rs +++ b/zstd/src/dictionary/finalize.rs @@ -16,8 +16,8 @@ use crate::encoding::blocks::{ }; use crate::encoding::workspace::Workspace; use crate::encoding::{ - CompressionLevel, CompressionParameters, DictionarySizes, EncoderDictionary, FrameCompressor, - HistoryBuf, LevelParameters, MatchGeneratorDriver, Matcher, Sequence, + CompressionLevel, DictionarySizes, EncoderDictionary, FrameCompressor, HistoryBuf, + MatchGeneratorDriver, Matcher, Sequence, }; use crate::fse::fse_encoder::{FSETable, write_ncount_at_log}; use crate::huff0::huff0_encoder::{HuffmanEncoder, HuffmanTable}; @@ -162,8 +162,8 @@ fn analyze_entropy( Ok(()) } -/// Compress the first block of each of the first `count` samples with -/// `content` as a raw dictionary and count what the compressed blocks hold. +/// Compress the first 128 KiB of each of the first `count` samples with +/// `content` as a raw dictionary and count what its compressed blocks hold. fn count_samples( counts: &mut EntropyCounts, content: &[u8], @@ -175,8 +175,7 @@ fn count_samples( .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?; let recorder = Recorder { inner: MatchGeneratorDriver::new(BLOCK_SIZE_MAX, 1), - armed: false, - literals: [0; 256], + blocks: Vec::new(), sequences: Vec::new(), }; let mut compressor: FrameCompressor<&[u8], Vec, Recorder> = @@ -184,61 +183,59 @@ fn count_samples( compressor .set_encoder_dictionary(EncoderDictionary::from_dictionary(dictionary)) .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?; - // Every sample runs the parameters of the average one with this content, - // as upstream's analysis does (`ZSTD_getParams(level, averageSampleSize, - // dictSize)`): one set of tables for the whole pass, which also keeps the - // dictionary resident from one sample to the next instead of indexing it - // again. + // Every sample runs the parameters of the average one, as upstream's + // analysis does (`ZSTD_getParams(level, averageSampleSize, dictSize)`): + // one set of tables for the whole pass, which also keeps the dictionary + // resident from one sample to the next instead of indexing it again. let average = samples.leading(count).len() / count.max(1); - let params = LevelParameters::for_level( - level, - Some(average as u64).filter(|&size| size != 0), - content.len(), - ); - compressor.set_parameters( - &CompressionParameters::builder(CompressionLevel::from_level(level)) - .window_log(params.window_log) - .hash_log(params.hash_log) - .chain_log(params.chain_log) - .search_log(params.search_log) - .min_match(params.min_match) - .target_length(params.target_length) - .strategy(params.strategy) - .build() - .expect("the level table's parameters are in range"), - ); - // Only a sample's first block is counted, `MIN(128 KiB, 1 << windowLog)` - // bytes compressed as one block (upstream zstd `ZDICT_countEStats`). The - // rest of a sample is not fed at all. - let block_size = BLOCK_SIZE_MAX.min(1 << params.window_log); - // A frame's window is fitted to its source, and upstream's analysis sizes - // it for the sample and the content together (`ZSTD_adjustCParams`, - // `tSize = srcSize + dictSize`), so that is the size every frame hints. - let sized_as = (average + content.len()) as u64; + // Every block of a sample is counted, not only the first as upstream's + // `ZDICT_countEStats` does: counting the first block alone, at our window + // or at upstream's, measured worse on the repository files (the samples + // compressed 0.013% and 0.05% larger in total across COVER, FastCOVER + // and the default search), so the rest of a sample is signal, not noise. let mut frame = Vec::new(); + let mut compressed = Vec::new(); for index in 0..count { let sample = samples.sample(index); - let block = &sample[..sample.len().min(block_size)]; + let block = &sample[..sample.len().min(BLOCK_SIZE_MAX)]; if block.is_empty() { continue; } frame.clear(); - compressor.matcher_mut().armed = true; compressor.set_source(block); compressor.set_drain(frame); - compressor.set_source_size_hint(sized_as); + compressor.set_source_size_hint(average as u64); compressor.compress(); frame = compressor.take_drain().expect("the drain was set above"); let fast_offset_codes = compressor.uses_fast_offset_codes(); + compressed_blocks(&frame, &mut compressed); let recorder = compressor.matcher_mut(); - // A block written raw holds no sequences to learn from, and its bytes - // are no literals of any compressed block (upstream skips it too). - if first_block_is_compressed(&frame) { - for (total, &seen) in counts.literals.iter_mut().zip(&recorder.literals) { + // A block written raw or as one repeated byte holds no sequences to + // learn from, and its bytes are no literals of any compressed block. + // Blocks are recorded one per block the frame holds; when one was + // split after matching the pieces cannot be told apart, and the + // frame's blocks count only if every piece was compressed. + let one_to_one = compressed.len() == recorder.blocks.len(); + let all_compressed = compressed.iter().all(|&kind| kind); + let mut reps = START_REPS; + let mut start = 0; + for (index, block) in recorder.blocks.iter().enumerate() { + let block_sequences = &recorder.sequences[start..block.sequences_end]; + start = block.sequences_end; + let counted = if one_to_one { + compressed[index] + } else { + all_compressed + }; + // A block left out does not advance the repeat offsets either: the + // encoder restores them when it writes a block raw. + if !counted { + continue; + } + for (total, &seen) in counts.literals.iter_mut().zip(&block.literals) { *total += seen as usize; } - let mut reps = START_REPS; - for &(ll, offset, ml) in &recorder.sequences { + for &(ll, offset, ml) in block_sequences { // The code the block wrote for this offset, under its policy. let off_base = if fast_offset_codes { encode_offset_with_history_fast(offset, ll, &mut reps) @@ -255,22 +252,33 @@ fn count_samples( } } } - recorder.literals = [0; 256]; + recorder.blocks.clear(); recorder.sequences.clear(); } Ok(()) } -/// Whether `frame`'s first block is a compressed one. -fn first_block_is_compressed(frame: &[u8]) -> bool { +/// Whether each block of `frame` is a compressed one, into `out`; nothing for +/// bytes that are not a frame. +fn compressed_blocks(frame: &[u8], out: &mut Vec) { + out.clear(); let Ok((_, header_len)) = crate::decoding::frame::read_frame_header_with_format(frame, false) else { - return false; + return; }; - // Block_Type is bits 1-2 of the block header (RFC 8878 3.1.1.2). - frame - .get(usize::from(header_len)) - .is_some_and(|&byte| (byte >> 1) & 3 == 2) + let mut at = usize::from(header_len); + while let Some(header) = frame.get(at..at + 3) { + // Block header (RFC 8878 3.1.1.2): Last_Block is bit 0, Block_Type + // bits 1-2, Block_Size the rest; an RLE block carries one byte. + let word = u32::from(header[0]) | u32::from(header[1]) << 8 | u32::from(header[2]) << 16; + let kind = (word >> 1) & 3; + out.push(kind == 2); + if word & 1 == 1 { + return; + } + let body = if kind == 1 { 1 } else { (word >> 3) as usize }; + at += 3 + body; + } } /// The literals code for `counts`, its weights scaled to its longest code @@ -290,18 +298,34 @@ fn literals_table(counts: &[usize; 256]) -> HuffmanTable { } } -/// The production matcher, recording the literals and sequences of the first -/// block it matches after being armed. Everything else is forwarded, so frames +/// The production matcher, recording the literals and sequences of every block +/// it matches, block by block. Everything else is forwarded, so frames /// compress as any compressor's would, dictionary reuse included. struct Recorder { inner: MatchGeneratorDriver, - /// Whether the next matched block is recorded; cleared once it is. - armed: bool, - literals: [u32; 256], + /// One entry per block of the frame, matched or skipped, in order. + blocks: Vec, /// Literal length, offset and match length of each sequence. sequences: Vec<(u32, u32, u32)>, } +/// What one block handed to the matcher produced. +struct RecordedBlock { + literals: [u32; 256], + /// Where the block's sequences end in [`Recorder::sequences`]. + sequences_end: usize, +} + +impl Recorder { + /// A block the matcher did not search: it yields nothing. + fn record_skipped(&mut self) { + self.blocks.push(RecordedBlock { + literals: [0; 256], + sequences_end: self.sequences.len(), + }); + } +} + impl Matcher for Recorder { fn get_last_space(&mut self) -> &[u8] { self.inner.get_last_space() @@ -325,22 +349,18 @@ impl Matcher for Recorder { fn skip_matching(&mut self) { self.inner.skip_matching(); + self.record_skipped(); } fn skip_matching_with_hint(&mut self, incompressible_hint: Option) { self.inner.skip_matching_with_hint(incompressible_hint); + self.record_skipped(); } fn start_matching(&mut self, mut handle_sequence: impl for<'a> FnMut(Sequence<'a>)) { - if !core::mem::take(&mut self.armed) { - self.inner.start_matching(handle_sequence); - return; - } + let mut literals = [0u32; 256]; let Self { - inner, - literals, - sequences, - .. + inner, sequences, .. } = self; inner.start_matching(|sequence| { match &sequence { @@ -363,6 +383,11 @@ impl Matcher for Recorder { } handle_sequence(sequence); }); + let sequences_end = self.sequences.len(); + self.blocks.push(RecordedBlock { + literals, + sequences_end, + }); } fn reset(&mut self, level: CompressionLevel) { @@ -389,10 +414,6 @@ impl Matcher for Recorder { self.inner.clear_param_overrides(); } - fn apply_parameters(&mut self, params: &CompressionParameters) { - self.inner.apply_parameters(params); - } - fn prime_with_dictionary(&mut self, dict_content: &[u8], offset_hist: [u32; 3]) { self.inner.prime_with_dictionary(dict_content, offset_hist); } diff --git a/zstd/src/dictionary/finalize/tests.rs b/zstd/src/dictionary/finalize/tests.rs index 5518b2d5e..f4003f6f1 100644 --- a/zstd/src/dictionary/finalize/tests.rs +++ b/zstd/src/dictionary/finalize/tests.rs @@ -124,68 +124,29 @@ fn literal_tables_take_the_depth_of_their_longest_code() { assert_eq!(parsed as usize, description.len()); } -/// Only a sample's first block is counted (upstream zstd `ZDICT_countEStats` -/// compresses `MIN(128 KiB, 1 << windowLog)` bytes as one block, the window of -/// the average sample plus the content). Samples averaging a line with half a -/// KiB of content keep that window at its 1 KiB floor; two versions of an -/// 8 KiB sample that differ only past 4 KiB must finalize to the same bytes. -#[test] -fn only_the_first_block_of_a_sample_is_counted() { - let (mut data, mut sizes) = log_samples(); - let content: Vec = data[..512].to_vec(); - let head = data[..4096].to_vec(); - let mut state = 0x2545_F491_4F6C_DD1Du64; - let noise: Vec = (0..4096) +/// Random bytes, different for each seed. +fn noise(seed: u64, len: usize) -> Vec { + let mut state = seed; + (0..len) .map(|_| { state ^= state << 13; state ^= state >> 7; state ^= state << 17; state as u8 }) - .collect(); - sizes.push(8192); - let finalized = |tail: &[u8], data: &mut Vec| { - let base = data.len(); - data.extend_from_slice(&head); - data.extend_from_slice(tail); - let samples = SampleSet::new(data, &sizes).unwrap(); - let dict = finalize( - &content, - &samples, - samples.count(), - 8192, - FinalizeOptions::default(), - ) - .unwrap(); - data.truncate(base); - dict - }; - let with_runs = finalized(&[b'z'; 4096], &mut data); - let with_noise = finalized(&noise, &mut data); - assert!( - with_runs == with_noise, - "a sample's second block reached the tables" - ); + .collect() } -/// The first block spans the window of the average sample plus the content, -/// not of the sample alone: with 4 KiB of content that window is 8 KiB, so the -/// second half of an 8 KiB sample is counted and two versions differing there -/// finalize differently. +/// A block written raw teaches nothing: its bytes are no literals of any +/// compressed block. Samples averaging a line keep the blocks at 1 KiB, so an +/// 8 KiB sample of 4 KiB of log lines and 4 KiB of noise is four compressed +/// blocks and four raw ones; two versions differing only in the noise must +/// finalize to the same bytes. #[test] -fn the_first_block_spans_the_window_of_sample_and_content() { +fn a_block_written_raw_is_not_counted() { let (mut data, mut sizes) = log_samples(); let content: Vec = data[..4096].to_vec(); let head = data[..4096].to_vec(); - let mut state = 0x2545_F491_4F6C_DD1Du64; - let noise: Vec = (0..4096) - .map(|_| { - state ^= state << 13; - state ^= state >> 7; - state ^= state << 17; - state as u8 - }) - .collect(); sizes.push(8192); let finalized = |tail: &[u8], data: &mut Vec| { let base = data.len(); @@ -203,12 +164,9 @@ fn the_first_block_spans_the_window_of_sample_and_content() { data.truncate(base); dict }; - let with_runs = finalized(&[b'z'; 4096], &mut data); - let with_noise = finalized(&noise, &mut data); - assert!( - with_runs != with_noise, - "the second half of the sample's first block was not counted" - ); + let one = finalized(&noise(0x2545_F491_4F6C_DD1D, 4096), &mut data); + let other = finalized(&noise(0x9E37_79B9_7F4A_7C15, 4096), &mut data); + assert!(one == other, "a raw block's bytes reached the tables"); } /// Offset codes are counted with the repeat policy the frame's blocks used. @@ -254,26 +212,25 @@ fn offset_codes_follow_the_frames_repeat_policy() { assert_eq!(offset_code_one(1), 1); } -/// Only a frame whose first block is compressed counts as one. +/// Each block's kind is read from the frame, one entry per block, and nothing +/// from bytes that are not a frame. #[test] -fn first_block_kind_is_read_from_the_frame() { +fn block_kinds_are_read_from_the_frame() { + let mut kinds = Vec::new(); let text: String = core::iter::repeat_n("compressible text ", 200).collect(); - assert!(first_block_is_compressed(&compress_slice_to_vec( - text.as_bytes(), - CompressionLevel::Default - ))); - let mut state = 1u32; - let noise: Vec = (0..4096) - .map(|_| { - state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); - (state >> 24) as u8 - }) - .collect(); - assert!(!first_block_is_compressed(&compress_slice_to_vec( - &noise, - CompressionLevel::Default - ))); - assert!(!first_block_is_compressed(b"not a frame")); + compressed_blocks( + &compress_slice_to_vec(text.as_bytes(), CompressionLevel::Default), + &mut kinds, + ); + assert_eq!(kinds, [true]); + // Past one full block of noise: two blocks, both written raw. + compressed_blocks( + &compress_slice_to_vec(&noise(1, 200_000), CompressionLevel::Default), + &mut kinds, + ); + assert_eq!(kinds, [false, false]); + compressed_blocks(b"not a frame", &mut kinds); + assert!(kinds.is_empty()); } /// The id is derived from the content when none is given, inside the range diff --git a/zstd/src/encoding/frame_compressor.rs b/zstd/src/encoding/frame_compressor.rs index 4b70fb32c..19b5c36a4 100644 --- a/zstd/src/encoding/frame_compressor.rs +++ b/zstd/src/encoding/frame_compressor.rs @@ -1802,6 +1802,59 @@ impl FrameCompressor { } } + /// Configure fine-grained compression parameters. + /// + /// Resets the base [`CompressionLevel`] to the parameters' level and installs the per-knob overrides + /// (window/hash/chain/search logs, strategy, LDM) applied at the next + /// frame. Pass `None`-equivalent (a builder that overrides nothing) + /// to fall back to plain level-based compression. + /// + /// ```rust + /// use structured_zstd::encoding::{ + /// CompressionLevel, CompressionParameters, FrameCompressor, Strategy, + /// }; + /// let params = CompressionParameters::builder(CompressionLevel::Level(19)) + /// .strategy(Strategy::Btultra2) + /// .enable_long_distance_matching(true) + /// .build() + /// .unwrap(); + /// let mut compressor: FrameCompressor = FrameCompressor::new(CompressionLevel::Default); + /// compressor.set_parameters(¶ms); + /// let compressed = compressor.compress_independent_frame(b"some data to compress"); + /// assert!(!compressed.is_empty()); + /// ``` + pub fn set_parameters(&mut self, params: &crate::encoding::CompressionParameters) { + self.compression_level = params.level(); + let overrides = params.overrides(); + self.tuning = FrameTuning::from_overrides(&overrides); + // Keep `state.strategy_tag` consistent immediately so the borrowed + // one-shot eligibility gate (`borrowed_eligible`) and literal gates + // are correct even before the next `compress()` re-sync. Resolve it + // size-adaptively (same `resolve_level_params` path `prepare_frame` + // uses) so a hint already set here yields the same strategy the matcher + // will run, not the bare level-only mapping. + // The dictionary counts only when the frame will prime it (same gate + // as `prepare_frame`'s `use_dictionary_state`): uncompressed mode + // ignores an attached dictionary and has no CDict tier to resolve. + let with_dictionary = !matches!(self.compression_level, CompressionLevel::Uncompressed) + && self.state.matcher.supports_dictionary_priming(); + let params = self.resolve_frame_params(self.source_size_hint, with_dictionary); + self.sync_effective_strategy(¶ms); + self.state.huf_optimal_search = + huf_search_enabled(self.state.strategy_tag, self.source_size_hint); + self.state.literal_compression_disabled = literal_compression_disabled( + self.state.strategy_tag, + self.compression_level, + gate_target_length( + self.compression_level, + &self.tuning, + self.dictionary.as_ref().filter(|_| with_dictionary), + ), + self.tuning.literal_compression, + ); + self.state.matcher.set_param_overrides(Some(overrides)); + } + /// Whether the borrowed (no per-block history copy) one-shot loop is /// valid for an `input_len`-byte slice under the resolved `prep`. /// @@ -2325,59 +2378,6 @@ impl FrameCompressor { ); } - /// Configure fine-grained compression parameters. - /// - /// Resets the base [`CompressionLevel`] to the parameters' level and installs the per-knob overrides - /// (window/hash/chain/search logs, strategy, LDM) applied at the next - /// frame. Pass `None`-equivalent (a builder that overrides nothing) - /// to fall back to plain level-based compression. A custom matcher - /// receives them through [`Matcher::apply_parameters`]. - /// - /// ```rust - /// use structured_zstd::encoding::{ - /// CompressionLevel, CompressionParameters, FrameCompressor, Strategy, - /// }; - /// let params = CompressionParameters::builder(CompressionLevel::Level(19)) - /// .strategy(Strategy::Btultra2) - /// .enable_long_distance_matching(true) - /// .build() - /// .unwrap(); - /// let mut compressor: FrameCompressor = FrameCompressor::new(CompressionLevel::Default); - /// compressor.set_parameters(¶ms); - /// let compressed = compressor.compress_independent_frame(b"some data to compress"); - /// assert!(!compressed.is_empty()); - /// ``` - pub fn set_parameters(&mut self, params: &crate::encoding::CompressionParameters) { - self.compression_level = params.level(); - self.tuning = FrameTuning::from_overrides(¶ms.overrides()); - // Keep `state.strategy_tag` consistent immediately so the borrowed - // one-shot eligibility gate (`borrowed_eligible`) and literal gates - // are correct even before the next `compress()` re-sync. Resolve it - // size-adaptively (same `resolve_level_params` path `prepare_frame` - // uses) so a hint already set here yields the same strategy the matcher - // will run, not the bare level-only mapping. - // The dictionary counts only when the frame will prime it (same gate - // as `prepare_frame`'s `use_dictionary_state`): uncompressed mode - // ignores an attached dictionary and has no CDict tier to resolve. - let with_dictionary = !matches!(self.compression_level, CompressionLevel::Uncompressed) - && self.state.matcher.supports_dictionary_priming(); - let level_params = self.resolve_frame_params(self.source_size_hint, with_dictionary); - self.sync_effective_strategy(&level_params); - self.state.huf_optimal_search = - huf_search_enabled(self.state.strategy_tag, self.source_size_hint); - self.state.literal_compression_disabled = literal_compression_disabled( - self.state.strategy_tag, - self.compression_level, - gate_target_length( - self.compression_level, - &self.tuning, - self.dictionary.as_ref().filter(|_| with_dictionary), - ), - self.tuning.literal_compression, - ); - self.state.matcher.apply_parameters(params); - } - /// Provide a hint about the total uncompressed size for the next frame. /// /// When set, the encoder selects smaller hash tables and windows for diff --git a/zstd/src/encoding/match_generator/mod.rs b/zstd/src/encoding/match_generator/mod.rs index dc82ddc21..a0bfe5831 100644 --- a/zstd/src/encoding/match_generator/mod.rs +++ b/zstd/src/encoding/match_generator/mod.rs @@ -1110,10 +1110,6 @@ impl Matcher for MatchGeneratorDriver { self.param_overrides = None; } - fn apply_parameters(&mut self, params: &super::CompressionParameters) { - self.param_overrides = Some(params.overrides()); - } - fn reset(&mut self, level: CompressionLevel) { // On its own the driver lays its tables out in a workspace of its own, // which is moved out for the call and back: moving it leaves the diff --git a/zstd/src/encoding/mod.rs b/zstd/src/encoding/mod.rs index aedfdb9bd..add752c44 100644 --- a/zstd/src/encoding/mod.rs +++ b/zstd/src/encoding/mod.rs @@ -562,12 +562,6 @@ pub trait Matcher { /// so switching back to a bare level after a customized frame does not /// keep the old overrides sticky. Default no-op for custom matchers. fn clear_param_overrides(&mut self) {} - /// Install the fine-grained overrides of `params` (window, table logs, - /// strategy, LDM) for the next frame. Called by - /// [`FrameCompressor::set_parameters`](crate::encoding::FrameCompressor::set_parameters); - /// consumed at the next [`reset`](Self::reset). Default no-op for custom - /// matchers, which size themselves. - fn apply_parameters(&mut self, _params: &CompressionParameters) {} /// Prime matcher state with dictionary history before compressing the next frame. /// Default implementation is a no-op for custom matchers that do not support this. fn prime_with_dictionary(&mut self, _dict_content: &[u8], _offset_hist: [u32; 3]) {} From f24bce8eefdb13bbbed9f5a087c61157a358423d Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Mon, 28 Sep 2026 20:35:24 +0300 Subject: [PATCH 09/20] perf(fse): describe a table without copying its probabilities Writing and sizing an NCount description copied the 256 probabilities of a built table into an array first, so every new sequence table a block emitted, and every header priced for the table-mode choice, paid a 1 KiB copy. The writer and the size count now read the probabilities through an accessor, from the table itself or from the normalized counts. --- zstd/src/fse/fse_encoder.rs | 71 +++++++++++++++++++++---------------- 1 file changed, 41 insertions(+), 30 deletions(-) diff --git a/zstd/src/fse/fse_encoder.rs b/zstd/src/fse/fse_encoder.rs index 807e1b043..060f5f5cc 100644 --- a/zstd/src/fse/fse_encoder.rs +++ b/zstd/src/fse/fse_encoder.rs @@ -287,20 +287,31 @@ impl FSETable { // single source of truth shared with the pre-build cost estimate in // `fse_header_bits_for_counts` (the table-mode selector prices a // candidate table from its normalized counts without building it). - let probs: [i32; 256] = core::array::from_fn(|i| self.states[i].probability); - ncount_header_bits(&probs, self.acc_log()) + let states = &self.states; + ncount_header_bits(states.len(), |i| states[i].probability, self.acc_log()) } pub(crate) fn write_table>>(&self, writer: &mut BitWriter) { - let probs: [i32; 256] = core::array::from_fn(|i| self.states[i].probability); - write_ncount(&probs, self.acc_log(), writer); + let states = &self.states; + write_ncount( + states.len(), + |i| states[i].probability, + self.acc_log(), + writer, + ); } } -/// Write the NCount description of the normalized distribution `probs` at -/// `acc_log` (RFC 8878 4.1.1), padded to a byte; what [`FSETable::write_table`] -/// writes for a built table, from the normalized counts alone. -fn write_ncount>>(probs: &[i32], acc_log: u8, writer: &mut BitWriter) { +/// Write the NCount description of the normalized distribution of `len` +/// symbols, `prob(i)` the probability of symbol `i`, at `acc_log` (RFC 8878 +/// 4.1.1), padded to a byte. The probabilities are read where they live, so a +/// built table and a bare distribution share the writer without a copy. +fn write_ncount>>( + len: usize, + prob: impl Fn(usize) -> i32 + Copy, + acc_log: u8, + writer: &mut BitWriter, +) { assert!( writer.index().is_multiple_of(8), "FSE table headers must start on a byte boundary" @@ -318,9 +329,9 @@ fn write_ncount>>(probs: &[i32], acc_log: u8, writer: &mut BitW let low_threshold = ((1 << bits_to_write) - 1) - (max_remaining_value); let mask = (1 << (bits_to_write - 1)) - 1; - let prob = probs[prob_idx]; + let p = prob(prob_idx); prob_idx += 1; - let value = (prob + 1) as u32; + let value = (p + 1) as u32; if value < low_threshold as u32 { writer.write_bits(value, bits_to_write as usize - 1); } else if value > mask { @@ -329,13 +340,13 @@ fn write_ncount>>(probs: &[i32], acc_log: u8, writer: &mut BitW writer.write_bits(value, bits_to_write as usize); } - if prob == -1 { + if p == -1 { probability_counter += 1; - } else if prob > 0 { - probability_counter += prob as usize; + } else if p > 0 { + probability_counter += p as usize; } else { let mut zeros = 0u8; - while prob_idx < probs.len() && probs[prob_idx] == 0 { + while prob_idx < len && prob(prob_idx) == 0 { zeros += 1; prob_idx += 1; if zeros == 3 { @@ -350,7 +361,7 @@ fn write_ncount>>(probs: &[i32], acc_log: u8, writer: &mut BitW #[cfg(debug_assertions)] { let written_bits = writer.index() - start_idx; - let computed = ncount_header_bits(probs, acc_log); + let computed = ncount_header_bits(len, prob, acc_log); debug_assert_eq!( written_bits, computed, "ncount_header_bits() mismatch: written={written_bits}, computed={computed}" @@ -533,7 +544,7 @@ pub(crate) fn write_ncount_at_log>>( max_symbol, true, ); - write_ncount(&probs[..=max_symbol], table_log, writer); + write_ncount(max_symbol + 1, |i| probs[i], table_log, writer); } #[cfg(any(test, feature = "fuzz-exports"))] @@ -578,13 +589,13 @@ pub(crate) fn build_table_from_counts_into( } /// Bit size of the serialized FSE NCount table header for the normalized -/// distribution `probs` at `acc_log`. This is exactly what -/// [`FSETable::table_header_bits`] returns (they share this implementation), -/// but computed from the normalized counts alone — no spread / state-table / -/// `symbolTT` construction. The table-mode selector uses it to price a -/// candidate custom table without building its (otherwise discarded) state -/// tables. -fn ncount_header_bits(probs: &[i32], acc_log: u8) -> usize { +/// distribution of `len` symbols, `prob(i)` the probability of symbol `i`, at +/// `acc_log`. This is exactly what [`FSETable::table_header_bits`] returns +/// (they share this implementation), but computed from the normalized counts +/// alone — no spread / state-table / `symbolTT` construction. The table-mode +/// selector uses it to price a candidate custom table without building its +/// (otherwise discarded) state tables. +fn ncount_header_bits(len: usize, prob: impl Fn(usize) -> i32, acc_log: u8) -> usize { let mut bits = 4; // acc_log - 5 let mut probability_counter = 0usize; let probability_sum = 1usize << acc_log; @@ -595,22 +606,22 @@ fn ncount_header_bits(probs: &[i32], acc_log: u8) -> usize { let bits_to_write = max_remaining_value.ilog2() + 1; let low_threshold = ((1 << bits_to_write) - 1) - max_remaining_value; - let prob = probs[prob_idx]; + let p = prob(prob_idx); prob_idx += 1; - let value = (prob + 1) as u32; + let value = (p + 1) as u32; if value < low_threshold as u32 { bits += bits_to_write as usize - 1; } else { bits += bits_to_write as usize; } - if prob == -1 { + if p == -1 { probability_counter += 1; - } else if prob > 0 { - probability_counter += prob as usize; + } else if p > 0 { + probability_counter += p as usize; } else { let mut zeros = 0u8; - while prob_idx < probs.len() && probs[prob_idx] == 0 { + while prob_idx < len && prob(prob_idx) == 0 { zeros += 1; prob_idx += 1; if zeros == 3 { @@ -667,7 +678,7 @@ pub(crate) fn fse_header_bits_for_counts( max_symbol, avoid_0_numbit, ); - ncount_header_bits(&probs[..counts.len()], table_log) + ncount_header_bits(counts.len(), |i| probs[i], table_log) } fn min_table_log(total: usize, max_symbol: usize) -> u8 { From 448d3eb075de7bb58590f86b39236c91799de192 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Mon, 28 Sep 2026 21:20:50 +0300 Subject: [PATCH 10/20] docs(fse): name the test-only table builders without linking them --- zstd/src/fse/fse_encoder.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/zstd/src/fse/fse_encoder.rs b/zstd/src/fse/fse_encoder.rs index 060f5f5cc..22d42a368 100644 --- a/zstd/src/fse/fse_encoder.rs +++ b/zstd/src/fse/fse_encoder.rs @@ -453,7 +453,7 @@ pub(crate) fn build_table_from_symbol_counts( build_table_from_counts(counts, max_log, avoid_0_numbit) } -/// [`build_table_from_symbol_counts`] writing into a caller-owned table. +/// `build_table_from_symbol_counts` writing into a caller-owned table. pub(crate) fn build_table_from_symbol_counts_into( counts: &[usize], max_log: u8, @@ -554,7 +554,7 @@ fn build_table_from_counts(counts: &[usize], max_log: u8, avoid_0_numbit: bool) out } -/// [`build_table_from_counts`] writing into a caller-owned table. +/// `build_table_from_counts` writing into a caller-owned table. pub(crate) fn build_table_from_counts_into( counts: &[usize], max_log: u8, @@ -641,7 +641,7 @@ fn ncount_header_bits(len: usize, prob: impl Fn(usize) -> i32, acc_log: u8) -> u } /// Header bit cost of the custom FSE table that -/// [`build_table_from_symbol_counts`] would produce for `counts`, without +/// `build_table_from_symbol_counts` would produce for `counts`, without /// building it. Mirrors the front of `build_table_from_counts` (table-log /// selection + count normalization) and then sizes the header via /// [`ncount_header_bits`]. Returns the exact value the built table's From dc2ffde98e68c1b45ebf7aac6e3ae9d8b5fa2a16 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Tue, 29 Sep 2026 01:32:08 +0300 Subject: [PATCH 11/20] fix(dict): read the recorded literals back by position --- zstd/src/dictionary/finalize.rs | 35 +++++++++++++++++++++------------ 1 file changed, 22 insertions(+), 13 deletions(-) diff --git a/zstd/src/dictionary/finalize.rs b/zstd/src/dictionary/finalize.rs index b7e4dc1e7..25e6b7fd5 100644 --- a/zstd/src/dictionary/finalize.rs +++ b/zstd/src/dictionary/finalize.rs @@ -357,32 +357,41 @@ impl Matcher for Recorder { self.record_skipped(); } - fn start_matching(&mut self, mut handle_sequence: impl for<'a> FnMut(Sequence<'a>)) { - let mut literals = [0u32; 256]; + fn start_matching(&mut self, mut handle_sequence: impl FnMut(Sequence)) { + let block_start = self.sequences.len(); + let mut tail = 0usize; let Self { inner, sequences, .. } = self; inner.start_matching(|sequence| { - match &sequence { - Sequence::Literals { literals: bytes } => { - for &byte in *bytes { - literals[usize::from(byte)] += 1; - } - } + match sequence { + Sequence::Literals { len } => tail += len, Sequence::Triple { - literals: bytes, + literal_len, offset, match_len, } => { - for &byte in *bytes { - literals[usize::from(byte)] += 1; - } // Lengths and offsets of one block, far below `u32::MAX`. - sequences.push((bytes.len() as u32, *offset as u32, *match_len as u32)); + sequences.push((literal_len as u32, offset as u32, match_len as u32)); } } handle_sequence(sequence); }); + // The literals are the block's bytes between the matches, read back by + // position once matching is done, as the encoder gathers them. + let block = self.inner.get_last_space(); + let mut literals = [0u32; 256]; + let mut pos = 0usize; + let mut count = |bytes: &[u8]| { + for &byte in bytes { + literals[usize::from(byte)] += 1; + } + }; + for &(ll, _, ml) in &self.sequences[block_start..] { + count(&block[pos..pos + ll as usize]); + pos += ll as usize + ml as usize; + } + count(&block[pos..pos + tail]); let sequences_end = self.sequences.len(); self.blocks.push(RecordedBlock { literals, From 44af562cfe9be767cacadb9223261b55e7dcee66 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Tue, 29 Sep 2026 01:51:10 +0300 Subject: [PATCH 12/20] perf(dict): finalize every candidate over the evaluator's buffers - The finalizer writes over a buffer the caller keeps, and the evaluator finalizes each candidate, and each shrunk one, over two buffers it holds; the best-so-far copies a candidate's dictionary and content only when it wins. Dictionaries are byte-identical. - The candidate copied into the encoder dictionary stays a copy: under callgrind the copy and the dictionary's preparation are 0.02% of a default training run, against 52.7% for compressing the samples with it. --- zstd/src/dictionary/finalize.rs | 27 ++++++++++++++++++++++++--- zstd/src/dictionary/selection.rs | 10 +++++----- zstd/src/dictionary/tests.rs | 2 +- 3 files changed, 30 insertions(+), 9 deletions(-) diff --git a/zstd/src/dictionary/finalize.rs b/zstd/src/dictionary/finalize.rs index 25e6b7fd5..1f98d21de 100644 --- a/zstd/src/dictionary/finalize.rs +++ b/zstd/src/dictionary/finalize.rs @@ -51,6 +51,21 @@ pub(super) fn finalize( dict_size: usize, options: FinalizeOptions, ) -> io::Result> { + let mut out = Vec::new(); + finalize_into(&mut out, content, samples, count, dict_size, options)?; + Ok(out) +} + +/// [`finalize`] over `out`, whose allocation a search keeps from one +/// candidate to the next. +pub(super) fn finalize_into( + out: &mut Vec, + content: &[u8], + samples: &SampleSet<'_>, + count: usize, + dict_size: usize, + options: FinalizeOptions, +) -> io::Result<()> { // Upstream zstd's order (`ZDICT_finalizeDictionary`): the size first. if dict_size < DICT_SIZE_MIN { return Err(too_small()); @@ -62,10 +77,11 @@ pub(super) fn finalize( if dict_id == 0 { return Err(invalid("dictionary id must be non-zero")); } - let mut out = Vec::with_capacity(dict_size); + out.clear(); + out.reserve(dict_size); out.extend_from_slice(&DICT_MAGIC_NUM); out.extend_from_slice(&dict_id.to_le_bytes()); - analyze_entropy(&mut out, content, samples, count, options.level)?; + analyze_entropy(out, content, samples, count, options.level)?; for rep in START_REPS { out.extend_from_slice(&rep.to_le_bytes()); } @@ -79,7 +95,7 @@ pub(super) fn finalize( out.resize(out.len() + MIN_CONTENT_SIZE - content.len(), 0); } out.extend_from_slice(content); - Ok(out) + Ok(()) } fn too_small() -> io::Error { @@ -171,6 +187,11 @@ fn count_samples( count: usize, level: i32, ) -> io::Result<()> { + // The candidate is copied into the encoder dictionary, where upstream zstd + // references it (`ZSTD_dlm_byRef`): under callgrind the copy and the + // dictionary's preparation are 0.02% of a default training run, against + // 52.7% for compressing the samples with it. Borrowing would put a lifetime + // on the encoder dictionary for none of it. let dictionary = crate::decoding::Dictionary::from_raw_content(0, content.to_vec()) .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?; let recorder = Recorder { diff --git a/zstd/src/dictionary/selection.rs b/zstd/src/dictionary/selection.rs index 819af5c6f..2bbbd6048 100644 --- a/zstd/src/dictionary/selection.rs +++ b/zstd/src/dictionary/selection.rs @@ -3,7 +3,7 @@ //! `COVER_selectDict` and `COVER_checkTotalCompressedSize`). use super::FinalizeOptions; -use super::finalize::finalize; +use super::finalize::finalize_into; use super::samples::SampleSet; use crate::encoding::{CompressionLevel, EncoderDictionary, FrameCompressor}; use core::ops::Range; @@ -70,16 +70,16 @@ impl<'s> Evaluator<'s> { Ok(out) } - /// [`Self::finalize`] over `out`. + /// [`Self::finalize`] over `out`, keeping its allocation. fn finalize_into(&self, content: &[u8], out: &mut Vec) -> io::Result<()> { - *out = finalize( + finalize_into( + out, content, self.samples, self.finalize_samples, self.capacity, self.finalize, - )?; - Ok(()) + ) } /// `content` finalized over `buffer` and priced. diff --git a/zstd/src/dictionary/tests.rs b/zstd/src/dictionary/tests.rs index 79c5e4608..4f4f7e93a 100644 --- a/zstd/src/dictionary/tests.rs +++ b/zstd/src/dictionary/tests.rs @@ -439,7 +439,7 @@ fn samples_too_short_for_a_dmer_are_refused() { #[test] fn the_preflight_names_the_cause_the_trainer_names() { let (data, sizes) = training_samples(); - let dict_size = TRAINER_DICT_SIZE_MIN - 1; + let dict_size = MIN_TRAINED_DICT_SIZE - 1; let cover = CoverOptions { k: 300, d: 8, From 6e9700283bb7d6177a92443a81cacb090dae3599 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Tue, 29 Sep 2026 02:18:16 +0300 Subject: [PATCH 13/20] perf(dict): keep the analysis compressor across candidates - The entropy analysis keeps its compressor, frame buffer and block kinds on the evaluator: between candidates only the dictionary changes, so the matcher and encoder scratch are no longer rebuilt for each one. The compressor is rebuilt when the level changes. - Each sample's frame is written straight into the kept buffer (compress_known_into: the header first, the length being known, then the blocks), without the streaming path's block accumulator and copy. - Offset codes are counted into a fixed array over every code, as upstream zstd's offcodeCount is, and described up to the alphabet's bound: no branch per sequence, no allocation per analysis. Dictionaries are byte-identical for --train and --train-cover on a frozen corpus. Time on an M1 moved within noise (--train -1.0%, --train-cover +1.3% by minimum, both under what two builds resolve). --- zstd/src/dictionary/finalize.rs | 112 +++++++++++++++++--------- zstd/src/dictionary/finalize/tests.rs | 17 +++- zstd/src/dictionary/selection.rs | 10 ++- zstd/src/encoding/frame_compressor.rs | 35 ++++++++ 4 files changed, 131 insertions(+), 43 deletions(-) diff --git a/zstd/src/dictionary/finalize.rs b/zstd/src/dictionary/finalize.rs index 1f98d21de..b6960e3ae 100644 --- a/zstd/src/dictionary/finalize.rs +++ b/zstd/src/dictionary/finalize.rs @@ -21,7 +21,7 @@ use crate::encoding::{ }; use crate::fse::fse_encoder::{FSETable, write_ncount_at_log}; use crate::huff0::huff0_encoder::{HuffmanEncoder, HuffmanTable}; -use std::{io, vec, vec::Vec}; +use std::{io, vec::Vec}; /// Smallest dictionary finalized (upstream zstd `ZDICT_DICTSIZE_MIN`). pub(super) const DICT_SIZE_MIN: usize = 256; @@ -44,24 +44,48 @@ const BLOCK_SIZE_MAX: usize = 128 << 10; /// Finalize `content` into a dictionary of at most `dict_size` bytes, drawing /// the entropy tables from the first `count` samples. Content that does not /// fit after the header is cut from the front, keeping the best segments. -pub(super) fn finalize( +pub(super) fn finalize<'s>( content: &[u8], - samples: &SampleSet<'_>, + samples: &SampleSet<'s>, count: usize, dict_size: usize, options: FinalizeOptions, ) -> io::Result> { let mut out = Vec::new(); - finalize_into(&mut out, content, samples, count, dict_size, options)?; + let mut analysis = Analysis::default(); + finalize_into( + &mut out, + &mut analysis, + content, + samples, + count, + dict_size, + options, + )?; Ok(out) } -/// [`finalize`] over `out`, whose allocation a search keeps from one -/// candidate to the next. -pub(super) fn finalize_into( +/// A compressor over one sample at a time, recording what it matches. +type AnalysisCompressor<'s> = FrameCompressor<&'s [u8], Vec, Recorder>; + +/// What the entropy analysis keeps from one candidate to the next: the +/// compressor it runs the samples through, whose matcher and scratch then +/// survive the change of dictionary, its frame buffer and its block kinds. +#[derive(Default)] +pub(super) struct Analysis<'s> { + /// The compressor and the level it was built for. + compressor: Option<(i32, AnalysisCompressor<'s>)>, + frame: Vec, + compressed: Vec, +} + +/// [`finalize`] over `out` and `analysis`, whose allocations a search keeps +/// from one candidate to the next. +pub(super) fn finalize_into<'s>( out: &mut Vec, + analysis: &mut Analysis<'s>, content: &[u8], - samples: &SampleSet<'_>, + samples: &SampleSet<'s>, count: usize, dict_size: usize, options: FinalizeOptions, @@ -81,7 +105,7 @@ pub(super) fn finalize_into( out.reserve(dict_size); out.extend_from_slice(&DICT_MAGIC_NUM); out.extend_from_slice(&dict_id.to_le_bytes()); - analyze_entropy(out, content, samples, count, options.level)?; + analyze_entropy(out, analysis, content, samples, count, options.level)?; for rep in START_REPS { out.extend_from_slice(&rep.to_le_bytes()); } @@ -124,7 +148,10 @@ fn derive_dict_id(raw_content: &[u8]) -> u32 { /// be described (upstream zstd `ZDICT_analyzeEntropy`). struct EntropyCounts { literals: [usize; 256], - offset_codes: Vec, + /// Every code a 32-bit offset base has, as upstream's `offcodeCount` + /// holds: only the codes up to the alphabet's bound are described, and a + /// sequence is counted without asking which side of it its code is on. + offset_codes: [usize; 32], match_lengths: [usize; 53], literal_lengths: [usize; 36], } @@ -132,10 +159,11 @@ struct EntropyCounts { /// Append the literals table, the three sequence tables and nothing else to /// `out`, from the statistics of the first `count` samples compressed with /// `content` as a raw dictionary at `level`. -fn analyze_entropy( +fn analyze_entropy<'s>( out: &mut Vec, + analysis: &mut Analysis<'s>, content: &[u8], - samples: &SampleSet<'_>, + samples: &SampleSet<'s>, count: usize, level: i32, ) -> io::Result<()> { @@ -149,11 +177,11 @@ fn analyze_entropy( } let mut counts = EntropyCounts { literals: [1; 256], - offset_codes: vec![1; offcode_max as usize + 1], + offset_codes: [1; 32], match_lengths: [1; 53], literal_lengths: [1; 36], }; - count_samples(&mut counts, content, samples, count, level)?; + count_samples(&mut counts, analysis, content, samples, count, level)?; let mut literals = literals_table(&counts.literals); if literals.table_log() == 8 { @@ -171,7 +199,11 @@ fn analyze_entropy( // are normalized and described without being built. let mut writer = BitWriter::from(&mut *out); HuffmanEncoder::new(&literals, &mut writer).write_table(); - write_ncount_at_log(&counts.offset_codes, OF_LOG, &mut writer); + write_ncount_at_log( + &counts.offset_codes[..=offcode_max as usize], + OF_LOG, + &mut writer, + ); write_ncount_at_log(&counts.match_lengths, ML_LOG, &mut writer); write_ncount_at_log(&counts.literal_lengths, LL_LOG, &mut writer); writer.flush(); @@ -180,10 +212,11 @@ fn analyze_entropy( /// Compress the first 128 KiB of each of the first `count` samples with /// `content` as a raw dictionary and count what its compressed blocks hold. -fn count_samples( +fn count_samples<'s>( counts: &mut EntropyCounts, + analysis: &mut Analysis<'s>, content: &[u8], - samples: &SampleSet<'_>, + samples: &SampleSet<'s>, count: usize, level: i32, ) -> io::Result<()> { @@ -194,13 +227,25 @@ fn count_samples( // on the encoder dictionary for none of it. let dictionary = crate::decoding::Dictionary::from_raw_content(0, content.to_vec()) .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?; - let recorder = Recorder { - inner: MatchGeneratorDriver::new(BLOCK_SIZE_MAX, 1), - blocks: Vec::new(), - sequences: Vec::new(), - }; - let mut compressor: FrameCompressor<&[u8], Vec, Recorder> = - FrameCompressor::new_with_matcher(recorder, CompressionLevel::from_level(level)); + let Analysis { + compressor, + frame, + compressed, + } = analysis; + // One compressor per level serves every candidate: the dictionary is all + // that changes, so its matcher and scratch are kept. + if compressor.as_ref().is_none_or(|(built, _)| *built != level) { + let recorder = Recorder { + inner: MatchGeneratorDriver::new(BLOCK_SIZE_MAX, 1), + blocks: Vec::new(), + sequences: Vec::new(), + }; + *compressor = Some(( + level, + FrameCompressor::new_with_matcher(recorder, CompressionLevel::from_level(level)), + )); + } + let (_, compressor) = compressor.as_mut().expect("built above"); compressor .set_encoder_dictionary(EncoderDictionary::from_dictionary(dictionary)) .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?; @@ -214,22 +259,17 @@ fn count_samples( // or at upstream's, measured worse on the repository files (the samples // compressed 0.013% and 0.05% larger in total across COVER, FastCOVER // and the default search), so the rest of a sample is signal, not noise. - let mut frame = Vec::new(); - let mut compressed = Vec::new(); for index in 0..count { let sample = samples.sample(index); let block = &sample[..sample.len().min(BLOCK_SIZE_MAX)]; if block.is_empty() { continue; } - frame.clear(); compressor.set_source(block); - compressor.set_drain(frame); compressor.set_source_size_hint(average as u64); - compressor.compress(); - frame = compressor.take_drain().expect("the drain was set above"); + compressor.compress_known_into(block.len() as u64, frame); let fast_offset_codes = compressor.uses_fast_offset_codes(); - compressed_blocks(&frame, &mut compressed); + compressed_blocks(frame, compressed); let recorder = compressor.matcher_mut(); // A block written raw or as one repeated byte holds no sequences to // learn from, and its bytes are no literals of any compressed block. @@ -265,12 +305,10 @@ fn count_samples( }; counts.literal_lengths[encode_literal_length(ll).0 as usize] += 1; counts.match_lengths[encode_match_len(ml).0 as usize] += 1; - if let Some(slot) = counts - .offset_codes - .get_mut(encode_offset(off_base).0 as usize) - { - *slot += 1; - } + // The code is the offset base's highest bit, below 32. + let code = encode_offset(off_base).0 as usize; + debug_assert!(code < counts.offset_codes.len()); + counts.offset_codes[code] += 1; } } recorder.blocks.clear(); diff --git a/zstd/src/dictionary/finalize/tests.rs b/zstd/src/dictionary/finalize/tests.rs index f4003f6f1..d6d32456a 100644 --- a/zstd/src/dictionary/finalize/tests.rs +++ b/zstd/src/dictionary/finalize/tests.rs @@ -198,14 +198,25 @@ fn offset_codes_follow_the_frames_repeat_policy() { } let samples = SampleSet::new(&data, &sizes).unwrap(); let content = samples.sample(0).to_vec(); - let offset_code_one = |level| { + // One analysis across both levels: the second must not run on the + // compressor the first built for another level. + let mut analysis = Analysis::default(); + let mut offset_code_one = |level| { let mut counts = EntropyCounts { literals: [1; 256], - offset_codes: vec![1; OFFCODE_MAX as usize + 1], + offset_codes: [1; 32], match_lengths: [1; 53], literal_lengths: [1; 36], }; - count_samples(&mut counts, &content, &samples, samples.count(), level).unwrap(); + count_samples( + &mut counts, + &mut analysis, + &content, + &samples, + samples.count(), + level, + ) + .unwrap(); counts.offset_codes[1] }; assert!(offset_code_one(3) > 1, "the input exercises deeper repeats"); diff --git a/zstd/src/dictionary/selection.rs b/zstd/src/dictionary/selection.rs index 2bbbd6048..3e9b54a9d 100644 --- a/zstd/src/dictionary/selection.rs +++ b/zstd/src/dictionary/selection.rs @@ -3,7 +3,7 @@ //! `COVER_selectDict` and `COVER_checkTotalCompressedSize`). use super::FinalizeOptions; -use super::finalize::finalize_into; +use super::finalize::{Analysis, finalize_into}; use super::samples::SampleSet; use crate::encoding::{CompressionLevel, EncoderDictionary, FrameCompressor}; use core::ops::Range; @@ -40,6 +40,8 @@ pub(super) struct Evaluator<'s> { /// only by a candidate that wins. dict: Vec, shrunk: Vec, + /// The entropy analysis every candidate is finalized through. + analysis: Analysis<'s>, } impl<'s> Evaluator<'s> { @@ -60,20 +62,22 @@ impl<'s> Evaluator<'s> { frame: Vec::new(), dict: Vec::new(), shrunk: Vec::new(), + analysis: Analysis::default(), } } /// `content` finalized into a dictionary of at most the capacity. - pub(super) fn finalize(&self, content: &[u8]) -> io::Result> { + pub(super) fn finalize(&mut self, content: &[u8]) -> io::Result> { let mut out = Vec::new(); self.finalize_into(content, &mut out)?; Ok(out) } /// [`Self::finalize`] over `out`, keeping its allocation. - fn finalize_into(&self, content: &[u8], out: &mut Vec) -> io::Result<()> { + fn finalize_into(&mut self, content: &[u8], out: &mut Vec) -> io::Result<()> { finalize_into( out, + &mut self.analysis, content, self.samples, self.finalize_samples, diff --git a/zstd/src/encoding/frame_compressor.rs b/zstd/src/encoding/frame_compressor.rs index 19b5c36a4..c1bb192e9 100644 --- a/zstd/src/encoding/frame_compressor.rs +++ b/zstd/src/encoding/frame_compressor.rs @@ -2483,6 +2483,41 @@ impl FrameCompressor { self.finish_frame(all_blocks, total_uncompressed, &prep); } + /// [`Self::compress`] for a source of `total` bytes, the frame written + /// over `out` rather than to the drain: with the content size known up + /// front the header goes first and the blocks straight after it, with no + /// block accumulator and no copy into the drain. The frame is the one + /// `compress` writes for the same source, hint and settings. + pub(crate) fn compress_known_into(&mut self, total: u64, out: &mut Vec) { + let prep = self.prepare_frame(crate::encoding::workspace::IngestPlan::Stream); + let mut source = self + .uncompressed_data + .take() + .expect("source must be set via set_source before compress_known_into()"); + out.clear(); + self.append_frame_header(total, &prep, out); + let header_len = out.len(); + let mut block_source = ReaderBlockSource::new(&mut source); + let read = self.run_owned_block_loop(&mut block_source, prep.initial_size_hint, false, out); + self.uncompressed_data = Some(source); + assert_eq!( + read, total, + "the source held another length than its frame header" + ); + #[cfg(feature = "hash")] + if self.content_checksum { + out.extend_from_slice(&(self.hasher.finish() as u32).to_le_bytes()); + } + #[cfg(feature = "lsm")] + { + let emit_checksum = cfg!(feature = "hash") && self.content_checksum; + let blocks_end = out.len() - if emit_checksum { 4 } else { 0 }; + self.populate_frame_emit_info(header_len, &out[header_len..blocks_end], emit_checksum); + } + #[cfg(not(feature = "lsm"))] + let _ = header_len; + } + /// Resets the compressor for the next frame, whose input reaches the /// match finder as `ingest` says. fn prepare_frame(&mut self, ingest: crate::encoding::workspace::IngestPlan) -> FramePrep { From 227b42d4cdefd5f927eb34a65335fcd1950eba0c Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Tue, 29 Sep 2026 03:01:12 +0300 Subject: [PATCH 14/20] perf(dict): build the literals code once, in kept buffers The finalizer built the literals code at the length limit and, when its longest code fell short of it, built it again at that length only to describe it there. Below the limit the height limiter does nothing, so both builds give the same code lengths and the second only lowers every weight by the same step: `build_limited_in` now does that step on the first build's weights, which is the table upstream zstd writes from the `maxNbBits` its build returns. The weights, the tree nodes and the table's buffers live in a `WeightScratch` the analysis keeps from one candidate to the next, so a parameter search allocates none of them per candidate. Dictionaries are byte-identical to the previous head for --train, COVER, FastCOVER and FastCOVER with shrink at levels 1, 3, 9 and 19. A test checks the code built under the limit against the code built at its longest length, over a scratch holding another table. --- zstd/src/dictionary/finalize.rs | 31 +++++++-------------- zstd/src/dictionary/finalize/tests.rs | 40 +++++++++++++++++++++++++-- zstd/src/huff0/huff0_encoder.rs | 33 ++++++++++++++++++---- 3 files changed, 75 insertions(+), 29 deletions(-) diff --git a/zstd/src/dictionary/finalize.rs b/zstd/src/dictionary/finalize.rs index b6960e3ae..65bc3b61d 100644 --- a/zstd/src/dictionary/finalize.rs +++ b/zstd/src/dictionary/finalize.rs @@ -20,7 +20,7 @@ use crate::encoding::{ MatchGeneratorDriver, Matcher, Sequence, }; use crate::fse::fse_encoder::{FSETable, write_ncount_at_log}; -use crate::huff0::huff0_encoder::{HuffmanEncoder, HuffmanTable}; +use crate::huff0::huff0_encoder::{HuffmanEncoder, HuffmanTable, WeightScratch}; use std::{io, vec::Vec}; /// Smallest dictionary finalized (upstream zstd `ZDICT_DICTSIZE_MIN`). @@ -70,13 +70,15 @@ type AnalysisCompressor<'s> = FrameCompressor<&'s [u8], Vec, Recorder>; /// What the entropy analysis keeps from one candidate to the next: the /// compressor it runs the samples through, whose matcher and scratch then -/// survive the change of dictionary, its frame buffer and its block kinds. +/// survive the change of dictionary, its frame buffer, its block kinds and +/// the buffers the literals code is built in. #[derive(Default)] pub(super) struct Analysis<'s> { /// The compressor and the level it was built for. compressor: Option<(i32, AnalysisCompressor<'s>)>, frame: Vec, compressed: Vec, + huffman: WeightScratch, } /// [`finalize`] over `out` and `analysis`, whose allocations a search keeps @@ -183,7 +185,8 @@ fn analyze_entropy<'s>( }; count_samples(&mut counts, analysis, content, samples, count, level)?; - let mut literals = literals_table(&counts.literals); + let huffman = &mut analysis.huffman; + let mut literals = HuffmanTable::build_limited_in(&counts.literals, HUF_MAX_BITS, huffman); if literals.table_log() == 8 { // Every symbol at eight bits describes nothing and cannot be written; // a mostly flat distribution that still compresses stands in @@ -192,13 +195,15 @@ fn analyze_entropy<'s>( counts.literals[0] = 4; counts.literals[253] = 1; counts.literals[254] = 1; - literals = literals_table(&counts.literals); + huffman.recycle(literals); + literals = HuffmanTable::build_limited_in(&counts.literals, HUF_MAX_BITS, huffman); debug_assert_eq!(literals.table_log(), 9); } // Each description goes straight into the dictionary; the sequence tables // are normalized and described without being built. let mut writer = BitWriter::from(&mut *out); HuffmanEncoder::new(&literals, &mut writer).write_table(); + huffman.recycle(literals); write_ncount_at_log( &counts.offset_codes[..=offcode_max as usize], OF_LOG, @@ -231,6 +236,7 @@ fn count_samples<'s>( compressor, frame, compressed, + .. } = analysis; // One compressor per level serves every candidate: the dictionary is all // that changes, so its matcher and scratch are kept. @@ -340,23 +346,6 @@ fn compressed_blocks(frame: &[u8], out: &mut Vec) { } } -/// The literals code for `counts`, its weights scaled to its longest code -/// rather than to the length limit, as upstream writes it (`HUF_writeCTable` -/// takes the `maxNbBits` its build returned): the same codes, and the table -/// log a decoder derives from the weights is the one the code needs. -fn literals_table(counts: &[usize; 256]) -> HuffmanTable { - let table = HuffmanTable::build_limited(counts, HUF_MAX_BITS); - let longest = (0..=255u8) - .filter_map(|symbol| table.num_bits_for_symbol(symbol)) - .max() - .map_or(HUF_MAX_BITS, usize::from); - if longest < HUF_MAX_BITS { - HuffmanTable::build_limited(counts, longest) - } else { - table - } -} - /// The production matcher, recording the literals and sequences of every block /// it matches, block by block. Everything else is forwarded, so frames /// compress as any compressor's would, dictionary reuse included. diff --git a/zstd/src/dictionary/finalize/tests.rs b/zstd/src/dictionary/finalize/tests.rs index d6d32456a..cac369a1f 100644 --- a/zstd/src/dictionary/finalize/tests.rs +++ b/zstd/src/dictionary/finalize/tests.rs @@ -108,12 +108,15 @@ fn samples_that_do_not_compress_leave_a_describable_flat_table() { /// can be written (every weight the same would have none). #[test] fn literal_tables_take_the_depth_of_their_longest_code() { - assert_eq!(literals_table(&[1; 256]).table_log(), 8); + let mut scratch = WeightScratch::default(); + let flat = HuffmanTable::build_limited_in(&[1; 256], HUF_MAX_BITS, &mut scratch); + assert_eq!(flat.table_log(), 8); + scratch.recycle(flat); let mut counts = [2usize; 256]; counts[0] = 4; counts[253] = 1; counts[254] = 1; - let table = literals_table(&counts); + let table = HuffmanTable::build_limited_in(&counts, HUF_MAX_BITS, &mut scratch); assert_eq!(table.table_log(), 9); let mut writer = BitWriter::new(); HuffmanEncoder::new(&table, &mut writer).write_table(); @@ -124,6 +127,39 @@ fn literal_tables_take_the_depth_of_their_longest_code() { assert_eq!(parsed as usize, description.len()); } +/// Below the limit a code is described at its longest length with the same +/// lengths the limit gave it: building under the limit equals building at the +/// longest code, over a scratch whose recycled table held another code. +#[test] +fn a_literal_table_under_the_limit_keeps_its_code_lengths() { + let mut counts = [0usize; 256]; + for (symbol, count) in counts.iter_mut().enumerate() { + *count = 10 + (symbol * symbol) % 7; + } + let mut scratch = WeightScratch::default(); + let stale = HuffmanTable::build_limited_in(&[1; 256], HUF_MAX_BITS, &mut scratch); + scratch.recycle(stale); + let under = HuffmanTable::build_limited_in(&counts, HUF_MAX_BITS, &mut scratch); + let longest = (0..=255u8) + .filter_map(|symbol| under.num_bits_for_symbol(symbol)) + .max() + .map(usize::from) + .unwrap(); + assert!( + longest < HUF_MAX_BITS, + "the code reaches the limit: {longest}" + ); + assert_eq!(under.table_log() as usize, longest); + let at = HuffmanTable::build_limited_in(&counts, longest, &mut WeightScratch::default()); + assert_eq!(at.table_log(), under.table_log()); + for symbol in 0..=255u8 { + assert_eq!( + under.num_bits_for_symbol(symbol), + at.num_bits_for_symbol(symbol) + ); + } +} + /// Random bytes, different for each seed. fn noise(seed: u64, len: usize) -> Vec { let mut state = seed; diff --git a/zstd/src/huff0/huff0_encoder.rs b/zstd/src/huff0/huff0_encoder.rs index d7e0009a3..0f0095608 100644 --- a/zstd/src/huff0/huff0_encoder.rs +++ b/zstd/src/huff0/huff0_encoder.rs @@ -995,13 +995,34 @@ impl HuffmanTable { Self::build_from_weights_reusing(weights, None) } - /// The optimal code for `counts` with no code longer than `max_bits` - /// (upstream zstd `HUF_buildCTable_wksp`, as the dictionary finalizer - /// calls it). Unlike [`Self::build_from_counts`] it runs no table-log - /// search: the table describes future literals, not these. + /// The optimal code for `counts` with no code longer than `max_bits`, + /// described at its longest code rather than at the limit (upstream zstd + /// `HUF_buildCTable_wksp`, whose returned `maxNbBits` the dictionary + /// finalizer hands to `HUF_writeCTable`). Unlike + /// [`Self::build_from_counts`] it runs no table-log search: the table + /// describes future literals, not these. The buffers come from + /// `scratch`, where [`WeightScratch::recycle`] returns the table. #[cfg(feature = "dict-builder")] - pub(crate) fn build_limited(counts: &[usize], max_bits: usize) -> Self { - Self::build_from_weights(&build_limited_weights(counts, max_bits)) + pub(crate) fn build_limited_in( + counts: &[usize], + max_bits: usize, + scratch: &mut WeightScratch, + ) -> Self { + build_limited_weights_into(counts, max_bits, scratch); + // The smallest weight belongs to the longest code. Lowering every + // weight until it is one keeps each code's length and makes the table + // log the longest code's. + let step = scratch + .weights + .iter() + .copied() + .filter(|&weight| weight > 0) + .min() + .map_or(0, |weight| weight - 1); + for weight in scratch.weights.iter_mut().filter(|weight| **weight > 0) { + *weight -= step; + } + Self::build_from_weights_reusing(&scratch.weights, scratch.spare_table.take()) } /// [`Self::build_from_weights`] filling a discarded table's buffers instead From f0e6d19539915d68acee37000fe2f26e621acb52 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Tue, 29 Sep 2026 03:33:41 +0300 Subject: [PATCH 15/20] fix(dict): fit summed literal counts to the Huffman tree The literal counts the finalizer sums over every sample fed the Huffman builder unchecked: past 2^32 literals, which a direct finalize over large samples can reach, the tree's u32 node counts wrapped in release and panicked in debug. The counts are now halved, rounding up so every symbol that occurred stays in the code, until they fit a node, and `build_limited_in` runs the entry check the other builders run. Regression test: literal_counts_past_a_tree_node_still_build_a_code. A search keeps the winner's content only when the winner will be shrunk; the shrink setting now lives in `Best`, which ranks candidates at full size and cuts only the winner, as `CoverOptions::shrink` documents. Dictionaries are byte-identical for --train, FastCOVER, COVER with a search and COVER with shrink. --- zstd/src/dictionary/finalize.rs | 42 ++++++++++++++++++--------- zstd/src/dictionary/finalize/tests.rs | 16 ++++++++++ zstd/src/dictionary/mod.rs | 13 ++++++--- zstd/src/dictionary/selection.rs | 32 +++++++++++--------- zstd/src/huff0/huff0_encoder.rs | 1 + 5 files changed, 74 insertions(+), 30 deletions(-) diff --git a/zstd/src/dictionary/finalize.rs b/zstd/src/dictionary/finalize.rs index 65bc3b61d..76cd3dd20 100644 --- a/zstd/src/dictionary/finalize.rs +++ b/zstd/src/dictionary/finalize.rs @@ -186,19 +186,7 @@ fn analyze_entropy<'s>( count_samples(&mut counts, analysis, content, samples, count, level)?; let huffman = &mut analysis.huffman; - let mut literals = HuffmanTable::build_limited_in(&counts.literals, HUF_MAX_BITS, huffman); - if literals.table_log() == 8 { - // Every symbol at eight bits describes nothing and cannot be written; - // a mostly flat distribution that still compresses stands in - // (upstream zstd `ZDICT_flatLit`). - counts.literals = [2; 256]; - counts.literals[0] = 4; - counts.literals[253] = 1; - counts.literals[254] = 1; - huffman.recycle(literals); - literals = HuffmanTable::build_limited_in(&counts.literals, HUF_MAX_BITS, huffman); - debug_assert_eq!(literals.table_log(), 9); - } + let literals = literals_table(&mut counts.literals, huffman); // Each description goes straight into the dictionary; the sequence tables // are normalized and described without being built. let mut writer = BitWriter::from(&mut *out); @@ -215,6 +203,34 @@ fn analyze_entropy<'s>( Ok(()) } +/// The literals code for `counts`, built in `scratch`. +fn literals_table(counts: &mut [usize; 256], scratch: &mut WeightScratch) -> HuffmanTable { + // Summed over every sample, the counts can pass what a tree node holds + // (`u32`); halving keeps their proportions, and rounding up keeps every + // symbol that occurred in the code. + // The literals counted are bytes held in memory, so the sum fits `u64`. + while counts.iter().map(|&count| count as u64).sum::() >= u64::from(u32::MAX) { + for count in counts.iter_mut() { + *count = count.div_ceil(2); + } + } + let literals = HuffmanTable::build_limited_in(counts, HUF_MAX_BITS, scratch); + if literals.table_log() != 8 { + return literals; + } + // Every symbol at eight bits describes nothing and cannot be written; a + // mostly flat distribution that still compresses stands in (upstream zstd + // `ZDICT_flatLit`). + *counts = [2; 256]; + counts[0] = 4; + counts[253] = 1; + counts[254] = 1; + scratch.recycle(literals); + let flat = HuffmanTable::build_limited_in(counts, HUF_MAX_BITS, scratch); + debug_assert_eq!(flat.table_log(), 9); + flat +} + /// Compress the first 128 KiB of each of the first `count` samples with /// `content` as a raw dictionary and count what its compressed blocks hold. fn count_samples<'s>( diff --git a/zstd/src/dictionary/finalize/tests.rs b/zstd/src/dictionary/finalize/tests.rs index cac369a1f..e85b10d3a 100644 --- a/zstd/src/dictionary/finalize/tests.rs +++ b/zstd/src/dictionary/finalize/tests.rs @@ -127,6 +127,22 @@ fn literal_tables_take_the_depth_of_their_longest_code() { assert_eq!(parsed as usize, description.len()); } +/// Literal counts summed over the samples can pass what a Huffman tree node +/// holds (more than 4 GiB of literals); the code is still built, from the same +/// proportions, instead of overflowing the tree. +#[test] +fn literal_counts_past_a_tree_node_still_build_a_code() { + let mut counts = [1usize; 256]; + counts[b'e' as usize] = 3 << 31; + counts[b't' as usize] = 1 << 31; + counts[b'a' as usize] = 1 << 28; + let table = literals_table(&mut counts, &mut WeightScratch::default()); + let bits = |symbol: u8| table.num_bits_for_symbol(symbol).unwrap(); + assert!(bits(b'e') <= bits(b't')); + assert!(bits(b't') <= bits(b'a')); + assert!(bits(b'a') < bits(b'z')); +} + /// Below the limit a code is described at its longest length with the same /// lengths the limit gave it: building under the limit equals building at the /// longest code, over a scratch whose recycled table held another code. diff --git a/zstd/src/dictionary/mod.rs b/zstd/src/dictionary/mod.rs index 432be6d5f..0965bc32e 100644 --- a/zstd/src/dictionary/mod.rs +++ b/zstd/src/dictionary/mod.rs @@ -754,7 +754,7 @@ fn run_cover( let scored = !plain || options.shrink.is_some() || space.split_point < 1.0; let mut evaluator = selection::Evaluator::new(&set, split.train, split.test.clone(), dict_size, finalize); - let mut best = selection::Best::new(); + let mut best = selection::Best::new(options.shrink); let mut state = Vec::new(); let mut content_scratch = Vec::new(); // `None` beside a `d` is a dmer size these samples cannot index. @@ -797,9 +797,13 @@ fn run_cover( if !scored { return Ok((evaluator.finalize(content)?, chosen)); } + // Ranked at full size; `finish` shrinks only the winner, as + // `CoverOptions::shrink` documents. Upstream's optimizer never applies + // its shrink at all, and cutting every candidate would multiply the + // search by the number of sizes tried. best.offer(evaluator.score(content), content, chosen); } - best.finish(&mut evaluator, options.shrink) + best.finish(&mut evaluator) } /// Train a FastCOVER dictionary of at most `dict_size` bytes with the `k` and @@ -909,7 +913,7 @@ fn run_fastcover( let mut window = fastcover::WindowCounts::default(); let mut freqs = Vec::new(); let mut content_scratch = Vec::new(); - let mut best = selection::Best::new(); + let mut best = selection::Best::new(options.cover.shrink); // `None` beside a `d` is a dmer size these samples cannot count. let mut context: Option<(usize, Option>)> = None; let mut evaluator = selection::Evaluator::new( @@ -958,9 +962,10 @@ fn run_fastcover( if !scored { return Ok((evaluator.finalize(content)?, chosen)); } + // Ranked at full size, the winner alone shrunk, as in the COVER search. best.offer(evaluator.score(content), content, chosen); } - best.finish(&mut evaluator, options.cover.shrink) + best.finish(&mut evaluator) } /// Train and finalize a dictionary with the reference's original trainer, the diff --git a/zstd/src/dictionary/selection.rs b/zstd/src/dictionary/selection.rs index 3e9b54a9d..68cde76cc 100644 --- a/zstd/src/dictionary/selection.rs +++ b/zstd/src/dictionary/selection.rs @@ -179,26 +179,35 @@ pub(super) struct Priced<'e> { pub(super) total: usize, } -/// The cheapest dictionary seen, the content it was finalized from and the -/// parameters that built it. A tie keeps the earlier one, as upstream's strict -/// comparison does. +/// The cheapest dictionary seen, the parameters that built it and, when the +/// winner is to be shrunk, the content it was finalized from. A tie keeps the +/// earlier one, as upstream's strict comparison does. +/// +/// Candidates are ranked at full size and only the winner is shrunk: the size +/// a dictionary is cut to is a separate choice from the `k` and `d` that built +/// it (see [`Evaluator::shrink`]). pub(super) struct Best

{ found: Option<(Scored, Vec, P)>, last_error: Option, + shrink: Option, } impl

Best

{ - pub(super) fn new() -> Self { + /// A search whose winner is cut down with `shrink` when given. + pub(super) fn new(shrink: Option) -> Self { Self { found: None, last_error: None, + shrink, } } - /// Keep `candidate` if it is the cheapest so far. Its dictionary and - /// `content` are copied only then, into the buffers the previous winner - /// held; a candidate that loses is not copied at all. + /// Keep `candidate` if it is the cheapest so far. Its dictionary, and its + /// `content` when the winner will be shrunk, are copied only then, into + /// the buffers the previous winner held; a candidate that loses is not + /// copied at all. pub(super) fn offer(&mut self, candidate: io::Result>, content: &[u8], params: P) { + let content = if self.shrink.is_some() { content } else { &[] }; match candidate { Ok(priced) => match &mut self.found { Some((best, _, _)) if priced.total >= best.total => {} @@ -222,13 +231,10 @@ impl

Best

{ } } - /// The winner's dictionary, cut down with `shrink` when given, and its + /// The winner's dictionary, cut down when the search shrinks, and its /// parameters; or why no candidate could be priced. - pub(super) fn finish( - self, - evaluator: &mut Evaluator<'_>, - shrink: Option, - ) -> io::Result<(Vec, P)> { + pub(super) fn finish(self, evaluator: &mut Evaluator<'_>) -> io::Result<(Vec, P)> { + let shrink = self.shrink; let (full, content, params) = match (self.found, self.last_error) { (Some(found), _) => found, (None, Some(err)) => return Err(err), diff --git a/zstd/src/huff0/huff0_encoder.rs b/zstd/src/huff0/huff0_encoder.rs index 0f0095608..e660d5b07 100644 --- a/zstd/src/huff0/huff0_encoder.rs +++ b/zstd/src/huff0/huff0_encoder.rs @@ -1008,6 +1008,7 @@ impl HuffmanTable { max_bits: usize, scratch: &mut WeightScratch, ) -> Self { + assert_histogram_fits_nodes::(counts); build_limited_weights_into(counts, max_bits, scratch); // The smallest weight belongs to the longest code. Lowering every // weight until it is one keeps each code's length and makes the table From 5771bb2ff856fe299d77b3f12eeadad469e31a86 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Tue, 29 Sep 2026 03:51:57 +0300 Subject: [PATCH 16/20] fix(dict): refuse an undersized finalize before the samples `finalize_raw_dict` built its sample set, walking every size and allocating the offsets, before it refused a dictionary under 256 bytes, and `ZDICT_finalizeDictionary` summed the sample sizes first too: an impossible request over a large corpus paid for the walk, and sizes that did not add up were reported instead of the size. Upstream zstd checks the capacity before anything else. The rule is the codec's `check_finalize_dict_size`, which `finalize_raw_dict` runs first and the C ABI runs before it reads any argument. Regression tests: finalize_raw_dict_refuses_an_undersized_dictionary_before_the_samples, and an overflowing size list in zdict_trainers_report_upstreams_error_codes. --- c-api/src/dict.rs | 5 +++++ c-api/src/tests.rs | 16 ++++++++++++++++ zstd/src/dictionary/finalize.rs | 15 +++++++++++---- zstd/src/dictionary/mod.rs | 25 +++++++++++++++++++++++++ zstd/src/dictionary/tests.rs | 20 ++++++++++++++++++++ 5 files changed, 77 insertions(+), 4 deletions(-) diff --git a/c-api/src/dict.rs b/c-api/src/dict.rs index 69476ecee..5b065045e 100644 --- a/c-api/src/dict.rs +++ b/c-api/src/dict.rs @@ -457,6 +457,11 @@ pub unsafe extern "C" fn ZDICT_finalizeDictionary( nb_samples: c_uint, parameters: ZDICT_params_t, ) -> usize { + // The finalizer's own first check, run before anything is read, the sample + // sizes included. + if let Err(err) = codec::dictionary::check_finalize_dict_size(max_dict_size) { + return encode(crate::error::code_for_training_error(&err)); + } // NULL + non-zero length is caller error, not a slice to build. if dict_content.is_null() && dict_content_size > 0 { return encode(ZSTD_ErrorCode::ZSTD_error_dictionaryCreation_failed); diff --git a/c-api/src/tests.rs b/c-api/src/tests.rs index 8a795c2e6..fccb6860a 100644 --- a/c-api/src/tests.rs +++ b/c-api/src/tests.rs @@ -549,6 +549,22 @@ fn zdict_trainers_report_upstreams_error_codes() { ) }; assert_eq!(code_from_result(n), ZSTD_error_dstSize_tooSmall); + // Before the sample sizes are summed, too: sizes that overflow are not + // what the refusal reports. + let overflowing = [usize::MAX, 2]; + let n = unsafe { + crate::dict::ZDICT_finalizeDictionary( + dict.as_mut_ptr(), + 255, + samples.as_ptr(), + 8, + samples.as_ptr(), + overflowing.as_ptr(), + overflowing.len() as u32, + zparams, + ) + }; + assert_eq!(code_from_result(n), ZSTD_error_dstSize_tooSmall); // A zero `k` asks the plain trainers for a search they do not run. let n = unsafe { ZDICT_trainFromBuffer_cover( diff --git a/zstd/src/dictionary/finalize.rs b/zstd/src/dictionary/finalize.rs index 76cd3dd20..eede80df2 100644 --- a/zstd/src/dictionary/finalize.rs +++ b/zstd/src/dictionary/finalize.rs @@ -92,10 +92,7 @@ pub(super) fn finalize_into<'s>( dict_size: usize, options: FinalizeOptions, ) -> io::Result<()> { - // Upstream zstd's order (`ZDICT_finalizeDictionary`): the size first. - if dict_size < DICT_SIZE_MIN { - return Err(too_small()); - } + check_size(dict_size)?; if content.is_empty() { return Err(invalid("raw dictionary content must not be empty")); } @@ -124,6 +121,16 @@ pub(super) fn finalize_into<'s>( Ok(()) } +/// Refuse a dictionary below the smallest one finalized. Upstream zstd's +/// order (`ZDICT_finalizeDictionary`): checked before anything else, the +/// samples included. +pub(super) fn check_size(dict_size: usize) -> io::Result<()> { + if dict_size < DICT_SIZE_MIN { + return Err(too_small()); + } + Ok(()) +} + fn too_small() -> io::Error { refuse( TrainingError::DictionaryTooSmall, diff --git a/zstd/src/dictionary/mod.rs b/zstd/src/dictionary/mod.rs index 0965bc32e..f9c75b3df 100644 --- a/zstd/src/dictionary/mod.rs +++ b/zstd/src/dictionary/mod.rs @@ -365,10 +365,35 @@ pub fn finalize_raw_dict( dict_size: usize, options: FinalizeOptions, ) -> io::Result> { + check_finalize_dict_size(dict_size)?; let set = samples::SampleSet::new(samples, sample_sizes)?; finalize::finalize(raw_content, &set, set.count(), dict_size, options) } +/// Refuse a `dict_size` [`finalize_raw_dict`] cannot fill, from the size +/// alone: below [`MIN_TRAINED_DICT_SIZE`]. Finalizing checks it first, before +/// the samples; a caller holding the samples somewhere costly to walk can run +/// it before that. +/// +/// # Errors +/// +/// `InvalidInput` carrying [`TrainingError::DictionaryTooSmall`]. +/// +/// # Examples +/// +/// ``` +/// use structured_zstd::dictionary::{ +/// MIN_TRAINED_DICT_SIZE, TrainingError, check_finalize_dict_size, +/// }; +/// +/// assert!(check_finalize_dict_size(MIN_TRAINED_DICT_SIZE).is_ok()); +/// let err = check_finalize_dict_size(MIN_TRAINED_DICT_SIZE - 1).unwrap_err(); +/// assert_eq!(TrainingError::of(&err), Some(TrainingError::DictionaryTooSmall)); +/// ``` +pub fn check_finalize_dict_size(dict_size: usize) -> io::Result<()> { + finalize::check_size(dict_size) +} + /// The `k` and `d` values a training run tries, resolved from the options the /// way the reference's optimizers resolve theirs. struct SearchSpace { diff --git a/zstd/src/dictionary/tests.rs b/zstd/src/dictionary/tests.rs index 4f4f7e93a..e9b1c8e13 100644 --- a/zstd/src/dictionary/tests.rs +++ b/zstd/src/dictionary/tests.rs @@ -804,6 +804,26 @@ fn finalize_raw_dict_rejects_zero_dict_id() { assert_eq!(err.to_string(), "dictionary id must be non-zero"); } +/// A dictionary too small is refused before the samples are walked, as +/// upstream's `ZDICT_finalizeDictionary` checks the capacity first: sizes that +/// do not describe the samples are not what the refusal reports. +#[test] +fn finalize_raw_dict_refuses_an_undersized_dictionary_before_the_samples() { + let (data, sizes) = training_samples(); + let err = finalize_raw_dict( + b"content", + &data, + &sizes[1..], + MIN_TRAINED_DICT_SIZE - 1, + FinalizeOptions::default(), + ) + .unwrap_err(); + assert_eq!( + TrainingError::of(&err), + Some(TrainingError::DictionaryTooSmall) + ); +} + /// The sample sizes must describe the samples exactly. #[test] fn finalize_raw_dict_rejects_sizes_that_do_not_add_up() { From 4bdcd0820c422d4217ea93e852e0bbcdd1a1e767 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Tue, 29 Sep 2026 05:27:41 +0300 Subject: [PATCH 17/20] docs(dict): name the raw content's size floor by its current constant --- zstd/src/dictionary/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zstd/src/dictionary/mod.rs b/zstd/src/dictionary/mod.rs index f9c75b3df..2f929c84c 100644 --- a/zstd/src/dictionary/mod.rs +++ b/zstd/src/dictionary/mod.rs @@ -227,7 +227,7 @@ pub fn create_raw_dict_from_source( /// of at most 128 KiB, and the content is what [`optimize_fastcover_dict`] /// picks over them at its defaults, as `zstd --train` does. A corpus no larger /// than `dict_size` is its own content; one the trainer refuses (a -/// `dict_size` under [`TRAINER_DICT_SIZE_MIN`], too little to search) gives +/// `dict_size` under [`MIN_TRAINED_DICT_SIZE`], too little to search) gives /// its last `dict_size` bytes. /// /// # Errors From 6fe11574102942ad75a4f549f13c1f0eab2c0515 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Tue, 29 Sep 2026 05:31:20 +0300 Subject: [PATCH 18/20] fix(dict): gate the known-size frame writer, place the split doc - `FrameCompressor::compress_known_into` has only the dictionary finalizer as a caller, so a build without `dict-builder` failed its dead-code lint; it is compiled with that feature only - the paragraph describing `SampleSet::split` sat above `check_holds_dmer`; it now documents `split` - the README and the loader say where the five-sample floor applies: to the samples the files make, as upstream's command checks it (dibio.c); COVER and FastCOVER check what the split and the memory limit keep, and the legacy trainer, as upstream's, trains on whatever the limit keeps Part of #128 --- README.md | 5 ++++- zstd/src/bin/structured-zstd/main.rs | 4 ++++ zstd/src/dictionary/samples.rs | 8 ++++---- zstd/src/encoding/frame_compressor.rs | 1 + 4 files changed, 13 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 77b755c85..ee61a5dad 100644 --- a/README.md +++ b/README.md @@ -125,7 +125,10 @@ upstream parses it and then never applies it, here it takes effect. `--train-legacy[=s=#]` (or `-s#`) runs upstream's original trainer. Every trainer counts samples, and they are loaded as upstream loads them (each file one sample of up to 128 KiB, or cut into `-B#` pieces, whole samples up to -2 GiB or `-M` when that is smaller); at least five are needed. For the same +2 GiB or `-M` when that is smaller). The files have to make at least five +samples; COVER and FastCOVER also need five to build from once the split and +the `-M` cap are applied, while the legacy trainer, as upstream's, trains on +whatever the cap keeps. For the same file list the legacy dictionary carries the same content as upstream's. `-D` takes either a dictionary produced by `--train` or any file at all, which is then used as raw content the way upstream does; such a dictionary has no ID, so the same bytes must be supplied when decoding. diff --git a/zstd/src/bin/structured-zstd/main.rs b/zstd/src/bin/structured-zstd/main.rs index d0c3aad76..c2c57e47c 100644 --- a/zstd/src/bin/structured-zstd/main.rs +++ b/zstd/src/bin/structured-zstd/main.rs @@ -3621,6 +3621,10 @@ fn load_training_samples( let budget = TRAINING_DATA_MAX.min(memory_limit.unwrap_or(u64::MAX)); let (plan, bytes) = plan_training_load(&file_sizes, block_size, budget, samples); // What the budget keeps, not what the files hold, is what the trainer gets. + // The five-sample floor above is on the files, as upstream's is + // (dibio.c, `fs.nbSamples < 5`); the samples kept under the budget go to + // the trainer's own check, and the legacy trainer, as upstream's, has + // none. let retained: u64 = plan.iter().sum(); enough_samples(usize::try_from(retained).unwrap_or(usize::MAX))?; // Both buffers are sized to exactly what the plan loads, which the budget diff --git a/zstd/src/dictionary/samples.rs b/zstd/src/dictionary/samples.rs index 8471c55c5..5af64a479 100644 --- a/zstd/src/dictionary/samples.rs +++ b/zstd/src/dictionary/samples.rs @@ -61,10 +61,6 @@ impl<'a> SampleSet<'a> { &self.data[..self.offsets[count]] } - /// Divide the samples into those a dictionary is built from and those it - /// is scored on (upstream zstd `COVER_ctx_init`): below 1 the leading - /// `split_point` share builds and the rest scores; at 1 every sample does - /// both. /// Refuse the first `train` samples when none of them is `span` bytes /// long: a trainer counts a dmer only inside one sample, so none would be /// counted however long the samples run together. Known from the sizes, @@ -82,6 +78,10 @@ impl<'a> SampleSet<'a> { )) } + /// Divide the samples into those a dictionary is built from and those it + /// is scored on (upstream zstd `COVER_ctx_init`): below 1 the leading + /// `split_point` share builds and the rest scores; at 1 every sample does + /// both. pub(super) fn split(&self, split_point: f64) -> io::Result { split_count(self.count(), split_point) } diff --git a/zstd/src/encoding/frame_compressor.rs b/zstd/src/encoding/frame_compressor.rs index c1bb192e9..b210347c4 100644 --- a/zstd/src/encoding/frame_compressor.rs +++ b/zstd/src/encoding/frame_compressor.rs @@ -2488,6 +2488,7 @@ impl FrameCompressor { /// front the header goes first and the blocks straight after it, with no /// block accumulator and no copy into the drain. The frame is the one /// `compress` writes for the same source, hint and settings. + #[cfg(feature = "dict-builder")] pub(crate) fn compress_known_into(&mut self, total: u64, out: &mut Vec) { let prep = self.prepare_frame(crate::encoding::workspace::IngestPlan::Stream); let mut source = self From 69168e23a7f7831deb7539676338b5d1b38a3d15 Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Tue, 29 Sep 2026 05:54:26 +0300 Subject: [PATCH 19/20] perf(dict): hand back the search winner's raw content The raw-dictionary path ran the FastCOVER search, then built a second context from the parameters it chose and selected the same segments again: another corpus-wide dmer count, another 2^20-entry frequency table and another selection, for content the search already held. The search now takes what to keep of its winner: its dictionary, cut down when shrinking, or the raw content it was finalized from. The raw path asks for the content and gets it straight from the search, and a content search no longer copies each winner's finalized dictionary. The content is unchanged; raw_content_is_the_fastcover_search_winner checks it against the dictionary the search finalizes. Part of #128 --- zstd/src/dictionary/mod.rs | 70 ++++++++++++---------- zstd/src/dictionary/selection.rs | 100 +++++++++++++++++++------------ 2 files changed, 103 insertions(+), 67 deletions(-) diff --git a/zstd/src/dictionary/mod.rs b/zstd/src/dictionary/mod.rs index 2f929c84c..c38f1af03 100644 --- a/zstd/src/dictionary/mod.rs +++ b/zstd/src/dictionary/mod.rs @@ -288,37 +288,19 @@ fn fastcover_raw_content( } /// The content of the dictionary [`optimize_fastcover_dict`] picks at its -/// defaults, rebuilt from the parameters it chose: a build is deterministic, -/// so this is the content the winner was finalized from. +/// defaults: the search keeps its winner's content and hands that back. fn search_raw_content(corpus: &[u8], sizes: &[usize], dict_size: usize) -> io::Result> { - let (_, chosen) = optimize_fastcover_dict( + let options = FastCoverOptions::default(); + let space = SearchSpace::optimizing(&options.cover, 0.75)?; + let (content, _) = run_fastcover( corpus, sizes, dict_size, - &FastCoverOptions::default(), + &options, FinalizeOptions::default(), + &space, + selection::Keep::Content, )?; - let set = samples::SampleSet::new(corpus, sizes)?; - let split = set.split(chosen.cover.split_point)?; - let ctx = fastcover::FastCoverContext::new( - &set, - split.train, - chosen.cover.d as usize, - chosen.f, - chosen.accel, - )?; - let mut content = Vec::new(); - // The build fills its buffer from the back; what it returns is the tail. - let kept = ctx - .build( - &mut Vec::new(), - &mut fastcover::WindowCounts::default(), - &mut content, - dict_size, - chosen.cover.k as usize, - )? - .len(); - content.drain(..content.len() - kept); Ok(content) } @@ -779,7 +761,9 @@ fn run_cover( let scored = !plain || options.shrink.is_some() || space.split_point < 1.0; let mut evaluator = selection::Evaluator::new(&set, split.train, split.test.clone(), dict_size, finalize); - let mut best = selection::Best::new(options.shrink); + let mut best = selection::Best::new(selection::Keep::Dictionary { + shrink: options.shrink, + }); let mut state = Vec::new(); let mut content_scratch = Vec::new(); // `None` beside a `d` is a dmer size these samples cannot index. @@ -872,7 +856,18 @@ pub fn train_fastcover_dict( finalize: FinalizeOptions, ) -> io::Result> { let space = SearchSpace::fixed(&options.cover)?; - let (dict, _) = run_fastcover(samples, sample_sizes, dict_size, options, finalize, &space)?; + let keep = selection::Keep::Dictionary { + shrink: options.cover.shrink, + }; + let (dict, _) = run_fastcover( + samples, + sample_sizes, + dict_size, + options, + finalize, + &space, + keep, + )?; Ok(dict) } @@ -918,7 +913,18 @@ pub fn optimize_fastcover_dict( finalize: FinalizeOptions, ) -> io::Result<(Vec, FastCoverOptions)> { let space = SearchSpace::optimizing(&options.cover, 0.75)?; - run_fastcover(samples, sample_sizes, dict_size, options, finalize, &space) + let keep = selection::Keep::Dictionary { + shrink: options.cover.shrink, + }; + run_fastcover( + samples, + sample_sizes, + dict_size, + options, + finalize, + &space, + keep, + ) } fn run_fastcover( @@ -928,6 +934,7 @@ fn run_fastcover( options: &FastCoverOptions, finalize: FinalizeOptions, space: &SearchSpace, + keep: selection::Keep, ) -> io::Result<(Vec, FastCoverOptions)> { let (f, accel) = fastcover_knobs(options, space)?; space.check_fits(dict_size)?; @@ -938,7 +945,7 @@ fn run_fastcover( let mut window = fastcover::WindowCounts::default(); let mut freqs = Vec::new(); let mut content_scratch = Vec::new(); - let mut best = selection::Best::new(options.cover.shrink); + let mut best = selection::Best::new(keep); // `None` beside a `d` is a dmer size these samples cannot count. let mut context: Option<(usize, Option>)> = None; let mut evaluator = selection::Evaluator::new( @@ -985,7 +992,10 @@ fn run_fastcover( }; let content = ctx.build(&mut freqs, &mut window, &mut content_scratch, dict_size, k)?; if !scored { - return Ok((evaluator.finalize(content)?, chosen)); + return match keep { + selection::Keep::Content => Ok((content.to_vec(), chosen)), + selection::Keep::Dictionary { .. } => Ok((evaluator.finalize(content)?, chosen)), + }; } // Ranked at full size, the winner alone shrunk, as in the COVER search. best.offer(evaluator.score(content), content, chosen); diff --git a/zstd/src/dictionary/selection.rs b/zstd/src/dictionary/selection.rs index 68cde76cc..e59a983a9 100644 --- a/zstd/src/dictionary/selection.rs +++ b/zstd/src/dictionary/selection.rs @@ -179,9 +179,30 @@ pub(super) struct Priced<'e> { pub(super) total: usize, } -/// The cheapest dictionary seen, the parameters that built it and, when the -/// winner is to be shrunk, the content it was finalized from. A tie keeps the -/// earlier one, as upstream's strict comparison does. +/// What a search hands back of its winner. +#[derive(Clone, Copy)] +pub(super) enum Keep { + /// The finalized dictionary, cut down with `shrink` when given. + Dictionary { shrink: Option }, + /// The raw content the dictionary was finalized from. + Content, +} + +impl Keep { + /// Whether the winner's content has to be kept through the search. + fn content(self) -> bool { + !matches!(self, Self::Dictionary { shrink: None }) + } + + /// Whether the winner's finalized dictionary is what is handed back. + fn dictionary(self) -> bool { + matches!(self, Self::Dictionary { .. }) + } +} + +/// The cheapest candidate seen, the parameters that built it, and what +/// [`Keep`] asks for of it: its dictionary, its content, or both. A tie keeps +/// the earlier one, as upstream's strict comparison does. /// /// Candidates are ranked at full size and only the winner is shrunk: the size /// a dictionary is cut to is a separate choice from the `k` and `d` that built @@ -189,52 +210,55 @@ pub(super) struct Priced<'e> { pub(super) struct Best

{ found: Option<(Scored, Vec, P)>, last_error: Option, - shrink: Option, + keep: Keep, } impl

Best

{ - /// A search whose winner is cut down with `shrink` when given. - pub(super) fn new(shrink: Option) -> Self { + /// A search handing back what `keep` names of its winner. + pub(super) fn new(keep: Keep) -> Self { Self { found: None, last_error: None, - shrink, + keep, } } - /// Keep `candidate` if it is the cheapest so far. Its dictionary, and its - /// `content` when the winner will be shrunk, are copied only then, into - /// the buffers the previous winner held; a candidate that loses is not - /// copied at all. + /// Keep `candidate` if it is the cheapest so far. What [`Keep`] asks for + /// of it is copied only then, into the buffers the previous winner held; + /// a candidate that loses is not copied at all. pub(super) fn offer(&mut self, candidate: io::Result>, content: &[u8], params: P) { - let content = if self.shrink.is_some() { content } else { &[] }; + let content = if self.keep.content() { content } else { &[] }; + let keep_dictionary = self.keep.dictionary(); match candidate { - Ok(priced) => match &mut self.found { - Some((best, _, _)) if priced.total >= best.total => {} - Some((best, kept, kept_params)) => { - best.total = priced.total; - best.dict.clear(); - best.dict.extend_from_slice(priced.dict); - kept.clear(); - kept.extend_from_slice(content); - *kept_params = params; - } - None => { - let best = Scored { - dict: priced.dict.to_vec(), - total: priced.total, - }; - self.found = Some((best, content.to_vec(), params)); + Ok(priced) => { + let dict = if keep_dictionary { priced.dict } else { &[] }; + match &mut self.found { + Some((best, _, _)) if priced.total >= best.total => {} + Some((best, kept, kept_params)) => { + best.total = priced.total; + best.dict.clear(); + best.dict.extend_from_slice(dict); + kept.clear(); + kept.extend_from_slice(content); + *kept_params = params; + } + None => { + let best = Scored { + dict: dict.to_vec(), + total: priced.total, + }; + self.found = Some((best, content.to_vec(), params)); + } } - }, + } Err(err) => self.last_error = Some(err), } } - /// The winner's dictionary, cut down when the search shrinks, and its - /// parameters; or why no candidate could be priced. + /// What [`Keep`] asked for of the winner, and its parameters; or why no + /// candidate could be priced. pub(super) fn finish(self, evaluator: &mut Evaluator<'_>) -> io::Result<(Vec, P)> { - let shrink = self.shrink; + let keep = self.keep; let (full, content, params) = match (self.found, self.last_error) { (Some(found), _) => found, (None, Some(err)) => return Err(err), @@ -245,10 +269,12 @@ impl

Best

{ )); } }; - let chosen = match shrink { - Some(regression) => evaluator.shrink(&content, full, regression)?, - None => full, - }; - Ok((chosen.dict, params)) + match keep { + Keep::Content => Ok((content, params)), + Keep::Dictionary { shrink: None } => Ok((full.dict, params)), + Keep::Dictionary { + shrink: Some(regression), + } => Ok((evaluator.shrink(&content, full, regression)?.dict, params)), + } } } From c9a7cc6a9f8a1f2c911fb6b89f413d861769a33f Mon Sep 17 00:00:00 2001 From: Dmitry Prudnikov Date: Tue, 29 Sep 2026 06:23:26 +0300 Subject: [PATCH 20/20] fix(dict): analyse samples one written block per matched block At btopt and above with a window of 128 KiB or more, the frame the analysis compresses could cut a matched block into several blocks after matching. The recorder sees one block per matcher call, so the pieces could not be told apart, and a frame with any raw piece had none of its blocks counted: a 128 KiB sample of 16 KiB of log lines and 112 KiB of noise at level 19 went out as two compressed pieces and a raw one and taught the tables nothing. The analysis compressor now never cuts after matching, as upstream's analysis compresses a single block (`ZDICT_countEStats`, `ZSTD_compressBlock`), so every written block is one recorded block and is counted by its own kind. Regression test: a_block_mixing_text_and_noise_is_counted. The literals table's description is encoded into the table's own buffer with the analysis' FSE table before it is written, both kept from one candidate to the next, instead of into a fresh buffer and table per candidate. Dictionaries on the repository files and the decodecorpus files are unchanged at levels 3 and 19, where no sample reaches a split. Part of #128 --- zstd/src/dictionary/finalize.rs | 30 ++++++------- zstd/src/dictionary/finalize/tests.rs | 47 ++++++++++++++++++++ zstd/src/encoding/blocks/compressed.rs | 1 + zstd/src/encoding/blocks/compressed/tests.rs | 6 +++ zstd/src/encoding/frame_compressor.rs | 15 +++++++ zstd/src/encoding/levels/fastest.rs | 3 +- zstd/src/encoding/levels/fastest/tests.rs | 1 + zstd/src/encoding/streaming_encoder.rs | 1 + zstd/src/huff0/huff0_encoder.rs | 8 +++- 9 files changed, 94 insertions(+), 18 deletions(-) diff --git a/zstd/src/dictionary/finalize.rs b/zstd/src/dictionary/finalize.rs index eede80df2..091e788df 100644 --- a/zstd/src/dictionary/finalize.rs +++ b/zstd/src/dictionary/finalize.rs @@ -193,7 +193,10 @@ fn analyze_entropy<'s>( count_samples(&mut counts, analysis, content, samples, count, level)?; let huffman = &mut analysis.huffman; - let literals = literals_table(&mut counts.literals, huffman); + let mut literals = literals_table(&mut counts.literals, huffman); + // Encoded in the table's own buffer and the scratch's FSE table, both kept + // from one candidate to the next, so the writer below reads it back. + literals.fill_weight_description_from_codes(huffman.weight_fse_table()); // Each description goes straight into the dictionary; the sequence tables // are normalized and described without being built. let mut writer = BitWriter::from(&mut *out); @@ -269,10 +272,12 @@ fn count_samples<'s>( blocks: Vec::new(), sequences: Vec::new(), }; - *compressor = Some(( - level, - FrameCompressor::new_with_matcher(recorder, CompressionLevel::from_level(level)), - )); + let mut built = + FrameCompressor::new_with_matcher(recorder, CompressionLevel::from_level(level)); + // One written block per matched one, so each is counted by its own + // kind; upstream's analysis compresses a single block, never split. + built.forbid_post_split(); + *compressor = Some((level, built)); } let (_, compressor) = compressor.as_mut().expect("built above"); compressor @@ -302,21 +307,14 @@ fn count_samples<'s>( let recorder = compressor.matcher_mut(); // A block written raw or as one repeated byte holds no sequences to // learn from, and its bytes are no literals of any compressed block. - // Blocks are recorded one per block the frame holds; when one was - // split after matching the pieces cannot be told apart, and the - // frame's blocks count only if every piece was compressed. - let one_to_one = compressed.len() == recorder.blocks.len(); - let all_compressed = compressed.iter().all(|&kind| kind); + // With no cut after matching, the frame holds one block per recorded + // one. + debug_assert_eq!(compressed.len(), recorder.blocks.len()); let mut reps = START_REPS; let mut start = 0; - for (index, block) in recorder.blocks.iter().enumerate() { + for (block, &counted) in recorder.blocks.iter().zip(compressed.iter()) { let block_sequences = &recorder.sequences[start..block.sequences_end]; start = block.sequences_end; - let counted = if one_to_one { - compressed[index] - } else { - all_compressed - }; // A block left out does not advance the repeat offsets either: the // encoder restores them when it writes a block raw. if !counted { diff --git a/zstd/src/dictionary/finalize/tests.rs b/zstd/src/dictionary/finalize/tests.rs index e85b10d3a..9f92570c2 100644 --- a/zstd/src/dictionary/finalize/tests.rs +++ b/zstd/src/dictionary/finalize/tests.rs @@ -221,6 +221,53 @@ fn a_block_written_raw_is_not_counted() { assert!(one == other, "a raw block's bytes reached the tables"); } +/// A block of text and noise still teaches its text. At level 19 a 128 KiB +/// sample of 16 KiB of log lines and 112 KiB of noise is one block for the +/// matcher; cut after matching, it went out as two compressed pieces and a raw +/// one, which could not be told apart from the one recorded block, so none of +/// the sample was counted. +#[test] +fn a_block_mixing_text_and_noise_is_counted() { + const TEXT: usize = 16 << 10; + let mut data = Vec::new(); + let mut line = 0u32; + while data.len() < TEXT { + data.extend_from_slice( + std::format!( + "ts={line} level=info msg=request served path=/api/v1/items/{}\n", + line % 97 + ) + .as_bytes(), + ); + line += 1; + } + data.truncate(TEXT); + data.extend_from_slice(&noise(0x2545_F491_4F6C_DD1D, (128 << 10) - TEXT)); + let sizes = [data.len()]; + let samples = SampleSet::new(&data, &sizes).unwrap(); + let content = data[..4096].to_vec(); + let mut counts = EntropyCounts { + literals: [1; 256], + offset_codes: [1; 32], + match_lengths: [1; 53], + literal_lengths: [1; 36], + }; + count_samples( + &mut counts, + &mut Analysis::default(), + &content, + &samples, + 1, + 19, + ) + .unwrap(); + let matches: usize = counts.match_lengths.iter().sum(); + assert!( + matches > 53 + 100, + "the text half's sequences were counted: {matches}" + ); +} + /// Offset codes are counted with the repeat policy the frame's blocks used. /// Two tokens alternating behind short varied gaps make an offset often equal /// the one before last: the full search writes it as repeat 2 or 3 (offset diff --git a/zstd/src/encoding/blocks/compressed.rs b/zstd/src/encoding/blocks/compressed.rs index b15a85f48..ef09a47f3 100644 --- a/zstd/src/encoding/blocks/compressed.rs +++ b/zstd/src/encoding/blocks/compressed.rs @@ -516,6 +516,7 @@ pub(crate) fn compress_block_with_post_split( pre_split: state.pre_split, huf_optimal_search: state.huf_optimal_search, literal_compression_disabled: state.literal_compression_disabled, + post_split_allowed: state.post_split_allowed, }, workspace, }; diff --git a/zstd/src/encoding/blocks/compressed/tests.rs b/zstd/src/encoding/blocks/compressed/tests.rs index 13eebf0a0..e0b798d17 100644 --- a/zstd/src/encoding/blocks/compressed/tests.rs +++ b/zstd/src/encoding/blocks/compressed/tests.rs @@ -399,6 +399,7 @@ fn estimator_literals_section_mirrors_emit_for_short_inputs() { pre_split: None, huf_optimal_search: true, literal_compression_disabled: false, + post_split_allowed: true, }; let mut emit_state = CompressState:: { matcher: EntropyOnlyMatcher, @@ -416,6 +417,7 @@ fn estimator_literals_section_mirrors_emit_for_short_inputs() { pre_split: None, huf_optimal_search: true, literal_compression_disabled: false, + post_split_allowed: true, }; let mut workspace = EstimatorWorkspace::default(); let est = estimate_block_parts_size(&mut est_state, &literals, &[], &mut workspace); @@ -481,6 +483,7 @@ fn a_section_with_flat_ends_costs_what_the_emitter_writes_for_it() { pre_split: None, huf_optimal_search: true, literal_compression_disabled: false, + post_split_allowed: true, }; let mut est_state = make_state(); let mut emit_state = make_state(); @@ -546,6 +549,7 @@ fn retained_heap_size_counts_the_block_buffers() { pre_split: None, huf_optimal_search: true, literal_compression_disabled: false, + post_split_allowed: true, }; let before = state.retained_scratch_heap_size(); let block = 64 * 1024; @@ -626,6 +630,7 @@ fn estimator_and_emitter_agree_on_a_block_with_sequences() { pre_split: None, huf_optimal_search: true, literal_compression_disabled: false, + post_split_allowed: true, }; let mut est_state = make_state(); let mut emit_state = make_state(); @@ -690,6 +695,7 @@ fn raw_partition_fallback_restores_repeat_offset_history() { pre_split: None, huf_optimal_search: true, literal_compression_disabled: false, + post_split_allowed: true, }; let source = [0xA5; 8]; let mut sequences = [RawSequence { diff --git a/zstd/src/encoding/frame_compressor.rs b/zstd/src/encoding/frame_compressor.rs index b210347c4..cbc02488a 100644 --- a/zstd/src/encoding/frame_compressor.rs +++ b/zstd/src/encoding/frame_compressor.rs @@ -1409,6 +1409,12 @@ pub(crate) struct CompressState { /// size and the encode cost in parity on the negative band. Set per frame /// alongside `strategy_tag`. pub(crate) literal_compression_disabled: bool, + /// Whether a matched block may be cut into several blocks after matching, + /// where its strategy and window call for it. `true` for every frame but a + /// dictionary's analysis, which has to see one written block per matched + /// block, as upstream's single-block analysis does (`ZDICT_countEStats` + /// compresses with `ZSTD_compressBlock`, which never splits). + pub(crate) post_split_allowed: bool, } /// Whether the HUF literal build should run the #167 table-log search for a @@ -1781,6 +1787,7 @@ impl FrameCompressor { compression_level, crate::encoding::CompressionLevel::Level(n) if n < 0 ), + post_split_allowed: true, }, magicless: false, content_checksum: false, @@ -2220,6 +2227,7 @@ impl FrameCompressor { compression_level, crate::encoding::CompressionLevel::Level(n) if n < 0 ), + post_split_allowed: true, }, compression_level, magicless: false, @@ -2483,6 +2491,13 @@ impl FrameCompressor { self.finish_frame(all_blocks, total_uncompressed, &prep); } + /// Write every matched block as one block, never cut after matching: what + /// a dictionary's analysis needs to tell each block's kind from the frame. + #[cfg(feature = "dict-builder")] + pub(crate) fn forbid_post_split(&mut self) { + self.state.post_split_allowed = false; + } + /// [`Self::compress`] for a source of `total` bytes, the frame written /// over `out` rather than to the drain: with the content size known up /// front the header goes first and the blocks straight after it, with no diff --git a/zstd/src/encoding/levels/fastest.rs b/zstd/src/encoding/levels/fastest.rs index 7d32a47a3..24ef1af5f 100644 --- a/zstd/src/encoding/levels/fastest.rs +++ b/zstd/src/encoding/levels/fastest.rs @@ -138,7 +138,8 @@ pub(crate) fn compress_block_encoded( // consume it: when no sink collects checksums (the common case), and when // the block is headed for the post-split helper, which emits several // physical blocks and records a checksum per partition of its own. - let post_split = post_split_enabled(state.strategy_tag, window_size); + let post_split = + state.post_split_allowed && post_split_enabled(state.strategy_tag, window_size); #[cfg(all(feature = "lsm", feature = "hash"))] let post_split_path = rle_byte_opt.is_none() && !raw_fast_path && post_split; #[cfg(all(feature = "lsm", feature = "hash"))] diff --git a/zstd/src/encoding/levels/fastest/tests.rs b/zstd/src/encoding/levels/fastest/tests.rs index fc66c24cc..510bdb358 100644 --- a/zstd/src/encoding/levels/fastest/tests.rs +++ b/zstd/src/encoding/levels/fastest/tests.rs @@ -83,6 +83,7 @@ fn rle_branch_passes_compressible_hint_to_skip_matching() { pre_split: None, huf_optimal_search: true, literal_compression_disabled: false, + post_split_allowed: true, }; let mut output = Vec::new(); let block = vec![0xAB; 1024]; diff --git a/zstd/src/encoding/streaming_encoder.rs b/zstd/src/encoding/streaming_encoder.rs index 002a0490c..1b4ce4891 100644 --- a/zstd/src/encoding/streaming_encoder.rs +++ b/zstd/src/encoding/streaming_encoder.rs @@ -413,6 +413,7 @@ impl CompressionContext { compression_level, CompressionLevel::Level(n) if n < 0 ), + post_split_allowed: true, }, encoded_scratch: Vec::new(), errored: false, diff --git a/zstd/src/huff0/huff0_encoder.rs b/zstd/src/huff0/huff0_encoder.rs index e660d5b07..d50939262 100644 --- a/zstd/src/huff0/huff0_encoder.rs +++ b/zstd/src/huff0/huff0_encoder.rs @@ -919,8 +919,14 @@ impl HuffmanTable { }; } + /// Encode the table's description into its own buffer, with `fse_table` + /// as the weights' FSE storage, so a writer that follows emits it without + /// encoding it again. #[cfg(feature = "std")] - fn fill_weight_description_from_codes(&mut self, fse_table: &mut fse_encoder::FSETable) { + pub(crate) fn fill_weight_description_from_codes( + &mut self, + fse_table: &mut fse_encoder::FSETable, + ) { // Before deriving the weights, not after: they cost a pass over the // whole alphabet, and a warm cache needs none of it. if self.cached_encoded_weight_description.state != DescriptionState::NotComputed {