featherdb_storage/

compression.rs

//! Adaptive Page Compression for FeatherDB
//!
//! This module provides transparent compression for database pages,
//! supporting multiple compression algorithms:
//! - LZ4: Fast compression/decompression, moderate compression ratio
//! - ZSTD: Higher compression ratio, configurable levels
//!
//! Compression is adaptive and skips pages smaller than a configurable threshold.

use featherdb_core::{Error, Result};
use std::sync::atomic::{AtomicU64, Ordering};

/// Compression algorithm type
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum CompressionType {
    /// No compression
    #[default]
    None,
    /// LZ4 - fast compression with moderate ratio
    Lz4,
    /// ZSTD - high compression ratio with configurable level
    /// Level ranges from 1 (fastest) to 22 (best compression)
    /// Default level is 3 for a good balance
    Zstd { level: i32 },
}

impl CompressionType {
    /// Create ZSTD compression with default level (3)
    pub fn zstd_default() -> Self {
        CompressionType::Zstd { level: 3 }
    }

    /// Create ZSTD compression with high ratio (level 9)
    pub fn zstd_high() -> Self {
        CompressionType::Zstd { level: 9 }
    }

    /// Create ZSTD compression with maximum ratio (level 19)
    pub fn zstd_max() -> Self {
        CompressionType::Zstd { level: 19 }
    }

    /// Get the byte identifier for this compression type
    pub fn to_byte(&self) -> u8 {
        match self {
            CompressionType::None => 0,
            CompressionType::Lz4 => 1,
            CompressionType::Zstd { .. } => 2,
        }
    }

    /// Create compression type from byte identifier
    pub fn from_byte(byte: u8, level: i32) -> Self {
        match byte {
            0 => CompressionType::None,
            1 => CompressionType::Lz4,
            2 => CompressionType::Zstd { level },
            _ => CompressionType::None,
        }
    }
}

/// Trait for compression implementations
pub trait Compressor: Send + Sync {
    /// Compress data
    fn compress(&self, data: &[u8]) -> Result<Vec<u8>>;

    /// Decompress data to expected size
    fn decompress(&self, data: &[u8], expected_size: usize) -> Result<Vec<u8>>;

    /// Get the compression type
    fn compression_type(&self) -> CompressionType;
}

/// No-op compressor (passthrough)
pub struct NoCompressor;

impl Compressor for NoCompressor {
    fn compress(&self, data: &[u8]) -> Result<Vec<u8>> {
        Ok(data.to_vec())
    }

    fn decompress(&self, data: &[u8], _expected_size: usize) -> Result<Vec<u8>> {
        Ok(data.to_vec())
    }

    fn compression_type(&self) -> CompressionType {
        CompressionType::None
    }
}

/// LZ4 compressor - optimized for speed
pub struct Lz4Compressor;

impl Compressor for Lz4Compressor {
    fn compress(&self, data: &[u8]) -> Result<Vec<u8>> {
        Ok(lz4_flex::compress_prepend_size(data))
    }

    fn decompress(&self, data: &[u8], expected_size: usize) -> Result<Vec<u8>> {
        lz4_flex::decompress_size_prepended(data).map_err(|e| Error::DecompressionError {
            message: format!("LZ4 decompression failed: {}", e),
            expected_size,
        })
    }

    fn compression_type(&self) -> CompressionType {
        CompressionType::Lz4
    }
}

/// ZSTD compressor - optimized for compression ratio
pub struct ZstdCompressor {
    level: i32,
}

impl ZstdCompressor {
    pub fn new(level: i32) -> Self {
        // Clamp level to valid range (1-22)
        let level = level.clamp(1, 22);
        ZstdCompressor { level }
    }
}

impl Compressor for ZstdCompressor {
    fn compress(&self, data: &[u8]) -> Result<Vec<u8>> {
        zstd::encode_all(data, self.level).map_err(|e| Error::CompressionError {
            message: format!("ZSTD compression failed: {}", e),
        })
    }

    fn decompress(&self, data: &[u8], expected_size: usize) -> Result<Vec<u8>> {
        zstd::decode_all(data).map_err(|e| Error::DecompressionError {
            message: format!("ZSTD decompression failed: {}", e),
            expected_size,
        })
    }

