Skip to main content

mediagit_compression/
zstd_compressor.rs

1// MediaGit - Git for Media Files
2// Copyright (C) 2025 MediaGit Contributors
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU Affero General Public License as published
6// by the Free Software Foundation, either version 3 of the License, or
7// (at your option) any later version.
8//
9// This program is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12// GNU Affero General Public License for more details.
13
14//! Zstd compression implementation
15//!
16//! Fast compression and decompression with good compression ratios.
17//! Ideal for frequently accessed data in MediaGit.
18
19use crate::error::{CompressionError, CompressionResult};
20use crate::{CompressionLevel, Compressor};
21use std::fmt;
22
23/// Zstd compressor implementation
24///
25/// Uses the Zstandard compression algorithm for balanced speed and ratio.
26/// Supports configurable compression levels.
27#[derive(Clone)]
28pub struct ZstdCompressor {
29    level: CompressionLevel,
30}
31
32impl ZstdCompressor {
33    /// Create a new Zstd compressor with the given compression level
34    pub fn new(level: CompressionLevel) -> Self {
35        ZstdCompressor { level }
36    }
37
38    /// Create a Zstd compressor with fast compression
39    pub fn fast() -> Self {
40        ZstdCompressor::new(CompressionLevel::Fast)
41    }
42
43    /// Create a Zstd compressor with default compression
44    pub fn default_level() -> Self {
45        ZstdCompressor::new(CompressionLevel::Default)
46    }
47
48    /// Create a Zstd compressor with best compression
49    pub fn best() -> Self {
50        ZstdCompressor::new(CompressionLevel::Best)
51    }
52}
53
54impl fmt::Debug for ZstdCompressor {
55    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
56        f.debug_struct("ZstdCompressor")
57            .field("level", &self.level)
58            .finish()
59    }
60}
61
62impl Compressor for ZstdCompressor {
63    fn compress(&self, data: &[u8]) -> CompressionResult<Vec<u8>> {
64        let level = self.level.to_zstd_level();
65
66        // For very small data, just return uncompressed to avoid overhead
67        if data.is_empty() {
68            return Ok(Vec::new());
69        }
70
71        match zstd::encode_all(data, level) {
72            Ok(compressed) => {
73                // Prepend zstd magic bytes (already in zstd output, but ensure it's there)
74                Ok(compressed)
75            }
76            Err(e) => Err(CompressionError::zstd_error(format!(
77                "zstd compression failed: {}",
78                e
79            ))),
80        }
81    }
82
83    fn decompress(&self, data: &[u8]) -> CompressionResult<Vec<u8>> {
84        if data.is_empty() {
85            return Ok(Vec::new());
86        }
87
88        // Check if this looks like zstd compressed data (has zstd magic bytes)
89        if data.len() >= 4 && data.starts_with(b"\x28\xb5\x2f\xfd") {
90            match zstd::decode_all(data) {
91                Ok(decompressed) => Ok(decompressed),
92                Err(e) => Err(CompressionError::decompression_failed(format!(
93                    "zstd decompression failed: {}",
94                    e
95                ))),
96            }
97        } else {
98            // Data is not zstd compressed, return as-is
99            // This handles the case where data was never compressed
100            Ok(data.to_vec())
101        }
102    }
103}
104
105#[cfg(test)]
106#[allow(clippy::unwrap_used)]
107mod tests {
108    use super::*;
109
110    #[test]
111    fn test_zstd_compress_decompress() {
112        let compressor = ZstdCompressor::new(CompressionLevel::Default);
113        let original = b"Hello, World! This is a test of zstd compression.";
114
115        let compressed = compressor.compress(original).unwrap();
116        let decompressed = compressor.decompress(&compressed).unwrap();
117
118        assert_eq!(original, &decompressed[..]);
119    }
120
121    #[test]
122    fn test_zstd_compress_empty() {
123        let compressor = ZstdCompressor::new(CompressionLevel::Default);
124        let compressed = compressor.compress(b"").unwrap();
125        assert_eq!(compressed.len(), 0);
126    }
127
128    #[test]
129    fn test_zstd_decompress_empty() {
130        let compressor = ZstdCompressor::new(CompressionLevel::Default);
131        let decompressed = compressor.decompress(b"").unwrap();
132        assert_eq!(decompressed.len(), 0);
133    }
134
135    #[test]
136    fn test_zstd_decompress_uncompressed_data() {
137        let compressor = ZstdCompressor::new(CompressionLevel::Default);
138        let data = b"This is not compressed";
139        let result = compressor.decompress(data).unwrap();
140        // Should return as-is since it doesn't have zstd magic bytes
141        assert_eq!(result, data);
142    }
143
144    #[test]
145    fn test_zstd_compression_levels() {
146        let data = b"This is test data that should compress. ".repeat(100);
147
148        let fast = ZstdCompressor::fast().compress(&data).unwrap();
149        let default = ZstdCompressor::default_level().compress(&data).unwrap();
150        let best = ZstdCompressor::best().compress(&data).unwrap();
151
152        // Best should compress better than default, which should be better than fast
153        // (though not always strictly true due to data characteristics)
154        assert!(best.len() <= default.len());
155        assert!(default.len() <= fast.len() || fast.len() < 100); // some tolerance
156
157        // All should decompress correctly
158        let decompressor = ZstdCompressor::default_level();
159        let fast_dec = decompressor.decompress(&fast).unwrap();
160        let default_dec = decompressor.decompress(&default).unwrap();
161        let best_dec = decompressor.decompress(&best).unwrap();
162
163        assert_eq!(fast_dec, data);
164        assert_eq!(default_dec, data);
165        assert_eq!(best_dec, data);
166    }
167
168    #[test]
169    fn test_zstd_large_data() {
170        let compressor = ZstdCompressor::new(CompressionLevel::Default);
171        let original = vec![0x42u8; 1024 * 1024]; // 1MB of repeated data
172
173        let compressed = compressor.compress(&original).unwrap();
174        let decompressed = compressor.decompress(&compressed).unwrap();
175
176        assert_eq!(original, decompressed);
177        // Repeated data should compress very well
178        assert!(compressed.len() < original.len() / 100);
179    }
180
181    #[test]
182    fn test_zstd_random_data() {
183        let compressor = ZstdCompressor::new(CompressionLevel::Default);
184        // Random data doesn't compress well
185        let original = (0..1000).map(|i| (i ^ 0xAA) as u8).collect::<Vec<_>>();
186
187        let compressed = compressor.compress(&original).unwrap();
188        let decompressed = compressor.decompress(&compressed).unwrap();
189
190        assert_eq!(original, decompressed);
191        // Random data might even expand slightly
192        assert!(compressed.len() <= original.len() + 1000);
193    }
194
195    #[test]
196    fn test_zstd_debug_format() {
197        let compressor = ZstdCompressor::new(CompressionLevel::Default);
198        let debug_str = format!("{:?}", compressor);
199        assert!(debug_str.contains("ZstdCompressor"));
200        assert!(debug_str.contains("Default"));
201    }
202}