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/c-api/src/dict.rs b/c-api/src/dict.rs index d5ede57b7..5b065045e 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,34 @@ 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); - }; + // 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 (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/c-api/src/tests.rs b/c-api/src/tests.rs index 4812b8f7c..fccb6860a 100644 --- a/c-api/src/tests.rs +++ b/c-api/src/tests.rs @@ -534,6 +534,37 @@ 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); + // 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/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/bin/structured-zstd/main.rs b/zstd/src/bin/structured-zstd/main.rs index 09463af60..c2c57e47c 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 @@ -3618,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 @@ -3722,17 +3729,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 +3746,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 +3757,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 +3785,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 +3796,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/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/fastcover.rs b/zstd/src/dictionary/fastcover.rs index c3f3cac1c..b53ff8d45 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,35 +120,17 @@ 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`). +/// 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> { @@ -161,7 +139,6 @@ pub(super) struct FastCoverContext<'s> { freqs: Vec, d: usize, f: u32, - finalize_percent: usize, } impl<'s> FastCoverContext<'s> { @@ -199,22 +176,16 @@ 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, 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 @@ -246,31 +217,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 +285,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 +326,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 +379,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 +420,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>); 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/finalize.rs b/zstd/src/dictionary/finalize.rs new file mode 100644 index 000000000..091e788df --- /dev/null +++ b/zstd/src/dictionary/finalize.rs @@ -0,0 +1,547 @@ +//! 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, + 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, write_ncount_at_log}; +use crate::huff0::huff0_encoder::{HuffmanEncoder, HuffmanTable, WeightScratch}; +use std::{io, 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<'s>( + content: &[u8], + samples: &SampleSet<'s>, + count: usize, + dict_size: usize, + options: FinalizeOptions, +) -> io::Result> { + let mut out = Vec::new(); + let mut analysis = Analysis::default(); + finalize_into( + &mut out, + &mut analysis, + content, + samples, + count, + dict_size, + options, + )?; + Ok(out) +} + +/// 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, 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 +/// from one candidate to the next. +pub(super) fn finalize_into<'s>( + out: &mut Vec, + analysis: &mut Analysis<'s>, + content: &[u8], + samples: &SampleSet<'s>, + count: usize, + dict_size: usize, + options: FinalizeOptions, +) -> io::Result<()> { + check_size(dict_size)?; + 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")); + } + out.clear(); + out.reserve(dict_size); + out.extend_from_slice(&DICT_MAGIC_NUM); + out.extend_from_slice(&dict_id.to_le_bytes()); + analyze_entropy(out, analysis, 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(()) +} + +/// 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, + "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], + /// 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], +} + +/// 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<'s>( + out: &mut Vec, + analysis: &mut Analysis<'s>, + content: &[u8], + samples: &SampleSet<'s>, + 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: [1; 32], + match_lengths: [1; 53], + literal_lengths: [1; 36], + }; + count_samples(&mut counts, analysis, content, samples, count, level)?; + + let huffman = &mut analysis.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); + HuffmanEncoder::new(&literals, &mut writer).write_table(); + huffman.recycle(literals); + 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(); + 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>( + counts: &mut EntropyCounts, + analysis: &mut Analysis<'s>, + content: &[u8], + samples: &SampleSet<'s>, + 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 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(), + }; + 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 + .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); + // 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. + for index in 0..count { + let sample = samples.sample(index); + let block = &sample[..sample.len().min(BLOCK_SIZE_MAX)]; + if block.is_empty() { + continue; + } + compressor.set_source(block); + compressor.set_source_size_hint(average as u64); + compressor.compress_known_into(block.len() as u64, frame); + let fast_offset_codes = compressor.uses_fast_offset_codes(); + 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. + // 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 (block, &counted) in recorder.blocks.iter().zip(compressed.iter()) { + let block_sequences = &recorder.sequences[start..block.sequences_end]; + start = block.sequences_end; + // 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; + } + 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) + } 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; + // 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(); + recorder.sequences.clear(); + } + Ok(()) +} + +/// 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; + }; + 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 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, + /// 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() + } + + 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(); + 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 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 { len } => tail += len, + Sequence::Triple { + literal_len, + offset, + match_len, + } => { + // Lengths and offsets of one block, far below `u32::MAX`. + 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, + sequences_end, + }); + } + + 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..9f92570c2 --- /dev/null +++ b/zstd/src/dictionary/finalize/tests.rs @@ -0,0 +1,353 @@ +use super::*; +use crate::encoding::compress_slice_to_vec; +use std::{format, 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() { + 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 = 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(); + 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()); +} + +/// 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. +#[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; + (0..len) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state as u8 + }) + .collect() +} + +/// 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 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(); + 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 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"); +} + +/// 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 +/// 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(); + // 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: [1; 32], + match_lengths: [1; 53], + literal_lengths: [1; 36], + }; + 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"); + assert_eq!(offset_code_one(1), 1); +} + +/// Each block's kind is read from the frame, one entry per block, and nothing +/// from bytes that are not a frame. +#[test] +fn block_kinds_are_read_from_the_frame() { + let mut kinds = Vec::new(); + let text: String = core::iter::repeat_n("compressible text ", 200).collect(); + 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 +/// 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/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/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/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..c38f1af03 100644 --- a/zstd/src/dictionary/mod.rs +++ b/zstd/src/dictionary/mod.rs @@ -17,61 +17,37 @@ //! 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 finalize; mod legacy; -mod lmc; -mod reservoir; 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::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}, 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`. @@ -94,10 +70,10 @@ 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, - /// Compression level candidates are scored at; zero is the default level. - pub level: i32, } impl Default for CoverOptions { @@ -108,7 +84,6 @@ impl Default for CoverOptions { steps: 4, split_point: 1.0, shrink: None, - level: 0, } } } @@ -121,8 +96,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. @@ -145,39 +120,29 @@ 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, } -/// 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 +174,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 +200,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 +220,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 [`MIN_TRAINED_DICT_SIZE`], 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,331 +249,133 @@ pub fn create_raw_dict_from_slice( output: &mut W, dict_size: usize, ) -> io::Result<()> { - if dict_size == 0 || all.is_empty() { - return Ok(()); - } - - if all.len() < K { - let keep = usize::min(all.len(), dict_size); - output.write_all(&all[all.len() - keep..])?; - return Ok(()); - } - - 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 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)?; - } - Ok(()) -} - -/// 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() + output.write_all(&fastcover_raw_content(all, None, dict_size)?) } -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, +/// 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> { - 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::() + if corpus.len() <= dict_size { + return Ok(corpus.to_vec()); } - - 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); + 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; + } } - if total <= 1 { - return Err(io::Error::new( - io::ErrorKind::InvalidInput, - "insufficient symbol statistics for FSE table", - )); + 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()); } - let table = build_table_from_symbol_counts(&counts, max_log, false); - Ok(serialize_fse_table(&table)) + trained } -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 +/// The content of the dictionary [`optimize_fastcover_dict`] picks at its +/// 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 options = FastCoverOptions::default(); + let space = SearchSpace::optimizing(&options.cover, 0.75)?; + let (content, _) = run_fastcover( + corpus, + sizes, + dict_size, + &options, + FinalizeOptions::default(), + &space, + selection::Keep::Content, + )?; + Ok(content) } -/// 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()`; a size too small carries +/// [`TrainingError::DictionaryTooSmall`]. +/// +/// # 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) + check_finalize_dict_size(dict_size)?; + let set = samples::SampleSet::new(samples, sample_sizes)?; + finalize::finalize(raw_content, &set, set.count(), dict_size, options) } -/// 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(()) +/// 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) } -/// 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 { @@ -733,7 +497,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 @@ -878,10 +642,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(()) @@ -898,7 +662,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. /// @@ -995,15 +759,11 @@ 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 best = selection::Best::new(); + let mut evaluator = + selection::Evaluator::new(&set, split.train, split.test.clone(), dict_size, finalize); + 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. @@ -1033,7 +793,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 } }; @@ -1046,9 +806,13 @@ fn run_cover( if !scored { return Ok((evaluator.finalize(content)?, chosen)); } - best.offer(evaluator.select(content, options.shrink), 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() + best.finish(&mut evaluator) } /// Train a FastCOVER dictionary of at most `dict_size` bytes with the `k` and @@ -1092,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) } @@ -1138,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( @@ -1148,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)?; @@ -1158,10 +945,16 @@ 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(keep); // `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; @@ -1188,7 +981,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 } }; @@ -1197,24 +990,17 @@ 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, - options.cover.level, - finalize, - ) - }); 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)), + }; } - best.offer(evaluator.select(content, options.cover.shrink), chosen); + // Ranked at full size, the winner alone shrunk, as in the COVER search. + best.offer(evaluator.score(content), content, chosen); } - best.finish() + best.finish(&mut evaluator) } /// Train and finalize a dictionary with the reference's original trainer, the @@ -1289,7 +1075,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/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/samples.rs b/zstd/src/dictionary/samples.rs index d6bbc6899..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) } @@ -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..e59a983a9 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::{Analysis, finalize_into}; 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: @@ -44,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> { @@ -52,7 +50,6 @@ impl<'s> Evaluator<'s> { finalize_samples: usize, scoring: Range, capacity: usize, - level: i32, finalize: FinalizeOptions, ) -> Self { Self { @@ -60,42 +57,33 @@ 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(), 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<()> { - 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( + fn finalize_into(&mut self, content: &[u8], out: &mut Vec) -> io::Result<()> { + finalize_into( + out, + &mut self.analysis, content, - self.samples.leading(self.finalize_samples), + self.samples, + self.finalize_samples, self.capacity, self.finalize, - )?; - Ok(()) + ) } /// `content` finalized over `buffer` and priced. @@ -114,7 +102,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))); @@ -129,42 +119,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) } } @@ -177,65 +173,108 @@ 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 +/// 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 +/// it (see [`Evaluator::shrink`]). pub(super) struct Best

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

