use super::types::{DxMachineError, Result};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum CompressionLevel {
#[default]
Fast,
Default,
High,
}
impl CompressionLevel {
pub fn to_zstd_level(self) -> i32 {
match self {
CompressionLevel::Fast => 1, CompressionLevel::Default => 3, CompressionLevel::High => 19, }
}
}
#[derive(Debug)]
pub struct DxCompressed {
compressed: Vec<u8>,
original_size: u32,
decompressed: Option<Vec<u8>>,
}
impl DxCompressed {
pub fn new() -> Self {
Self {
compressed: Vec::new(),
original_size: 0,
decompressed: None,
}
}
pub fn compress(data: &[u8]) -> Self {
let compressed = lz4_compress_fast(data);
let original_size = data.len() as u32;
Self {
compressed,
original_size,
decompressed: None,
}
}
pub fn compress_level(data: &[u8], _level: CompressionLevel) -> Self {
Self::compress(data)
}
#[inline(always)]
pub fn compressed_size(&self) -> usize {
self.compressed.len()
}
#[inline(always)]
pub fn original_size(&self) -> usize {
self.original_size as usize
}
#[inline(always)]
pub fn ratio(&self) -> f64 {
if self.original_size == 0 {
return 1.0;
}
self.compressed.len() as f64 / self.original_size as f64
}
#[inline(always)]
pub fn savings(&self) -> f64 {
1.0 - self.ratio()
}
#[inline(always)]
pub fn as_compressed(&self) -> &[u8] {
&self.compressed
}
pub fn decompress(&mut self) -> Result<&[u8]> {
if self.decompressed.is_none() {
let data = lz4_decompress_fast(&self.compressed)?;
self.decompressed = Some(data);
}
Ok(self.decompressed.as_ref().unwrap_or_else(|| unreachable!()))
}
pub fn decompress_owned(&self) -> Result<Vec<u8>> {
lz4_decompress_fast(&self.compressed)
}
#[inline(always)]
pub fn is_cached(&self) -> bool {
self.decompressed.is_some()
}
pub fn clear_cache(&mut self) {
self.decompressed = None;
}
pub fn from_compressed(compressed: Vec<u8>, original_size: u32) -> Self {
Self {
compressed,
original_size,
decompressed: None,
}
}
pub fn to_wire(&self) -> Vec<u8> {
self.compressed.clone()
}
pub fn from_wire(data: &[u8]) -> Result<Self> {
if data.len() < 4 {
return Err(DxMachineError::BufferTooSmall {
required: 4,
actual: data.len(),
});
}
let original_size = u32::from_le_bytes([data[0], data[1], data[2], data[3]]);
let compressed = data.to_vec();
Ok(Self {
compressed,
original_size,
decompressed: None,
})
}
}
impl Default for DxCompressed {
fn default() -> Self {
Self::new()
}
}
#[cfg(feature = "compression-lz4")]
fn lz4_compress_fast(input: &[u8]) -> Vec<u8> {
lz4_flex::compress_prepend_size(input)
}
#[cfg(not(feature = "compression-lz4"))]
fn lz4_compress_fast(input: &[u8]) -> Vec<u8> {
lz4_compress(input)
}
#[cfg(feature = "compression-lz4")]
pub(crate) fn lz4_decompress_fast(input: &[u8]) -> Result<Vec<u8>> {
lz4_flex::decompress_size_prepended(input)
.map_err(|e| DxMachineError::DecompressionFailed(e.to_string()))
}
#[cfg(not(feature = "compression-lz4"))]
fn lz4_decompress_fast(input: &[u8]) -> Result<Vec<u8>> {
if input.len() < 4 {
return Err(DxMachineError::DecompressionFailed(
"Input too short".into(),
));
}
let size = u32::from_le_bytes([input[0], input[1], input[2], input[3]]) as usize;
lz4_decompress(&input[4..], size)
}
#[allow(dead_code)]
#[cfg(feature = "compression")]
fn zstd_compress(input: &[u8], level: i32) -> Vec<u8> {
zstd::encode_all(input, level).unwrap_or_else(|_| input.to_vec())
}
#[allow(dead_code)]
#[cfg(not(feature = "compression"))]
fn zstd_compress(input: &[u8], _level: i32) -> Vec<u8> {
lz4_compress(input)
}
#[allow(dead_code)]
#[cfg(feature = "compression")]
fn zstd_decompress(input: &[u8]) -> Result<Vec<u8>> {
zstd::decode_all(input).map_err(|e| DxMachineError::DecompressionFailed(e.to_string()))
}
#[allow(dead_code)]
#[cfg(not(feature = "compression"))]
fn zstd_decompress(input: &[u8]) -> Result<Vec<u8>> {
if input.len() < 4 {
return Err(DxMachineError::DecompressionFailed(
"Input too short".into(),
));
}
let size = u32::from_le_bytes([input[0], input[1], input[2], input[3]]) as usize;
lz4_decompress(&input[4..], size)
}
#[allow(dead_code)]
fn lz4_compress(input: &[u8]) -> Vec<u8> {
if input.is_empty() {
return Vec::new();
}
let mut output = Vec::with_capacity(input.len());
let mut i = 0;
while i < input.len() {
let byte = input[i];
let mut run_len = 1;
while i + run_len < input.len() && input[i + run_len] == byte && run_len < 255 {
run_len += 1;
}
if run_len >= 4 {
output.push(0xFF);
output.push(run_len as u8);
output.push(byte);
i += run_len;
} else {
let lit_start = i;
while i < input.len() {
if i + 4 <= input.len() {
let b = input[i];
if input[i + 1] == b && input[i + 2] == b && input[i + 3] == b {
break;
}
}
i += 1;
if i - lit_start >= 254 {
break;
}
}
let lit_len = i - lit_start;
if lit_len > 0 {
output.push(lit_len as u8);
output.extend_from_slice(&input[lit_start..i]);
}
}
}
output
}
#[allow(dead_code)] fn lz4_decompress(input: &[u8], expected_size: usize) -> Result<Vec<u8>> {
if input.is_empty() {
return Ok(Vec::new());
}
let mut output = Vec::with_capacity(expected_size);
let mut i = 0;
while i < input.len() {
let marker = input[i];
i += 1;
if marker == 0xFF {
if i + 2 > input.len() {
return Err(DxMachineError::InvalidData("Truncated RLE sequence".into()));
}
let run_len = input[i] as usize;
let byte = input[i + 1];
i += 2;
output.extend(std::iter::repeat_n(byte, run_len));
} else {
let lit_len = marker as usize;
if i + lit_len > input.len() {
return Err(DxMachineError::InvalidData(
"Truncated literal sequence".into(),
));
}
output.extend_from_slice(&input[i..i + lit_len]);
i += lit_len;
}
}
Ok(output)
}
pub struct StreamCompressor {
chunk_size: usize,
chunks: Vec<DxCompressed>,
buffer: Vec<u8>,
}
impl StreamCompressor {
pub fn new(chunk_size: usize) -> Self {
Self {
chunk_size,
chunks: Vec::new(),
buffer: Vec::with_capacity(chunk_size),
}
}
pub fn default_chunk() -> Self {
Self::new(64 * 1024)
}
pub fn write(&mut self, data: &[u8]) {
let mut remaining = data;
while !remaining.is_empty() {
let space = self.chunk_size - self.buffer.len();
let take = remaining.len().min(space);
self.buffer.extend_from_slice(&remaining[..take]);
remaining = &remaining[take..];
if self.buffer.len() >= self.chunk_size {
self.flush_chunk();
}
}
}
fn flush_chunk(&mut self) {
if !self.buffer.is_empty() {
let chunk = DxCompressed::compress(&self.buffer);
self.chunks.push(chunk);
self.buffer.clear();
}
}
pub fn finish(mut self) -> Vec<DxCompressed> {
self.flush_chunk();
self.chunks
}
pub fn chunk_count(&self) -> usize {
self.chunks.len()
}
pub fn total_compressed_size(&self) -> usize {
self.chunks
.iter()
.map(|c| c.compressed_size())
.sum::<usize>()
+ self.buffer.len()
}
}
pub struct StreamDecompressor {
chunks: Vec<DxCompressed>,
current_chunk: usize,
current_offset: usize,
}
impl StreamDecompressor {
pub fn new(chunks: Vec<DxCompressed>) -> Self {
Self {
chunks,
current_chunk: 0,
current_offset: 0,
}
}
pub fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
if self.current_chunk >= self.chunks.len() {
return Ok(0);
}
let mut written = 0;
while written < buf.len() && self.current_chunk < self.chunks.len() {
let chunk = &mut self.chunks[self.current_chunk];
let data = chunk.decompress()?;
let remaining_in_chunk = data.len() - self.current_offset;
let to_copy = (buf.len() - written).min(remaining_in_chunk);
buf[written..written + to_copy]
.copy_from_slice(&data[self.current_offset..self.current_offset + to_copy]);
written += to_copy;
self.current_offset += to_copy;
if self.current_offset >= data.len() {
self.current_chunk += 1;
self.current_offset = 0;
}
}
Ok(written)
}
pub fn decompress_all(&mut self) -> Result<Vec<u8>> {
let total_size: usize = self.chunks.iter().map(|c| c.original_size()).sum();
let mut output = Vec::with_capacity(total_size);
for chunk in &mut self.chunks {
let data = chunk.decompress()?;
output.extend_from_slice(data);
}
Ok(output)
}
}
#[cfg(feature = "compression-lz4")]
pub fn compress_lz4(data: &[u8]) -> Result<Vec<u8>> {
if data.is_empty() {
return Ok(Vec::new());
}
Ok(lz4_flex::compress_prepend_size(data))
}
#[cfg(not(feature = "compression-lz4"))]
pub fn compress_lz4(_data: &[u8]) -> Result<Vec<u8>> {
Err(DxMachineError::InvalidData(
"LZ4 compression not available (enable 'compression-lz4' feature)".into(),
))
}
#[cfg(feature = "compression-lz4")]
pub fn decompress_lz4(data: &[u8]) -> Result<Vec<u8>> {
if data.is_empty() {
return Ok(Vec::new());
}
lz4_flex::decompress_size_prepended(data).map_err(|e| {
DxMachineError::DecompressionFailed(format!("LZ4 decompression failed: {}", e))
})
}
#[cfg(not(feature = "compression-lz4"))]
pub fn decompress_lz4(_data: &[u8]) -> Result<Vec<u8>> {
Err(DxMachineError::InvalidData(
"LZ4 decompression not available (enable 'compression-lz4' feature)".into(),
))
}
#[cfg(feature = "compression")]
pub fn compress_zstd(data: &[u8]) -> Result<Vec<u8>> {
compress_zstd_level(data, CompressionLevel::Default)
}
#[cfg(not(feature = "compression"))]
pub fn compress_zstd(_data: &[u8]) -> Result<Vec<u8>> {
Err(DxMachineError::InvalidData(
"Zstd compression not available (enable 'compression' feature)".into(),
))
}
#[cfg(feature = "compression")]
pub fn compress_zstd_level(data: &[u8], level: CompressionLevel) -> Result<Vec<u8>> {
if data.is_empty() {
return Ok(Vec::new());
}
let zstd_level = level.to_zstd_level();
zstd::encode_all(data, zstd_level)
.map_err(|e| DxMachineError::InvalidData(format!("Zstd compression failed: {}", e)))
}
#[cfg(not(feature = "compression"))]
pub fn compress_zstd_level(_data: &[u8], _level: CompressionLevel) -> Result<Vec<u8>> {
Err(DxMachineError::InvalidData(
"Zstd compression not available (enable 'compression' feature)".into(),
))
}
#[cfg(feature = "compression")]
pub fn decompress_zstd(data: &[u8]) -> Result<Vec<u8>> {
if data.is_empty() {
return Ok(Vec::new());
}
zstd::decode_all(data)
.map_err(|e| DxMachineError::InvalidData(format!("Zstd decompression failed: {}", e)))
}
#[cfg(not(feature = "compression"))]
pub fn decompress_zstd(_data: &[u8]) -> Result<Vec<u8>> {
Err(DxMachineError::InvalidData(
"Zstd decompression not available (enable 'compression' feature)".into(),
))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
#[cfg(feature = "compression-lz4")]
fn test_compress_decompress() {
let original = b"Hello, World! This is a test of the compression system.";
let mut compressed = DxCompressed::compress(original);
println!(
"Original: {} bytes, Compressed: {} bytes, Ratio: {:.2}",
original.len(),
compressed.compressed_size(),
compressed.ratio()
);
let decompressed = compressed.decompress().unwrap();
assert_eq!(decompressed, original);
}
#[test]
#[cfg(feature = "compression-lz4")]
fn test_compress_repetitive_data() {
let original: Vec<u8> = std::iter::repeat_n(b'A', 1000).collect();
let compressed = DxCompressed::compress(&original);
println!(
"Repetitive: {} bytes -> {} bytes ({:.1}% savings)",
original.len(),
compressed.compressed_size(),
compressed.savings() * 100.0
);
assert!(compressed.ratio() < 0.1); }
#[test]
#[cfg(feature = "compression-lz4")]
fn test_compress_random_data() {
let original: Vec<u8> = (0..1000).map(|i| (i % 256) as u8).collect();
let compressed = DxCompressed::compress(&original);
println!(
"Sequential: {} bytes -> {} bytes ({:.1}% savings)",
original.len(),
compressed.compressed_size(),
compressed.savings() * 100.0
);
}
#[test]
#[cfg(feature = "compression-lz4")]
fn test_wire_format() {
let original = b"Test data for wire format";
let compressed = DxCompressed::compress(original);
let wire = compressed.to_wire();
let restored = DxCompressed::from_wire(&wire).unwrap();
assert_eq!(restored.original_size(), original.len());
assert_eq!(restored.compressed_size(), compressed.compressed_size());
}
#[test]
#[cfg(feature = "compression-lz4")]
fn test_streaming_compressor() {
let mut compressor = StreamCompressor::new(32);
for i in 0..10 {
let data: Vec<u8> = (0..20).map(|j| ((i * 20 + j) % 256) as u8).collect();
compressor.write(&data);
}
let chunks = compressor.finish();
println!("Produced {} chunks", chunks.len());
let mut decompressor = StreamDecompressor::new(chunks);
let output = decompressor.decompress_all().unwrap();
let expected: Vec<u8> = (0..200).map(|i| (i % 256) as u8).collect();
assert_eq!(output, expected);
}
#[test]
#[cfg(feature = "compression-lz4")]
fn test_cache() {
let original = b"Cache test data";
let mut compressed = DxCompressed::compress(original);
assert!(!compressed.is_cached());
compressed.decompress().unwrap();
assert!(compressed.is_cached());
compressed.clear_cache();
assert!(!compressed.is_cached());
}
#[test]
#[cfg(feature = "compression-lz4")]
fn test_empty_data() {
let original: &[u8] = &[];
let mut compressed = DxCompressed::compress(original);
assert_eq!(compressed.original_size(), 0);
let decompressed = compressed.decompress().unwrap();
assert!(decompressed.is_empty());
}
#[test]
#[cfg(feature = "compression")]
fn test_zstd_compress_decompress() {
let original = b"Hello, World! This is a test of Zstd compression.";
let compressed = compress_zstd(original).unwrap();
println!(
"Zstd: Original {} bytes -> Compressed {} bytes ({:.1}% of original)",
original.len(),
compressed.len(),
(compressed.len() as f64 / original.len() as f64) * 100.0
);
let decompressed = decompress_zstd(&compressed).unwrap();
assert_eq!(decompressed, original);
}
#[test]
#[cfg(feature = "compression")]
fn test_zstd_compression_levels() {
let original: Vec<u8> = std::iter::repeat_n(b'A', 1000).collect();
let fast = compress_zstd_level(&original, CompressionLevel::Fast).unwrap();
let default = compress_zstd_level(&original, CompressionLevel::Default).unwrap();
let high = compress_zstd_level(&original, CompressionLevel::High).unwrap();
println!("Zstd compression levels on 1000 bytes of 'A':");
println!(" Fast (level 1): {} bytes", fast.len());
println!(" Default (level 3): {} bytes", default.len());
println!(" High (level 19): {} bytes", high.len());
assert!(high.len() <= default.len());
assert!(default.len() <= fast.len());
assert_eq!(decompress_zstd(&fast).unwrap(), original);
assert_eq!(decompress_zstd(&default).unwrap(), original);
assert_eq!(decompress_zstd(&high).unwrap(), original);
}
#[test]
#[cfg(feature = "compression")]
fn test_zstd_empty_data() {
let original: &[u8] = &[];
let compressed = compress_zstd(original).unwrap();
assert!(compressed.is_empty());
let decompressed = decompress_zstd(&compressed).unwrap();
assert!(decompressed.is_empty());
}
#[test]
#[cfg(feature = "compression")]
fn test_zstd_vs_lz4() {
let original: Vec<u8> = std::iter::repeat_n(b'X', 10000).collect();
let lz4_compressed = lz4_compress(&original);
let zstd_compressed = compress_zstd(&original).unwrap();
println!("Compression comparison on 10KB repetitive data:");
println!(
" LZ4: {} bytes ({:.1}% of original)",
lz4_compressed.len(),
(lz4_compressed.len() as f64 / original.len() as f64) * 100.0
);
println!(
" Zstd: {} bytes ({:.1}% of original)",
zstd_compressed.len(),
(zstd_compressed.len() as f64 / original.len() as f64) * 100.0
);
let lz4_decompressed = lz4_decompress(&lz4_compressed, original.len()).unwrap();
let zstd_decompressed = decompress_zstd(&zstd_compressed).unwrap();
assert_eq!(lz4_decompressed, original);
assert_eq!(zstd_decompressed, original);
}
}