heddle_format/compression/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2#![deny(clippy::cast_possible_truncation)]
3
4//! Compression utilities for Heddle storage.
5//!
6//! Provides configurable compression with support for:
7//! - zstd: High compression ratio, good speed
8//! - Delta encoding: For similar versions of the same file
9
10mod dictionaries;
11mod frame;
12mod zstd_codec;
13
14pub use dictionaries::CompressionDictionary;
15use frame::{
16    DICTIONARY_HEADER_LEN as DICTIONARY_COMPRESSED_HEADER_LEN, HEADER_LEN as COMPRESSED_HEADER_LEN,
17};
18
19/// Compression algorithm selection.
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21#[repr(u8)]
22enum CompressionType {
23    /// Zstandard compression.
24    Zstd = 1,
25}
26
27impl CompressionType {
28    /// Convert from byte value.
29    fn from_u8(value: u8) -> Option<Self> {
30        match value {
31            1 => Some(CompressionType::Zstd),
32            _ => None,
33        }
34    }
35}
36
37/// Compression configuration.
38#[derive(Debug, Clone, Copy)]
39pub struct CompressionConfig {
40    /// Whether compression is enabled.
41    pub enabled: bool,
42    /// Compression level (algorithm-specific).
43    /// For zstd: 1-22 (1=fast, 22=best, 3=default)
44    pub level: i32,
45    /// Minimum size to compress (smaller objects aren't worth it).
46    pub min_size: usize,
47    /// Maximum size for delta compression base.
48    pub max_delta_size: usize,
49}
50
51impl Default for CompressionConfig {
52    fn default() -> Self {
53        Self {
54            enabled: cfg!(feature = "zstd"),
55            level: 3,                   // zstd default
56            min_size: 256,              // Don't compress tiny objects
57            max_delta_size: 10_000_000, // 10MB max for delta base
58        }
59    }
60}
61
62impl CompressionConfig {
63    /// Create configuration from environment variables.
64    pub fn from_env() -> Self {
65        let mut config = Self::default();
66
67        if let Ok(val) = std::env::var("HEDDLE_COMPRESSION") {
68            let requested = val != "0" && val.to_lowercase() != "false";
69            config.enabled = requested && cfg!(feature = "zstd");
70        }
71
72        if let Ok(val) = std::env::var("HEDDLE_COMPRESSION_LEVEL")
73            && let Ok(level) = val.parse::<i32>()
74        {
75            config.level = level.clamp(1, 22);
76        }
77
78        if let Ok(val) = std::env::var("HEDDLE_COMPRESSION_MIN_SIZE")
79            && let Ok(size) = val.parse::<usize>()
80        {
81            config.min_size = size;
82        }
83
84        config
85    }
86
87    /// Disable compression.
88    pub fn disabled() -> Self {
89        Self {
90            enabled: false,
91            level: 0,
92            min_size: usize::MAX,
93            max_delta_size: 0,
94        }
95    }
96}
97
98/// Compression error type.
99#[derive(Debug, thiserror::Error)]
100pub enum CompressionError {
101    #[error("decompression failed: {0}")]
102    DecompressionFailed(String),
103    #[error("compression failed: {0}")]
104    CompressionFailed(String),
105    #[error("invalid compression type: {0}")]
106    InvalidType(u8),
107    #[error("corrupted data: {0}")]
108    CorruptedData(String),
109    #[error("invalid operation: {0}")]
110    InvalidOperation(String),
111    #[error("unknown compression dictionary id: {0}")]
112    UnknownDictionary(u32),
113    #[error("object size {size} exceeds maximum {max}")]
114    SizeLimitExceeded { size: u64, max: u64 },
115}
116
117#[cfg(feature = "bench")]
118/// Compress data using zstd.
119pub fn compress_zstd(data: &[u8], level: i32) -> Result<Vec<u8>, CompressionError> {
120    zstd_codec::compress(data, level)
121}
122
123#[cfg(feature = "bench")]
124/// Decompress zstd data while enforcing the recorded output size.
125pub fn decompress_zstd(data: &[u8], expected_size: u64) -> Result<Vec<u8>, CompressionError> {
126    zstd_codec::decompress(data, expected_size)
127}
128
129/// Compress data with automatic algorithm selection.
130///
131/// Returns the compressed data with header, or None if compression
132/// doesn't help (compressed would be larger).
133pub fn compress(
134    data: &[u8],
135    config: &CompressionConfig,
136) -> Result<Option<Vec<u8>>, CompressionError> {
137    if !config.enabled || data.len() < config.min_size {
138        return Ok(None);
139    }
140
141    zstd_codec::validate_size(data.len() as u64)?;
142    let compressed = zstd_codec::compress(data, config.level)?;
143
144    if compressed.len() >= data.len() {
145        return Ok(None);
146    }
147
148    let mut result = Vec::with_capacity(COMPRESSED_HEADER_LEN + compressed.len());
149    result.push(CompressionType::Zstd as u8);
150    result.extend_from_slice(&(data.len() as u64).to_be_bytes());
151    result.extend_from_slice(&compressed);
152    Ok(Some(result))
153}
154
155/// Compress data with a durable, versioned dictionary.
156///
157/// The dictionary ID is embedded in the compression wrapper so the matching
158/// dictionary decoder can select the exact bundled bytes.
159pub fn compress_with_dictionary(
160    data: &[u8],
161    config: &CompressionConfig,
162    dictionary: CompressionDictionary,
163) -> Result<Option<Vec<u8>>, CompressionError> {
164    if !config.enabled || data.len() < config.min_size {
165        return Ok(None);
166    }
167
168    zstd_codec::validate_size(data.len() as u64)?;
169    let compressed = zstd_codec::compress_with_dictionary(data, config.level, dictionary.bytes())?;
170
171    // Only use compression if it actually helps
172    if compressed.len() >= data.len() {
173        return Ok(None);
174    }
175
176    let mut result = Vec::with_capacity(DICTIONARY_COMPRESSED_HEADER_LEN + compressed.len());
177    result.push(CompressionType::Zstd as u8);
178    result.extend_from_slice(&(data.len() as u64).to_be_bytes());
179    result.extend_from_slice(&dictionary.id().to_be_bytes());
180    result.extend_from_slice(&compressed);
181
182    Ok(Some(result))
183}
184
185/// Decompress data based on header.
186///
187/// Returns the decompressed data, or original data if uncompressed.
188pub fn decompress(data: &[u8]) -> Result<Vec<u8>, CompressionError> {
189    if data.len() < COMPRESSED_HEADER_LEN {
190        // Too short for header, assume uncompressed
191        return Ok(data.to_vec());
192    }
193
194    let compression_type =
195        CompressionType::from_u8(data[0]).ok_or_else(|| CompressionError::InvalidType(data[0]))?;
196
197    match compression_type {
198        CompressionType::Zstd if frame::parse_zstd(data).is_some() => {
199            decompress_zstd_with_header(data)
200        }
201        CompressionType::Zstd => Ok(data.to_vec()),
202    }
203}
204
205/// Check if data is compressed (has compression header).
206pub fn is_compressed(data: &[u8]) -> bool {
207    if data.len() < COMPRESSED_HEADER_LEN {
208        return false;
209    }
210
211    matches!(
212        CompressionType::from_u8(data[0]),
213        Some(CompressionType::Zstd)
214    ) && frame::parse_zstd(data).is_some()
215}
216
217/// Peek at the recorded *uncompressed* size in a header-prefixed blob,
218/// without decompressing the payload. Returns `None` for short or
219/// unprefixed inputs (the caller can then fall back to the file length).
220///
221/// Used by header-only size queries (e.g. [`ObjectStore::blob_size`])
222/// where reading the full blob would dominate. Only the first 9 bytes
223/// of the input are consulted.
224pub fn header_uncompressed_size(data: &[u8]) -> Option<u64> {
225    if data.len() < COMPRESSED_HEADER_LEN {
226        return None;
227    }
228    let CompressionType::Zstd = CompressionType::from_u8(data[0])?;
229    Some(frame::parse_zstd(data)?.uncompressed_size)
230}
231
232#[cfg(test)]
233/// Get compression info from header.
234fn compression_info(data: &[u8]) -> Option<(CompressionType, u64)> {
235    if data.len() < COMPRESSED_HEADER_LEN {
236        return None;
237    }
238
239    let compression_type = CompressionType::from_u8(data[0])?;
240    let uncompressed_size = u64::from_be_bytes(data[1..COMPRESSED_HEADER_LEN].try_into().ok()?);
241
242    Some((compression_type, uncompressed_size))
243}
244
245fn decompress_zstd_with_header(data: &[u8]) -> Result<Vec<u8>, CompressionError> {
246    let header = frame::parse_zstd(data).ok_or_else(|| {
247        CompressionError::CorruptedData("zstd compression header is invalid".to_string())
248    })?;
249    zstd_codec::decompress(&data[header.len..], header.uncompressed_size)
250}
251
252/// Decompress data that uses the dictionary compression wrapper.
253///
254/// Returns the original data when it does not have a dictionary compression
255/// header.
256pub fn decompress_with_dictionary(data: &[u8]) -> Result<Vec<u8>, CompressionError> {
257    if data.len() < DICTIONARY_COMPRESSED_HEADER_LEN {
258        return Ok(data.to_vec());
259    }
260
261    let Some(header) = frame::parse_dictionary_zstd(data) else {
262        return Ok(data.to_vec());
263    };
264    let dictionary = dictionaries::lookup(header.dictionary_id)
265        .ok_or(CompressionError::UnknownDictionary(header.dictionary_id))?;
266    zstd_codec::decompress_with_dictionary(
267        &data[DICTIONARY_COMPRESSED_HEADER_LEN..],
268        header.uncompressed_size,
269        dictionary,
270    )
271}
272
273#[cfg(test)]
274mod compression_tests;
heddle_format/compression/mod.rs

heddle_format/compression/
mod.rs