Best

{ - pub(super) fn new() -> Self { + /// A search handing back what `keep` names of its winner. + pub(super) fn new(keep: Keep) -> Self { Self { found: None, last_error: None, + keep, } } - /// 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. 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.keep.content() { content } else { &[] }; + let keep_dictionary = self.keep.dictionary(); 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, + 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, - }, - params, - )); + }; + 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", - )), + /// 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 keep = self.keep; + 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", + )); + } + }; + 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)), } } } diff --git a/zstd/src/dictionary/tests.rs b/zstd/src/dictionary/tests.rs index 8d831b623..e9b1c8e13 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(); @@ -294,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] @@ -438,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, @@ -471,7 +472,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 +550,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 +562,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"); @@ -678,60 +681,104 @@ fn create_raw_dict_from_source_never_exceeds_requested_size() { ); } -/// 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. +/// 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 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 [ +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(), - 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() - ); + ) + .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 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() - ); + 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)); } #[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); @@ -740,15 +787,54 @@ 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"); } + +/// 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() { + 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..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, }; @@ -773,10 +774,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 +1264,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, @@ -2932,7 +2930,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 +2947,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 +2966,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], @@ -3033,6 +3031,14 @@ pub(in crate::encoding) 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 +3051,7 @@ pub(in crate::encoding) 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], @@ -3075,7 +3081,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/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/blocks/mod.rs b/zstd/src/encoding/blocks/mod.rs index 671ff6fb8..149f9c23b 100644 --- a/zstd/src/encoding/blocks/mod.rs +++ b/zstd/src/encoding/blocks/mod.rs @@ -6,3 +6,9 @@ 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, + 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 08c43e3ed..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,49 @@ 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 + /// 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 + .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 { @@ -3231,6 +3282,19 @@ 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 + } + + /// 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/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(), 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/fse/fse_encoder.rs b/zstd/src/fse/fse_encoder.rs index e8cbbc347..22d42a368 100644 --- a/zstd/src/fse/fse_encoder.rs +++ b/zstd/src/fse/fse_encoder.rs @@ -287,67 +287,86 @@ 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) { - assert!( - writer.index().is_multiple_of(8), - "FSE table headers must start on a byte boundary" + let states = &self.states; + write_ncount( + states.len(), + |i| states[i].probability, + self.acc_log(), + writer, ); - #[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); - } + } +} + +/// 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" + ); + #[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; - 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; - } + let p = prob(prob_idx); + prob_idx += 1; + 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 { + writer.write_bits(value + low_threshold as u32, bits_to_write as usize); + } else { + writer.write_bits(value, bits_to_write as usize); + } + + if p == -1 { + probability_counter += 1; + } else if p > 0 { + probability_counter += p as usize; + } else { + let mut zeros = 0u8; + while prob_idx < len && prob(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(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(zeros, 2); } } + writer.write_bits(0u8, writer.misaligned()); + #[cfg(debug_assertions)] + { + let written_bits = writer.index() - start_idx; + let computed = ncount_header_bits(len, prob, acc_log); + debug_assert_eq!( + written_bits, computed, + "ncount_header_bits() mismatch: written={written_bits}, computed={computed}" + ); + } } #[derive(Debug, Clone, Default)] @@ -425,7 +444,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, @@ -434,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, @@ -499,14 +518,43 @@ 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"))] +/// 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 write_ncount_at_log>>( + counts: &[usize], + table_log: u8, + writer: &mut BitWriter, +) { + 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, + ); + write_ncount(max_symbol + 1, |i| probs[i], table_log, writer); +} + +#[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); 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, @@ -541,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; @@ -558,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 { @@ -593,7 +641,7 @@ fn ncount_header_bits(probs: &[i32], acc_log: u8) -> usize { } /// 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 @@ -630,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 { 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 9ed0897ca..d50939262 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 @@ -918,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 { @@ -994,6 +1001,37 @@ impl HuffmanTable { Self::build_from_weights_reusing(weights, None) } + /// 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_in( + counts: &[usize], + 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 + // 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 /// 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