tx2_pack/
compression.rs

1use crate::error::{PackError, Result};
2use crate::format::CompressionType;
3
4#[derive(Debug, Clone, Copy)]
5pub enum CompressionCodec {
6    None,
7    Zstd(i32),
8    Lz4,
9}
10
11impl CompressionCodec {
12    pub fn none() -> Self {
13        CompressionCodec::None
14    }
15
16    pub fn zstd_default() -> Self {
17        CompressionCodec::Zstd(3)
18    }
19
20    pub fn zstd_fast() -> Self {
21        CompressionCodec::Zstd(1)
22    }
23
24    pub fn zstd_best() -> Self {
25        CompressionCodec::Zstd(19)
26    }
27
28    pub fn lz4_default() -> Self {
29        CompressionCodec::Lz4
30    }
31}
32
33impl From<CompressionCodec> for CompressionType {
34    fn from(codec: CompressionCodec) -> Self {
35        match codec {
36            CompressionCodec::None => CompressionType::None,
37            CompressionCodec::Zstd(_) => CompressionType::Zstd,
38            CompressionCodec::Lz4 => CompressionType::Lz4,
39        }
40    }
41}
42
43pub fn compress(data: &[u8], codec: CompressionCodec) -> Result<Vec<u8>> {
44    match codec {
45        CompressionCodec::None => Ok(data.to_vec()),
46
47        CompressionCodec::Zstd(level) => {
48            zstd::bulk::compress(data, level)
49                .map_err(|e| PackError::Compression(e.to_string()))
50        }
51
52        CompressionCodec::Lz4 => {
53            let mut encoder = lz4::EncoderBuilder::new()
54                .level(4)
55                .build(Vec::new())
56                .map_err(|e| PackError::Compression(e.to_string()))?;
57
58            std::io::copy(&mut &data[..], &mut encoder)
59                .map_err(|e| PackError::Compression(e.to_string()))?;
60
61            let (compressed, result) = encoder.finish();
62            result.map_err(|e| PackError::Compression(e.to_string()))?;
63
64            Ok(compressed)
65        }
66    }
67}
68
69pub fn decompress(data: &[u8], compression_type: CompressionType) -> Result<Vec<u8>> {
70    match compression_type {
71        CompressionType::None => Ok(data.to_vec()),
72
73        CompressionType::Zstd => {
74            zstd::bulk::decompress(data, 100 * 1024 * 1024)
75                .map_err(|e| PackError::Decompression(e.to_string()))
76        }
77
78        CompressionType::Lz4 => {
79            let mut decoder = lz4::Decoder::new(data)
80                .map_err(|e| PackError::Decompression(e.to_string()))?;
81
82            let mut decompressed = Vec::new();
83            std::io::copy(&mut decoder, &mut decompressed)
84                .map_err(|e| PackError::Decompression(e.to_string()))?;
85
86            Ok(decompressed)
87        }
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94
95    #[test]
96    fn test_zstd_compression() {
97        let data = b"Hello, World! This is a test of ZSTD compression.".repeat(100);
98
99        let compressed = compress(&data, CompressionCodec::zstd_default()).unwrap();
100        assert!(compressed.len() < data.len());
101
102        let decompressed = decompress(&compressed, CompressionType::Zstd).unwrap();
103        assert_eq!(data, decompressed);
104    }
105
106    #[test]
107    fn test_lz4_compression() {
108        let data = b"Hello, World! This is a test of LZ4 compression.".repeat(100);
109
110        let compressed = compress(&data, CompressionCodec::Lz4).unwrap();
111        assert!(compressed.len() < data.len());
112
113        let decompressed = decompress(&compressed, CompressionType::Lz4).unwrap();
114        assert_eq!(data, decompressed);
115    }
116
117    #[test]
118    fn test_no_compression() {
119        let data = b"Hello, World!";
120
121        let compressed = compress(data, CompressionCodec::None).unwrap();
122        assert_eq!(data, compressed.as_slice());
123
124        let decompressed = decompress(&compressed, CompressionType::None).unwrap();
125        assert_eq!(data, decompressed.as_slice());
126    }
127}