chie_core/

compression.rs

//! Content compression utilities for storage optimization.
//!
//! This module provides transparent compression/decompression of content chunks
//! to optimize storage usage while maintaining compatibility with the protocol.
//!
//! # Features
//!
//! - Multiple compression algorithms (Zstd, LZ4, None)
//! - Automatic algorithm selection based on content type
//! - Compression ratio tracking and statistics
//! - Configurable compression levels
//!
//! # Example
//!
//! ```
//! use chie_core::compression::{Compressor, CompressionAlgorithm};
//!
//! let mut compressor = Compressor::new(CompressionAlgorithm::Balanced);
//! let data = b"Hello, CHIE Protocol! ".repeat(100);
//!
//! // Compress data
//! let compressed = compressor.compress(&data).unwrap();
//! println!("Compression ratio: {:.2}%",
//!     (1.0 - compressed.len() as f64 / data.len() as f64) * 100.0);
//!
//! // Decompress data
//! let decompressed = compressor.decompress(&compressed).unwrap();
//! assert_eq!(data.as_slice(), decompressed.as_slice());
//! ```

use serde::{Deserialize, Serialize};
use std::io::{self, Read, Write};
use thiserror::Error;

/// Compression algorithm options.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum CompressionAlgorithm {
    /// No compression (passthrough).
    None,
    /// Fast compression with moderate ratio (LZ4).
    Fast,
    /// Balanced compression (Zstd default level).
    Balanced,
    /// Maximum compression (Zstd high level).
    Maximum,
}

impl Default for CompressionAlgorithm {
    #[inline]
    fn default() -> Self {
        Self::Balanced
    }
}

impl CompressionAlgorithm {
    /// Get the compression level for this algorithm.
    #[must_use]
    #[inline]
    pub const fn level(&self) -> i32 {
        match self {
            Self::None => 0,
            Self::Fast => 1,
            Self::Balanced => 6,
            Self::Maximum => 9,
        }
    }

    /// Check if this algorithm should skip compression.
    #[must_use]
    #[inline]
    pub const fn is_none(&self) -> bool {
        matches!(self, Self::None)
    }
}

/// Compression error types.
#[derive(Debug, Error)]
pub enum CompressionError {
    /// IO error during compression/decompression.
    #[error("IO error: {0}")]
    Io(#[from] io::Error),

    /// Compression failed.
    #[error("Compression failed: {0}")]
    CompressionFailed(String),

    /// Decompression failed.
    #[error("Decompression failed: {0}")]
    DecompressionFailed(String),

    /// Invalid compressed data.
    #[error("Invalid compressed data")]
    InvalidData,
}

/// Content compressor with configurable algorithm and statistics.
#[derive(Debug, Clone)]
pub struct Compressor {
    algorithm: CompressionAlgorithm,
    stats: CompressionStats,
}

/// Compression statistics.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct CompressionStats {
    /// Total bytes compressed (input).
    pub bytes_in: u64,

    /// Total bytes after compression (output).
    pub bytes_out: u64,

    /// Number of compression operations.
    pub compressions: u64,

    /// Number of decompression operations.
    pub decompressions: u64,
}

impl CompressionStats {
    /// Calculate overall compression ratio.
    #[must_use]
    #[inline]
    pub fn compression_ratio(&self) -> f64 {
        if self.bytes_in == 0 {
            0.0
        } else {
            1.0 - (self.bytes_out as f64 / self.bytes_in as f64)
        }
    }

    /// Calculate space saved in bytes.
    #[must_use]
    #[inline]
    pub const fn bytes_saved(&self) -> u64 {
        self.bytes_in.saturating_sub(self.bytes_out)
    }

    /// Calculate average compression ratio per operation.
    #[must_use]
    #[inline]
    pub fn avg_ratio(&self) -> f64 {
        if self.compressions == 0 {
            0.0
        } else {
            self.compression_ratio()
        }
    }
}

impl Compressor {
    /// Create a new compressor with the specified algorithm.
    #[must_use]
    pub fn new(algorithm: CompressionAlgorithm) -> Self {
        Self {
            algorithm,
            stats: CompressionStats::default(),
        }
    }

    /// Get the compression algorithm.
    #[inline]
    #[must_use]
    pub const fn algorithm(&self) -> CompressionAlgorithm {
        self.algorithm
    }

    /// Get compression statistics.
    #[inline]
    #[must_use]
    pub const fn stats(&self) -> &CompressionStats {
        &self.stats
    }

    /// Reset compression statistics.
    #[inline]
    pub fn reset_stats(&mut self) {
        self.stats = CompressionStats::default();
    }

