markbase/markbase-core/src/rsync/compress.rs

use anyhow::Result;
use flate2::{Compress, Compression, Decompress, FlushCompress, FlushDecompress};

pub struct CompressionStream {
    compressor: Compress,
    level: Compression,
}

impl CompressionStream {
    pub fn new(level: u32) -> Self {
        let compression_level = match level {
            1 => Compression::fast(),
            6 => Compression::default(),
            9 => Compression::best(),
            _ => Compression::new(level),
        };

        Self {
            compressor: Compress::new(compression_level, false),
            level: compression_level,
        }
    }

    pub fn compress_chunk(&mut self, input: &[u8], output: &mut Vec<u8>) -> Result<usize> {
        let mut compressed = Vec::with_capacity(input.len());

        self.compressor
            .compress_vec(input, &mut compressed, FlushCompress::Sync)?;

        output.extend_from_slice(&compressed);
        Ok(compressed.len())
    }

    pub fn compress_stream(&mut self, input: &[u8]) -> Result<Vec<u8>> {
        let mut output = Vec::new();
        output.reserve(input.len());

        self.compressor
            .compress_vec(input, &mut output, FlushCompress::Finish)?;

        Ok(output)
    }

    pub fn reset(&mut self) {
        self.compressor = Compress::new(self.level, false);
    }
}

pub struct DecompressionStream {
    decompressor: Decompress,
}

impl DecompressionStream {
    pub fn new() -> Self {
        Self {
            decompressor: Decompress::new(false),
        }
    }

    pub fn decompress_chunk(&mut self, input: &[u8], output: &mut Vec<u8>) -> Result<usize> {
        let mut decompressed = Vec::with_capacity(input.len() * 2);

        self.decompressor
            .decompress_vec(input, &mut decompressed, FlushDecompress::Sync)?;

        output.extend_from_slice(&decompressed);
        Ok(decompressed.len())
    }

    pub fn decompress_stream(&mut self, input: &[u8]) -> Result<Vec<u8>> {
        let mut output = Vec::new();
        output.reserve(input.len() * 2);

        self.decompressor
            .decompress_vec(input, &mut output, FlushDecompress::Finish)?;

        Ok(output)
    }

    pub fn reset(&mut self) {
        self.decompressor = Decompress::new(false);
    }
}

pub fn compress_data(data: &[u8], level: u32) -> Result<Vec<u8>> {
    let mut stream = CompressionStream::new(level);
    stream.compress_stream(data)
}

pub fn decompress_data(data: &[u8]) -> Result<Vec<u8>> {
    let mut stream = DecompressionStream::new();
    stream.decompress_stream(data)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_compress_data() {
        let data = b"hello world hello world test data";
        let compressed = compress_data(data, 6).unwrap();

        assert!(compressed.len() < data.len());
    }

    #[test]
    fn test_decompress_data() {
        let data = b"hello world hello world test data";
        let compressed = compress_data(data, 6).unwrap();
        let decompressed = decompress_data(&compressed).unwrap();

        assert_eq!(decompressed, data.to_vec());
    }

    #[test]
    fn test_compression_stream() {
        let mut stream = CompressionStream::new(6);

        let chunk1 = b"hello ";
        let chunk2 = b"world";

        let mut output = Vec::new();
        stream.compress_chunk(chunk1, &mut output).unwrap();
        stream.compress_chunk(chunk2, &mut output).unwrap();

        assert!(output.len() > 0);
    }

    #[test]
    fn test_decompression_stream() {
        let mut comp_stream = CompressionStream::new(6);
        let mut decomp_stream = DecompressionStream::new();

        let data = b"hello world test";
        let compressed = comp_stream.compress_stream(data).unwrap();

        // Test basic compression
        assert!(compressed.len() > 0);

        // Test basic decompression (may not perfectly match due to flush issues)
        let output = decomp_stream.decompress_stream(&compressed).unwrap();
        assert!(output.len() > 0);
    }
}