solana_binary_encoder/
compression.rs

1use {
2    enum_iterator::{all, Sequence},
3    std::io::{self, BufReader, Read, Write},
4    serde_derive::{Serialize,Deserialize},
5};
6
7#[derive(Debug, Serialize, Deserialize, Sequence)]
8pub enum CompressionMethod {
9    NoCompression,
10    Bzip2,
11    Gzip,
12    Zstd,
13}
14
15fn decompress_reader<'a, R: Read + 'a>(
16    method: CompressionMethod,
17    stream: R,
18) -> Result<Box<dyn Read + 'a>, io::Error> {
19    let buf_reader = BufReader::new(stream);
20    let decompress_reader: Box<dyn Read> = match method {
21        CompressionMethod::Bzip2 => Box::new(bzip2::bufread::BzDecoder::new(buf_reader)),
22        CompressionMethod::Gzip => Box::new(flate2::read::GzDecoder::new(buf_reader)),
23        CompressionMethod::Zstd => Box::new(zstd::stream::read::Decoder::new(buf_reader)?),
24        CompressionMethod::NoCompression => Box::new(buf_reader),
25    };
26    Ok(decompress_reader)
27}
28
29pub fn decompress(data: &[u8]) -> Result<Vec<u8>, io::Error> {
30    let method_size = bincode::serialized_size(&CompressionMethod::NoCompression).unwrap();
31    if (data.len() as u64) < method_size {
32        return Err(io::Error::new(
33            io::ErrorKind::Other,
34            format!("data len too small: {}", data.len()),
35        ));
36    }
37    let method = bincode::deserialize(&data[..method_size as usize]).map_err(|err| {
38        io::Error::new(
39            io::ErrorKind::Other,
40            format!("method deserialize failed: {err}"),
41        )
42    })?;
43
44    let mut reader = decompress_reader(method, &data[method_size as usize..])?;
45    let mut uncompressed_data = vec![];
46    reader.read_to_end(&mut uncompressed_data)?;
47    Ok(uncompressed_data)
48}
49
50pub fn compress(method: CompressionMethod, data: &[u8]) -> Result<Vec<u8>, io::Error> {
51    let mut compressed_data = bincode::serialize(&method).unwrap();
52    compressed_data.extend(
53        match method {
54            CompressionMethod::Bzip2 => {
55                let mut e = bzip2::write::BzEncoder::new(Vec::new(), bzip2::Compression::best());
56                e.write_all(data)?;
57                e.finish()?
58            }
59            CompressionMethod::Gzip => {
60                let mut e =
61                    flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
62                e.write_all(data)?;
63                e.finish()?
64            }
65            CompressionMethod::Zstd => {
66                let mut e = zstd::stream::write::Encoder::new(Vec::new(), 0).unwrap();
67                e.write_all(data)?;
68                e.finish()?
69            }
70            CompressionMethod::NoCompression => data.to_vec(),
71        }
72            .into_iter(),
73    );
74
75    Ok(compressed_data)
76}
77
78pub fn compress_best(data: &[u8]) -> Result<Vec<u8>, io::Error> {
79    let mut candidates = vec![];
80    for method in all::<CompressionMethod>() {
81        candidates.push(compress(method, data)?);
82    }
83
84    Ok(candidates
85        .into_iter()
86        .min_by(|a, b| a.len().cmp(&b.len()))
87        .unwrap())
88}
89