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