mars_xlog_core/
compress.rs

1use std::io::{Read, Write};
2
3use flate2::Compression;
4use thiserror::Error;
5
6#[derive(Debug, Error)]
7/// Errors raised by the compression and decompression helpers.
8pub enum CompressError {
9    /// Raw-deflate compression failed.
10    #[error("zlib compress failed: {0}")]
11    Zlib(String),
12    /// zstd compression failed.
13    #[error("zstd compress failed: {0}")]
14    Zstd(String),
15    /// Raw-deflate decompression failed.
16    #[error("zlib decompress failed: {0}")]
17    ZlibDecompress(String),
18    /// zstd frame decompression failed.
19    #[error("zstd decompress failed: {0}")]
20    ZstdDecompress(String),
21}
22
23/// Stateful compressor used by the async writer to emit one logical stream.
24pub trait StreamCompressor {
25    /// Compresses one input chunk and appends newly emitted bytes to `output`.
26    fn compress_chunk(&mut self, input: &[u8], output: &mut Vec<u8>) -> Result<(), CompressError>;
27    /// Finalizes the stream and appends any remaining encoded bytes to `output`.
28    fn flush(&mut self, output: &mut Vec<u8>) -> Result<(), CompressError>;
29}
30
31/// Raw-deflate stream compressor compatible with Mars zlib settings.
32pub struct ZlibStreamCompressor {
33    inner: flate2::write::DeflateEncoder<Vec<u8>>,
34}
35
36impl ZlibStreamCompressor {
37    /// Creates a raw-deflate stream compressor with the requested compression level.
38    pub fn new(level: i32) -> Self {
39        let level = level.clamp(0, 9) as u32;
40        Self {
41            inner: flate2::write::DeflateEncoder::new(Vec::new(), Compression::new(level)),
42        }
43    }
44}
45
46impl Default for ZlibStreamCompressor {
47    fn default() -> Self {
48        Self::new(9)
49    }
50}
51
52impl StreamCompressor for ZlibStreamCompressor {
53    fn compress_chunk(&mut self, input: &[u8], output: &mut Vec<u8>) -> Result<(), CompressError> {
54        self.inner
55            .write_all(input)
56            .map_err(|e| CompressError::Zlib(e.to_string()))?;
57        self.inner
58            .flush()
59            .map_err(|e| CompressError::Zlib(e.to_string()))?;
60        drain_encoded_bytes(self.inner.get_mut(), output);
61        Ok(())
62    }
63
64    fn flush(&mut self, output: &mut Vec<u8>) -> Result<(), CompressError> {
65        self.inner
66            .try_finish()
67            .map_err(|e| CompressError::Zlib(e.to_string()))?;
68        drain_encoded_bytes(self.inner.get_mut(), output);
69        Ok(())
70    }
71}
72
73/// zstd chunk compressor.
74///
75/// Each chunk is encoded as a zstd frame and can be decoded from concatenated frames.
76pub struct ZstdChunkCompressor {
77    level: i32,
78}
79
80impl ZstdChunkCompressor {
81    /// Creates a compressor that encodes each input chunk as a separate zstd frame.
82    pub fn new(level: i32) -> Self {
83        Self { level }
84    }
85}
86
87impl StreamCompressor for ZstdChunkCompressor {
88    fn compress_chunk(&mut self, input: &[u8], output: &mut Vec<u8>) -> Result<(), CompressError> {
89        let compressed = zstd::bulk::compress(input, self.level)
90            .map_err(|e| CompressError::Zstd(e.to_string()))?;
91        output.extend_from_slice(&compressed);
92        Ok(())
93    }
94
95    fn flush(&mut self, _output: &mut Vec<u8>) -> Result<(), CompressError> {
96        Ok(())
97    }
98}
99
100/// Streaming zstd compressor compatible with Mars async path.
101///
102/// Uses one compression stream with `windowLog=16`, flushing each chunk and
103/// emitting the final frame epilogue in `flush()`.
104pub struct ZstdStreamCompressor {
105    inner: Option<zstd::stream::write::Encoder<'static, Vec<u8>>>,
106}
107
108impl ZstdStreamCompressor {
109    /// Creates a streaming zstd compressor configured to match Mars async settings.
110    pub fn new(level: i32) -> Result<Self, CompressError> {
111        let mut encoder = zstd::stream::write::Encoder::new(Vec::new(), level)
112            .map_err(|e| CompressError::Zstd(e.to_string()))?;
113        encoder
114            .window_log(16)
115            .map_err(|e| CompressError::Zstd(e.to_string()))?;
116        Ok(Self {
117            inner: Some(encoder),
118        })
119    }
120}
121
122impl StreamCompressor for ZstdStreamCompressor {
123    fn compress_chunk(&mut self, input: &[u8], output: &mut Vec<u8>) -> Result<(), CompressError> {
124        let Some(encoder) = self.inner.as_mut() else {
125            return Err(CompressError::Zstd(
126                "zstd stream compressor already finished".to_string(),
127            ));
128        };
129        encoder
130            .write_all(input)
131            .map_err(|e| CompressError::Zstd(e.to_string()))?;
132        encoder
133            .flush()
134            .map_err(|e| CompressError::Zstd(e.to_string()))?;
135        drain_encoded_bytes(encoder.get_mut(), output);
136        Ok(())
137    }
138
139    fn flush(&mut self, output: &mut Vec<u8>) -> Result<(), CompressError> {
140        let Some(encoder) = self.inner.take() else {
141            return Ok(());
142        };
143        let mut encoded = encoder
144            .finish()
145            .map_err(|e| CompressError::Zstd(e.to_string()))?;
146        drain_encoded_bytes(&mut encoded, output);
147        Ok(())
148    }
149}
150
151fn drain_encoded_bytes(encoded: &mut Vec<u8>, output: &mut Vec<u8>) {
152    if encoded.is_empty() {
153        return;
154    }
155    output.extend_from_slice(encoded);
156    encoded.clear();
157}
158
159/// Decompresses raw-deflate data produced by the zlib stream compressor.
160pub fn decompress_raw_zlib(input: &[u8]) -> Result<Vec<u8>, CompressError> {
161    let mut decoder = flate2::read::DeflateDecoder::new(input);
162    let mut out = Vec::new();
163    decoder
164        .read_to_end(&mut out)
165        .map_err(|e| CompressError::ZlibDecompress(e.to_string()))?;
166    Ok(out)
167}
168
169/// Decompresses one or more concatenated zstd frames into a single byte buffer.
170pub fn decompress_zstd_frames(input: &[u8]) -> Result<Vec<u8>, CompressError> {
171    let mut reader = std::io::Cursor::new(input);
172    let mut decoder = zstd::stream::Decoder::new(&mut reader)
173        .map_err(|e| CompressError::ZstdDecompress(e.to_string()))?;
174    let mut out = Vec::new();
175    decoder
176        .read_to_end(&mut out)
177        .map_err(|e| CompressError::ZstdDecompress(e.to_string()))?;
178    Ok(out)
179}
180
181#[cfg(test)]
182mod tests {
183    use super::{
184        decompress_raw_zlib, decompress_zstd_frames, CompressError, StreamCompressor,
185        ZlibStreamCompressor, ZstdChunkCompressor, ZstdStreamCompressor,
186    };
187
188    #[test]
189    fn zlib_stream_roundtrips_incremental_chunks() {
190        let mut compressor = ZlibStreamCompressor::new(6);
191        let mut encoded = Vec::new();
192
193        compressor.compress_chunk(b"mars", &mut encoded).unwrap();
194        compressor.compress_chunk(b" xlog", &mut encoded).unwrap();
195        compressor.flush(&mut encoded).unwrap();
196
197        assert_eq!(decompress_raw_zlib(&encoded).unwrap(), b"mars xlog");
198    }
199
200    #[test]
201    fn stream_compressors_release_emitted_bytes_after_each_chunk() {
202        let mut zlib = ZlibStreamCompressor::new(6);
203        let mut zlib_out = Vec::new();
204        zlib.compress_chunk(b"mars", &mut zlib_out).unwrap();
205        assert!(zlib.inner.get_ref().is_empty());
206        zlib.compress_chunk(b" xlog", &mut zlib_out).unwrap();
207        assert!(zlib.inner.get_ref().is_empty());
208
209        let mut zstd = ZstdStreamCompressor::new(3).unwrap();
210        let mut zstd_out = Vec::new();
211        zstd.compress_chunk(b"mars", &mut zstd_out).unwrap();
212        assert!(zstd.inner.as_ref().unwrap().get_ref().is_empty());
213        zstd.compress_chunk(b" xlog", &mut zstd_out).unwrap();
214        assert!(zstd.inner.as_ref().unwrap().get_ref().is_empty());
215    }
216
217    #[test]
218    fn zstd_chunk_roundtrips_concatenated_frames() {
219        let mut compressor = ZstdChunkCompressor::new(3);
220        let mut encoded = Vec::new();
221
222        compressor.compress_chunk(b"mars", &mut encoded).unwrap();
223        compressor.compress_chunk(b" xlog", &mut encoded).unwrap();
224        compressor.flush(&mut encoded).unwrap();
225
226        assert_eq!(decompress_zstd_frames(&encoded).unwrap(), b"mars xlog");
227    }
228
229    #[test]
230    fn zstd_stream_rejects_compress_after_finish_and_double_flush_is_noop() {
231        let mut compressor = ZstdStreamCompressor::new(3).unwrap();
232        let mut encoded = Vec::new();
233
234        compressor.compress_chunk(b"hello", &mut encoded).unwrap();
235        compressor.flush(&mut encoded).unwrap();
236        let finished = encoded.clone();
237
238        let err = compressor
239            .compress_chunk(b"world", &mut encoded)
240            .unwrap_err();
241        assert!(
242            matches!(err, CompressError::Zstd(message) if message.contains("already finished"))
243        );
244
245        compressor.flush(&mut encoded).unwrap();
246        assert_eq!(encoded, finished);
247    }
248
249    #[test]
250    fn invalid_decompression_maps_to_specific_error_variants() {
251        assert!(matches!(
252            decompress_raw_zlib(b"not-a-deflate-stream"),
253            Err(CompressError::ZlibDecompress(_))
254        ));
255        assert!(matches!(
256            decompress_zstd_frames(b"not-a-zstd-frame"),
257            Err(CompressError::ZstdDecompress(_))
258        ));
259    }
260}
mars_xlog_core/compress.rs

mars_xlog_core/
compress.rs
