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}