use crate::error::{EdgeError, Result};
use bytes::Bytes;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum CompressionLevel {
Fast,
Balanced,
Best,
}
impl CompressionLevel {
pub fn lz4_level(&self) -> i32 {
match self {
Self::Fast => 1,
Self::Balanced => 4,
Self::Best => 9,
}
}
pub fn deflate_level_u8(&self) -> u8 {
match self {
Self::Fast => 1,
Self::Balanced => 6,
Self::Best => 9,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum CompressionStrategy {
Lz4,
Snappy,
Deflate,
None,
}
impl CompressionStrategy {
pub fn auto_select(data: &[u8]) -> Self {
if data.len() < 1024 {
Self::None
} else if Self::estimate_entropy(data) > 0.9 {
Self::None
} else if data.len() < 10 * 1024 {
Self::Snappy
} else {
Self::Lz4
}
}
fn estimate_entropy(data: &[u8]) -> f64 {
if data.is_empty() {
return 0.0;
}
let mut counts = [0u32; 256];
for &byte in data.iter().take(1024.min(data.len())) {
counts[byte as usize] = counts[byte as usize].saturating_add(1);
}
let len = data.len().min(1024) as f64;
let mut entropy = 0.0;
for &count in &counts {
if count > 0 {
let p = count as f64 / len;
entropy -= p * p.log2();
}
}
entropy / 8.0 }
}
pub struct EdgeCompressor {
strategy: CompressionStrategy,
level: CompressionLevel,
}
impl EdgeCompressor {
pub fn new(strategy: CompressionStrategy, level: CompressionLevel) -> Self {
Self { strategy, level }
}
pub fn auto() -> Self {
Self {
strategy: CompressionStrategy::Lz4,
level: CompressionLevel::Balanced,
}
}
pub fn fast() -> Self {
Self {
strategy: CompressionStrategy::Snappy,
level: CompressionLevel::Fast,
}
}
pub fn best() -> Self {
Self {
strategy: CompressionStrategy::Deflate,
level: CompressionLevel::Best,
}
}
pub fn compress(&self, data: &[u8]) -> Result<Bytes> {
match self.strategy {
CompressionStrategy::None => Ok(Bytes::copy_from_slice(data)),
CompressionStrategy::Lz4 => self.compress_lz4(data),
CompressionStrategy::Snappy => self.compress_snappy(data),
CompressionStrategy::Deflate => self.compress_deflate(data),
}
}
pub fn decompress(&self, data: &[u8]) -> Result<Bytes> {
match self.strategy {
CompressionStrategy::None => Ok(Bytes::copy_from_slice(data)),
CompressionStrategy::Lz4 => self.decompress_lz4(data),
CompressionStrategy::Snappy => self.decompress_snappy(data),
CompressionStrategy::Deflate => self.decompress_deflate(data),
}
}
fn compress_lz4(&self, data: &[u8]) -> Result<Bytes> {
let compressed = oxiarc_lz4::compress_block_with_accel(data, self.level.lz4_level())
.map_err(|e| EdgeError::compression(e.to_string()))?;
let orig_size = data.len() as i32;
let mut result = Vec::with_capacity(4 + compressed.len());
result.extend_from_slice(&orig_size.to_le_bytes());
result.extend_from_slice(&compressed);
Ok(Bytes::from(result))
}
fn decompress_lz4(&self, data: &[u8]) -> Result<Bytes> {
if data.len() < 4 {
return Err(EdgeError::decompression("LZ4 data too short".to_string()));
}
let orig_size = i32::from_le_bytes([data[0], data[1], data[2], data[3]]) as usize;
let decompressed = oxiarc_lz4::decompress_block(&data[4..], orig_size)
.map_err(|e| EdgeError::decompression(e.to_string()))?;
Ok(Bytes::from(decompressed))
}
fn compress_snappy(&self, data: &[u8]) -> Result<Bytes> {
Ok(Bytes::from(oxiarc_snappy::compress(data)))
}
fn decompress_snappy(&self, data: &[u8]) -> Result<Bytes> {
oxiarc_snappy::decompress(data)
.map(Bytes::from)
.map_err(|e| EdgeError::decompression(e.to_string()))
}
fn compress_deflate(&self, data: &[u8]) -> Result<Bytes> {
oxiarc_deflate::deflate(data, self.level.deflate_level_u8())
.map(Bytes::from)
.map_err(|e| EdgeError::compression(e.to_string()))
}
fn decompress_deflate(&self, data: &[u8]) -> Result<Bytes> {
oxiarc_deflate::inflate(data)
.map(Bytes::from)
.map_err(|e| EdgeError::decompression(e.to_string()))
}
pub fn compression_ratio(&self, original_size: usize, compressed_size: usize) -> f64 {
if original_size == 0 {
return 0.0;
}
compressed_size as f64 / original_size as f64
}
pub fn estimate_compressed_size(&self, data: &[u8]) -> usize {
match self.strategy {
CompressionStrategy::None => data.len(),
CompressionStrategy::Snappy => {
(data.len() as f64 * 1.5) as usize
}
CompressionStrategy::Lz4 => {
data.len() + (data.len() / 255) + 16
}
CompressionStrategy::Deflate => {
(data.len() as f64 * 1.1) as usize
}
}
}
}
pub struct AdaptiveCompressor {
level: CompressionLevel,
}
impl AdaptiveCompressor {
pub fn new(level: CompressionLevel) -> Self {
Self { level }
}
pub fn compress(&self, data: &[u8]) -> Result<(Bytes, CompressionStrategy)> {
let strategy = CompressionStrategy::auto_select(data);
let compressor = EdgeCompressor::new(strategy, self.level);
let compressed = compressor.compress(data)?;
Ok((compressed, strategy))
}
pub fn decompress(&self, data: &[u8], strategy: CompressionStrategy) -> Result<Bytes> {
let compressor = EdgeCompressor::new(strategy, self.level);
compressor.decompress(data)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompressedData {
pub strategy: CompressionStrategy,
pub original_size: usize,
pub compressed_size: usize,
pub data: Vec<u8>,
}
impl CompressedData {
pub fn new(strategy: CompressionStrategy, original_size: usize, data: Bytes) -> Self {
let compressed_size = data.len();
Self {
strategy,
original_size,
compressed_size,
data: data.to_vec(),
}
}
pub fn ratio(&self) -> f64 {
if self.original_size == 0 {
return 0.0;
}
self.compressed_size as f64 / self.original_size as f64
}
pub fn space_saved(&self) -> usize {
self.original_size.saturating_sub(self.compressed_size)
}
pub fn space_saved_percent(&self) -> f64 {
if self.original_size == 0 {
return 0.0;
}
(self.space_saved() as f64 / self.original_size as f64) * 100.0
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_compression_lz4() -> Result<()> {
let compressor = EdgeCompressor::new(CompressionStrategy::Lz4, CompressionLevel::Balanced);
let data = b"Hello, World! This is a test message for compression. \
Repeat this several times to make it worth compressing. \
Hello, World! This is a test message for compression.";
let compressed = compressor.compress(data)?;
let decompressed = compressor.decompress(&compressed)?;
assert_eq!(&decompressed[..], &data[..]);
Ok(())
}
#[test]
fn test_compression_snappy() -> Result<()> {
let compressor = EdgeCompressor::new(CompressionStrategy::Snappy, CompressionLevel::Fast);
let data = b"Hello, World! This is a test message for compression.";
let compressed = compressor.compress(data)?;
let decompressed = compressor.decompress(&compressed)?;
assert_eq!(&decompressed[..], &data[..]);
Ok(())
}
#[test]
fn test_compression_deflate() -> Result<()> {
let compressor = EdgeCompressor::new(CompressionStrategy::Deflate, CompressionLevel::Best);
let data = b"Hello, World! This is a test message for compression.";
let compressed = compressor.compress(data)?;
let decompressed = compressor.decompress(&compressed)?;
assert_eq!(&decompressed[..], &data[..]);
Ok(())
}
#[test]
fn test_compression_none() -> Result<()> {
let compressor = EdgeCompressor::new(CompressionStrategy::None, CompressionLevel::Fast);
let data = b"Hello, World!";
let compressed = compressor.compress(data)?;
assert_eq!(&compressed[..], &data[..]);
Ok(())
}
#[test]
fn test_adaptive_compression() -> Result<()> {
let compressor = AdaptiveCompressor::new(CompressionLevel::Balanced);
let data = b"Hello, World! This is a test message for adaptive compression.";
let (compressed, strategy) = compressor.compress(data)?;
let decompressed = compressor.decompress(&compressed, strategy)?;
assert_eq!(&decompressed[..], &data[..]);
Ok(())
}
#[test]
fn test_auto_select_strategy() {
let small_data = b"Hi";
let strategy = CompressionStrategy::auto_select(small_data);
assert_eq!(strategy, CompressionStrategy::None);
let medium_data = vec![0u8; 5000];
let strategy = CompressionStrategy::auto_select(&medium_data);
assert!(matches!(
strategy,
CompressionStrategy::Snappy | CompressionStrategy::Lz4
));
}
#[test]
fn test_compression_ratio() {
let compressor = EdgeCompressor::fast();
let ratio = compressor.compression_ratio(1000, 500);
assert_eq!(ratio, 0.5);
}
#[test]
fn test_compressed_data_metadata() -> Result<()> {
let original = b"Test data for compression. This message repeats. \
Test data for compression. This message repeats. \
Test data for compression. This message repeats.";
let compressor = EdgeCompressor::fast();
let compressed = compressor.compress(original)?;
let metadata = CompressedData::new(CompressionStrategy::Snappy, original.len(), compressed);
assert_eq!(metadata.original_size, original.len());
assert!(metadata.ratio() > 0.0);
assert!(
metadata.space_saved_percent() >= -100.0 && metadata.space_saved_percent() <= 100.0
);
Ok(())
}
}