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}