llm_optimizer_processor/state/
compression.rs

1//! Compression algorithms for state backends
2//!
3//! This module provides transparent compression and decompression for state backends,
4//! supporting multiple algorithms with configurable compression levels and thresholds.
5//!
6//! ## Supported Algorithms
7//!
8//! - **LZ4**: Fast compression with good compression ratio (default)
9//! - **Snappy**: Very fast compression, moderate compression ratio
10//! - **Zstd**: Best compression ratio, moderate speed
11//! - **None**: No compression (passthrough)
12//!
13//! ## Example
14//!
15//! ```rust
16//! use processor::state::compression::{CompressionConfig, CompressionAlgorithm, compress, decompress};
17//!
18//! let config = CompressionConfig::default()
19//!     .with_algorithm(CompressionAlgorithm::Lz4)
20//!     .with_level(4)
21//!     .with_threshold(1024); // Only compress data > 1KB
22//!
23//! let data = b"Hello, World!".repeat(100);
24//!
25//! // Compress
26//! let compressed = compress(&data, &config).unwrap();
27//! println!("Original: {} bytes, Compressed: {} bytes", data.len(), compressed.len());
28//!
29//! // Decompress
30//! let decompressed = decompress(&compressed, &config).unwrap();
31//! assert_eq!(data, decompressed);
32//! ```
33
34use serde::{Deserialize, Serialize};
35use std::io::{self, Read, Write};
36use thiserror::Error;
37
38/// Compression algorithm selection
39#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
40pub enum CompressionAlgorithm {
41    /// No compression (passthrough)
42    None,
43    /// LZ4 compression - fast with good compression ratio
44    Lz4,
45    /// Snappy compression - very fast with moderate compression
46    Snappy,
47    /// Zstd compression - best ratio with moderate speed
48    Zstd,
49}
50
51impl CompressionAlgorithm {
52    /// Get the algorithm name as a string
53    pub fn as_str(&self) -> &'static str {
54        match self {
55            Self::None => "none",
56            Self::Lz4 => "lz4",
57            Self::Snappy => "snappy",
58            Self::Zstd => "zstd",
59        }
60    }
61
62    /// Get the magic byte prefix for this algorithm
63    ///
64    /// This prefix is prepended to compressed data to identify the algorithm used.
65    fn magic_byte(&self) -> u8 {
66        match self {
67            Self::None => 0x00,
68            Self::Lz4 => 0x01,
69            Self::Snappy => 0x02,
70            Self::Zstd => 0x03,
71        }
72    }
73
74    /// Parse algorithm from magic byte
75    fn from_magic_byte(byte: u8) -> Result<Self, CompressionError> {
76        match byte {
77            0x00 => Ok(Self::None),
78            0x01 => Ok(Self::Lz4),
79            0x02 => Ok(Self::Snappy),
80            0x03 => Ok(Self::Zstd),
81            _ => Err(CompressionError::InvalidMagicByte(byte)),
82        }
83    }
84}
85
86impl Default for CompressionAlgorithm {
87    fn default() -> Self {
88        Self::Lz4 // Fast and good compression ratio
89    }
90}
91
92impl std::fmt::Display for CompressionAlgorithm {
93    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
94        write!(f, "{}", self.as_str())
95    }
96}
97
98/// Compression configuration
99#[derive(Debug, Clone, Serialize, Deserialize)]
100pub struct CompressionConfig {
101    /// Compression algorithm to use
102    pub algorithm: CompressionAlgorithm,
103
104    /// Compression level (algorithm-specific):
105    /// - LZ4: 0 (default) to 16 (max)
106    /// - Snappy: Ignored (no levels)
107    /// - Zstd: 1 to 22 (default: 3)
108    pub level: i32,
109
110    /// Minimum size in bytes before compression is applied
111    /// Data smaller than this threshold is stored uncompressed
112    pub threshold: usize,
113
114    /// Enable statistics tracking
115    pub enable_stats: bool,
116}
117
118impl Default for CompressionConfig {
119    fn default() -> Self {
120        Self {
121            algorithm: CompressionAlgorithm::Lz4,
122            level: 4, // Good balance for LZ4
123            threshold: 256, // Compress anything >= 256 bytes
124            enable_stats: false,
125        }
126    }
127}
128
129impl CompressionConfig {
130    /// Create a new configuration with default values
131    pub fn new() -> Self {
132        Self::default()
133    }
134
135    /// Set the compression algorithm
136    pub fn with_algorithm(mut self, algorithm: CompressionAlgorithm) -> Self {
137        self.algorithm = algorithm;
138        self
139    }
140
141    /// Set the compression level
142    pub fn with_level(mut self, level: i32) -> Self {
143        self.level = level;
144        self
145    }
146
147    /// Set the compression threshold
148    pub fn with_threshold(mut self, threshold: usize) -> Self {
149        self.threshold = threshold;
150        self
151    }
152
153    /// Enable statistics tracking
154    pub fn with_stats(mut self, enable: bool) -> Self {
155        self.enable_stats = enable;
156        self
157    }
158
159    /// Create a configuration with no compression
160    pub fn none() -> Self {
161        Self {
162            algorithm: CompressionAlgorithm::None,
163            ..Default::default()
164        }
165    }
166
167    /// Create a configuration optimized for speed
168    pub fn fast() -> Self {
169        Self {
170            algorithm: CompressionAlgorithm::Snappy,
171            level: 0,
172            threshold: 512,
173            enable_stats: false,
174        }
175    }
176
177    /// Create a configuration balanced for speed and size
178    pub fn balanced() -> Self {
179        Self {
180            algorithm: CompressionAlgorithm::Lz4,
181            level: 4,
182            threshold: 256,
183            enable_stats: false,
184        }
185    }
186
187    /// Create a configuration optimized for size
188    pub fn best_size() -> Self {
189        Self {
190            algorithm: CompressionAlgorithm::Zstd,
191            level: 15,
192            threshold: 128,
193            enable_stats: false,
194        }
195    }
196}
197
198/// Compression statistics
199#[derive(Debug, Clone, Default)]
200pub struct CompressionStats {
201    /// Number of compression operations
202    pub compress_count: u64,
203    /// Number of decompression operations
204    pub decompress_count: u64,
205    /// Total bytes input to compression
206    pub bytes_in: u64,
207    /// Total bytes output from compression
208    pub bytes_out: u64,
209    /// Total bytes input to decompression
210    pub decompress_bytes_in: u64,
211    /// Total bytes output from decompression
212    pub decompress_bytes_out: u64,
213    /// Number of times threshold prevented compression
214    pub threshold_skipped: u64,
215}
216
217impl CompressionStats {
218    /// Get the compression ratio (0.0 to 1.0, lower is better)
219    pub fn compression_ratio(&self) -> f64 {
220        if self.bytes_in == 0 {
221            1.0
222        } else {
223            self.bytes_out as f64 / self.bytes_in as f64
224        }
225    }
226
227    /// Get the space saved in bytes
228    pub fn space_saved(&self) -> i64 {
229        self.bytes_in as i64 - self.bytes_out as i64
230    }
231
232    /// Get the space saved as a percentage
233    pub fn space_saved_pct(&self) -> f64 {
234        if self.bytes_in == 0 {
235            0.0
236        } else {
237            (1.0 - self.compression_ratio()) * 100.0
238        }
239    }
240}
241
242/// Compression errors
243#[derive(Debug, Error)]
244pub enum CompressionError {
245    /// LZ4 compression error
246    #[error("LZ4 compression error: {0}")]
247    Lz4Error(String),
248
249    /// Snappy compression error
250    #[error("Snappy compression error: {0}")]
251    SnappyError(String),
252
253    /// Zstd compression error
254    #[error("Zstd compression error: {0}")]
255    ZstdError(String),
256
257    /// Invalid magic byte in compressed data
258    #[error("Invalid compression magic byte: {0:#x}")]
259    InvalidMagicByte(u8),
260
261    /// Data too short to contain magic byte
262    #[error("Data too short (must be at least 1 byte)")]
263    DataTooShort,
264
265    /// IO error during compression/decompression
266    #[error("IO error: {0}")]
267    IoError(#[from] io::Error),
268}
269
270pub type CompressionResult<T> = Result<T, CompressionError>;
271
272/// Compress data according to the configuration
273///
274/// The compressed data will have a 1-byte magic header prepended to identify
275/// the compression algorithm used. This allows for transparent decompression.
276///
277/// # Arguments
278///
279/// * `data` - Input data to compress
280/// * `config` - Compression configuration
281///
282/// # Returns
283///
284/// Compressed data with magic byte prefix
285///
286/// # Example
287///
288/// ```
289/// # use processor::state::compression::{compress, CompressionConfig};
290/// let data = b"Hello, World!".repeat(100);
291/// let config = CompressionConfig::balanced();
292/// let compressed = compress(&data, &config).unwrap();
293/// assert!(compressed.len() < data.len());
294/// ```
295pub fn compress(data: &[u8], config: &CompressionConfig) -> CompressionResult<Vec<u8>> {
296    // Check threshold - don't compress small data
297    if data.len() < config.threshold {
298        let mut result = Vec::with_capacity(1 + data.len());
299        result.push(CompressionAlgorithm::None.magic_byte());
300        result.extend_from_slice(data);
301        return Ok(result);
302    }
303
304    match config.algorithm {
305        CompressionAlgorithm::None => {
306            let mut result = Vec::with_capacity(1 + data.len());
307            result.push(CompressionAlgorithm::None.magic_byte());
308            result.extend_from_slice(data);
309            Ok(result)
310        }
311
312        CompressionAlgorithm::Lz4 => {
313            let compressed = lz4::block::compress(data, Some(lz4::block::CompressionMode::HIGHCOMPRESSION(config.level)), false)
314                .map_err(|e| CompressionError::Lz4Error(e.to_string()))?;
315
316            let mut result = Vec::with_capacity(1 + compressed.len());
317            result.push(CompressionAlgorithm::Lz4.magic_byte());
318            result.extend_from_slice(&compressed);
319            Ok(result)
320        }
321
322        CompressionAlgorithm::Snappy => {
323            let mut encoder = snap::raw::Encoder::new();
324            let compressed = encoder
325                .compress_vec(data)
326                .map_err(|e| CompressionError::SnappyError(e.to_string()))?;
327
328            let mut result = Vec::with_capacity(1 + compressed.len());
329            result.push(CompressionAlgorithm::Snappy.magic_byte());
330            result.extend_from_slice(&compressed);
331            Ok(result)
332        }
333
334        CompressionAlgorithm::Zstd => {
335            let mut compressed = Vec::new();
336            let mut encoder = zstd::Encoder::new(&mut compressed, config.level)?;
337            encoder.write_all(data)?;
338            encoder.finish()?;
339
340            let mut result = Vec::with_capacity(1 + compressed.len());
341            result.push(CompressionAlgorithm::Zstd.magic_byte());
342            result.extend_from_slice(&compressed);
343            Ok(result)
344        }
345    }
346}
347
348/// Decompress data that was compressed with `compress()`
349///
350/// The function automatically detects the compression algorithm from the magic byte
351/// and decompresses accordingly.
352///
353/// # Arguments
354///
355/// * `data` - Compressed data with magic byte prefix
356/// * `config` - Compression configuration (not currently used but kept for API consistency)
357///
358/// # Returns
359///
360/// Decompressed data
361///
362/// # Example
363///
364/// ```
365/// # use processor::state::compression::{compress, decompress, CompressionConfig};
366/// let original = b"Hello, World!".repeat(100);
367/// let config = CompressionConfig::balanced();
368/// let compressed = compress(&original, &config).unwrap();
369/// let decompressed = decompress(&compressed, &config).unwrap();
370/// assert_eq!(original.to_vec(), decompressed);
371/// ```
372pub fn decompress(data: &[u8], _config: &CompressionConfig) -> CompressionResult<Vec<u8>> {
373    if data.is_empty() {
374        return Err(CompressionError::DataTooShort);
375    }
376
377    let magic_byte = data[0];
378    let compressed_data = &data[1..];
379
380    let algorithm = CompressionAlgorithm::from_magic_byte(magic_byte)?;
381
382    match algorithm {
383        CompressionAlgorithm::None => {
384            Ok(compressed_data.to_vec())
385        }
386
387        CompressionAlgorithm::Lz4 => {
388            // LZ4 needs to know the decompressed size. We'll use a heuristic of 4x the compressed size
389            // and let LZ4 tell us the actual size
390            let decompressed = lz4::block::decompress(compressed_data, None)
391                .map_err(|e| CompressionError::Lz4Error(e.to_string()))?;
392            Ok(decompressed)
393        }
394
395        CompressionAlgorithm::Snappy => {
396            let mut decoder = snap::raw::Decoder::new();
397            let decompressed = decoder
398                .decompress_vec(compressed_data)
399                .map_err(|e| CompressionError::SnappyError(e.to_string()))?;
400            Ok(decompressed)
401        }
402
403        CompressionAlgorithm::Zstd => {
404            let mut decompressed = Vec::new();
405            let mut decoder = zstd::Decoder::new(compressed_data)?;
406            decoder.read_to_end(&mut decompressed)?;
407            Ok(decompressed)
408        }
409    }
410}
411
412#[cfg(test)]
413mod tests {
414    use super::*;
415
416    const TEST_DATA_SMALL: &[u8] = b"Hello, World!";
417    const TEST_DATA_MEDIUM: &[u8] = b"This is a longer test string that should definitely be compressed because it has enough repetition. This is a longer test string that should definitely be compressed because it has enough repetition.";
418
419    fn generate_large_data(size: usize) -> Vec<u8> {
420        (0..size).map(|i| (i % 256) as u8).collect()
421    }
422
423    #[test]
424    fn test_algorithm_display() {
425        assert_eq!(CompressionAlgorithm::None.to_string(), "none");
426        assert_eq!(CompressionAlgorithm::Lz4.to_string(), "lz4");
427        assert_eq!(CompressionAlgorithm::Snappy.to_string(), "snappy");
428        assert_eq!(CompressionAlgorithm::Zstd.to_string(), "zstd");
429    }
430
431    #[test]
432    fn test_magic_byte_roundtrip() {
433        for algo in &[
434            CompressionAlgorithm::None,
435            CompressionAlgorithm::Lz4,
436            CompressionAlgorithm::Snappy,
437            CompressionAlgorithm::Zstd,
438        ] {
439            let byte = algo.magic_byte();
440            let parsed = CompressionAlgorithm::from_magic_byte(byte).unwrap();
441            assert_eq!(*algo, parsed);
442        }
443    }
444
445    #[test]
446    fn test_invalid_magic_byte() {
447        let result = CompressionAlgorithm::from_magic_byte(0xFF);
448        assert!(result.is_err());
449    }
450
451    #[test]
452    fn test_compression_config_presets() {
453        let fast = CompressionConfig::fast();
454        assert_eq!(fast.algorithm, CompressionAlgorithm::Snappy);
455
456        let balanced = CompressionConfig::balanced();
457        assert_eq!(balanced.algorithm, CompressionAlgorithm::Lz4);
458
459        let best = CompressionConfig::best_size();
460        assert_eq!(best.algorithm, CompressionAlgorithm::Zstd);
461
462        let none = CompressionConfig::none();
463        assert_eq!(none.algorithm, CompressionAlgorithm::None);
464    }
465
466    #[test]
467    fn test_compression_none() {
468        let config = CompressionConfig::none();
469        let compressed = compress(TEST_DATA_MEDIUM, &config).unwrap();
470        let decompressed = decompress(&compressed, &config).unwrap();
471
472        assert_eq!(decompressed, TEST_DATA_MEDIUM);
473        // No compression should add only magic byte
474        assert_eq!(compressed.len(), TEST_DATA_MEDIUM.len() + 1);
475    }
476
477    #[test]
478    fn test_compression_lz4() {
479        let config = CompressionConfig::default().with_algorithm(CompressionAlgorithm::Lz4);
480        let compressed = compress(TEST_DATA_MEDIUM, &config).unwrap();
481        let decompressed = decompress(&compressed, &config).unwrap();
482
483        assert_eq!(decompressed, TEST_DATA_MEDIUM);
484        // Should achieve some compression on repetitive data
485        assert!(compressed.len() < TEST_DATA_MEDIUM.len());
486    }
487
488    #[test]
489    fn test_compression_snappy() {
490        let config = CompressionConfig::default().with_algorithm(CompressionAlgorithm::Snappy);
491        let compressed = compress(TEST_DATA_MEDIUM, &config).unwrap();
492        let decompressed = decompress(&compressed, &config).unwrap();
493
494        assert_eq!(decompressed, TEST_DATA_MEDIUM);
495        assert!(compressed.len() < TEST_DATA_MEDIUM.len());
496    }
497
498    #[test]
499    fn test_compression_zstd() {
500        let config = CompressionConfig::default().with_algorithm(CompressionAlgorithm::Zstd);
501        let compressed = compress(TEST_DATA_MEDIUM, &config).unwrap();
502        let decompressed = decompress(&compressed, &config).unwrap();
503
504        assert_eq!(decompressed, TEST_DATA_MEDIUM);
505        assert!(compressed.len() < TEST_DATA_MEDIUM.len());
506    }
507
508    #[test]
509    fn test_threshold_prevents_compression() {
510        let config = CompressionConfig::default()
511            .with_algorithm(CompressionAlgorithm::Lz4)
512            .with_threshold(1000); // Higher than test data
513
514        let compressed = compress(TEST_DATA_SMALL, &config).unwrap();
515
516        // Should not be compressed (only magic byte added)
517        assert_eq!(compressed[0], CompressionAlgorithm::None.magic_byte());
518        assert_eq!(compressed.len(), TEST_DATA_SMALL.len() + 1);
519
520        let decompressed = decompress(&compressed, &config).unwrap();
521        assert_eq!(decompressed, TEST_DATA_SMALL);
522    }
523
524    #[test]
525    fn test_large_data_compression() {
526        let large_data = generate_large_data(1024 * 1024); // 1MB
527        let config = CompressionConfig::balanced();
528
529        let compressed = compress(&large_data, &config).unwrap();
530        let decompressed = decompress(&compressed, &config).unwrap();
531
532        assert_eq!(decompressed, large_data);
533        println!("Original: {} bytes, Compressed: {} bytes, Ratio: {:.2}%",
534            large_data.len(),
535            compressed.len(),
536            (compressed.len() as f64 / large_data.len() as f64) * 100.0
537        );
538    }
539
540    #[test]
541    fn test_empty_data() {
542        let config = CompressionConfig::default();
543        let compressed = compress(&[], &config).unwrap();
544
545        // Empty data below threshold, should not be compressed
546        assert_eq!(compressed.len(), 1); // Just magic byte
547
548        let decompressed = decompress(&compressed, &config).unwrap();
549        assert_eq!(decompressed.len(), 0);
550    }
551
552    #[test]
553    fn test_decompress_empty_fails() {
554        let config = CompressionConfig::default();
555        let result = decompress(&[], &config);
556        assert!(result.is_err());
557    }
558
559    #[test]
560    fn test_all_algorithms_roundtrip() {
561        let test_data = generate_large_data(10000);
562
563        for algo in &[
564            CompressionAlgorithm::None,
565            CompressionAlgorithm::Lz4,
566            CompressionAlgorithm::Snappy,
567            CompressionAlgorithm::Zstd,
568        ] {
569            let config = CompressionConfig::default()
570                .with_algorithm(*algo)
571                .with_threshold(0); // Force compression
572
573            let compressed = compress(&test_data, &config).unwrap();
574            let decompressed = decompress(&compressed, &config).unwrap();
575
576            assert_eq!(
577                decompressed, test_data,
578                "Roundtrip failed for algorithm: {:?}",
579                algo
580            );
581        }
582    }
583
584    #[test]
585    fn test_compression_levels() {
586        let test_data = generate_large_data(10000);
587
588        // Test different Zstd levels
589        for level in [1, 3, 10, 15] {
590            let config = CompressionConfig::default()
591                .with_algorithm(CompressionAlgorithm::Zstd)
592                .with_level(level)
593                .with_threshold(0);
594
595            let compressed = compress(&test_data, &config).unwrap();
596            let decompressed = decompress(&compressed, &config).unwrap();
597
598            assert_eq!(decompressed, test_data);
599            println!("Zstd level {}: {} bytes", level, compressed.len());
600        }
601    }
602
603    #[test]
604    fn test_stats() {
605        let mut stats = CompressionStats::default();
606
607        stats.compress_count = 100;
608        stats.bytes_in = 10000;
609        stats.bytes_out = 3000;
610
611        assert_eq!(stats.compression_ratio(), 0.3);
612        assert_eq!(stats.space_saved(), 7000);
613        assert_eq!(stats.space_saved_pct(), 70.0);
614    }
615}
llm_optimizer_processor/state/compression.rs

llm_optimizer_processor/state/
compression.rs
