limnifs_write/
dictionary.rs

1//! ZSTD dictionary training — writer-side API.
2//!
3//! Wraps `limnifs_core::codec::zstd_dict` (which wraps
4//! `omnizip_zstd::train_dictionary`) at the writer layer. The
5//! codec-layer API takes raw samples; this layer adds:
6//!
7//! - Per-content-class sample collection (text, binary, etc.).
8//! - Configurable trainer (FrequencyTrainer today; FastCover via
9//!   `omnizip_zstd::FastCoverTrainer` when needed).
10//! - Dictionary id allocation (0x00..=0xFE; 0xFF is `NO_DICT`).
11//! - Integration with `WriteConfig::dictionaries`.
12//!
13//! ## Pipeline integration (planned)
14//!
15//! Today this module exposes the trainer; the writer pipeline does
16//! NOT yet call it. The plan, filed in
17//! `TODO.impl/04-writer-pipeline/04-zstd-dictionary-training.md`:
18//!
19//! 1. Walk + parallel chunk + compress (existing pipeline).
20//! 2. Collect unique plaintext drops per content class during
21//!    `merge_chunked_file`.
22//! 3. After parallel phase, train one dict per class with ≥
23//!    `min_class_size` drops.
24//! 4. Re-compress eligible drops with `compress_with_dict`; keep
25//!    the smaller of (original, dict-compressed).
26//! 5. Drop records carry `dict_id`; manifest emits
27//!    `dictionary_section`.
28//!
29//! This module is the public API for steps 3–4. The pipeline glue
30//! lands in a follow-up PR.
31
32use limnifs_core::codec::zstd_dict::{
33    compress_with_dict, decompress_with_dict, train_dictionary as core_train,
34    train_dictionary_fastcover,
35};
36
37/// Default target dictionary size (64 KiB). Matches `DictionaryConfig::max_dict_size` default.
38pub const DEFAULT_TARGET_SIZE: usize = 65_536;
39
40/// Minimum sample count before training is worthwhile. Below this,
41/// the trainer returns empty (not enough signal). Matches
42/// `DictionaryConfig::min_class_size` default.
43pub const DEFAULT_MIN_SAMPLES: usize = 100;
44
45/// Trained dictionary ready for use with `compress_with_dict`.
46#[derive(Clone, Debug)]
47pub struct TrainedDictionary {
48    /// Allocated id (0x00..=0xFE). Stored in the manifest's
49    /// `dictionary_section` and referenced by `DropRecord::dict_id`.
50    pub id: u8,
51    /// Codec id this dictionary targets (today always CODEC_ZSTD).
52    pub codec: u8,
53    /// Raw dictionary bytes (omnizip's serialized form).
54    pub content: Vec<u8>,
55}
56
57impl TrainedDictionary {
58    /// Compress `plaintext` with this dictionary.
59    ///
60    /// # Errors
61    /// Returns [`crate::WriteError`] on compression failure.
62    pub fn compress(&self, plaintext: &[u8]) -> Result<Vec<u8>, crate::WriteError> {
63        compress_with_dict(plaintext, &self.content).map_err(|e| {
64            crate::WriteError::Io(std::io::Error::other(format!(
65                "dict compress (id {}): {e}",
66                self.id
67            )))
68        })
69    }
70
71    /// Decompress `compressed` with this dictionary.
72    ///
73    /// # Errors
74    /// Returns [`crate::WriteError`] on decompression failure.
75    pub fn decompress(
76        &self,
77        compressed: &[u8],
78        expected_len: u32,
79    ) -> Result<Vec<u8>, crate::WriteError> {
80        decompress_with_dict(compressed, expected_len, &self.content).map_err(|e| {
81            crate::WriteError::Io(std::io::Error::other(format!(
82                "dict decompress (id {}): {e}",
83                self.id
84            )))
85        })
86    }
87}
88
89/// Adopt the ZSTD dictionaries carried in a base image's
90/// `dictionary_section` for reuse by a layer write. Non-ZSTD
91/// entries and unknown class ids are skipped (forward
92/// compatibility: ids 0 = text, 1 = binary today).
93#[must_use]
94pub fn adopt_from_section(
95    section: limnifs_core::dictionary_section::DictionarySection,
96) -> Vec<TrainedDictionary> {
97    section
98        .dicts
99        .into_iter()
100        .filter(|d| d.codec_id == limnifs_core::codec::CODEC_ZSTD && matches!(d.class_id, 0 | 1))
101        .map(|d| TrainedDictionary {
102            id: d.class_id,
103            codec: d.codec_id,
104            content: d.data,
105        })
106        .collect()
107}
108
109/// Train a ZSTD dictionary from `samples` using the default
110/// FrequencyTrainer. Returns `None` if `samples` is empty, target
111/// size is 0, or the trainer produces an empty dictionary (not
112/// enough signal).
113///
114/// `id` is the caller-allocated dictionary id (0x00..=0xFE).
115#[must_use]
116pub fn train_zstd(id: u8, samples: &[&[u8]], target_size: usize) -> Option<TrainedDictionary> {
117    train_zstd_with_trainer(id, samples, target_size, TrainerKind::Frequency)
118}
119
120/// Trainer algorithm selection. See [`train_zstd_with_trainer`].
121#[derive(Clone, Copy, Debug, Eq, PartialEq)]
122pub enum TrainerKind {
123    /// Top-K substrings by frequency × length. Default. Wins on
124    /// corpora with strong common substrings.
125    Frequency,
126    /// Dmer-frequency scoring per FastCover (Facebook 2018). Wins on
127    /// corpora with distributed redundancy (mixed JSON, source files,
128    /// log lines).
129    FastCover,
130}
131
132impl TrainerKind {
133    /// Parse from a config string. Unknown values fall back to
134    /// `Frequency` (the default).
135    #[must_use]
136    pub fn from_config_str(s: &str) -> Self {
137        match s.to_ascii_lowercase().as_str() {
138            "fastcover" => Self::FastCover,
139            _ => Self::Frequency,
140        }
141    }
142}
143
144/// Train with explicit trainer selection. See [`train_zstd`] for the
145/// default-FrequencyTrainer shortcut.
146#[must_use]
147pub fn train_zstd_with_trainer(
148    id: u8,
149    samples: &[&[u8]],
150    target_size: usize,
151    trainer: TrainerKind,
152) -> Option<TrainedDictionary> {
153    if samples.is_empty() || target_size == 0 {
154        return None;
155    }
156    let content = match trainer {
157        TrainerKind::Frequency => core_train(samples, target_size),
158        TrainerKind::FastCover => train_dictionary_fastcover(samples, target_size),
159    };
160    if content.is_empty() {
161        return None;
162    }
163    Some(TrainedDictionary {
164        id,
165        codec: limnifs_core::codec::CODEC_ZSTD,
166        content,
167    })
168}
169
170/// Allocate dictionary ids 0x00..=0xFE for a set of trained dicts.
171/// The 0xFF slot is reserved as `NO_DICT` sentinel.
172///
173/// Returns a map from class name to allocated id. Errors if more
174/// than 254 classes need ids.
175pub fn allocate_ids<'a>(class_names: &'a [&'a str]) -> Result<Vec<(&'a str, u8)>, &'static str> {
176    if class_names.len() > 254 {
177        return Err("dictionary id space exhausted (max 254 classes)");
178    }
179    Ok(class_names
180        .iter()
181        .enumerate()
182        .map(|(i, name)| (*name, u8::try_from(i).expect("≤ 254")))
183        .collect())
184}
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189
190    fn synthetic_text_samples(n: usize) -> Vec<Vec<u8>> {
191        // Repetitive source-code-like content. The trainer should
192        // find common substrings like "function", "return", "const".
193        (0..n)
194            .map(|i| format!("function test_case_{i}() {{ return {i}; }}\n").into_bytes())
195            .collect()
196    }
197
198    #[test]
199    fn train_zstd_returns_dict_for_repetitive_samples() {
200        let samples_vec = synthetic_text_samples(50);
201        let samples: Vec<&[u8]> = samples_vec.iter().map(Vec::as_slice).collect();
202        let dict = train_zstd(0, &samples, 4096);
203        // FrequencyTrainer may return empty on some inputs; assert
204        // at minimum that the function ran without panicking.
205        if let Some(d) = &dict {
206            assert!(!d.content.is_empty(), "trained dict content non-empty");
207            assert_eq!(d.id, 0);
208            assert_eq!(d.codec, limnifs_core::codec::CODEC_ZSTD);
209        }
210    }
211
212    #[test]
213    fn train_zstd_returns_none_for_empty_samples() {
214        assert!(train_zstd(0, &[], 4096).is_none());
215    }
216
217    #[test]
218    fn train_zstd_returns_none_for_zero_target_size() {
219        let samples_vec = synthetic_text_samples(10);
220        let samples: Vec<&[u8]> = samples_vec.iter().map(Vec::as_slice).collect();
221        assert!(train_zstd(0, &samples, 0).is_none());
222    }
223
224    #[test]
225    fn dict_round_trips_when_trained() {
226        let samples_vec = synthetic_text_samples(50);
227        let samples: Vec<&[u8]> = samples_vec.iter().map(Vec::as_slice).collect();
228        let Some(dict) = train_zstd(0, &samples, 4096) else {
229            return; // Trainer may legitimately return None
230        };
231        let plaintext = b"function test_case_99() { return 99; }\n";
232        let compressed = dict.compress(plaintext).expect("compress");
233        let recovered = dict
234            .decompress(&compressed, plaintext.len() as u32)
235            .expect("decompress");
236        assert_eq!(recovered.as_slice(), &plaintext[..]);
237    }
238
239    #[test]
240    fn allocate_ids_assigns_sequential_ids() {
241        let names = vec!["text", "binary", "source"];
242        let allocated = allocate_ids(&names).expect("allocate");
243        assert_eq!(allocated.len(), 3);
244        assert_eq!(allocated[0], ("text", 0));
245        assert_eq!(allocated[1], ("binary", 1));
246        assert_eq!(allocated[2], ("source", 2));
247    }
248
249    #[test]
250    fn allocate_ids_rejects_more_than_254_classes() {
251        let names: Vec<&str> = (0..255).map(|_| "x").collect();
252        assert!(allocate_ids(&names).is_err());
253    }
254    #[test]
255    fn adopt_from_section_maps_ids_and_filters_codecs() {
256        let section = limnifs_core::dictionary_section::DictionarySection {
257            version: limnifs_core::dictionary_section::DICTIONARY_SECTION_VERSION,
258            dicts: vec![
259                limnifs_core::dictionary_section::Dictionary {
260                    codec_id: limnifs_core::codec::CODEC_ZSTD,
261                    class_id: 0,
262                    data: b"text-dict".to_vec(),
263                },
264                limnifs_core::dictionary_section::Dictionary {
265                    codec_id: limnifs_core::codec::CODEC_ZSTD,
266                    class_id: 1,
267                    data: b"binary-dict".to_vec(),
268                },
269                limnifs_core::dictionary_section::Dictionary {
270                    codec_id: limnifs_core::codec::CODEC_LZ4,
271                    class_id: 0,
272                    data: b"wrong-codec".to_vec(),
273                },
274                limnifs_core::dictionary_section::Dictionary {
275                    codec_id: limnifs_core::codec::CODEC_ZSTD,
276                    class_id: 7,
277                    data: b"unknown-class".to_vec(),
278                },
279            ],
280        };
281        let adopted = adopt_from_section(section);
282        assert_eq!(
283            adopted.len(),
284            2,
285            "non-zstd and unknown-class entries dropped"
286        );
287        assert_eq!(adopted[0].id, 0);
288        assert_eq!(adopted[0].content, b"text-dict");
289        assert_eq!(adopted[1].id, 1);
290        assert_eq!(adopted[1].content, b"binary-dict");
291    }
292}
limnifs_write/dictionary.rs

limnifs_write/
dictionary.rs
