use std::io::{self, Write};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
#[repr(u8)]
pub enum Compression {
Unknown = 0,
None = 1,
#[default]
Gzip = 2,
Brotli = 3,
Zstd = 4,
}
impl Compression {
#[allow(clippy::should_implement_trait)]
pub fn from_str(s: &str) -> Option<Self> {
match s.to_lowercase().as_str() {
"none" => Some(Compression::None),
"gzip" => Some(Compression::Gzip),
"brotli" => Some(Compression::Brotli),
"zstd" => Some(Compression::Zstd),
_ => Option::None,
}
}
pub fn code(&self) -> u8 {
*self as u8
}
pub fn from_code(code: u8) -> Option<Self> {
match code {
0 => Some(Compression::Unknown),
1 => Some(Compression::None),
2 => Some(Compression::Gzip),
3 => Some(Compression::Brotli),
4 => Some(Compression::Zstd),
_ => Option::None,
}
}
pub fn name(&self) -> &'static str {
match self {
Compression::Unknown => "unknown",
Compression::None => "none",
Compression::Gzip => "gzip",
Compression::Brotli => "brotli",
Compression::Zstd => "zstd",
}
}
}
struct Codec {
verb: &'static str,
gzip: fn(&[u8]) -> io::Result<Vec<u8>>,
brotli: fn(&[u8]) -> io::Result<Vec<u8>>,
zstd: fn(&[u8]) -> io::Result<Vec<u8>>,
}
const COMPRESS: &Codec = &Codec {
verb: "compress",
gzip: compress_gzip,
brotli: compress_brotli,
zstd: compress_zstd,
};
const DECOMPRESS: &Codec = &Codec {
verb: "decompress",
gzip: decompress_gzip,
brotli: decompress_brotli,
zstd: decompress_zstd,
};
fn dispatch(data: &[u8], compression: Compression, codec: &Codec) -> io::Result<Vec<u8>> {
match compression {
Compression::Unknown => Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!("Cannot {} with unknown compression type", codec.verb),
)),
Compression::None => Ok(data.to_vec()),
Compression::Gzip => (codec.gzip)(data),
Compression::Brotli => (codec.brotli)(data),
Compression::Zstd => (codec.zstd)(data),
}
}
pub fn compress(data: &[u8], compression: Compression) -> io::Result<Vec<u8>> {
dispatch(data, compression, COMPRESS)
}
pub fn decompress(data: &[u8], compression: Compression) -> io::Result<Vec<u8>> {
dispatch(data, compression, DECOMPRESS)
}
fn decompress_gzip(data: &[u8]) -> io::Result<Vec<u8>> {
use std::io::Read;
let mut out = Vec::new();
flate2::read::GzDecoder::new(data).read_to_end(&mut out)?;
Ok(out)
}
fn decompress_brotli(data: &[u8]) -> io::Result<Vec<u8>> {
use std::io::Read;
let mut out = Vec::new();
brotli::Decompressor::new(data, 4096).read_to_end(&mut out)?;
Ok(out)
}
fn decompress_zstd(data: &[u8]) -> io::Result<Vec<u8>> {
zstd::decode_all(data)
}
fn compress_gzip(data: &[u8]) -> io::Result<Vec<u8>> {
use flate2::write::GzEncoder;
use flate2::Compression as GzCompression;
let mut encoder = GzEncoder::new(Vec::new(), GzCompression::default());
encoder.write_all(data)?;
encoder.finish()
}
fn compress_brotli(data: &[u8]) -> io::Result<Vec<u8>> {
use brotli::enc::BrotliEncoderParams;
use brotli::CompressorWriter;
let params = BrotliEncoderParams {
quality: 4,
..Default::default()
};
let mut output = Vec::new();
{
let mut writer = CompressorWriter::with_params(&mut output, 4096, ¶ms);
writer.write_all(data)?;
}
Ok(output)
}
fn compress_zstd(data: &[u8]) -> io::Result<Vec<u8>> {
zstd::encode_all(data, 3)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_compression_codes_match_pmtiles_spec() {
assert_eq!(Compression::Unknown.code(), 0);
assert_eq!(Compression::None.code(), 1);
assert_eq!(Compression::Gzip.code(), 2);
assert_eq!(Compression::Brotli.code(), 3);
assert_eq!(Compression::Zstd.code(), 4);
}
#[test]
fn test_compression_default_is_gzip() {
assert_eq!(Compression::default(), Compression::Gzip);
}
#[test]
fn test_compression_from_str() {
assert_eq!(Compression::from_str("none"), Some(Compression::None));
assert_eq!(Compression::from_str("gzip"), Some(Compression::Gzip));
assert_eq!(Compression::from_str("brotli"), Some(Compression::Brotli));
assert_eq!(Compression::from_str("zstd"), Some(Compression::Zstd));
assert_eq!(Compression::from_str("GZIP"), Some(Compression::Gzip)); assert_eq!(Compression::from_str("invalid"), Option::None);
}
#[test]
fn test_compression_names() {
assert_eq!(Compression::None.name(), "none");
assert_eq!(Compression::Gzip.name(), "gzip");
assert_eq!(Compression::Brotli.name(), "brotli");
assert_eq!(Compression::Zstd.name(), "zstd");
}
#[test]
fn test_compress_none_returns_original() {
let data = b"Hello, PMTiles!";
let compressed = compress(data, Compression::None).unwrap();
assert_eq!(compressed, data);
}
#[test]
fn test_compress_unknown_returns_error() {
let data = b"Hello, PMTiles!";
let result = compress(data, Compression::Unknown);
assert!(result.is_err());
}
#[test]
fn test_compress_gzip_produces_smaller_output() {
let data = "Hello, PMTiles! ".repeat(100);
let compressed = compress(data.as_bytes(), Compression::Gzip).unwrap();
assert!(
compressed.len() < data.len(),
"Gzip should compress repetitive data: {} < {}",
compressed.len(),
data.len()
);
}
#[test]
fn test_compress_brotli_produces_smaller_output() {
let data = "Hello, PMTiles! ".repeat(100);
let compressed = compress(data.as_bytes(), Compression::Brotli).unwrap();
assert!(
compressed.len() < data.len(),
"Brotli should compress repetitive data: {} < {}",
compressed.len(),
data.len()
);
}
#[test]
fn test_compress_zstd_produces_smaller_output() {
let data = "Hello, PMTiles! ".repeat(100);
let compressed = compress(data.as_bytes(), Compression::Zstd).unwrap();
assert!(
compressed.len() < data.len(),
"Zstd should compress repetitive data: {} < {}",
compressed.len(),
data.len()
);
}
#[test]
fn test_gzip_roundtrip() {
use flate2::read::GzDecoder;
use std::io::Read;
let original = b"Hello, PMTiles! This is test data for compression roundtrip.";
let compressed = compress(original, Compression::Gzip).unwrap();
let mut decoder = GzDecoder::new(&compressed[..]);
let mut decompressed = Vec::new();
decoder.read_to_end(&mut decompressed).unwrap();
assert_eq!(decompressed, original);
}
#[test]
fn test_brotli_roundtrip() {
use brotli::Decompressor;
use std::io::Read;
let original = b"Hello, PMTiles! This is test data for compression roundtrip.";
let compressed = compress(original, Compression::Brotli).unwrap();
let mut decompressor = Decompressor::new(&compressed[..], 4096);
let mut decompressed = Vec::new();
decompressor.read_to_end(&mut decompressed).unwrap();
assert_eq!(decompressed, original);
}
#[test]
fn test_zstd_roundtrip() {
let original = b"Hello, PMTiles! This is test data for compression roundtrip.";
let compressed = compress(original, Compression::Zstd).unwrap();
let decompressed = zstd::decode_all(&compressed[..]).unwrap();
assert_eq!(decompressed, original);
}
#[test]
fn test_compress_empty_data() {
for compression in [
Compression::None,
Compression::Gzip,
Compression::Brotli,
Compression::Zstd,
] {
let result = compress(&[], compression);
assert!(
result.is_ok(),
"{} should handle empty data",
compression.name()
);
}
}
#[test]
fn test_compress_large_data() {
let data = vec![0x42u8; 1_000_000];
for compression in [Compression::Gzip, Compression::Brotli, Compression::Zstd] {
let result = compress(&data, compression);
assert!(
result.is_ok(),
"{} should handle large data",
compression.name()
);
let compressed = result.unwrap();
assert!(
compressed.len() < data.len() / 10,
"{} should achieve >10x compression on uniform data",
compression.name()
);
}
}
#[test]
fn test_from_code_inverts_code() {
for compression in [
Compression::Unknown,
Compression::None,
Compression::Gzip,
Compression::Brotli,
Compression::Zstd,
] {
assert_eq!(
Compression::from_code(compression.code()),
Some(compression)
);
}
assert_eq!(Compression::from_code(5), None);
assert_eq!(Compression::from_code(255), None);
}
#[test]
fn test_decompress_roundtrips_every_codec() {
let original = b"PMTiles round-trip payload \x00\x01\x02 with some repetition repetition";
for compression in [
Compression::None,
Compression::Gzip,
Compression::Brotli,
Compression::Zstd,
] {
let compressed = compress(original, compression).unwrap();
let decompressed = decompress(&compressed, compression).unwrap();
assert_eq!(
decompressed,
original.to_vec(),
"{} round-trip",
compression.name()
);
}
}
#[test]
fn test_decompress_empty_payload() {
for compression in [
Compression::None,
Compression::Gzip,
Compression::Brotli,
Compression::Zstd,
] {
let compressed = compress(&[], compression).unwrap();
let decompressed = decompress(&compressed, compression).unwrap();
assert!(decompressed.is_empty(), "{}", compression.name());
}
}
#[test]
fn test_decompress_unknown_is_error() {
assert!(decompress(b"anything", Compression::Unknown).is_err());
}
#[test]
fn test_decompress_corrupt_gzip_is_error() {
assert!(decompress(b"not gzip at all", Compression::Gzip).is_err());
}
}