solana_storage_utils/
compression.rs

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

solana_storage_utils/
compression.rs
