use runmat_time::Instant;
use std::collections::HashMap;
#[cfg(all(feature = "compression", not(target_arch = "wasm32")))]
use std::convert::TryFrom;
#[cfg(all(feature = "compression", target_arch = "wasm32"))]
use std::io::{Cursor, Read};
use crate::format::{CompressionAlgorithm, CompressionInfo};
use crate::{SnapshotError, SnapshotResult};
#[cfg(all(feature = "compression", target_arch = "wasm32"))]
use ruzstd::decoding::StreamingDecoder;
pub struct CompressionEngine {
config: CompressionConfig,
stats: CompressionStats,
}
#[derive(Debug, Clone)]
pub struct CompressionConfig {
pub default_level: u32,
pub adaptive_selection: bool,
pub size_threshold: usize,
pub target_ratio: f64,
pub max_compression_time: std::time::Duration,
pub prefer_speed: bool,
}
#[derive(Debug, Clone, Default)]
pub struct CompressionStats {
pub attempts: HashMap<String, u64>,
pub ratios: HashMap<String, f64>,
pub times: HashMap<String, std::time::Duration>,
pub total_bytes: u64,
pub total_time: std::time::Duration,
}
#[derive(Debug)]
pub struct CompressionResult {
pub data: Vec<u8>,
pub info: CompressionInfo,
pub metrics: CompressionMetrics,
}
#[derive(Debug)]
pub struct CompressionMetrics {
pub compression_time: std::time::Duration,
pub compression_ratio: f64,
pub throughput: f64,
pub memory_usage: usize,
}
impl CompressionEngine {
pub fn new(config: CompressionConfig) -> Self {
Self {
config,
stats: CompressionStats::default(),
}
}
pub fn compress(&mut self, data: &[u8]) -> SnapshotResult<CompressionResult> {
if data.len() < self.config.size_threshold {
return Ok(CompressionResult {
data: data.to_vec(),
info: CompressionInfo {
algorithm: CompressionAlgorithm::None,
level: 0,
parameters: HashMap::new(),
},
metrics: CompressionMetrics {
compression_time: std::time::Duration::ZERO,
compression_ratio: 1.0,
throughput: 0.0,
memory_usage: data.len(),
},
});
}
let algorithm = if self.config.adaptive_selection {
self.select_optimal_algorithm(data)?
} else {
CompressionAlgorithm::Lz4 {
fast: self.config.prefer_speed,
}
};
self.compress_with_algorithm(data, algorithm)
}
pub fn decompress(&self, data: &[u8], info: &CompressionInfo) -> SnapshotResult<Vec<u8>> {
let start = Instant::now();
let result = match &info.algorithm {
CompressionAlgorithm::None => data.to_vec(),
#[cfg(feature = "compression")]
CompressionAlgorithm::Lz4 { .. } => {
let original_size = info
.parameters
.get("original_size")
.and_then(|s| s.parse::<usize>().ok())
.unwrap_or_else(|| {
data.len() * 4
});
decompress_lz4_block(data, original_size)?
}
#[cfg(feature = "compression")]
CompressionAlgorithm::Zstd { dictionary } => {
if dictionary.is_some() {
log::debug!(
"Dictionary decompression not available, using standard decompression"
);
}
decompress_zstd_block(data)?
}
#[cfg(not(feature = "compression"))]
_ => {
return Err(SnapshotError::Configuration {
message: "Compression feature not enabled".to_string(),
});
}
};
let duration = start.elapsed();
log::debug!(
"Decompressed {} bytes to {} bytes in {:?}",
data.len(),
result.len(),
duration
);
Ok(result)
}
pub fn select_optimal_algorithm(&self, data: &[u8]) -> SnapshotResult<CompressionAlgorithm> {
let characteristics = self.analyze_data(data);
if characteristics.entropy > 0.9 {
return Ok(CompressionAlgorithm::None);
}
if self.config.prefer_speed {
Ok(CompressionAlgorithm::Lz4 { fast: true })
} else if characteristics.repetition_ratio > 0.7 {
Ok(CompressionAlgorithm::Zstd { dictionary: None })
} else {
Ok(CompressionAlgorithm::Lz4 { fast: false })
}
}
pub fn compress_with_algorithm(
&mut self,
data: &[u8],
algorithm: CompressionAlgorithm,
) -> SnapshotResult<CompressionResult> {
let start = Instant::now();
let start_memory = self.estimate_memory_usage();
let (compressed_data, algorithm_used) = match algorithm {
CompressionAlgorithm::None => (data.to_vec(), algorithm),
#[cfg(feature = "compression")]
CompressionAlgorithm::Lz4 { fast } => {
let compressed = compress_lz4_block(data, fast)?;
(compressed, CompressionAlgorithm::Lz4 { fast })
}
#[cfg(feature = "compression")]
CompressionAlgorithm::Zstd { dictionary } => {
let level = if self.config.prefer_speed {
1
} else {
self.config.default_level as i32
};
if dictionary.is_some() {
log::debug!("Dictionary compression not available, using standard compression");
}
let compressed = compress_zstd_block(data, level)?;
(compressed, CompressionAlgorithm::Zstd { dictionary })
}
#[cfg(not(feature = "compression"))]
_ => {
return Err(SnapshotError::Configuration {
message: "Compression feature not enabled".to_string(),
});
}
};
let compression_time = start.elapsed();
let end_memory = self.estimate_memory_usage();
let compression_ratio = compressed_data.len() as f64 / data.len() as f64;
if compression_ratio > 0.95 && !matches!(algorithm_used, CompressionAlgorithm::None) {
log::debug!(
"Compression ratio {compression_ratio:.3} not effective, using uncompressed data"
);
return Ok(CompressionResult {
data: data.to_vec(),
info: CompressionInfo {
algorithm: CompressionAlgorithm::None,
level: 0,
parameters: HashMap::new(),
},
metrics: CompressionMetrics {
compression_time,
compression_ratio: 1.0,
throughput: data.len() as f64 / compression_time.as_secs_f64(),
memory_usage: end_memory.saturating_sub(start_memory),
},
});
}
self.update_stats(
&algorithm_used,
data.len(),
compressed_data.len(),
compression_time,
);
let metrics = CompressionMetrics {
compression_time,
compression_ratio,
throughput: data.len() as f64 / compression_time.as_secs_f64(),
memory_usage: end_memory.saturating_sub(start_memory),
};
let mut parameters = HashMap::new();
parameters.insert("original_size".to_string(), data.len().to_string());
parameters.insert(
"compressed_size".to_string(),
compressed_data.len().to_string(),
);
Ok(CompressionResult {
data: compressed_data,
info: CompressionInfo {
algorithm: algorithm_used,
level: self.config.default_level,
parameters,
},
metrics,
})
}
pub fn analyze_data(&self, data: &[u8]) -> DataCharacteristics {
let mut characteristics = DataCharacteristics {
entropy: 0.0,
repetition_ratio: 0.0,
pattern_density: 0.0,
ascii_ratio: 0.0,
};
if data.is_empty() {
return characteristics;
}
let mut byte_counts = [0u32; 256];
let mut ascii_count = 0;
for &byte in data {
byte_counts[byte as usize] += 1;
if byte.is_ascii() {
ascii_count += 1;
}
}
let len = data.len() as f64;
let mut entropy = 0.0;
for &count in &byte_counts {
if count > 0 {
let p = count as f64 / len;
entropy -= p * p.log2();
}
}
characteristics.entropy = entropy / 8.0; characteristics.ascii_ratio = ascii_count as f64 / len;
let mut repetition_count = 0;
let window_size = 64.min(data.len() / 2);
if data.len() > window_size * 2 {
for i in 0..data.len() - window_size {
for j in (i + window_size)..data.len() - window_size {
if data[i..i + window_size] == data[j..j + window_size] {
repetition_count += 1;
break;
}
}
}
characteristics.repetition_ratio =
repetition_count as f64 / (data.len() - window_size) as f64;
}
characteristics
}
fn update_stats(
&mut self,
algorithm: &CompressionAlgorithm,
original_size: usize,
compressed_size: usize,
time: std::time::Duration,
) {
let algo_name = match algorithm {
CompressionAlgorithm::None => "none",
CompressionAlgorithm::Lz4 { fast } => {
if *fast {
"lz4-fast"
} else {
"lz4"
}
}
CompressionAlgorithm::Zstd { .. } => "zstd",
};
let ratio = compressed_size as f64 / original_size as f64;
*self
.stats
.attempts
.entry(algo_name.to_string())
.or_insert(0) += 1;
let attempts = self.stats.attempts[algo_name];
if let Some(existing_ratio) = self.stats.ratios.get_mut(algo_name) {
*existing_ratio = (*existing_ratio * (attempts - 1) as f64 + ratio) / attempts as f64;
} else {
self.stats.ratios.insert(algo_name.to_string(), ratio);
}
if let Some(existing_time) = self.stats.times.get_mut(algo_name) {
*existing_time = (*existing_time * (attempts as u32 - 1) + time) / attempts as u32;
} else {
self.stats.times.insert(algo_name.to_string(), time);
}
self.stats.total_bytes += original_size as u64;
self.stats.total_time += time;
}
fn estimate_memory_usage(&self) -> usize {
std::mem::size_of::<Self>()
+ self.stats.attempts.len() * 64
+ self.stats.ratios.len() * 64
+ self.stats.times.len() * 64
}
pub fn stats(&self) -> &CompressionStats {
&self.stats
}
pub fn reset_stats(&mut self) {
self.stats = CompressionStats::default();
}
}
#[derive(Debug)]
pub struct DataCharacteristics {
pub entropy: f64,
pub repetition_ratio: f64,
pub pattern_density: f64,
pub ascii_ratio: f64,
}
impl Default for CompressionConfig {
fn default() -> Self {
Self {
default_level: 6,
adaptive_selection: true,
size_threshold: 1024, target_ratio: 0.7,
max_compression_time: std::time::Duration::from_secs(30),
prefer_speed: false,
}
}
}
impl CompressionStats {
pub fn overall_ratio(&self) -> f64 {
if self.ratios.is_empty() {
1.0
} else {
self.ratios.values().sum::<f64>() / self.ratios.len() as f64
}
}
pub fn overall_throughput(&self) -> f64 {
if self.total_time.as_secs_f64() > 0.0 {
self.total_bytes as f64 / self.total_time.as_secs_f64()
} else {
0.0
}
}
pub fn best_algorithm(&self) -> Option<String> {
self.ratios
.iter()
.min_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal))
.map(|(name, _)| name.clone())
}
}
#[cfg(feature = "compression")]
fn decompress_lz4_block(data: &[u8], original_size: usize) -> SnapshotResult<Vec<u8>> {
let size_hint = original_size.max(1);
#[cfg(target_arch = "wasm32")]
{
let min_uncompressed = size_hint.max(data.len().saturating_mul(4));
lz4_flex::block::decompress(data, min_uncompressed).map_err(|e| {
SnapshotError::Compression {
message: format!("LZ4 decompression failed: {e}"),
}
})
}
#[cfg(not(target_arch = "wasm32"))]
{
let clamped = size_hint.min(i32::MAX as usize);
let size_i32 = i32::try_from(clamped).unwrap_or(i32::MAX);
lz4::block::decompress(data, Some(size_i32)).map_err(|e| SnapshotError::Compression {
message: format!("LZ4 decompression failed: {e}"),
})
}
}
#[cfg(all(feature = "compression", target_arch = "wasm32"))]
fn decompress_zstd_block(data: &[u8]) -> SnapshotResult<Vec<u8>> {
let cursor = Cursor::new(data);
let mut decoder = StreamingDecoder::new(cursor).map_err(|e| SnapshotError::Compression {
message: format!("ZSTD decompression failed: {e}"),
})?;
let mut output = Vec::new();
decoder
.read_to_end(&mut output)
.map_err(|e| SnapshotError::Compression {
message: format!("ZSTD decompression failed: {e}"),
})?;
Ok(output)
}
#[cfg(all(feature = "compression", not(target_arch = "wasm32")))]
fn decompress_zstd_block(data: &[u8]) -> SnapshotResult<Vec<u8>> {
zstd::decode_all(data).map_err(|e| SnapshotError::Compression {
message: format!("ZSTD decompression failed: {e}"),
})
}
#[cfg(feature = "compression")]
fn compress_lz4_block(data: &[u8], fast: bool) -> SnapshotResult<Vec<u8>> {
#[cfg(target_arch = "wasm32")]
{
let compressed = lz4_flex::block::compress(data);
if !fast {
log::trace!("High-compression LZ4 mode not available on wasm; using fast path");
}
Ok(compressed)
}
#[cfg(not(target_arch = "wasm32"))]
{
let result = if fast {
lz4::block::compress(data, None, false)
} else {
lz4::block::compress(
data,
Some(lz4::block::CompressionMode::HIGHCOMPRESSION(12)),
false,
)
};
result.map_err(|e| SnapshotError::Compression {
message: format!("LZ4 compression failed: {e}"),
})
}
}
#[cfg(all(feature = "compression", target_arch = "wasm32"))]
fn compress_zstd_block(_data: &[u8], _level: i32) -> SnapshotResult<Vec<u8>> {
Err(SnapshotError::Configuration {
message: "ZSTD compression is not supported on wasm targets".to_string(),
})
}
#[cfg(all(feature = "compression", not(target_arch = "wasm32")))]
fn compress_zstd_block(data: &[u8], level: i32) -> SnapshotResult<Vec<u8>> {
zstd::encode_all(data, level).map_err(|e| SnapshotError::Compression {
message: format!("ZSTD compression failed: {e}"),
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_compression_config_default() {
let config = CompressionConfig::default();
assert_eq!(config.default_level, 6);
assert!(config.adaptive_selection);
assert_eq!(config.size_threshold, 1024);
}
#[test]
fn test_compression_engine_creation() {
let config = CompressionConfig::default();
let engine = CompressionEngine::new(config);
assert_eq!(engine.stats.total_bytes, 0);
}
#[test]
fn test_small_data_compression() {
let mut engine = CompressionEngine::new(CompressionConfig::default());
let small_data = vec![1, 2, 3];
let result = engine.compress(&small_data).unwrap();
assert_eq!(result.data, small_data);
assert!(matches!(result.info.algorithm, CompressionAlgorithm::None));
}
#[cfg(feature = "compression")]
#[test]
fn test_lz4_compression() {
let mut engine = CompressionEngine::new(CompressionConfig {
adaptive_selection: false,
prefer_speed: true,
size_threshold: 10, ..CompressionConfig::default()
});
let data = b"Hello, World! This is a longer test string for LZ4 compression.".repeat(50);
let result = engine.compress(&data).unwrap();
assert!(
result.data.len() <= data.len(),
"Compression did not reduce size"
);
assert!(matches!(
result.info.algorithm,
CompressionAlgorithm::Lz4 { .. } | CompressionAlgorithm::None
));
match engine.decompress(&result.data, &result.info) {
Ok(decompressed) => {
assert_eq!(decompressed, data);
}
Err(e) => {
if matches!(result.info.algorithm, CompressionAlgorithm::None) {
assert_eq!(result.data, data);
} else {
panic!("Decompression failed: {e:?}");
}
}
}
}
#[test]
fn test_compression_stats() {
let stats = CompressionStats::default();
assert_eq!(stats.overall_ratio(), 1.0);
assert_eq!(stats.overall_throughput(), 0.0);
assert!(stats.best_algorithm().is_none());
}
#[test]
fn test_data_characteristics() {
let engine = CompressionEngine::new(CompressionConfig::default());
let ascii_data = b"Hello, World!".to_vec();
let characteristics = engine.analyze_data(&ascii_data);
assert!(characteristics.ascii_ratio > 0.9);
let binary_data = vec![0u8, 1u8, 2u8, 255u8];
let characteristics = engine.analyze_data(&binary_data);
assert!(characteristics.ascii_ratio < 1.0);
}
}