    fn compression_type(&self) -> CompressionType {
        CompressionType::Zstd { level: self.level }
    }
}

/// Factory to create compressors based on type
pub fn create_compressor(compression_type: CompressionType) -> Box<dyn Compressor> {
    match compression_type {
        CompressionType::None => Box::new(NoCompressor),
        CompressionType::Lz4 => Box::new(Lz4Compressor),
        CompressionType::Zstd { level } => Box::new(ZstdCompressor::new(level)),
    }
}

/// Statistics for compression operations
#[derive(Debug, Default)]
pub struct CompressionStats {
    /// Total bytes before compression
    pub bytes_before_compression: AtomicU64,
    /// Total bytes after compression
    pub bytes_after_compression: AtomicU64,
    /// Number of pages compressed
    pub pages_compressed: AtomicU64,
    /// Number of pages skipped (below threshold or incompressible)
    pub pages_skipped: AtomicU64,
    /// Number of pages decompressed
    pub pages_decompressed: AtomicU64,
}

impl CompressionStats {
    /// Create new compression stats
    pub fn new() -> Self {
        Self::default()
    }

    /// Get the overall compression ratio (0.0-1.0, lower is better)
    /// Returns 1.0 if no compression has occurred
    pub fn compression_ratio(&self) -> f64 {
        let before = self.bytes_before_compression.load(Ordering::Relaxed);
        let after = self.bytes_after_compression.load(Ordering::Relaxed);

        if before == 0 {
            1.0
        } else {
            after as f64 / before as f64
        }
    }

    /// Get the space savings percentage (0-100%, higher is better)
    pub fn space_savings_percent(&self) -> f64 {
        (1.0 - self.compression_ratio()) * 100.0
    }

    /// Get a snapshot of current statistics
    pub fn snapshot(&self) -> CompressionStatsSnapshot {
        CompressionStatsSnapshot {
            bytes_before_compression: self.bytes_before_compression.load(Ordering::Relaxed),
            bytes_after_compression: self.bytes_after_compression.load(Ordering::Relaxed),
            pages_compressed: self.pages_compressed.load(Ordering::Relaxed),
            pages_skipped: self.pages_skipped.load(Ordering::Relaxed),
            pages_decompressed: self.pages_decompressed.load(Ordering::Relaxed),
        }
    }

    /// Record a compression operation
    pub fn record_compression(&self, original_size: usize, compressed_size: usize) {
        self.bytes_before_compression
            .fetch_add(original_size as u64, Ordering::Relaxed);
        self.bytes_after_compression
            .fetch_add(compressed_size as u64, Ordering::Relaxed);
        self.pages_compressed.fetch_add(1, Ordering::Relaxed);
    }

    /// Record a skipped page (not compressed)
    pub fn record_skipped(&self, size: usize) {
        self.bytes_before_compression
            .fetch_add(size as u64, Ordering::Relaxed);
        self.bytes_after_compression
            .fetch_add(size as u64, Ordering::Relaxed);
        self.pages_skipped.fetch_add(1, Ordering::Relaxed);
    }

    /// Record a decompression operation
    pub fn record_decompression(&self) {
        self.pages_decompressed.fetch_add(1, Ordering::Relaxed);
    }

    /// Reset all statistics
    pub fn reset(&self) {
        self.bytes_before_compression.store(0, Ordering::Relaxed);
        self.bytes_after_compression.store(0, Ordering::Relaxed);
        self.pages_compressed.store(0, Ordering::Relaxed);
        self.pages_skipped.store(0, Ordering::Relaxed);
        self.pages_decompressed.store(0, Ordering::Relaxed);
    }
}

/// Immutable snapshot of compression statistics
#[derive(Debug, Clone, Copy)]
pub struct CompressionStatsSnapshot {
    pub bytes_before_compression: u64,
    pub bytes_after_compression: u64,
    pub pages_compressed: u64,
    pub pages_skipped: u64,
    pub pages_decompressed: u64,
}

impl CompressionStatsSnapshot {
    /// Get the overall compression ratio (0.0-1.0, lower is better)
    /// Returns 1.0 if no compression has occurred
    pub fn compression_ratio(&self) -> f64 {
        if self.bytes_before_compression == 0 {
            1.0
        } else {
            self.bytes_after_compression as f64 / self.bytes_before_compression as f64
        }
    }

