hermes_core/compression/
zstd.rs

1//! Zstd compression backend with dictionary support
2//!
3////! For static indexes, we use:
4//! - Maximum compression level (22) for best compression ratio
5//! - Trained dictionaries for even better compression of similar documents
6//! - Larger block sizes to improve compression efficiency
7
8use std::io::{self, Read, Write};
9
10/// Compression level (1-22 for zstd)
11#[derive(Debug, Clone, Copy)]
12pub struct CompressionLevel(pub i32);
13
14impl CompressionLevel {
15    /// Fast compression (level 1)
16    pub const FAST: Self = Self(1);
17    /// Default compression (level 3)
18    pub const DEFAULT: Self = Self(3);
19    /// Better compression (level 9)
20    pub const BETTER: Self = Self(9);
21    /// Best compression (level 19)
22    pub const BEST: Self = Self(19);
23    /// Maximum compression (level 22) - slowest but smallest
24    pub const MAX: Self = Self(22);
25}
26
27impl Default for CompressionLevel {
28    fn default() -> Self {
29        Self::MAX // Use maximum compression for static indexes
30    }
31}
32
33/// Trained Zstd dictionary for improved compression
34#[derive(Clone)]
35pub struct CompressionDict {
36    raw_dict: Vec<u8>,
37}
38
39impl CompressionDict {
40    /// Train a dictionary from sample data
41    ///
42    /// For best results, provide many small samples (e.g., serialized documents)
43    /// The dictionary size should typically be 16KB-112KB
44    pub fn train(samples: &[&[u8]], dict_size: usize) -> io::Result<Self> {
45        let raw_dict = zstd::dict::from_samples(samples, dict_size).map_err(io::Error::other)?;
46        Ok(Self { raw_dict })
47    }
48
49    /// Create dictionary from raw bytes (for loading saved dictionaries)
50    pub fn from_bytes(bytes: Vec<u8>) -> Self {
51        Self { raw_dict: bytes }
52    }
53
54    /// Get raw dictionary bytes (for saving)
55    pub fn as_bytes(&self) -> &[u8] {
56        &self.raw_dict
57    }
58
59    /// Dictionary size in bytes
60    pub fn len(&self) -> usize {
61        self.raw_dict.len()
62    }
63
64    /// Check if dictionary is empty
65    pub fn is_empty(&self) -> bool {
66        self.raw_dict.is_empty()
67    }
68}
69
70/// Compress data using Zstd
71pub fn compress(data: &[u8], level: CompressionLevel) -> io::Result<Vec<u8>> {
72    zstd::encode_all(data, level.0).map_err(io::Error::other)
73}
74
75/// Compress data using Zstd with a trained dictionary
76pub fn compress_with_dict(
77    data: &[u8],
78    level: CompressionLevel,
79    dict: &CompressionDict,
80) -> io::Result<Vec<u8>> {
81    let mut encoder = zstd::Encoder::with_dictionary(Vec::new(), level.0, &dict.raw_dict)
82        .map_err(io::Error::other)?;
83    encoder.write_all(data)?;
84    encoder.finish().map_err(io::Error::other)
85}
86
87/// Decompress data using Zstd
88pub fn decompress(data: &[u8]) -> io::Result<Vec<u8>> {
89    zstd::decode_all(data).map_err(io::Error::other)
90}
91
92/// Decompress data using Zstd with a trained dictionary
93pub fn decompress_with_dict(data: &[u8], dict: &CompressionDict) -> io::Result<Vec<u8>> {
94    let mut decoder =
95        zstd::Decoder::with_dictionary(data, &dict.raw_dict).map_err(io::Error::other)?;
96    let mut result = Vec::new();
97    decoder.read_to_end(&mut result)?;
98    Ok(result)
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    #[test]
106    fn test_roundtrip() {
107        let data = b"Hello, World! This is a test of compression.".repeat(100);
108        let compressed = compress(&data, CompressionLevel::default()).unwrap();
109        let decompressed = decompress(&compressed).unwrap();
110        assert_eq!(data, decompressed.as_slice());
111        assert!(compressed.len() < data.len());
112    }
113
114    #[test]
115    fn test_empty_data() {
116        let data: &[u8] = &[];
117        let compressed = compress(data, CompressionLevel::default()).unwrap();
118        let decompressed = decompress(&compressed).unwrap();
119        assert!(decompressed.is_empty());
120    }
121
122    #[test]
123    fn test_compression_levels() {
124        let data = b"Test data for compression levels".repeat(100);
125        for level in [1, 3, 9, 19] {
126            let compressed = compress(&data, CompressionLevel(level)).unwrap();
127            let decompressed = decompress(&compressed).unwrap();
128            assert_eq!(data.as_slice(), decompressed.as_slice());
129        }
130    }
131}
hermes_core/compression/zstd.rs

hermes_core/compression/
zstd.rs
