zstd_rs/dict/
mod.rs

1//! Dictionaries: zstd-format (magic, ID, entropy tables, repeat offsets, content,
2//! as produced by `zstd --train` / ZDICT) and raw-content dictionaries.
3use crate::error::{Error, Result};
4use crate::fse::{ncount, Norm};
5use crate::huf::{weights, HufWeights};
6use alloc::vec::Vec;
7
8pub const DICT_MAGIC: u32 = 0xEC30_A437;
9
10/// Entropy tables carried by a zstd-format dictionary.
11#[derive(Clone, Debug)]
12pub struct DictTables {
13    pub huf: HufWeights,
14    pub of: Norm,
15    pub ml: Norm,
16    pub ll: Norm,
17}
18
19/// A parsed dictionary, usable by both the decoder and (via `EncoderDictionary`) the encoder.
20#[derive(Clone, Debug)]
21pub struct Dictionary {
22    pub(crate) id: u32,
23    pub(crate) content: Vec<u8>,
24    pub(crate) tables: Option<DictTables>,
25    pub(crate) reps: [u32; 3],
26}
27
28/// How to interpret dictionary bytes. There is deliberately no auto-detection default.
29#[derive(Clone, Copy, Debug, PartialEq, Eq)]
30pub enum DictionaryFormat {
31    /// Must start with the zstd dictionary magic; entropy tables and repeat offsets are loaded.
32    Zstd,
33    /// The bytes are history only; frames reference it with dictionary ID 0 unless `id` is given.
34    Raw {
35        /// ID to write in frame headers (0 = none).
36        id: u32,
37    },
38    /// Zstd format if the magic is present, else raw with ID 0 (the reference library's behaviour).
39    Detect,
40}
41
42impl Dictionary {
43    /// Parses `bytes` in the given format.
44    pub fn new(bytes: &[u8], format: DictionaryFormat) -> Result<Dictionary> {
45        let has_magic = bytes.len() >= 8 && u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]) == DICT_MAGIC;
46        match format {
47            DictionaryFormat::Zstd if has_magic => Self::parse_zstd(bytes),
48            DictionaryFormat::Zstd => Err(Error::BadDictionary("missing zstd dictionary magic")),
49            DictionaryFormat::Raw { id } => Ok(Self::raw(bytes, id)),
50            DictionaryFormat::Detect if has_magic => Self::parse_zstd(bytes),
51            DictionaryFormat::Detect => Ok(Self::raw(bytes, 0)),
52        }
53    }
54
55    fn raw(bytes: &[u8], id: u32) -> Dictionary {
56        Dictionary { id, content: bytes.to_vec(), tables: None, reps: [1, 4, 8] }
57    }
58
59    fn parse_zstd(bytes: &[u8]) -> Result<Dictionary> {
60        let id = u32::from_le_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]);
61        let mut at = 8;
62        let bad = |e: Error| match e {
63            Error::Corrupt(why) => Error::BadDictionary(why),
64            Error::Truncated => Error::BadDictionary("truncated entropy tables"),
65            other => other,
66        };
67        let (huf, used) = weights::read(&bytes[at..]).map_err(bad)?;
68        at += used;
69        let (of, used) = ncount::read(&bytes[at..], 32, 8).map_err(bad)?;
70        at += used;
71        let (ml, used) = ncount::read(&bytes[at..], 53, 9).map_err(bad)?;
72        at += used;
73        let (ll, used) = ncount::read(&bytes[at..], 36, 9).map_err(bad)?;
74        at += used;
75        let rep_bytes = bytes.get(at..at + 12).ok_or(Error::BadDictionary("truncated repeat offsets"))?;
76        let mut reps = [0u32; 3];
77        for (i, r) in reps.iter_mut().enumerate() {
78            let b = &rep_bytes[4 * i..4 * i + 4];
79            *r = u32::from_le_bytes([b[0], b[1], b[2], b[3]]);
80        }
81        let content = bytes[at + 12..].to_vec();
82        for &r in &reps {
83            if r == 0 || r as usize > content.len() {
84                return Err(Error::BadDictionary("repeat offset outside content"));
85            }
86        }
87        Ok(Dictionary { id, content, tables: Some(DictTables { huf, of, ml, ll }), reps })
88    }
89
90    /// Dictionary ID written to / expected in frame headers (0 = none).
91    pub fn id(&self) -> u32 {
92        self.id
93    }
94
95    /// The history bytes matches may reference.
96    pub fn content(&self) -> &[u8] {
97        &self.content
98    }
99
100    /// True when the dictionary carries entropy tables.
101    pub fn has_entropy_tables(&self) -> bool {
102        self.tables.is_some()
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109
110    #[test]
111    fn raw_and_detect() {
112        let d = Dictionary::new(b"hello world", DictionaryFormat::Detect).unwrap();
113        assert_eq!((d.id(), d.content(), d.has_entropy_tables()), (0, &b"hello world"[..], false));
114        assert!(Dictionary::new(b"hello world", DictionaryFormat::Zstd).is_err());
115        let r = Dictionary::new(b"x", DictionaryFormat::Raw { id: 9 }).unwrap();
116        assert_eq!(r.id(), 9);
117    }
118
119    #[test]
120    fn truncated_zstd_dict_is_error() {
121        let mut b = DICT_MAGIC.to_le_bytes().to_vec();
122        b.extend_from_slice(&[1, 0, 0, 0, 0x80]);
123        assert!(matches!(Dictionary::new(&b, DictionaryFormat::Zstd), Err(Error::BadDictionary(_))));
124    }
125}
zstd_rs/dict/mod.rs

zstd_rs/dict/
mod.rs