    /// Get the space savings percentage (0-100%, higher is better)
    pub fn space_savings_percent(&self) -> f64 {
        (1.0 - self.compression_ratio()) * 100.0
    }
}

impl Clone for CompressionStats {
    fn clone(&self) -> Self {
        CompressionStats {
            bytes_before_compression: AtomicU64::new(
                self.bytes_before_compression.load(Ordering::Relaxed),
            ),
            bytes_after_compression: AtomicU64::new(
                self.bytes_after_compression.load(Ordering::Relaxed),
            ),
            pages_compressed: AtomicU64::new(self.pages_compressed.load(Ordering::Relaxed)),
            pages_skipped: AtomicU64::new(self.pages_skipped.load(Ordering::Relaxed)),
            pages_decompressed: AtomicU64::new(self.pages_decompressed.load(Ordering::Relaxed)),
        }
    }
}

/// Compressed page header stored at the beginning of compressed pages
/// This is stored within the page data on disk
#[derive(Debug, Clone, Copy)]
pub struct CompressedPageHeader {
    /// Magic byte to identify compressed pages (0xC0 = "CO" for compressed)
    pub magic: u8,
    /// Compression type (0=none, 1=lz4, 2=zstd)
    pub compression_type: u8,
    /// Original uncompressed size (u32 to support up to 4GB pages)
    pub original_size: u32,
    /// Compressed data size (excluding this header)
    pub compressed_size: u32,
    /// ZSTD compression level (only meaningful for ZSTD)
    pub compression_level: u8,
    /// Reserved for future use
    pub reserved: u8,
}

impl CompressedPageHeader {
    /// Magic byte for compressed pages
    pub const MAGIC: u8 = 0xC0;

    /// Header size in bytes
    pub const SIZE: usize = 12;

    /// Create a new compressed page header
    pub fn new(
        compression_type: CompressionType,
        original_size: usize,
        compressed_size: usize,
    ) -> Self {
        let (comp_type, level) = match compression_type {
            CompressionType::None => (0, 0),
            CompressionType::Lz4 => (1, 0),
            CompressionType::Zstd { level } => (2, level as u8),
        };

        CompressedPageHeader {
            magic: Self::MAGIC,
            compression_type: comp_type,
            original_size: original_size as u32,
            compressed_size: compressed_size as u32,
            compression_level: level,
            reserved: 0,
        }
    }

    /// Serialize header to bytes
    pub fn to_bytes(&self) -> [u8; Self::SIZE] {
        let mut buf = [0u8; Self::SIZE];
        buf[0] = self.magic;
        buf[1] = self.compression_type;
        buf[2..6].copy_from_slice(&self.original_size.to_le_bytes());
        buf[6..10].copy_from_slice(&self.compressed_size.to_le_bytes());
        buf[10] = self.compression_level;
        buf[11] = self.reserved;
        buf
    }

    /// Deserialize header from bytes
    pub fn from_bytes(data: &[u8]) -> Option<Self> {
        if data.len() < Self::SIZE {
            return None;
        }

        if data[0] != Self::MAGIC {
            return None;
        }

        Some(CompressedPageHeader {
            magic: data[0],
            compression_type: data[1],
            original_size: u32::from_le_bytes(data[2..6].try_into().ok()?),
            compressed_size: u32::from_le_bytes(data[6..10].try_into().ok()?),
            compression_level: data[10],
            reserved: data[11],
        })
    }

    /// Get the compression type enum
    pub fn get_compression_type(&self) -> CompressionType {
        CompressionType::from_byte(self.compression_type, self.compression_level as i32)
    }

    /// Check if this is a valid compressed page header
    pub fn is_valid(&self) -> bool {
        self.magic == Self::MAGIC && self.compression_type <= 2
    }
}

/// Page compression manager
/// Handles compression/decompression of pages with configurable settings
pub struct PageCompressor {
    /// The compressor implementation
    compressor: Box<dyn Compressor>,
    /// Minimum page size to compress (pages smaller than this are stored uncompressed)
    threshold: usize,
    /// Compression statistics
    pub stats: CompressionStats,
}

impl PageCompressor {
    /// Default compression threshold (512 bytes)
    pub const DEFAULT_THRESHOLD: usize = 512;

