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

base_d/
compression.rs
