zrip_core/dict/
mod.rs

1#![forbid(unsafe_code)]
2
3#[cfg(feature = "dict_builder")]
4pub mod cover;
5#[cfg(feature = "dict_builder")]
6pub mod fastcover;
7#[cfg(feature = "dict_builder")]
8pub mod finalize;
9
10#[cfg(feature = "alloc")]
11use alloc::vec::Vec;
12
13use crate::bitstream::reader::BitReader;
14use crate::error::DecompressError;
15use crate::fse::table_builder::{build_decode_table, parse_fse_table_description};
16use crate::fse::{
17    FseDecodeEntry, LL_MAX_ACCURACY_LOG, LL_MAX_SYMBOL, ML_MAX_ACCURACY_LOG, ML_MAX_SYMBOL,
18    OF_MAX_ACCURACY_LOG, OF_MAX_SYMBOL,
19};
20use crate::huffman::HuffmanDecodeEntry;
21use crate::huffman::weights::{build_huffman_decode_table, parse_huffman_weights};
22
23pub const DICT_MAGIC: u32 = 0xEC30_A437;
24
25/// A pre-trained zstd dictionary for improved compression of small data.
26///
27/// Load from raw bytes with [`Dictionary::from_bytes`], or train with
28/// [`train_dict_fastcover`] (requires `dict_builder` feature).
29#[cfg(feature = "alloc")]
30#[derive(Clone)]
31pub struct Dictionary {
32    id: u32,
33    content: Vec<u8>,
34    huf_table: Option<(Vec<HuffmanDecodeEntry>, u8)>,
35    of_table: Option<(Vec<FseDecodeEntry>, u8)>,
36    ml_table: Option<(Vec<FseDecodeEntry>, u8)>,
37    ll_table: Option<(Vec<FseDecodeEntry>, u8)>,
38    rep_offsets: [u32; 3],
39}
40
41#[cfg(feature = "alloc")]
42impl Dictionary {
43    /// Parses a dictionary from its raw byte representation.
44    pub fn from_bytes(data: &[u8]) -> Result<Self, DecompressError> {
45        if data.len() < 8 {
46            return Err(DecompressError::InvalidDictionary);
47        }
48
49        let magic = u32::from_le_bytes([data[0], data[1], data[2], data[3]]);
50        if magic != DICT_MAGIC {
51            return Err(DecompressError::InvalidDictionary);
52        }
53
54        let id = u32::from_le_bytes([data[4], data[5], data[6], data[7]]);
55        if id == 0 {
56            return Err(DecompressError::InvalidDictionary);
57        }
58        let mut pos = 8;
59
60        let huf_table = parse_dict_huffman(&data[pos..])?;
61        pos += huf_table.1;
62        let huf_decode = if huf_table.0.is_some() {
63            huf_table.0
64        } else {
65            None
66        };
67
68        let (of_table, of_consumed) =
69            parse_dict_fse_checked(&data[pos..], OF_MAX_SYMBOL, OF_MAX_ACCURACY_LOG)?;
70        pos += of_consumed;
71
72        let (ml_table, ml_consumed) =
73            parse_dict_fse_checked(&data[pos..], ML_MAX_SYMBOL, ML_MAX_ACCURACY_LOG)?;
74        pos += ml_consumed;
75
76        let (ll_table, ll_consumed) =
77            parse_dict_fse_checked(&data[pos..], LL_MAX_SYMBOL, LL_MAX_ACCURACY_LOG)?;
78        pos += ll_consumed;
79
80        if pos + 12 > data.len() {
81            return Err(DecompressError::InvalidDictionary);
82        }
83
84        let rep1 = u32::from_le_bytes([data[pos], data[pos + 1], data[pos + 2], data[pos + 3]]);
85        let rep2 = u32::from_le_bytes([data[pos + 4], data[pos + 5], data[pos + 6], data[pos + 7]]);
86        let rep3 =
87            u32::from_le_bytes([data[pos + 8], data[pos + 9], data[pos + 10], data[pos + 11]]);
88        pos += 12;
89
90        if rep1 == 0 || rep2 == 0 || rep3 == 0 {
91            return Err(DecompressError::InvalidDictionary);
92        }
93
94        let content = data[pos..].to_vec();
95
96        Ok(Self {
97            id,
98            content,
99            huf_table: huf_decode,
100            of_table,
101            ml_table,
102            ll_table,
103            rep_offsets: [rep1, rep2, rep3],
104        })
105    }
106
107    /// Returns the dictionary ID embedded in the header.
108    pub fn id(&self) -> u32 {
109        self.id
110    }
111
112    /// Returns the raw content segment used as match-finding history prefix.
113    pub fn content(&self) -> &[u8] {
114        &self.content
115    }
116
117    /// Returns the three initial repeat offsets stored in the dictionary.
118    pub fn rep_offsets(&self) -> &[u32; 3] {
119        &self.rep_offsets
120    }
121
122    /// Returns the Huffman decode table and its log2 size, if present.
123    pub fn huf_table(&self) -> Option<(&[HuffmanDecodeEntry], u8)> {
124        self.huf_table.as_ref().map(|(t, l)| (t.as_slice(), *l))
125    }
126
127    /// Returns the offset-code FSE decode table and accuracy log, if present.
128    pub fn of_table(&self) -> Option<(&[FseDecodeEntry], u8)> {
129        self.of_table.as_ref().map(|(t, l)| (t.as_slice(), *l))
130    }
131
132    /// Returns the match-length FSE decode table and accuracy log, if present.
133    pub fn ml_table(&self) -> Option<(&[FseDecodeEntry], u8)> {
134        self.ml_table.as_ref().map(|(t, l)| (t.as_slice(), *l))
135    }
136
137    /// Returns the literal-length FSE decode table and accuracy log, if present.
138    pub fn ll_table(&self) -> Option<(&[FseDecodeEntry], u8)> {
139        self.ll_table.as_ref().map(|(t, l)| (t.as_slice(), *l))
140    }
141}
142
143#[cfg(feature = "alloc")]
144#[allow(clippy::type_complexity)]
145fn parse_dict_huffman(
146    data: &[u8],
147) -> Result<(Option<(Vec<HuffmanDecodeEntry>, u8)>, usize), DecompressError> {
148    if data.is_empty() {
149        return Err(DecompressError::InvalidDictionary);
150    }
151
152    let (weights, consumed) = parse_huffman_weights(data)?;
153    if weights.is_empty() {
154        return Ok((None, consumed));
155    }
156    let (table, table_log) = build_huffman_decode_table(&weights)?;
157    Ok((Some((table, table_log)), consumed))
158}
159
160#[cfg(feature = "alloc")]
161#[allow(clippy::type_complexity)]
162fn parse_dict_fse_checked(
163    data: &[u8],
164    max_symbol: u8,
165    max_accuracy_log: u8,
166) -> Result<(Option<(Vec<FseDecodeEntry>, u8)>, usize), DecompressError> {
167    if data.is_empty() {
168        return Err(DecompressError::InvalidDictionary);
169    }
170
171    let mut reader = BitReader::new(data);
172    let (distribution, accuracy_log) = parse_fse_table_description(&mut reader, max_symbol)?;
173    if accuracy_log > max_accuracy_log {
174        return Err(DecompressError::InvalidDictionary);
175    }
176    let consumed = reader.bytes_consumed();
177    let table = build_decode_table(&distribution, accuracy_log)
178        .map_err(|_| DecompressError::InvalidDictionary)?;
179    Ok((Some((table, accuracy_log)), consumed))
180}
181
182#[cfg(all(test, feature = "alloc"))]
183mod tests {
184    use super::*;
185    use crate::fse::table_builder::serialize_fse_table_description;
186
187    fn build_dict_with_of_accuracy(accuracy_log: u8) -> Vec<u8> {
188        let mut dict = Vec::new();
189        dict.extend_from_slice(&DICT_MAGIC.to_le_bytes());
190        dict.extend_from_slice(&1u32.to_le_bytes());
191        // Huffman table: header byte 0 triggers FSE path with compressed_size=0 → error.
192        // Use header byte 128 (direct mode, 1 symbol with weight in upper nibble).
193        dict.push(128);
194        dict.push(0x10); // symbol 0, weight = 1
195
196        // OF FSE table with the specified accuracy_log.
197        // Distribution: symbol 0 gets all probability.
198        let table_size = 1i16 << accuracy_log;
199        let mut dist = vec![0i16; 32];
200        dist[0] = table_size;
201        let of_bytes = serialize_fse_table_description(&dist, accuracy_log);
202        dict.extend_from_slice(&of_bytes);
203
204        // ML FSE table: accuracy_log = 6, symbol 0 gets all probability.
205        let mut ml_dist = vec![0i16; 53];
206        ml_dist[0] = 1 << 6;
207        let ml_bytes = serialize_fse_table_description(&ml_dist, 6);
208        dict.extend_from_slice(&ml_bytes);
209
210        // LL FSE table: accuracy_log = 6, symbol 0 gets all probability.
211        let mut ll_dist = vec![0i16; 36];
212        ll_dist[0] = 1 << 6;
213        let ll_bytes = serialize_fse_table_description(&ll_dist, 6);
214        dict.extend_from_slice(&ll_bytes);
215
216        // Rep offsets: [1, 4, 8]
217        dict.extend_from_slice(&1u32.to_le_bytes());
218        dict.extend_from_slice(&4u32.to_le_bytes());
219        dict.extend_from_slice(&8u32.to_le_bytes());
220
221        // Content (at least 1 byte)
222        dict.extend_from_slice(b"content");
223        dict
224    }
225
226    #[test]
227    fn dict_rejects_oversized_of_accuracy() {
228        let dict_bytes = build_dict_with_of_accuracy(10);
229        let result = Dictionary::from_bytes(&dict_bytes);
230        assert!(
231            result.is_err(),
232            "accuracy_log=10 should be rejected for OF (max=8)"
233        );
234    }
235
236    #[test]
237    fn dict_accepts_valid_of_accuracy() {
238        let dict_bytes = build_dict_with_of_accuracy(8);
239        let result = Dictionary::from_bytes(&dict_bytes);
240        assert!(result.is_ok(), "accuracy_log=8 should be accepted for OF");
241    }
242}
243
244#[cfg(feature = "dict_builder")]
245/// Trains a dictionary from sample data using the FastCOVER algorithm.
246pub fn train_dict_fastcover(
247    samples: &[&[u8]],
248    dict_size: usize,
249    params: fastcover::FastCoverParams,
250) -> Dictionary {
251    let content = fastcover::select_segments(samples, dict_size, &params);
252    let dict_bytes = finalize::finalize_dictionary(&content, samples, dict_size);
253    Dictionary::from_bytes(&dict_bytes).expect("finalized dictionary must be valid")
254}
zrip_core/dict/mod.rs

zrip_core/dict/
mod.rs
