base_d/features/
compression.rs

1use std::io::{Read, Write};
2
3/// Supported compression algorithms.
4#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5pub enum CompressionAlgorithm {
6    Gzip,
7    Zstd,
8    Brotli,
9    Lz4,
10    Snappy,
11    Lzma,
12}
13
14impl CompressionAlgorithm {
15    /// Returns all available compression algorithms.
16    pub fn all() -> Vec<CompressionAlgorithm> {
17        vec![
18            CompressionAlgorithm::Gzip,
19            CompressionAlgorithm::Zstd,
20            CompressionAlgorithm::Brotli,
21            CompressionAlgorithm::Lz4,
22            CompressionAlgorithm::Snappy,
23            CompressionAlgorithm::Lzma,
24        ]
25    }
26
27    /// Select a random compression algorithm.
28    pub fn random() -> CompressionAlgorithm {
29        use rand::prelude::IndexedRandom;
30        let all = Self::all();
31        *all.choose(&mut rand::rng()).unwrap()
32    }
33
34    /// Get default compression level for this algorithm.
35    pub fn default_level(&self) -> u32 {
36        match self {
37            CompressionAlgorithm::Gzip => 6,
38            CompressionAlgorithm::Zstd => 3,
39            CompressionAlgorithm::Brotli => 6,
40            CompressionAlgorithm::Lz4 => 0,    // LZ4 ignores level
41            CompressionAlgorithm::Snappy => 0, // Snappy ignores level
42            CompressionAlgorithm::Lzma => 6,
43        }
44    }
45
46    /// Parse compression algorithm from string.
47    #[allow(clippy::should_implement_trait)]
48    pub fn from_str(s: &str) -> Result<Self, String> {
49        match s.to_lowercase().as_str() {
50            "gzip" | "gz" => Ok(CompressionAlgorithm::Gzip),
51            "zstd" | "zst" => Ok(CompressionAlgorithm::Zstd),
52            "brotli" | "br" => Ok(CompressionAlgorithm::Brotli),
53            "lz4" => Ok(CompressionAlgorithm::Lz4),
54            "snappy" | "snap" => Ok(CompressionAlgorithm::Snappy),
55            "lzma" | "xz" => Ok(CompressionAlgorithm::Lzma),
56            _ => Err(format!("Unknown compression algorithm: {}", s)),
57        }
58    }
59
60    pub fn as_str(&self) -> &str {
61        match self {
62            CompressionAlgorithm::Gzip => "gzip",
63            CompressionAlgorithm::Zstd => "zstd",
64            CompressionAlgorithm::Brotli => "brotli",
65            CompressionAlgorithm::Lz4 => "lz4",
66            CompressionAlgorithm::Snappy => "snappy",
67            CompressionAlgorithm::Lzma => "lzma",
68        }
69    }
70}
71
72/// Compress data using the specified algorithm and level.
73pub fn compress(
74    data: &[u8],
75    algorithm: CompressionAlgorithm,
76    level: u32,
77) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
78    match algorithm {
79        CompressionAlgorithm::Gzip => compress_gzip(data, level),
80        CompressionAlgorithm::Zstd => compress_zstd(data, level),
81        CompressionAlgorithm::Brotli => compress_brotli(data, level),
82        CompressionAlgorithm::Lz4 => compress_lz4(data, level),
83        CompressionAlgorithm::Snappy => compress_snappy(data, level),
84        CompressionAlgorithm::Lzma => compress_lzma(data, level),
85    }
86}
87
88/// Decompress data using the specified algorithm.
89pub fn decompress(
90    data: &[u8],
91    algorithm: CompressionAlgorithm,
92) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
93    match algorithm {
94        CompressionAlgorithm::Gzip => decompress_gzip(data),
95        CompressionAlgorithm::Zstd => decompress_zstd(data),
96        CompressionAlgorithm::Brotli => decompress_brotli(data),
97        CompressionAlgorithm::Lz4 => decompress_lz4(data),
98        CompressionAlgorithm::Snappy => decompress_snappy(data),
99        CompressionAlgorithm::Lzma => decompress_lzma(data),
100    }
101}
102
103fn compress_gzip(data: &[u8], level: u32) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
104    use flate2::Compression;
105    use flate2::write::GzEncoder;
106
107    let mut encoder = GzEncoder::new(Vec::new(), Compression::new(level));
108    encoder.write_all(data)?;
109    Ok(encoder.finish()?)
110}
111
112fn decompress_gzip(data: &[u8]) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
113    use flate2::read::GzDecoder;
114
115    let mut decoder = GzDecoder::new(data);
116    let mut result = Vec::new();
117    decoder.read_to_end(&mut result)?;
118    Ok(result)
119}
120
121fn compress_zstd(data: &[u8], level: u32) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
122    Ok(zstd::encode_all(data, level as i32)?)
123}
124
125fn decompress_zstd(data: &[u8]) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
126    Ok(zstd::decode_all(data)?)
127}
128
129fn compress_brotli(data: &[u8], level: u32) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
130    let mut result = Vec::new();
131    let mut reader = brotli::CompressorReader::new(data, 4096, level, 22);
132    reader.read_to_end(&mut result)?;
133    Ok(result)
134}
135
136fn decompress_brotli(data: &[u8]) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
137    let mut result = Vec::new();
138    let mut reader = brotli::Decompressor::new(data, 4096);
139    reader.read_to_end(&mut result)?;
140    Ok(result)
141}
142
143fn compress_lz4(data: &[u8], _level: u32) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
144    // LZ4 doesn't use compression levels in the same way
145    Ok(lz4::block::compress(data, None, false)?)
146}
147
148fn decompress_lz4(data: &[u8]) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
149    // We need to know the uncompressed size for LZ4, but we don't have it
150    // Use a reasonable max size (100MB)
151    Ok(lz4::block::decompress(data, Some(100 * 1024 * 1024))?)
152}
153
154fn compress_snappy(data: &[u8], _level: u32) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
155    // Snappy doesn't support compression levels
156    let mut encoder = snap::raw::Encoder::new();
157    Ok(encoder.compress_vec(data)?)
158}
159
160fn decompress_snappy(data: &[u8]) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
161    let mut decoder = snap::raw::Decoder::new();
162    Ok(decoder.decompress_vec(data)?)
163}
164
165fn compress_lzma(data: &[u8], level: u32) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
166    use xz2::write::XzEncoder;
167
168    let mut encoder = XzEncoder::new(Vec::new(), level);
169    encoder.write_all(data)?;
170    Ok(encoder.finish()?)
171}
172
173fn decompress_lzma(data: &[u8]) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
174    use xz2::read::XzDecoder;
175
176    let mut decoder = XzDecoder::new(data);
177    let mut result = Vec::new();
178    decoder.read_to_end(&mut result)?;
179    Ok(result)
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185
186    #[test]
187    fn test_gzip_roundtrip() {
188        let data = b"Hello, world! This is a test of gzip compression.";
189        let compressed = compress(data, CompressionAlgorithm::Gzip, 6).unwrap();
190        let decompressed = decompress(&compressed, CompressionAlgorithm::Gzip).unwrap();
191        assert_eq!(data.as_ref(), decompressed.as_slice());
192    }
193
194    #[test]
195    fn test_zstd_roundtrip() {
196        let data = b"Hello, world! This is a test of zstd compression.";
197        let compressed = compress(data, CompressionAlgorithm::Zstd, 3).unwrap();
198        let decompressed = decompress(&compressed, CompressionAlgorithm::Zstd).unwrap();
199        assert_eq!(data.as_ref(), decompressed.as_slice());
200    }
201
202    #[test]
203    fn test_brotli_roundtrip() {
204        let data = b"Hello, world! This is a test of brotli compression.";
205        let compressed = compress(data, CompressionAlgorithm::Brotli, 6).unwrap();
206        let decompressed = decompress(&compressed, CompressionAlgorithm::Brotli).unwrap();
207        assert_eq!(data.as_ref(), decompressed.as_slice());
208    }
209
210    #[test]
211    fn test_lz4_roundtrip() {
212        let data = b"Hello, world! This is a test of lz4 compression.";
213        let compressed = compress(data, CompressionAlgorithm::Lz4, 0).unwrap();
214        let decompressed = decompress(&compressed, CompressionAlgorithm::Lz4).unwrap();
215        assert_eq!(data.as_ref(), decompressed.as_slice());
216    }
217
218    #[test]
219    fn test_snappy_roundtrip() {
220        let data = b"Hello, world! This is a test of snappy compression.";
221        let compressed = compress(data, CompressionAlgorithm::Snappy, 0).unwrap();
222        let decompressed = decompress(&compressed, CompressionAlgorithm::Snappy).unwrap();
223        assert_eq!(data.as_ref(), decompressed.as_slice());
224    }
225
226    #[test]
227    fn test_lzma_roundtrip() {
228        let data = b"Hello, world! This is a test of lzma compression.";
229        let compressed = compress(data, CompressionAlgorithm::Lzma, 6).unwrap();
230        let decompressed = decompress(&compressed, CompressionAlgorithm::Lzma).unwrap();
231        assert_eq!(data.as_ref(), decompressed.as_slice());
232    }
233}