agcodex_persistence/
compression.rs

1//! Compression utilities using Zstd
2
3use crate::error::PersistenceError;
4use crate::error::Result;
5use std::io::Read;
6use std::io::Write;
7
8/// Compression level for Zstd
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum CompressionLevel {
11    /// Fast compression (level 1)
12    Fast,
13    /// Balanced compression (level 3)
14    Balanced,
15    /// Maximum compression (level 9)
16    Maximum,
17    /// Custom level (1-22)
18    Custom(i32),
19}
20
21impl CompressionLevel {
22    /// Convert to Zstd compression level
23    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
39/// Zstd compressor for session data
40pub struct Compressor {
41    level: CompressionLevel,
42}
43
44impl Compressor {
45    /// Create a new compressor with the specified level
46    pub const fn new(level: CompressionLevel) -> Self {
47        Self { level }
48    }
49
50    /// Compress data using Zstd
51    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    /// Decompress data using Zstd
65    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    /// Compress data with dictionary for better compression of similar data
78    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    /// Calculate compression ratio
84    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
92/// Stream compressor for large sessions
93pub struct StreamCompressor {
94    level: CompressionLevel,
95}
96
97impl StreamCompressor {
98    pub const fn new(level: CompressionLevel) -> Self {
99        Self { level }
100    }
101
102    /// Compress from reader to writer
103    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    /// Decompress from reader to writer
117    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        // Use larger, repetitive data that compresses well
134        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]; // Highly compressible data
147
148        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        // Maximum should compress better than fast
157        assert!(maximum_compressed.len() <= fast_compressed.len());
158
159        // All should decompress to the same data
160        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}