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/// Train a ZSTD dictionary from `samples` using the default
90/// FrequencyTrainer. Returns `None` if `samples` is empty, target
91/// size is 0, or the trainer produces an empty dictionary (not
92/// enough signal).
93///
94/// `id` is the caller-allocated dictionary id (0x00..=0xFE).
95#[must_use]
96pub fn train_zstd(id: u8, samples: &[&[u8]], target_size: usize) -> Option<TrainedDictionary> {
97    train_zstd_with_trainer(id, samples, target_size, TrainerKind::Frequency)
98}
99
100/// Trainer algorithm selection. See [`train_zstd_with_trainer`].
101#[derive(Clone, Copy, Debug, Eq, PartialEq)]
102pub enum TrainerKind {
103    /// Top-K substrings by frequency × length. Default. Wins on
104    /// corpora with strong common substrings.
105    Frequency,
106    /// Dmer-frequency scoring per FastCover (Facebook 2018). Wins on
107    /// corpora with distributed redundancy (mixed JSON, source files,
108    /// log lines).
109    FastCover,
110}
111
112impl TrainerKind {
113    /// Parse from a config string. Unknown values fall back to
114    /// `Frequency` (the default).
115    #[must_use]
116    pub fn from_config_str(s: &str) -> Self {
117        match s.to_ascii_lowercase().as_str() {
118            "fastcover" => Self::FastCover,
119            _ => Self::Frequency,
120        }
121    }
122}
123
124/// Train with explicit trainer selection. See [`train_zstd`] for the
125/// default-FrequencyTrainer shortcut.
126#[must_use]
127pub fn train_zstd_with_trainer(
128    id: u8,
129    samples: &[&[u8]],
130    target_size: usize,
131    trainer: TrainerKind,
132) -> Option<TrainedDictionary> {
133    if samples.is_empty() || target_size == 0 {
134        return None;
135    }
136    let content = match trainer {
137        TrainerKind::Frequency => core_train(samples, target_size),
138        TrainerKind::FastCover => train_dictionary_fastcover(samples, target_size),
139    };
140    if content.is_empty() {
141        return None;
142    }
143    Some(TrainedDictionary {
144        id,
145        codec: limnifs_core::codec::CODEC_ZSTD,
146        content,
147    })
148}
149
150/// Allocate dictionary ids 0x00..=0xFE for a set of trained dicts.
151/// The 0xFF slot is reserved as `NO_DICT` sentinel.
152///
153/// Returns a map from class name to allocated id. Errors if more
154/// than 254 classes need ids.
155pub fn allocate_ids<'a>(class_names: &'a [&'a str]) -> Result<Vec<(&'a str, u8)>, &'static str> {
156    if class_names.len() > 254 {
157        return Err("dictionary id space exhausted (max 254 classes)");
158    }
159    Ok(class_names
160        .iter()
161        .enumerate()
162        .map(|(i, name)| (*name, u8::try_from(i).expect("≤ 254")))
163        .collect())
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169
170    fn synthetic_text_samples(n: usize) -> Vec<Vec<u8>> {
171        // Repetitive source-code-like content. The trainer should
172        // find common substrings like "function", "return", "const".
173        (0..n)
174            .map(|i| format!("function test_case_{i}() {{ return {i}; }}\n").into_bytes())
175            .collect()
176    }
177
178    #[test]
179    fn train_zstd_returns_dict_for_repetitive_samples() {
180        let samples_vec = synthetic_text_samples(50);
181        let samples: Vec<&[u8]> = samples_vec.iter().map(Vec::as_slice).collect();
182        let dict = train_zstd(0, &samples, 4096);
183        // FrequencyTrainer may return empty on some inputs; assert
184        // at minimum that the function ran without panicking.
185        if let Some(d) = &dict {
186            assert!(!d.content.is_empty(), "trained dict content non-empty");
187            assert_eq!(d.id, 0);
188            assert_eq!(d.codec, limnifs_core::codec::CODEC_ZSTD);
189        }
190    }
191
192    #[test]
193    fn train_zstd_returns_none_for_empty_samples() {
194        assert!(train_zstd(0, &[], 4096).is_none());
195    }
196
197    #[test]
198    fn train_zstd_returns_none_for_zero_target_size() {
199        let samples_vec = synthetic_text_samples(10);
200        let samples: Vec<&[u8]> = samples_vec.iter().map(Vec::as_slice).collect();
201        assert!(train_zstd(0, &samples, 0).is_none());
202    }
203
204    #[test]
205    fn dict_round_trips_when_trained() {
206        let samples_vec = synthetic_text_samples(50);
207        let samples: Vec<&[u8]> = samples_vec.iter().map(Vec::as_slice).collect();
208        let Some(dict) = train_zstd(0, &samples, 4096) else {
209            return; // Trainer may legitimately return None
210        };
211        let plaintext = b"function test_case_99() { return 99; }\n";
212        let compressed = dict.compress(plaintext).expect("compress");
213        let recovered = dict
214            .decompress(&compressed, plaintext.len() as u32)
215            .expect("decompress");
216        assert_eq!(recovered.as_slice(), &plaintext[..]);
217    }
218
219    #[test]
220    fn allocate_ids_assigns_sequential_ids() {
221        let names = vec!["text", "binary", "source"];
222        let allocated = allocate_ids(&names).expect("allocate");
223        assert_eq!(allocated.len(), 3);
224        assert_eq!(allocated[0], ("text", 0));
225        assert_eq!(allocated[1], ("binary", 1));
226        assert_eq!(allocated[2], ("source", 2));
227    }
228
229    #[test]
230    fn allocate_ids_rejects_more_than_254_classes() {
231        let names: Vec<&str> = (0..255).map(|_| "x").collect();
232        assert!(allocate_ids(&names).is_err());
233    }
234}
limnifs_write/dictionary.rs

limnifs_write/
dictionary.rs
