Skip to main content

batuta_common/
compression.rs

1//! Shared compression utilities for the Batuta stack.
2//!
3//! Provides LZ4 (fast, real-time) and ZSTD (better ratio) compression
4//! with a common enum interface. Used by trueno-db and trueno-rag.
5
6/// Compression algorithm selector.
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
8pub enum Compression {
9    /// LZ4 - Fast compression with prepended size, good for real-time
10    #[default]
11    Lz4,
12    /// ZSTD - Better compression ratio, slower
13    Zstd,
14}
15
16impl Compression {
17    /// Get algorithm name as string
18    #[must_use]
19    pub const fn as_str(&self) -> &'static str {
20        match self {
21            Self::Lz4 => "lz4",
22            Self::Zstd => "zstd",
23        }
24    }
25
26    /// Compress data using this algorithm.
27    ///
28    /// Returns empty vec for empty input (short-circuit).
29    ///
30    /// # Errors
31    /// Returns [`CompressionError`] if the compression algorithm fails.
32    pub fn compress(&self, data: &[u8]) -> Result<Vec<u8>, CompressionError> {
33        if data.is_empty() {
34            return Ok(Vec::new());
35        }
36        match self {
37            Self::Lz4 => Ok(lz4_flex::compress_prepend_size(data)),
38            Self::Zstd => zstd::encode_all(data, 3)
39                .map_err(|e| CompressionError(format!("ZSTD compression failed: {e}"))),
40        }
41    }
42
43    /// Decompress data using this algorithm.
44    ///
45    /// Returns empty vec for empty input (short-circuit).
46    ///
47    /// # Errors
48    /// Returns [`CompressionError`] if the decompression algorithm fails.
49    pub fn decompress(&self, data: &[u8]) -> Result<Vec<u8>, CompressionError> {
50        if data.is_empty() {
51            return Ok(Vec::new());
52        }
53        match self {
54            Self::Lz4 => lz4_flex::decompress_size_prepended(data)
55                .map_err(|e| CompressionError(format!("LZ4 decompression failed: {e}"))),
56            Self::Zstd => zstd::decode_all(data)
57                .map_err(|e| CompressionError(format!("ZSTD decompression failed: {e}"))),
58        }
59    }
60}
61
62/// Error type for compression operations.
63#[derive(Debug, Clone)]
64pub struct CompressionError(pub String);
65
66impl std::fmt::Display for CompressionError {
67    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
68        write!(f, "{}", self.0)
69    }
70}
71
72impl std::error::Error for CompressionError {}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77
78    #[test]
79    fn test_lz4_roundtrip() {
80        let data = b"hello world hello world hello world";
81        let compressed = Compression::Lz4.compress(data).unwrap();
82        let decompressed = Compression::Lz4.decompress(&compressed).unwrap();
83        assert_eq!(data.as_slice(), decompressed.as_slice());
84    }
85
86    #[test]
87    fn test_zstd_roundtrip() {
88        let data = b"hello world hello world hello world";
89        let compressed = Compression::Zstd.compress(data).unwrap();
90        let decompressed = Compression::Zstd.decompress(&compressed).unwrap();
91        assert_eq!(data.as_slice(), decompressed.as_slice());
92    }
93
94    #[test]
95    fn test_as_str() {
96        assert_eq!(Compression::Lz4.as_str(), "lz4");
97        assert_eq!(Compression::Zstd.as_str(), "zstd");
98    }
99
100    #[test]
101    fn test_default_is_lz4() {
102        assert_eq!(Compression::default(), Compression::Lz4);
103    }
104
105    #[test]
106    fn test_empty_data_lz4() {
107        let empty: &[u8] = &[];
108        let compressed = Compression::Lz4.compress(empty).unwrap();
109        assert!(compressed.is_empty());
110        let decompressed = Compression::Lz4.decompress(&compressed).unwrap();
111        assert!(decompressed.is_empty());
112    }
113
114    #[test]
115    fn test_empty_data_zstd() {
116        let empty: &[u8] = &[];
117        let compressed = Compression::Zstd.compress(empty).unwrap();
118        assert!(compressed.is_empty());
119        let decompressed = Compression::Zstd.decompress(&compressed).unwrap();
120        assert!(decompressed.is_empty());
121    }
122
123    #[test]
124    fn test_lz4_compresses_repeated_data() {
125        let data = vec![0u8; 10000];
126        let compressed = Compression::Lz4.compress(&data).unwrap();
127        assert!(compressed.len() < data.len() / 10);
128    }
129
130    #[test]
131    fn test_zstd_compresses_repeated_data() {
132        let data = vec![0u8; 10000];
133        let compressed = Compression::Zstd.compress(&data).unwrap();
134        assert!(compressed.len() < data.len() / 10);
135    }
136
137    #[test]
138    fn test_compression_error_display() {
139        let err = CompressionError("test error".to_string());
140        assert_eq!(err.to_string(), "test error");
141    }
142}