    /// Create a new page compressor
    pub fn new(compression_type: CompressionType, threshold: usize) -> Self {
        PageCompressor {
            compressor: create_compressor(compression_type),
            threshold,
            stats: CompressionStats::new(),
        }
    }

    /// Create a page compressor with default settings (LZ4, 512 byte threshold)
    pub fn default_lz4() -> Self {
        Self::new(CompressionType::Lz4, Self::DEFAULT_THRESHOLD)
    }

    /// Create a page compressor with ZSTD and default settings
    pub fn default_zstd() -> Self {
        Self::new(CompressionType::zstd_default(), Self::DEFAULT_THRESHOLD)
    }

    /// Compress page data for writing to disk
    /// Returns the compressed data with header, or original data if compression is not beneficial
    pub fn compress_page(&self, data: &[u8]) -> Result<Vec<u8>> {
        let compression_type = self.compressor.compression_type();

        // Skip compression for small pages or if compression is disabled
        if data.len() < self.threshold || matches!(compression_type, CompressionType::None) {
            self.stats.record_skipped(data.len());
            return Ok(data.to_vec());
        }

        // Compress the data
        let compressed = self.compressor.compress(data)?;

        // Check if compression is beneficial (at least 10% smaller)
        // Also ensure compressed + header is smaller than original
        let total_compressed_size = CompressedPageHeader::SIZE + compressed.len();
        if total_compressed_size >= data.len() || compressed.len() >= (data.len() * 9 / 10) {
            // Compression not beneficial, store uncompressed
            self.stats.record_skipped(data.len());
            return Ok(data.to_vec());
        }

        // Create header and prepend to compressed data
        let header = CompressedPageHeader::new(compression_type, data.len(), compressed.len());
        let mut result = Vec::with_capacity(total_compressed_size);
        result.extend_from_slice(&header.to_bytes());
        result.extend_from_slice(&compressed);

        self.stats.record_compression(data.len(), result.len());
        Ok(result)
    }

    /// Decompress page data read from disk
    /// Automatically detects if the page is compressed and decompresses if needed
    pub fn decompress_page(&self, data: &[u8], page_size: usize) -> Result<Vec<u8>> {
        // Check if this is compressed data by looking for the magic byte
        if let Some(header) = CompressedPageHeader::from_bytes(data) {
            if header.is_valid() {
                // Extract compressed data (skip header)
                let compressed_data = &data[CompressedPageHeader::SIZE..];
                if compressed_data.len() < header.compressed_size as usize {
                    return Err(Error::DecompressionError {
                        message: "Compressed data truncated".into(),
                        expected_size: header.compressed_size as usize,
                    });
                }

                let compressed_data = &compressed_data[..header.compressed_size as usize];

                // Create appropriate decompressor
                let decompressor = create_compressor(header.get_compression_type());
                let decompressed =
                    decompressor.decompress(compressed_data, header.original_size as usize)?;

                // Verify size
                if decompressed.len() != header.original_size as usize {
                    return Err(Error::DecompressionError {
                        message: format!(
                            "Decompressed size mismatch: expected {}, got {}",
                            header.original_size,
                            decompressed.len()
                        ),
                        expected_size: header.original_size as usize,
                    });
                }

                self.stats.record_decompression();
                return Ok(decompressed);
            }
        }

        // Not compressed, return as-is (padded to page size if needed)
        if data.len() < page_size {
            let mut result = data.to_vec();
            result.resize(page_size, 0);
            Ok(result)
        } else {
            Ok(data[..page_size].to_vec())
        }
    }

    /// Get the compression type
    pub fn compression_type(&self) -> CompressionType {
        self.compressor.compression_type()
    }

    /// Get the compression threshold
    pub fn threshold(&self) -> usize {
        self.threshold
    }

