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}
11
12impl CompressionAlgorithm {
13    /// Parse compression algorithm from string.
14    pub fn from_str(s: &str) -> Result<Self, String> {
15        match s.to_lowercase().as_str() {
16            "gzip" | "gz" => Ok(CompressionAlgorithm::Gzip),
17            "zstd" | "zst" => Ok(CompressionAlgorithm::Zstd),
18            "brotli" | "br" => Ok(CompressionAlgorithm::Brotli),
19            "lz4" => Ok(CompressionAlgorithm::Lz4),
20            _ => Err(format!("Unknown compression algorithm: {}", s)),
21        }
22    }
23    
24    pub fn as_str(&self) -> &str {
25        match self {
26            CompressionAlgorithm::Gzip => "gzip",
27            CompressionAlgorithm::Zstd => "zstd",
28            CompressionAlgorithm::Brotli => "brotli",
29            CompressionAlgorithm::Lz4 => "lz4",
30        }
31    }
32}
33
34/// Compress data using the specified algorithm and level.
35pub fn compress(data: &[u8], algorithm: CompressionAlgorithm, level: u32) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
36    match algorithm {
37        CompressionAlgorithm::Gzip => compress_gzip(data, level),
38        CompressionAlgorithm::Zstd => compress_zstd(data, level),
39        CompressionAlgorithm::Brotli => compress_brotli(data, level),
40        CompressionAlgorithm::Lz4 => compress_lz4(data, level),
41    }
42}
43
44/// Decompress data using the specified algorithm.
45pub fn decompress(data: &[u8], algorithm: CompressionAlgorithm) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
46    match algorithm {
47        CompressionAlgorithm::Gzip => decompress_gzip(data),
48        CompressionAlgorithm::Zstd => decompress_zstd(data),
49        CompressionAlgorithm::Brotli => decompress_brotli(data),
50        CompressionAlgorithm::Lz4 => decompress_lz4(data),
51    }
52}
53
54fn compress_gzip(data: &[u8], level: u32) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
55    use flate2::Compression;
56    use flate2::write::GzEncoder;
57    
58    let mut encoder = GzEncoder::new(Vec::new(), Compression::new(level));
59    encoder.write_all(data)?;
60    Ok(encoder.finish()?)
61}
62
63fn decompress_gzip(data: &[u8]) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
64    use flate2::read::GzDecoder;
65    
66    let mut decoder = GzDecoder::new(data);
67    let mut result = Vec::new();
68    decoder.read_to_end(&mut result)?;
69    Ok(result)
70}
71
72fn compress_zstd(data: &[u8], level: u32) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
73    Ok(zstd::encode_all(data, level as i32)?)
74}
75
76fn decompress_zstd(data: &[u8]) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
77    Ok(zstd::decode_all(data)?)
78}
79
80fn compress_brotli(data: &[u8], level: u32) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
81    let mut result = Vec::new();
82    let mut reader = brotli::CompressorReader::new(data, 4096, level, 22);
83    reader.read_to_end(&mut result)?;
84    Ok(result)
85}
86
87fn decompress_brotli(data: &[u8]) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
88    let mut result = Vec::new();
89    let mut reader = brotli::Decompressor::new(data, 4096);
90    reader.read_to_end(&mut result)?;
91    Ok(result)
92}
93
94fn compress_lz4(data: &[u8], _level: u32) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
95    // LZ4 doesn't use compression levels in the same way
96    Ok(lz4::block::compress(data, None, false)?)
97}
98
99fn decompress_lz4(data: &[u8]) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
100    // We need to know the uncompressed size for LZ4, but we don't have it
101    // Use a reasonable max size (100MB)
102    Ok(lz4::block::decompress(data, Some(100 * 1024 * 1024))?)
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108    
109    #[test]
110    fn test_gzip_roundtrip() {
111        let data = b"Hello, world! This is a test of gzip compression.";
112        let compressed = compress(data, CompressionAlgorithm::Gzip, 6).unwrap();
113        let decompressed = decompress(&compressed, CompressionAlgorithm::Gzip).unwrap();
114        assert_eq!(data.as_ref(), decompressed.as_slice());
115    }
116    
117    #[test]
118    fn test_zstd_roundtrip() {
119        let data = b"Hello, world! This is a test of zstd compression.";
120        let compressed = compress(data, CompressionAlgorithm::Zstd, 3).unwrap();
121        let decompressed = decompress(&compressed, CompressionAlgorithm::Zstd).unwrap();
122        assert_eq!(data.as_ref(), decompressed.as_slice());
123    }
124    
125    #[test]
126    fn test_brotli_roundtrip() {
127        let data = b"Hello, world! This is a test of brotli compression.";
128        let compressed = compress(data, CompressionAlgorithm::Brotli, 6).unwrap();
129        let decompressed = decompress(&compressed, CompressionAlgorithm::Brotli).unwrap();
130        assert_eq!(data.as_ref(), decompressed.as_slice());
131    }
132    
133    #[test]
134    fn test_lz4_roundtrip() {
135        let data = b"Hello, world! This is a test of lz4 compression.";
136        let compressed = compress(data, CompressionAlgorithm::Lz4, 0).unwrap();
137        let decompressed = decompress(&compressed, CompressionAlgorithm::Lz4).unwrap();
138        assert_eq!(data.as_ref(), decompressed.as_slice());
139    }
140}