zstd/bulk/
compressor.rs

1use crate::map_error_code;
2
3use std::io;
4use zstd_safe;
5
6/// Allows to compress independently multiple chunks of data.
7///
8/// Each job will be processed entirely in-memory without streaming, so this
9/// is most fitting for many small jobs. To compress larger volume that don't
10/// easily fit in memory, a streaming compression may be more appropriate.
11///
12/// It is more efficient than a streaming compressor for 2 reasons:
13/// * It re-uses the zstd context between jobs to avoid re-allocations
14/// * It avoids copying data from a `Read` into a temporary buffer before compression.
15#[derive(Default)]
16pub struct Compressor<'a> {
17    context: zstd_safe::CCtx<'a>,
18}
19
20impl Compressor<'static> {
21    /// Creates a new zstd compressor
22    pub fn new(level: i32) -> io::Result<Self> {
23        Self::with_dictionary(level, &[])
24    }
25
26    /// Creates a new zstd compressor, using the given dictionary.
27    ///
28    /// Note that using a dictionary means that decompression will need to use
29    /// the same dictionary.
30    pub fn with_dictionary(level: i32, dictionary: &[u8]) -> io::Result<Self> {
31        let mut compressor = Self::default();
32
33        compressor.set_dictionary(level, dictionary)?;
34
35        Ok(compressor)
36    }
37}
38
39impl<'a> Compressor<'a> {
40    /// Creates a new compressor using an existing `EncoderDictionary`.
41    ///
42    /// The compression level will be the one specified when creating the dictionary.
43    ///
44    /// Note that using a dictionary means that decompression will need to use
45    /// the same dictionary.
46    pub fn with_prepared_dictionary<'b>(
47        dictionary: &'a crate::dict::EncoderDictionary<'b>,
48    ) -> io::Result<Self>
49    where
50        'b: 'a,
51    {
52        let mut compressor = Self::default();
53
54        compressor.set_prepared_dictionary(dictionary)?;
55
56        Ok(compressor)
57    }
58
59    /// Changes the compression level used by this compressor.
60    ///
61    /// *This will clear any dictionary previously registered.*
62    ///
63    /// If you want to keep the existing dictionary, you will need to pass it again to
64    /// `Self::set_dictionary` instead of using this method.
65    pub fn set_compression_level(&mut self, level: i32) -> io::Result<()> {
66        self.set_dictionary(level, &[])
67    }
68
69    /// Changes the dictionary and compression level used by this compressor.
70    ///
71    /// Will affect future compression jobs.
72    ///
73    /// Note that using a dictionary means that decompression will need to use
74    /// the same dictionary.
75    pub fn set_dictionary(
76        &mut self,
77        level: i32,
78        dictionary: &[u8],
79    ) -> io::Result<()> {
80        self.context
81            .set_parameter(zstd_safe::CParameter::CompressionLevel(level))
82            .map_err(map_error_code)?;
83
84        self.context
85            .load_dictionary(dictionary)
86            .map_err(map_error_code)?;
87
88        Ok(())
89    }
90
91    /// Changes the dictionary used by this compressor.
92    ///
93    /// The compression level used when preparing the dictionary will be used.
94    ///
95    /// Note that using a dictionary means that decompression will need to use
96    /// the same dictionary.
97    pub fn set_prepared_dictionary<'b>(
98        &mut self,
99        dictionary: &'a crate::dict::EncoderDictionary<'b>,
100    ) -> io::Result<()>
101    where
102        'b: 'a,
103    {
104        self.context
105            .ref_cdict(dictionary.as_cdict())
106            .map_err(map_error_code)?;
107
108        Ok(())
109    }
110
111    /// Compress a single block of data to the given destination buffer.
112    ///
113    /// Returns the number of bytes written, or an error if something happened
114    /// (for instance if the destination buffer was too small).
115    ///
116    /// A level of `0` uses zstd's default (currently `3`).
117    pub fn compress_to_buffer<C: zstd_safe::WriteBuf + ?Sized>(
118        &mut self,
119        source: &[u8],
120        destination: &mut C,
121    ) -> io::Result<usize> {
122        self.context
123            .compress2(destination, source)
124            .map_err(map_error_code)
125    }
126
127    /// Compresses a block of data and returns the compressed result.
128    ///
129    /// A level of `0` uses zstd's default (currently `3`).
130    pub fn compress(&mut self, data: &[u8]) -> io::Result<Vec<u8>> {
131        // We allocate a big buffer, slightly larger than the input data.
132        let buffer_len = zstd_safe::compress_bound(data.len());
133        let mut buffer = Vec::with_capacity(buffer_len);
134
135        self.compress_to_buffer(data, &mut buffer)?;
136
137        // Should we shrink the vec? Meh, let the user do it if he wants.
138        Ok(buffer)
139    }
140
141    /// Gives mutable access to the internal context.
142    pub fn context_mut(&mut self) -> &mut zstd_safe::CCtx<'a> {
143        &mut self.context
144    }
145
146    /// Sets a compression parameter for this compressor.
147    pub fn set_parameter(
148        &mut self,
149        parameter: zstd_safe::CParameter,
150    ) -> io::Result<()> {
151        self.context
152            .set_parameter(parameter)
153            .map_err(map_error_code)?;
154        Ok(())
155    }
156
157    crate::encoder_parameters!();
158}
159
160fn _assert_traits() {
161    fn _assert_send<T: Send>(_: T) {}
162
163    _assert_send(Compressor::new(0));
164}