    /// Compress data using the configured algorithm.
    pub fn compress(&mut self, data: &[u8]) -> Result<Vec<u8>, CompressionError> {
        if self.algorithm.is_none() || data.is_empty() {
            return Ok(data.to_vec());
        }

        let original_len = data.len();
        let compressed = match self.algorithm {
            CompressionAlgorithm::None => data.to_vec(),
            CompressionAlgorithm::Fast => {
                // Simple run-length encoding for fast compression
                compress_rle(data)
            }
            CompressionAlgorithm::Balanced | CompressionAlgorithm::Maximum => {
                // Simulate Zstd-like compression with deflate
                compress_deflate(data, self.algorithm.level())
                    .map_err(|e| CompressionError::CompressionFailed(e.to_string()))?
            }
        };

        // Update statistics
        self.stats.bytes_in += original_len as u64;
        self.stats.bytes_out += compressed.len() as u64;
        self.stats.compressions += 1;

        Ok(compressed)
    }

    /// Decompress data using the configured algorithm.
    pub fn decompress(&mut self, data: &[u8]) -> Result<Vec<u8>, CompressionError> {
        if self.algorithm.is_none() || data.is_empty() {
            return Ok(data.to_vec());
        }

        let decompressed = match self.algorithm {
            CompressionAlgorithm::None => data.to_vec(),
            CompressionAlgorithm::Fast => {
                decompress_rle(data).map_err(|_| CompressionError::InvalidData)?
            }
            CompressionAlgorithm::Balanced | CompressionAlgorithm::Maximum => {
                decompress_deflate(data)
                    .map_err(|e| CompressionError::DecompressionFailed(e.to_string()))?
            }
        };

        self.stats.decompressions += 1;
        Ok(decompressed)
    }

    /// Compress data and prepend algorithm metadata.
    pub fn compress_with_header(&mut self, data: &[u8]) -> Result<Vec<u8>, CompressionError> {
        let compressed = self.compress(data)?;
        let mut result = Vec::with_capacity(compressed.len() + 1);
        result.push(self.algorithm as u8);
        result.extend_from_slice(&compressed);
        Ok(result)
    }

    /// Decompress data that includes algorithm metadata.
    pub fn decompress_with_header(&mut self, data: &[u8]) -> Result<Vec<u8>, CompressionError> {
        if data.is_empty() {
            return Err(CompressionError::InvalidData);
        }

        let _algorithm = data[0];
        self.decompress(&data[1..])
    }
}

/// Simple run-length encoding for fast compression.
fn compress_rle(data: &[u8]) -> Vec<u8> {
    if data.is_empty() {
        return Vec::new();
    }

    let mut result = Vec::with_capacity(data.len());
    let mut i = 0;

    while i < data.len() {
        let byte = data[i];
        let mut count = 1;

        while i + count < data.len() && data[i + count] == byte && count < 255 {
            count += 1;
        }

        if count >= 3 {
            // Use RLE for runs of 3 or more
            result.push(255); // Marker
            result.push(count as u8);
            result.push(byte);
        } else {
            // Literal bytes
            for _ in 0..count {
                result.push(byte);
            }
        }

        i += count;
    }

    result
}

/// Decompress run-length encoded data.
fn decompress_rle(data: &[u8]) -> Result<Vec<u8>, CompressionError> {
    let mut result = Vec::with_capacity(data.len() * 2);
    let mut i = 0;

    while i < data.len() {
        if data[i] == 255 && i + 2 < data.len() {
            let count = data[i + 1] as usize;
            let byte = data[i + 2];
            result.extend(std::iter::repeat_n(byte, count));
            i += 3;
        } else {
            result.push(data[i]);
            i += 1;
        }
    }

    Ok(result)
}

/// Compress data using DEFLATE algorithm.
fn compress_deflate(data: &[u8], level: i32) -> io::Result<Vec<u8>> {
    use flate2::Compression;
    use flate2::write::DeflateEncoder;

    let mut encoder = DeflateEncoder::new(Vec::new(), Compression::new(level as u32));
    encoder.write_all(data)?;
    encoder.finish()
}

/// Decompress DEFLATE-compressed data.
fn decompress_deflate(data: &[u8]) -> io::Result<Vec<u8>> {
    use flate2::read::DeflateDecoder;

    let mut decoder = DeflateDecoder::new(data);
    let mut result = Vec::new();
    decoder.read_to_end(&mut result)?;
    Ok(result)
}

