agcodex_persistence/
compression.rs1use crate::error::PersistenceError;
4use crate::error::Result;
5use std::io::Read;
6use std::io::Write;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum CompressionLevel {
11 Fast,
13 Balanced,
15 Maximum,
17 Custom(i32),
19}
20
21impl CompressionLevel {
22 pub fn to_level(self) -> i32 {
24 match self {
25 Self::Fast => 1,
26 Self::Balanced => 3,
27 Self::Maximum => 9,
28 Self::Custom(level) => level.clamp(1, 22),
29 }
30 }
31}
32
33impl Default for CompressionLevel {
34 fn default() -> Self {
35 Self::Balanced
36 }
37}
38
39pub struct Compressor {
41 level: CompressionLevel,
42}
43
44impl Compressor {
45 pub const fn new(level: CompressionLevel) -> Self {
47 Self { level }
48 }
49
50 pub fn compress(&self, data: &[u8]) -> Result<Vec<u8>> {
52 let mut encoder = zstd::Encoder::new(Vec::new(), self.level.to_level())
53 .map_err(|e| PersistenceError::Compression(e.to_string()))?;
54
55 encoder
56 .write_all(data)
57 .map_err(|e| PersistenceError::Compression(e.to_string()))?;
58
59 encoder
60 .finish()
61 .map_err(|e| PersistenceError::Compression(e.to_string()))
62 }
63
64 pub fn decompress(&self, compressed: &[u8]) -> Result<Vec<u8>> {
66 let mut decoder = zstd::Decoder::new(compressed)
67 .map_err(|e| PersistenceError::Compression(e.to_string()))?;
68
69 let mut decompressed = Vec::new();
70 decoder
71 .read_to_end(&mut decompressed)
72 .map_err(|e| PersistenceError::Compression(e.to_string()))?;
73
74 Ok(decompressed)
75 }
76
77 pub fn compress_with_dict(&self, data: &[u8], _dict: &[u8]) -> Result<Vec<u8>> {
79 zstd::encode_all(std::io::Cursor::new(data), self.level.to_level())
80 .map_err(|e| PersistenceError::Compression(e.to_string()))
81 }
82
83 pub fn compression_ratio(original_size: usize, compressed_size: usize) -> f32 {
85 if compressed_size == 0 {
86 return 0.0;
87 }
88 1.0 - (compressed_size as f32 / original_size as f32)
89 }
90}
91
92pub struct StreamCompressor {
94 level: CompressionLevel,
95}
96
97impl StreamCompressor {
98 pub const fn new(level: CompressionLevel) -> Self {
99 Self { level }
100 }
101
102 pub fn compress_stream<R: Read, W: Write>(&self, mut reader: R, writer: W) -> Result<u64> {
104 let mut encoder = zstd::Encoder::new(writer, self.level.to_level())
105 .map_err(|e| PersistenceError::Compression(e.to_string()))?;
106
107 let bytes_written = std::io::copy(&mut reader, &mut encoder)?;
108
109 encoder
110 .finish()
111 .map_err(|e| PersistenceError::Compression(e.to_string()))?;
112
113 Ok(bytes_written)
114 }
115
116 pub fn decompress_stream<R: Read, W: Write>(&self, reader: R, mut writer: W) -> Result<u64> {
118 let mut decoder =
119 zstd::Decoder::new(reader).map_err(|e| PersistenceError::Compression(e.to_string()))?;
120
121 let bytes_written = std::io::copy(&mut decoder, &mut writer)?;
122 Ok(bytes_written)
123 }
124}
125
126#[cfg(test)]
127mod tests {
128 use super::*;
129
130 #[test]
131 fn test_compression_roundtrip() {
132 let compressor = Compressor::new(CompressionLevel::Balanced);
133 let data = "Hello, AGCodex! This is a test of the compression system. ".repeat(100);
135 let data = data.as_bytes();
136
137 let compressed = compressor.compress(data).unwrap();
138 assert!(compressed.len() < data.len());
139
140 let decompressed = compressor.decompress(&compressed).unwrap();
141 assert_eq!(decompressed, data);
142 }
143
144 #[test]
145 fn test_compression_levels() {
146 let data = vec![b'A'; 10000]; let fast = Compressor::new(CompressionLevel::Fast);
149 let balanced = Compressor::new(CompressionLevel::Balanced);
150 let maximum = Compressor::new(CompressionLevel::Maximum);
151
152 let fast_compressed = fast.compress(&data).unwrap();
153 let balanced_compressed = balanced.compress(&data).unwrap();
154 let maximum_compressed = maximum.compress(&data).unwrap();
155
156 assert!(maximum_compressed.len() <= fast_compressed.len());
158
159 assert_eq!(fast.decompress(&fast_compressed).unwrap(), data);
161 assert_eq!(balanced.decompress(&balanced_compressed).unwrap(), data);
162 assert_eq!(maximum.decompress(&maximum_compressed).unwrap(), data);
163 }
164
165 #[test]
166 fn test_compression_ratio() {
167 let ratio = Compressor::compression_ratio(1000, 100);
168 assert!((ratio - 0.9).abs() < 0.001);
169
170 let ratio = Compressor::compression_ratio(1000, 500);
171 assert!((ratio - 0.5).abs() < 0.001);
172 }
173}