    /// Get compression statistics
    pub fn stats(&self) -> &CompressionStats {
        &self.stats
    }
}

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

    #[test]
    fn test_compression_type_byte_conversion() {
        assert_eq!(CompressionType::None.to_byte(), 0);
        assert_eq!(CompressionType::Lz4.to_byte(), 1);
        assert_eq!(CompressionType::Zstd { level: 5 }.to_byte(), 2);

        assert_eq!(CompressionType::from_byte(0, 0), CompressionType::None);
        assert_eq!(CompressionType::from_byte(1, 0), CompressionType::Lz4);
        assert_eq!(
            CompressionType::from_byte(2, 5),
            CompressionType::Zstd { level: 5 }
        );
    }

    #[test]
    fn test_lz4_compression() {
        let compressor = Lz4Compressor;
        let data = b"Hello, World! This is a test of LZ4 compression. ".repeat(100);

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

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

    #[test]
    fn test_zstd_compression() {
        let compressor = ZstdCompressor::new(3);
        let data = b"Hello, World! This is a test of ZSTD compression. ".repeat(100);

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

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

    #[test]
    fn test_zstd_levels() {
        let data = b"Test data for compression level comparison. ".repeat(100);

        // Higher levels should produce smaller output (usually)
        let low = ZstdCompressor::new(1).compress(&data).unwrap();
        let high = ZstdCompressor::new(19).compress(&data).unwrap();

        // Both should decompress correctly
        let dec_low = ZstdCompressor::new(1).decompress(&low, data.len()).unwrap();
        let dec_high = ZstdCompressor::new(19)
            .decompress(&high, data.len())
            .unwrap();

        assert_eq!(dec_low, data.to_vec());
        assert_eq!(dec_high, data.to_vec());
    }

    #[test]
    fn test_compressed_page_header() {
        let header = CompressedPageHeader::new(CompressionType::Lz4, 4096, 2048);

        let bytes = header.to_bytes();
        let parsed = CompressedPageHeader::from_bytes(&bytes).unwrap();

        assert_eq!(parsed.magic, CompressedPageHeader::MAGIC);
        assert_eq!(parsed.compression_type, 1); // LZ4
        assert_eq!(parsed.original_size, 4096);
        assert_eq!(parsed.compressed_size, 2048);
        assert!(parsed.is_valid());
    }

    #[test]
    fn test_page_compressor_small_page() {
        let compressor = PageCompressor::new(CompressionType::Lz4, 512);
        let small_data = vec![0u8; 256]; // Below threshold

        let result = compressor.compress_page(&small_data).unwrap();
        assert_eq!(result, small_data); // Should be unchanged
        assert_eq!(compressor.stats.pages_skipped.load(Ordering::Relaxed), 1);
    }

    #[test]
    fn test_page_compressor_compressible_data() {
        let compressor = PageCompressor::new(CompressionType::Lz4, 512);
        let data = vec![0xABu8; 4096]; // Highly compressible

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

        // Should have compression header
        let header = CompressedPageHeader::from_bytes(&compressed).unwrap();
        assert!(header.is_valid());
        assert_eq!(header.original_size, 4096);

        // Decompress and verify
        let decompressed = compressor.decompress_page(&compressed, 4096).unwrap();
        assert_eq!(decompressed, data);
    }

    #[test]
    fn test_page_compressor_incompressible_data() {
        let compressor = PageCompressor::new(CompressionType::Lz4, 512);

        // Random-ish data that doesn't compress well
        let data: Vec<u8> = (0..4096).map(|i| (i * 17 + i / 3) as u8).collect();

        let result = compressor.compress_page(&data).unwrap();

        // Decompress should work regardless
        let decompressed = compressor.decompress_page(&result, 4096).unwrap();
        assert_eq!(decompressed, data);
    }

    #[test]
    fn test_compression_stats() {
        let stats = CompressionStats::new();

        stats.record_compression(4096, 2048);
        stats.record_compression(4096, 1024);

        assert_eq!(stats.bytes_before_compression.load(Ordering::Relaxed), 8192);
        assert_eq!(stats.bytes_after_compression.load(Ordering::Relaxed), 3072);
        assert_eq!(stats.pages_compressed.load(Ordering::Relaxed), 2);

        // 3072 / 8192 = 0.375
        assert!((stats.compression_ratio() - 0.375).abs() < 0.001);

        // 62.5% space savings
        assert!((stats.space_savings_percent() - 62.5).abs() < 0.1);
    }

    #[test]
    fn test_no_compression() {
        let compressor = PageCompressor::new(CompressionType::None, 512);
        let data = vec![0xABu8; 4096];

        let result = compressor.compress_page(&data).unwrap();
        assert_eq!(result, data); // Should be unchanged

        let decompressed = compressor.decompress_page(&result, 4096).unwrap();
        assert_eq!(decompressed, data);
    }
}