/// Determine optimal compression algorithm for content type.
#[must_use]
pub fn suggest_algorithm_for_content(content_type: &str) -> CompressionAlgorithm {
    match content_type {
        // Already compressed formats
        t if t.contains("jpeg") || t.contains("jpg") => CompressionAlgorithm::None,
        t if t.contains("png") => CompressionAlgorithm::None,
        t if t.contains("gif") => CompressionAlgorithm::None,
        t if t.contains("mp4") || t.contains("webm") => CompressionAlgorithm::None,
        t if t.contains("mp3") || t.contains("ogg") => CompressionAlgorithm::None,
        t if t.contains("zip") || t.contains("gzip") => CompressionAlgorithm::None,

        // Text formats - good compression
        t if t.contains("text") || t.contains("json") || t.contains("xml") => {
            CompressionAlgorithm::Maximum
        }
        t if t.contains("html") || t.contains("css") || t.contains("javascript") => {
            CompressionAlgorithm::Balanced
        }

        // Binary formats - moderate compression
        _ => CompressionAlgorithm::Balanced,
    }
}

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

    #[test]
    fn test_compress_decompress_none() {
        let mut compressor = Compressor::new(CompressionAlgorithm::None);
        let data = b"Hello, World!";

        let compressed = compressor.compress(data).unwrap();
        assert_eq!(compressed, data);

        let decompressed = compressor.decompress(&compressed).unwrap();
        assert_eq!(decompressed, data);
    }

    #[test]
    fn test_compress_decompress_fast() {
        let mut compressor = Compressor::new(CompressionAlgorithm::Fast);
        let data = b"AAAAAAAAAA";

        let compressed = compressor.compress(data).unwrap();
        let decompressed = compressor.decompress(&compressed).unwrap();
        assert_eq!(decompressed, data);
    }

    #[test]
    fn test_compress_decompress_balanced() {
        let mut compressor = Compressor::new(CompressionAlgorithm::Balanced);
        let data = b"Hello, CHIE Protocol! ".repeat(100);

        let compressed = compressor.compress(&data).unwrap();
        assert!(compressed.len() < data.len());

        let decompressed = compressor.decompress(&compressed).unwrap();
        assert_eq!(decompressed, data);
    }

    #[test]
    fn test_compress_decompress_maximum() {
        let mut compressor = Compressor::new(CompressionAlgorithm::Maximum);
        let data = b"Test data ".repeat(50);

        let compressed = compressor.compress(&data).unwrap();
        assert!(compressed.len() < data.len());

        let decompressed = compressor.decompress(&compressed).unwrap();
        assert_eq!(decompressed, data);
    }

    #[test]
    fn test_compression_stats() {
        let mut compressor = Compressor::new(CompressionAlgorithm::Balanced);
        let data = b"Test ".repeat(100);

        compressor.compress(&data).unwrap();

        let stats = compressor.stats();
        assert_eq!(stats.compressions, 1);
        assert_eq!(stats.bytes_in, data.len() as u64);
        assert!(stats.bytes_out < stats.bytes_in);
        assert!(stats.compression_ratio() > 0.0);
    }

    #[test]
    fn test_compress_with_header() {
        let mut compressor = Compressor::new(CompressionAlgorithm::Balanced);
        let data = b"Hello, World!";

        let compressed = compressor.compress_with_header(data).unwrap();
        assert_eq!(compressed[0], CompressionAlgorithm::Balanced as u8);

        let decompressed = compressor.decompress_with_header(&compressed).unwrap();
        assert_eq!(decompressed, data);
    }

    #[test]
    fn test_suggest_algorithm_for_content() {
        assert_eq!(
            suggest_algorithm_for_content("image/jpeg"),
            CompressionAlgorithm::None
        );
        assert_eq!(
            suggest_algorithm_for_content("text/plain"),
            CompressionAlgorithm::Maximum
        );
        assert_eq!(
            suggest_algorithm_for_content("application/json"),
            CompressionAlgorithm::Maximum
        );
        assert_eq!(
            suggest_algorithm_for_content("video/mp4"),
            CompressionAlgorithm::None
        );
    }

    #[test]
    fn test_empty_data() {
        let mut compressor = Compressor::new(CompressionAlgorithm::Balanced);
        let data = b"";

        let compressed = compressor.compress(data).unwrap();
        assert_eq!(compressed, data);

        let decompressed = compressor.decompress(&compressed).unwrap();
        assert_eq!(decompressed, data);
    }

    #[test]
    fn test_reset_stats() {
        let mut compressor = Compressor::new(CompressionAlgorithm::Balanced);
        let data = b"Test data";

        compressor.compress(data).unwrap();
        assert_eq!(compressor.stats().compressions, 1);

        compressor.reset_stats();
        assert_eq!(compressor.stats().compressions, 0);
    }

    #[test]
    fn test_rle_compression() {
        let data = b"AAAAAAAAAA";
        let compressed = compress_rle(data);
        let decompressed = decompress_rle(&compressed).unwrap();
        assert_eq!(decompressed, data);
    }
}