structured_zstd/dictionary/samples.rs
1//! Training samples: every sample back to back in one buffer, and where each
2//! one ends.
3
4use core::ops::Range;
5use std::{io, string::String, vec::Vec};
6
7/// Fewest samples the segment trainers build from (upstream zstd `cover.c`,
8/// `COVER_ctx_init`): a dmer's frequency is the number of samples holding it,
9/// which says nothing when there are only a handful.
10pub(super) const MIN_TRAIN_SAMPLES: usize = 5;
11
12/// Samples a trainer reads: sample `i` is `data[offsets[i]..offsets[i + 1]]`.
13pub(super) struct SampleSet<'a> {
14 data: &'a [u8],
15 offsets: Vec<usize>,
16}
17
18impl<'a> SampleSet<'a> {
19 /// `data` cut into consecutive samples of `sizes` bytes, which must cover it
20 /// exactly.
21 pub(super) fn new(data: &'a [u8], sizes: &[usize]) -> io::Result<Self> {
22 let mut offsets = Vec::with_capacity(sizes.len() + 1);
23 offsets.push(0);
24 let mut end = 0usize;
25 for &size in sizes {
26 end = end
27 .checked_add(size)
28 .filter(|&end| end <= data.len())
29 .ok_or_else(|| {
30 refuse(
31 TrainingError::Samples,
32 "the sample sizes add up to more than the samples",
33 )
34 })?;
35 offsets.push(end);
36 }
37 if end != data.len() {
38 return Err(refuse(
39 TrainingError::Samples,
40 "the sample sizes do not add up to the samples",
41 ));
42 }
43 Ok(Self { data, offsets })
44 }
45
46 pub(super) fn count(&self) -> usize {
47 self.offsets.len() - 1
48 }
49
50 /// Start of every sample, then the end of the last.
51 pub(super) fn offsets(&self) -> &[usize] {
52 &self.offsets
53 }
54
55 pub(super) fn sample(&self, index: usize) -> &'a [u8] {
56 &self.data[self.offsets[index]..self.offsets[index + 1]]
57 }
58
59 /// The first `count` samples, back to back.
60 pub(super) fn leading(&self, count: usize) -> &'a [u8] {
61 &self.data[..self.offsets[count]]
62 }
63
64 /// Refuse the first `train` samples when none of them is `span` bytes
65 /// long: a trainer counts a dmer only inside one sample, so none would be
66 /// counted however long the samples run together. Known from the sizes,
67 /// so it is refused before the dmers are indexed.
68 pub(super) fn check_holds_dmer(&self, train: usize, span: usize) -> io::Result<()> {
69 if self.offsets[..=train]
70 .windows(2)
71 .any(|bounds| bounds[1] - bounds[0] >= span)
72 {
73 return Ok(());
74 }
75 Err(refuse(
76 TrainingError::Samples,
77 &std::format!("no training sample is {span} bytes long, so none holds a dmer to count"),
78 ))
79 }
80
81 /// Divide the samples into those a dictionary is built from and those it
82 /// is scored on (upstream zstd `COVER_ctx_init`): below 1 the leading
83 /// `split_point` share builds and the rest scores; at 1 every sample does
84 /// both.
85 pub(super) fn split(&self, split_point: f64) -> io::Result<Split> {
86 split_count(self.count(), split_point)
87 }
88}
89
90/// [`SampleSet::split`] of `count` samples, which needs only their number.
91pub(super) fn split_count(count: usize, split_point: f64) -> io::Result<Split> {
92 let split = if split_point < 1.0 {
93 let train = (count as f64 * split_point) as usize;
94 Split {
95 train,
96 test: train..count,
97 }
98 } else {
99 Split {
100 train: count,
101 test: 0..count,
102 }
103 };
104 if split.train < MIN_TRAIN_SAMPLES {
105 return Err(refuse(
106 TrainingError::Samples,
107 &std::format!(
108 "{} training sample(s) is too few; at least {MIN_TRAIN_SAMPLES} are needed",
109 split.train
110 ),
111 ));
112 }
113 if split.test.is_empty() {
114 return Err(refuse(
115 TrainingError::Samples,
116 "the split leaves no sample to score the dictionary on",
117 ));
118 }
119 Ok(split)
120}
121
122/// How the samples divide between building a dictionary and scoring it.
123#[derive(Clone, Debug, PartialEq, Eq)]
124pub(super) struct Split {
125 /// The leading samples a dictionary is built from.
126 pub(super) train: usize,
127 /// The samples a dictionary is scored on.
128 pub(super) test: Range<usize>,
129}
130
131/// Why a segment trainer or the finalizer refused its input. It rides inside
132/// the `InvalidInput` error returned; [`TrainingError::of`] reads it back.
133#[derive(Clone, Copy, Debug, PartialEq, Eq)]
134#[non_exhaustive]
135pub enum TrainingError {
136 /// A training parameter lies outside its range.
137 Parameter,
138 /// The samples cannot be trained on: none, too few, or too few or too many
139 /// bytes.
140 Samples,
141 /// The dictionary asked for is smaller than
142 /// [`MIN_TRAINED_DICT_SIZE`](super::MIN_TRAINED_DICT_SIZE).
143 DictionaryTooSmall,
144}
145
146impl TrainingError {
147 /// The cause carried by an error a trainer returned, or `None` for one it
148 /// did not raise itself (a failed read, write or allocation).
149 ///
150 /// # Examples
151 ///
152 /// ```
153 /// use structured_zstd::dictionary::{
154 /// CoverOptions, FinalizeOptions, TrainingError, train_cover_dict,
155 /// };
156 ///
157 /// let options = CoverOptions { k: 64, d: 8, ..CoverOptions::default() };
158 /// let err = train_cover_dict(&[0; 64], &[64], 100, &options, FinalizeOptions::default())
159 /// .unwrap_err();
160 /// assert_eq!(TrainingError::of(&err), Some(TrainingError::DictionaryTooSmall));
161 /// ```
162 pub fn of(error: &io::Error) -> Option<Self> {
163 error
164 .get_ref()?
165 .downcast_ref::<Refusal>()
166 .map(|refusal| refusal.cause)
167 }
168}
169
170/// A trainer's refusal: its cause, and the reason shown to a person.
171#[derive(Debug)]
172struct Refusal {
173 cause: TrainingError,
174 reason: String,
175}
176
177impl core::fmt::Display for Refusal {
178 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
179 f.write_str(&self.reason)
180 }
181}
182
183impl std::error::Error for Refusal {}
184
185/// An `InvalidInput` error carrying `cause`.
186pub(super) fn refuse(cause: TrainingError, reason: &str) -> io::Error {
187 io::Error::new(
188 io::ErrorKind::InvalidInput,
189 Refusal {
190 cause,
191 reason: reason.into(),
192 },
193 )
194}