feat(dict)!: entropy tables measured from the samples

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.

c-api/src/dict.rs

@@ -163,7 +163,6 @@ impl ZDICT_cover_params_t {
             self.splitPoint,
             self.shrinkDict,
             self.shrinkDictMaxRegression,
-            self.zParams.compressionLevel,
         )
     }
 
@@ -181,21 +180,22 @@ fn cover_options(
     split_point: f64,
     shrink: c_uint,
     max_regression: c_uint,
-    level: c_int,
 ) -> CoverOptions {
     CoverOptions {
         k,
         d,
         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(&parameters),
+                )
+            },
+        )
     }
-    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

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(

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(

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<structured_zstd::dictionary::CoverOptions> {
-    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<structured_zstd::dictionary::CoverOptions> {
+    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<structured_zstd::dictionary::CoverOptions> {
     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<structured_zstd::dictionary::FastCoverOptions> {
     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 {

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(&params, 3, "--train-cover").unwrap();
+        let options = cover_tuning(&params, "--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(&params, 5).unwrap();
+    let fixed = fastcover_options(&params).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();