feat(rust-core): add zstd chunk compression support and rewrite partially referenced packs during prune

rust/src/compression.rs

@@ -3,8 +3,63 @@ use flate2::read::{GzDecoder, GzEncoder};
 use std::io::Read;
 use crate::error::ArchiveError;
 
-/// Gzip compress data.
-pub fn compress(data: &[u8]) -> Result<Vec<u8>, ArchiveError> {
+/// Supported compression algorithms.
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum CompressionAlgorithm {
+    Gzip,
+    Zstd,
+}
+
+impl CompressionAlgorithm {
+    pub fn from_str(s: &str) -> Self {
+        match s {
+            "zstd" => Self::Zstd,
+            _ => Self::Gzip,
+        }
+    }
+
+    /// Map to the flag bits stored in IDX entries (bits 1-2).
+    pub fn to_flags(self) -> u32 {
+        match self {
+            Self::Gzip => 0x02,  // FLAG_GZIP
+            Self::Zstd => 0x04,  // FLAG_ZSTD
+        }
+    }
+
+    /// Determine algorithm from IDX flags.
+    pub fn from_flags(flags: u32) -> Self {
+        if flags & 0x04 != 0 {
+            Self::Zstd
+        } else {
+            Self::Gzip
+        }
+    }
+}
+
+/// Compress data with the specified algorithm.
+pub fn compress(data: &[u8], algo: CompressionAlgorithm) -> Result<Vec<u8>, ArchiveError> {
+    match algo {
+        CompressionAlgorithm::Gzip => compress_gzip(data),
+        CompressionAlgorithm::Zstd => compress_zstd(data),
+    }
+}
+
+/// Decompress data with the specified algorithm.
+pub fn decompress(data: &[u8], algo: CompressionAlgorithm) -> Result<Vec<u8>, ArchiveError> {
+    match algo {
+        CompressionAlgorithm::Gzip => decompress_gzip(data),
+        CompressionAlgorithm::Zstd => decompress_zstd(data),
+    }
+}
+
+/// Decompress data by detecting algorithm from flags.
+pub fn decompress_by_flags(data: &[u8], flags: u32) -> Result<Vec<u8>, ArchiveError> {
+    decompress(data, CompressionAlgorithm::from_flags(flags))
+}
+
+// ==================== Gzip ====================
+
+fn compress_gzip(data: &[u8]) -> Result<Vec<u8>, ArchiveError> {
     let mut encoder = GzEncoder::new(data, Compression::default());
     let mut compressed = Vec::new();
     encoder.read_to_end(&mut compressed)
@@ -12,8 +67,7 @@ pub fn compress(data: &[u8]) -> Result<Vec<u8>, ArchiveError> {
     Ok(compressed)
 }
 
-/// Gzip decompress data.
-pub fn decompress(data: &[u8]) -> Result<Vec<u8>, ArchiveError> {
+fn decompress_gzip(data: &[u8]) -> Result<Vec<u8>, ArchiveError> {
     let mut decoder = GzDecoder::new(data);
     let mut decompressed = Vec::new();
     decoder.read_to_end(&mut decompressed)
@@ -21,23 +75,57 @@ pub fn decompress(data: &[u8]) -> Result<Vec<u8>, ArchiveError> {
     Ok(decompressed)
 }
 
+// ==================== Zstd ====================
+
+fn compress_zstd(data: &[u8]) -> Result<Vec<u8>, ArchiveError> {
+    zstd::encode_all(data, 3) // level 3 = good balance of speed/ratio
+        .map_err(|e| ArchiveError::Io(e))
+}
+
+fn decompress_zstd(data: &[u8]) -> Result<Vec<u8>, ArchiveError> {
+    zstd::decode_all(data)
+        .map_err(|e| ArchiveError::Io(e))
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
-    fn test_roundtrip() {
+    fn test_gzip_roundtrip() {
         let data = b"Hello, this is test data for compression!";
-        let compressed = compress(data).unwrap();
-        let decompressed = decompress(&compressed).unwrap();
+        let compressed = compress(data, CompressionAlgorithm::Gzip).unwrap();
+        let decompressed = decompress(&compressed, CompressionAlgorithm::Gzip).unwrap();
+        assert_eq!(data.as_slice(), decompressed.as_slice());
+    }
+
+    #[test]
+    fn test_zstd_roundtrip() {
+        let data = b"Hello, this is test data for zstd compression!";
+        let compressed = compress(data, CompressionAlgorithm::Zstd).unwrap();
+        let decompressed = decompress(&compressed, CompressionAlgorithm::Zstd).unwrap();
         assert_eq!(data.as_slice(), decompressed.as_slice());
     }
 
     #[test]
     fn test_compression_reduces_size() {
-        // Highly compressible data
         let data = vec![b'A'; 10000];
-        let compressed = compress(&data).unwrap();
-        assert!(compressed.len() < data.len());
+        let gzip = compress(&data, CompressionAlgorithm::Gzip).unwrap();
+        let zstd = compress(&data, CompressionAlgorithm::Zstd).unwrap();
+        assert!(gzip.len() < data.len());
+        assert!(zstd.len() < data.len());
+    }
+
+    #[test]
+    fn test_decompress_by_flags() {
+        let data = b"flag-based decompression test";
+
+        let gzip_compressed = compress(data, CompressionAlgorithm::Gzip).unwrap();
+        let result = decompress_by_flags(&gzip_compressed, 0x02).unwrap();
+        assert_eq!(data.as_slice(), result.as_slice());
+
+        let zstd_compressed = compress(data, CompressionAlgorithm::Zstd).unwrap();
+        let result = decompress_by_flags(&zstd_compressed, 0x04).unwrap();
+        assert_eq!(data.as_slice(), result.as_slice());
     }
 }