shard_core/
compression.rs

1use anyhow::Result;
2use std::str::FromStr;
3
4#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5pub enum Compression {
6    None,
7    Zstd,
8    Zlib,
9}
10
11impl FromStr for Compression {
12    type Err = anyhow::Error;
13
14    fn from_str(s: &str) -> Result<Self, Self::Err> {
15        match s {
16            "none" => Ok(Compression::None),
17            "zstd" => Ok(Compression::Zstd),
18            "zlib" => Ok(Compression::Zlib),
19            other => anyhow::bail!("Unknown compression algorithm: {other}"),
20        }
21    }
22}
23
24impl Compression {
25    pub fn as_str(&self) -> &'static str {
26        match self {
27            Compression::None => "none",
28            Compression::Zstd => "zstd",
29            Compression::Zlib => "zlib",
30        }
31    }
32
33    pub fn compress(&self, data: &[u8]) -> Result<Vec<u8>> {
34        match self {
35            Compression::None => Ok(data.to_vec()),
36            Compression::Zstd => {
37                let compressed = zstd::bulk::compress(data, 3)?;
38                Ok(compressed)
39            }
40            Compression::Zlib => {
41                use std::io::Write;
42                let mut encoder =
43                    flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::default());
44                encoder.write_all(data)?;
45                let compressed = encoder.finish()?;
46                Ok(compressed)
47            }
48        }
49    }
50
51    pub fn decompress(&self, data: &[u8]) -> Result<Vec<u8>> {
52        match self {
53            Compression::None => Ok(data.to_vec()),
54            Compression::Zstd => {
55                let decompressed = zstd::bulk::decompress(data, 1024 * 1024 * 1024)?;
56                Ok(decompressed)
57            }
58            Compression::Zlib => {
59                use std::io::Read;
60                let mut decoder = flate2::read::ZlibDecoder::new(data);
61                let mut buf = Vec::new();
62                decoder.read_to_end(&mut buf)?;
63                Ok(buf)
64            }
65        }
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use super::*;
72
73    #[test]
74    fn test_compress_decompress_none() {
75        let data = b"hello world";
76        let compressed = Compression::None.compress(data).unwrap();
77        assert_eq!(compressed, data);
78        let decompressed = Compression::None.decompress(&compressed).unwrap();
79        assert_eq!(decompressed, data);
80    }
81
82    #[test]
83    fn test_compress_decompress_zstd() {
84        let data = vec![0u8; 4096];
85        let compressed = Compression::Zstd.compress(&data).unwrap();
86        assert!(compressed.len() < data.len(), "zstd should compress zeros");
87        let decompressed = Compression::Zstd.decompress(&compressed).unwrap();
88        assert_eq!(decompressed, data);
89    }
90
91    #[test]
92    fn test_compress_decompress_zlib() {
93        let data = vec![0u8; 4096];
94        let compressed = Compression::Zlib.compress(&data).unwrap();
95        assert!(compressed.len() < data.len(), "zlib should compress zeros");
96        let decompressed = Compression::Zlib.decompress(&compressed).unwrap();
97        assert_eq!(decompressed, data);
98    }
99
100    #[test]
101    fn test_roundtrip_all() {
102        let data = b"The quick brown fox jumps over the lazy dog";
103        for algo in &[Compression::None, Compression::Zstd, Compression::Zlib] {
104            let compressed = algo.compress(data).unwrap();
105            let decompressed = algo.decompress(&compressed).unwrap();
106            assert_eq!(decompressed, data, "roundtrip failed for {:?}", algo);
107        }
108    }
109
110    #[test]
111    fn test_from_str() {
112        assert_eq!("none".parse::<Compression>().unwrap(), Compression::None);
113        assert_eq!("zstd".parse::<Compression>().unwrap(), Compression::Zstd);
114        assert_eq!("zlib".parse::<Compression>().unwrap(), Compression::Zlib);
115        assert!("invalid".parse::<Compression>().is_err());
116    }
117
118    #[test]
119    fn test_as_str() {
120        assert_eq!(Compression::None.as_str(), "none");
121        assert_eq!(Compression::Zstd.as_str(), "zstd");
122        assert_eq!(Compression::Zlib.as_str(), "zlib");
123    }
124}
shard_core/compression.rs

shard_core/
compression.rs
