ailake_vec/
compress.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2#[derive(Debug, Clone, Copy, PartialEq, Eq)]
3pub enum CompressionCodec {
4    None,
5    Lz4,
6    Zstd,
7}
8
9pub struct BlockCompressor {
10    codec: CompressionCodec,
11    zstd_level: i32,
12}
13
14impl BlockCompressor {
15    pub fn none() -> Self {
16        Self {
17            codec: CompressionCodec::None,
18            zstd_level: 3,
19        }
20    }
21
22    pub fn lz4() -> Self {
23        Self {
24            codec: CompressionCodec::Lz4,
25            zstd_level: 3,
26        }
27    }
28
29    pub fn zstd(level: i32) -> Self {
30        Self {
31            codec: CompressionCodec::Zstd,
32            zstd_level: level,
33        }
34    }
35
36    pub fn codec(&self) -> CompressionCodec {
37        self.codec
38    }
39
40    pub fn compress(&self, data: &[u8]) -> Vec<u8> {
41        match self.codec {
42            CompressionCodec::None => data.to_vec(),
43            CompressionCodec::Lz4 => lz4_flex::compress_prepend_size(data),
44            CompressionCodec::Zstd => {
45                zstd::bulk::compress(data, self.zstd_level).unwrap_or_else(|_| data.to_vec())
46            }
47        }
48    }
49
50    /// Decompresses `data` written by [`compress`](Self::compress) with the same codec.
51    ///
52    /// Returns an error on truncated/corrupted input instead of silently substituting
53    /// the still-compressed bytes as if they were the decompressed payload — a caller
54    /// has no way to detect corruption if a failed decompress looks like success.
55    pub fn decompress(&self, data: &[u8]) -> ailake_core::AilakeResult<Vec<u8>> {
56        match self.codec {
57            CompressionCodec::None => Ok(data.to_vec()),
58            CompressionCodec::Lz4 => lz4_flex::decompress_size_prepended(data).map_err(|e| {
59                ailake_core::AilakeError::Io(std::io::Error::other(format!(
60                    "ailake: LZ4 block decompression failed ({} bytes input): {e}",
61                    data.len()
62                )))
63            }),
64            CompressionCodec::Zstd => zstd::bulk::decompress(data, 64 * 1024 * 1024).map_err(|e| {
65                ailake_core::AilakeError::Io(std::io::Error::other(format!(
66                    "ailake: Zstd block decompression failed ({} bytes input): {e}",
67                    data.len()
68                )))
69            }),
70        }
71    }
72}
73
74impl Default for BlockCompressor {
75    fn default() -> Self {
76        Self::zstd(3)
77    }
78}
79
80#[cfg(test)]
81mod tests {
82    use super::*;
83
84    fn roundtrip(codec: BlockCompressor, data: &[u8]) {
85        let compressed = codec.compress(data);
86        let decompressed = codec.decompress(&compressed).unwrap();
87        assert_eq!(decompressed, data);
88    }
89
90    #[test]
91    fn lz4_roundtrip() {
92        let data: Vec<u8> = (0u8..200).cycle().take(4096).collect();
93        roundtrip(BlockCompressor::lz4(), &data);
94    }
95
96    #[test]
97    fn zstd_roundtrip() {
98        let data: Vec<u8> = (0u8..200).cycle().take(4096).collect();
99        roundtrip(BlockCompressor::zstd(3), &data);
100    }
101
102    #[test]
103    fn none_passthrough() {
104        let data = b"hello ailake";
105        roundtrip(BlockCompressor::none(), data);
106    }
107
108    /// Regression: `decompress()` used to swallow the error and return the still-compressed
109    /// bytes verbatim on truncated/corrupted input, masking corruption as if it were a
110    /// successful decompress instead of surfacing it to the caller.
111    #[test]
112    fn lz4_decompress_of_corrupt_data_errors() {
113        let c = BlockCompressor::lz4();
114        let garbage = vec![0xFFu8; 8];
115        assert!(c.decompress(&garbage).is_err());
116    }
117
118    #[test]
119    fn zstd_decompress_of_corrupt_data_errors() {
120        let c = BlockCompressor::zstd(3);
121        let garbage = vec![0xFFu8; 8];
122        assert!(c.decompress(&garbage).is_err());
123    }
124
125    #[test]
126    fn zstd_compresses_repetitive_data() {
127        // Repetitive float data (like zero vectors) should compress well
128        let data = vec![0u8; 8192];
129        let c = BlockCompressor::zstd(3);
130        let compressed = c.compress(&data);
131        assert!(
132            compressed.len() < data.len() / 4,
133            "expected >4x compression ratio"
134        );
135    }
136}
ailake_vec/compress.rs

ailake_vec/
compress.rs
