bitar/
compression.rs

1use bytes::Bytes;
2use std::fmt;
3
4use crate::chunk_dictionary as dict;
5
6#[derive(Debug)]
7pub enum CompressionError {
8    Io(std::io::Error),
9    #[cfg(feature = "lzma-compression")]
10    LZMA(lzma::LzmaError),
11}
12impl std::error::Error for CompressionError {
13    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
14        match self {
15            CompressionError::Io(err) => Some(err),
16            #[cfg(feature = "lzma-compression")]
17            CompressionError::LZMA(err) => Some(err),
18        }
19    }
20}
21impl fmt::Display for CompressionError {
22    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
23        match self {
24            Self::Io(_) => write!(f, "i/o error"),
25            #[cfg(feature = "lzma-compression")]
26            Self::LZMA(_) => write!(f, "LZMA error"),
27        }
28    }
29}
30impl From<std::io::Error> for CompressionError {
31    fn from(e: std::io::Error) -> Self {
32        Self::Io(e)
33    }
34}
35#[cfg(feature = "lzma-compression")]
36impl From<lzma::LzmaError> for CompressionError {
37    fn from(e: lzma::LzmaError) -> Self {
38        Self::LZMA(e)
39    }
40}
41
42#[derive(Debug)]
43pub struct CompressionLevelOutOfRangeError(CompressionAlgorithm);
44impl std::error::Error for CompressionLevelOutOfRangeError {}
45impl fmt::Display for CompressionLevelOutOfRangeError {
46    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
47        write!(
48            f,
49            "{} compression level out of range (valid range is 1-{})",
50            self.0,
51            self.0.max_level()
52        )
53    }
54}
55
56#[derive(Debug, Clone, Copy, PartialEq, Eq)]
57pub enum CompressionAlgorithm {
58    #[cfg(feature = "lzma-compression")]
59    Lzma,
60    #[cfg(feature = "zstd-compression")]
61    Zstd,
62    Brotli,
63}
64
65impl CompressionAlgorithm {
66    /// Get the compression algorithm's max level.
67    pub fn max_level(self) -> u32 {
68        match self {
69            #[cfg(feature = "lzma-compression")]
70            CompressionAlgorithm::Lzma => 9,
71            #[cfg(feature = "zstd-compression")]
72            CompressionAlgorithm::Zstd => {
73                u32::try_from(*zstd::compression_level_range().end()).unwrap()
74            }
75            CompressionAlgorithm::Brotli => 11,
76        }
77    }
78    /// Decompress a block of data using the set compression.
79    pub(crate) fn decompress(
80        self,
81        compressed: Bytes,
82        size_hint: usize,
83    ) -> Result<Bytes, CompressionError> {
84        let mut output = Vec::with_capacity(size_hint);
85        match self {
86            #[cfg(feature = "lzma-compression")]
87            CompressionAlgorithm::Lzma => {
88                use lzma::LzmaWriter;
89                use std::io::prelude::*;
90                let mut f = LzmaWriter::new_decompressor(&mut output)?;
91                f.write_all(&compressed)?;
92                f.finish()?;
93            }
94            #[cfg(feature = "zstd-compression")]
95            CompressionAlgorithm::Zstd => {
96                zstd::stream::copy_decode(&compressed[..], &mut output)?;
97            }
98            CompressionAlgorithm::Brotli => {
99                let mut input_slice = &compressed[..];
100                brotli_decompressor::BrotliDecompress(&mut input_slice, &mut output)?;
101            }
102        }
103        Ok(Bytes::from(output))
104    }
105}
106
107impl fmt::Display for CompressionAlgorithm {
108    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
109        let algorithm_name = match self {
110            #[cfg(feature = "lzma-compression")]
111            CompressionAlgorithm::Lzma => "LZMA",
112            #[cfg(feature = "zstd-compression")]
113            CompressionAlgorithm::Zstd => "zstd",
114            CompressionAlgorithm::Brotli => "Brotli",
115        };
116        write!(f, "{}", algorithm_name)
117    }
118}
119
120/// Compression.
121#[derive(Debug, Clone, Copy, PartialEq, Eq)]
122pub struct Compression {
123    pub(crate) algorithm: CompressionAlgorithm,
124    pub(crate) level: u32,
125}
126
127impl Compression {
128    /// Create a new compression of given algorithm and level.
129    pub fn try_new(
130        algorithm: CompressionAlgorithm,
131        level: u32,
132    ) -> Result<Compression, CompressionLevelOutOfRangeError> {
133        if level < 1 || level > algorithm.max_level() {
134            return Err(CompressionLevelOutOfRangeError(algorithm));
135        }
136        Ok(Compression { algorithm, level })
137    }
138    /// Create a new brotli compression of given level.
139    pub fn brotli(level: u32) -> Result<Compression, CompressionLevelOutOfRangeError> {
140        Self::try_new(CompressionAlgorithm::Brotli, level)
141    }
142    #[cfg(feature = "lzma-compression")]
143    /// Create a new lzma compression of given level.
144    pub fn lzma(level: u32) -> Result<Compression, CompressionLevelOutOfRangeError> {
145        Self::try_new(CompressionAlgorithm::Lzma, level)
146    }
147    #[cfg(feature = "zstd-compression")]
148    /// Create a new zstd compression of given level.
149    pub fn zstd(level: u32) -> Result<Compression, CompressionLevelOutOfRangeError> {
150        Self::try_new(CompressionAlgorithm::Zstd, level)
151    }
152    /// Compress a block of data with set compression.
153    #[cfg(feature = "compress")]
154    pub(crate) fn compress(self, chunk: Bytes) -> Result<Bytes, CompressionError> {
155        use brotli::enc::backward_references::BrotliEncoderParams;
156        use std::io::Write;
157        let mut output = Vec::with_capacity(chunk.len());
158        match self.algorithm {
159            #[cfg(feature = "lzma-compression")]
160            CompressionAlgorithm::Lzma => {
161                use lzma::LzmaWriter;
162                use std::io::prelude::*;
163                let mut f = LzmaWriter::new_compressor(&mut output, self.level)?;
164                f.write_all(&chunk)?;
165                f.finish()?;
166            }
167            #[cfg(feature = "zstd-compression")]
168            CompressionAlgorithm::Zstd => {
169                zstd::stream::copy_encode(&chunk[..], &mut output, self.level as i32)?;
170            }
171            CompressionAlgorithm::Brotli => {
172                let params = BrotliEncoderParams {
173                    quality: self.level as i32,
174                    magic_number: false,
175                    ..Default::default()
176                };
177                let mut writer =
178                    brotli::CompressorWriter::with_params(&mut output, 1024 * 1024, &params);
179                writer.write_all(&chunk)?;
180            }
181        }
182        Ok(Bytes::from(output))
183    }
184}
185
186impl fmt::Display for Compression {
187    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
188        write!(f, "{} (level {})", self.algorithm, self.level)
189    }
190}
191
192impl From<Option<Compression>> for dict::ChunkCompression {
193    fn from(c: Option<Compression>) -> Self {
194        let (compression, compression_level) = match c {
195            #[cfg(feature = "lzma-compression")]
196            Some(Compression {
197                algorithm: CompressionAlgorithm::Lzma,
198                level,
199            }) => (dict::chunk_compression::CompressionType::Lzma, level),
200            #[cfg(feature = "zstd-compression")]
201            Some(Compression {
202                algorithm: CompressionAlgorithm::Zstd,
203                level,
204            }) => (dict::chunk_compression::CompressionType::Zstd, level),
205            Some(Compression {
206                algorithm: CompressionAlgorithm::Brotli,
207                level,
208            }) => (dict::chunk_compression::CompressionType::Brotli, level),
209            None => (dict::chunk_compression::CompressionType::None, 0),
210        };
211        Self {
212            compression: compression as i32,
213            compression_level,
214        }
215    }
216}