assetpack_core/
codec.rs

1use std::io::Cursor;
2
3use brotli::enc::backward_references::BrotliEncoderParams;
4
5use crate::error::{Error, Result};
6
7pub const SAMPLE_LIMIT: usize = 64 * 1024;
8const COMPRESSED_WHITELIST: &[&str] = &[
9  "jpg", "jpeg", "png", "webp", "zip", "epub", "gz", "tgz", "bz2", "xz", "br", "zst", "mp4", "mov", "pdf",
10];
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum Codec {
14  Raw,
15  Zstd,
16  Brotli,
17}
18
19pub struct EntropyEstimator;
20
21impl EntropyEstimator {
22  pub fn shannon_entropy(bytes: &[u8]) -> f64 {
23    let sample = if bytes.len() > SAMPLE_LIMIT {
24      &bytes[..SAMPLE_LIMIT]
25    } else {
26      bytes
27    };
28    let mut counts = [0usize; 256];
29    for b in sample {
30      counts[*b as usize] += 1;
31    }
32    let total = sample.len() as f64;
33    counts
34      .iter()
35      .filter(|&&c| c > 0)
36      .map(|&c| {
37        let p = c as f64 / total;
38        -p * p.log2()
39      })
40      .sum()
41  }
42}
43
44fn is_compressed_extension(ext: &str) -> bool {
45  COMPRESSED_WHITELIST.contains(&ext)
46}
47
48pub fn decide_codec(data: &[u8], extension: Option<&str>) -> Codec {
49  if let Some(ext) = extension.map(|s| s.to_ascii_lowercase())
50    && is_compressed_extension(&ext)
51  {
52    return Codec::Raw;
53  }
54
55  let entropy = EntropyEstimator::shannon_entropy(data);
56  if entropy >= 7.7 {
57    Codec::Raw
58  } else if entropy >= 6.9 {
59    Codec::Zstd
60  } else {
61    Codec::Brotli
62  }
63}
64
65pub fn compress(codec: Codec, data: &[u8]) -> Result<Vec<u8>> {
66  match codec {
67    Codec::Raw => Ok(data.to_vec()),
68    Codec::Zstd => zstd::encode_all(Cursor::new(data), 5).map_err(|e| Error::Compress(e.to_string())),
69    Codec::Brotli => {
70      let params = BrotliEncoderParams {
71        quality: 5,
72        lgwin: 20,
73        ..Default::default()
74      };
75      let mut output = Vec::new();
76      brotli::BrotliCompress(&mut Cursor::new(data), &mut output, &params).map_err(|e| Error::Compress(e.to_string()))?;
77      Ok(output)
78    }
79  }
80}
81
82pub fn decompress(codec: Codec, data: &[u8]) -> Result<Vec<u8>> {
83  match codec {
84    Codec::Raw => Ok(data.to_vec()),
85    Codec::Zstd => {
86      let mut out = Vec::new();
87      zstd::stream::copy_decode(Cursor::new(data), &mut out).map_err(|e| Error::Decompress(e.to_string()))?;
88      Ok(out)
89    }
90    Codec::Brotli => {
91      let mut out = Vec::new();
92      brotli::BrotliDecompress(&mut Cursor::new(data), &mut out).map_err(|e| Error::Decompress(e.to_string()))?;
93      Ok(out)
94    }
95  }
96}
97
98pub fn compress_auto(data: &[u8], extension: Option<&str>) -> Result<(Codec, Vec<u8>)> {
99  let choice = decide_codec(data, extension);
100  let compressed = compress(choice, data)?;
101  if choice == Codec::Raw {
102    return Ok((choice, compressed));
103  }
104  let ratio = compressed.len() as f64 / data.len().max(1) as f64;
105  if ratio >= 0.98 {
106    Ok((Codec::Raw, data.to_vec()))
107  } else {
108    Ok((choice, compressed))
109  }
110}
111
112#[cfg(test)]
113mod tests {
114  use rand::Rng;
115
116  use super::*;
117
118  #[test]
119  fn entropy_detects_noise_and_structure() {
120    let mut random = [0u8; 1024];
121    rand::rng().fill_bytes(&mut random);
122    let random_h = EntropyEstimator::shannon_entropy(&random);
123    let structured = vec![0u8; 1024];
124    let structured_h = EntropyEstimator::shannon_entropy(&structured);
125    assert!(random_h > 7.5);
126    assert!(structured_h < 1.0);
127  }
128
129  #[test]
130  fn roundtrip_for_all_codecs() {
131    let data = b"The quick brown fox jumps over the lazy dog".repeat(64);
132    for codec in [Codec::Raw, Codec::Zstd, Codec::Brotli] {
133      let compressed = compress(codec, &data).unwrap();
134      let restored = decompress(codec, &compressed).unwrap();
135      assert_eq!(restored, data);
136    }
137  }
138
139  #[test]
140  fn auto_prefers_brotli_for_low_entropy() {
141    let data = vec![1u8; 4096];
142    let (codec, compressed) = compress_auto(&data, None).unwrap();
143    assert_eq!(codec, Codec::Brotli);
144    assert!(compressed.len() < data.len());
145  }
146
147  #[test]
148  fn auto_falls_back_to_raw_when_not_smaller() {
149    let data = vec![0u8; 8];
150    let (codec, compressed) = compress_auto(&data, None).unwrap();
151    assert_eq!(codec, Codec::Raw);
152    assert_eq!(compressed, data);
153  }
154
155  #[test]
156  fn auto_respects_compressed_whitelist_extensions() {
157    let mut data = vec![0u8; 4096];
158    rand::rng().fill_bytes(&mut data);
159    let (codec, compressed) = compress_auto(&data, Some("jpg")).unwrap();
160    assert_eq!(codec, Codec::Raw);
161    assert_eq!(compressed, data);
162  }
163
164  #[test]
165  fn auto_skips_compression_for_high_entropy_data() {
166    let data = (0u8..=255).cycle().take(64 * 1024).collect::<Vec<_>>();
167    let (codec, compressed) = compress_auto(&data, None).unwrap();
168    assert_eq!(codec, Codec::Raw);
169    assert_eq!(compressed, data);
170  }
171}
assetpack_core/codec.rs

assetpack_core/
codec.rs
