zstd/stream/write/
mod.rs

Help
1//! Implement push-based [`Write`] trait for both compressing and decompressing.
2use std::io::{self, Write};
3
4use zstd_safe;
5
6use crate::dict::{DecoderDictionary, EncoderDictionary};
7use crate::stream::{raw, zio};
8
9#[cfg(test)]
10mod tests;
11
12/// An encoder that compress and forward data to another writer.
13///
14/// This allows to compress a stream of data
15/// (good for files or heavy network stream).
16///
17/// Don't forget to call [`finish()`] before dropping it!
18///
19/// Alternatively, you can call [`auto_finish()`] to use an
20/// [`AutoFinishEncoder`] that will finish on drop.
21///
22/// Note: The zstd library has its own internal input buffer (~128kb).
23///
24/// [`finish()`]: #method.finish
25/// [`auto_finish()`]: #method.auto_finish
26/// [`AutoFinishEncoder`]: AutoFinishEncoder
27pub struct Encoder<'a, W: Write> {
28    // output writer (compressed data)
29    writer: zio::Writer<W, raw::Encoder<'a>>,
30}
31
32/// A decoder that decompress and forward data to another writer.
33///
34/// Note that you probably want to `flush()` after writing your stream content.
35/// You can use [`auto_flush()`] to automatically flush the writer on drop.
36///
37/// [`auto_flush()`]: Decoder::auto_flush
38pub struct Decoder<'a, W: Write> {
39    // output writer (decompressed data)
40    writer: zio::Writer<W, raw::Decoder<'a>>,
41}
42
43/// A wrapper around an `Encoder<W>` that finishes the stream on drop.
44///
45/// This can be created by the [`auto_finish()`] method on the [`Encoder`].
46///
47/// [`auto_finish()`]: Encoder::auto_finish
48/// [`Encoder`]: Encoder
49pub struct AutoFinishEncoder<
50    'a,
51    W: Write,
52    F: FnMut(io::Result<W>) = Box<dyn Send + FnMut(io::Result<W>)>,
53> {
54    // We wrap this in an option to take it during drop.
55    encoder: Option<Encoder<'a, W>>,
56
57    on_finish: Option<F>,
58}
59
60/// A wrapper around a `Decoder<W>` that flushes the stream on drop.
61///
62/// This can be created by the [`auto_flush()`] method on the [`Decoder`].
63///
64/// [`auto_flush()`]: Decoder::auto_flush
65/// [`Decoder`]: Decoder
66pub struct AutoFlushDecoder<
67    'a,
68    W: Write,
69    F: FnMut(io::Result<()>) = Box<dyn Send + FnMut(io::Result<()>)>,
70> {
71    // We wrap this in an option to take it during drop.
72    decoder: Option<Decoder<'a, W>>,
73
74    on_flush: Option<F>,
75}
76
77impl<'a, W: Write, F: FnMut(io::Result<()>)> AutoFlushDecoder<'a, W, F> {
78    fn new(decoder: Decoder<'a, W>, on_flush: F) -> Self {
79        AutoFlushDecoder {
80            decoder: Some(decoder),
81            on_flush: Some(on_flush),
82        }
83    }
84
85    /// Acquires a reference to the underlying writer.
86    pub fn get_ref(&self) -> &W {
87        self.decoder.as_ref().unwrap().get_ref()
88    }
89
90    /// Acquires a mutable reference to the underlying writer.
91    ///
92    /// Note that mutation of the writer may result in surprising results if
93    /// this decoder is continued to be used.
94    ///
95    /// Mostly used for testing purposes.
96    pub fn get_mut(&mut self) -> &mut W {
97        self.decoder.as_mut().unwrap().get_mut()
98    }
99}
100
101impl<W, F> Drop for AutoFlushDecoder<'_, W, F>
102where
103    W: Write,
104    F: FnMut(io::Result<()>),
105{
106    fn drop(&mut self) {
107        let mut decoder = self.decoder.take().unwrap();
108        let result = decoder.flush();
109        if let Some(mut on_finish) = self.on_flush.take() {
110            on_finish(result);
111        }
112    }
113}
114
115impl<W: Write, F: FnMut(io::Result<()>)> Write for AutoFlushDecoder<'_, W, F> {
116    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
117        self.decoder.as_mut().unwrap().write(buf)
118    }
119
120    fn flush(&mut self) -> io::Result<()> {
121        self.decoder.as_mut().unwrap().flush()
122    }
123}
124
125impl<'a, W: Write, F: FnMut(io::Result<W>)> AutoFinishEncoder<'a, W, F> {
126    fn new(encoder: Encoder<'a, W>, on_finish: F) -> Self {
127        AutoFinishEncoder {
128            encoder: Some(encoder),
129            on_finish: Some(on_finish),
130        }
131    }
132
133    /// Acquires a reference to the underlying writer.
134    pub fn get_ref(&self) -> &W {
135        self.encoder.as_ref().unwrap().get_ref()
136    }
137
138    /// Acquires a mutable reference to the underlying writer.
139    ///
140    /// Note that mutation of the writer may result in surprising results if
141    /// this encoder is continued to be used.
142    ///
143    /// Mostly used for testing purposes.
144    pub fn get_mut(&mut self) -> &mut W {
145        self.encoder.as_mut().unwrap().get_mut()
146    }
147}
148
149impl<W: Write, F: FnMut(io::Result<W>)> Drop for AutoFinishEncoder<'_, W, F> {
150    fn drop(&mut self) {
151        let result = self.encoder.take().unwrap().finish();
152        if let Some(mut on_finish) = self.on_finish.take() {
153            on_finish(result);
154        }
155    }
156}
157
158impl<W: Write, F: FnMut(io::Result<W>)> Write for AutoFinishEncoder<'_, W, F> {
159    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
160        self.encoder.as_mut().unwrap().write(buf)
161    }
162
163    fn flush(&mut self) -> io::Result<()> {
164        self.encoder.as_mut().unwrap().flush()
165    }
166}
167
168impl<W: Write> Encoder<'static, W> {
169    /// Creates a new encoder.
170    ///
171    /// `level`: compression level (1-21).
172    ///
173    /// A level of `0` uses zstd's default (currently `3`).
174    pub fn new(writer: W, level: i32) -> io::Result<Self> {
175        Self::with_dictionary(writer, level, &[])
176    }
177
178    /// Creates a new encoder, using an existing dictionary.
179    ///
180    /// (Provides better compression ratio for small files,
181    /// but requires the dictionary to be present during decompression.)
182    ///
183    /// A level of `0` uses zstd's default (currently `3`).
184    pub fn with_dictionary(
185        writer: W,
186        level: i32,
187        dictionary: &[u8],
188    ) -> io::Result<Self> {
189        let encoder = raw::Encoder::with_dictionary(level, dictionary)?;
190        let writer = zio::Writer::new(writer, encoder);
191        Ok(Encoder { writer })
192    }
193}
194
195impl<'a, W: Write> Encoder<'a, W> {
196    /// Creates a new encoder, using an existing prepared `EncoderDictionary`.
197    ///
198    /// (Provides better compression ratio for small files,
199    /// but requires the dictionary to be present during decompression.)
200    pub fn with_prepared_dictionary<'b>(
201        writer: W,
202        dictionary: &EncoderDictionary<'b>,
203    ) -> io::Result<Self>
204    where
205        'b: 'a,
206    {
207        let encoder = raw::Encoder::with_prepared_dictionary(dictionary)?;
208        let writer = zio::Writer::new(writer, encoder);
209        Ok(Encoder { writer })
210    }
211
212    /// Returns a wrapper around `self` that will finish the stream on drop.
213    ///
214    /// # Panic
215    ///
216    /// Panics on drop if an error happens when finishing the stream.
217    pub fn auto_finish(self) -> AutoFinishEncoder<'a, W> {
218        self.on_finish(Box::new(|result| {
219            result.unwrap();
220        }))
221    }
222
223    /// Returns an encoder that will finish the stream on drop.
224    ///
225    /// Calls the given callback with the result from `finish()`.
226    pub fn on_finish<F: FnMut(io::Result<W>)>(
227        self,
228        f: F,
229    ) -> AutoFinishEncoder<'a, W, F> {
230        AutoFinishEncoder::new(self, f)
231    }
232
233    /// Acquires a reference to the underlying writer.
234    pub fn get_ref(&self) -> &W {
235        self.writer.writer()
236    }
237
238    /// Acquires a mutable reference to the underlying writer.
239    ///
240    /// Note that mutation of the writer may result in surprising results if
241    /// this encoder is continued to be used.
242    pub fn get_mut(&mut self) -> &mut W {
243        self.writer.writer_mut()
244    }
245
246    /// **Required**: Finishes the stream.
247    ///
248    /// You *need* to finish the stream when you're done writing, either with
249    /// this method or with [`try_finish(self)`](#method.try_finish).
250    ///
251    /// This returns the inner writer in case you need it.
252    ///
253    /// To get back `self` in case an error happened, use `try_finish`.
254    ///
255    /// **Note**: If you don't want (or can't) call `finish()` manually after
256    ///           writing your data, consider using `auto_finish()` to get an
257    ///           `AutoFinishEncoder`.
258    pub fn finish(self) -> io::Result<W> {
259        self.try_finish().map_err(|(_, err)| err)
260    }
261
262    /// **Required**: Attempts to finish the stream.
263    ///
264    /// You *need* to finish the stream when you're done writing, either with
265    /// this method or with [`finish(self)`](#method.finish).
266    ///
267    /// This returns the inner writer if the finish was successful, or the
268    /// object plus an error if it wasn't.
269    ///
270    /// `write` on this object will panic after `try_finish` has been called,
271    /// even if it fails.
272    pub fn try_finish(mut self) -> Result<W, (Self, io::Error)> {
273        match self.writer.finish() {
274            // Return the writer, because why not
275            Ok(()) => Ok(self.writer.into_inner().0),
276            Err(e) => Err((self, e)),
277        }
278    }
279
280    /// Attemps to finish the stream.
281    ///
282    /// You *need* to finish the stream when you're done writing, either with
283    /// this method or with [`finish(self)`](#method.finish).
284    pub fn do_finish(&mut self) -> io::Result<()> {
285        self.writer.finish()
286    }
287
288    /// Return a recommendation for the size of data to write at once.
289    pub fn recommended_input_size() -> usize {
290        zstd_safe::CCtx::in_size()
291    }
292
293    crate::encoder_common!(writer);
294}
295
296impl<'a, W: Write> Write for Encoder<'a, W> {
297    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
298        self.writer.write(buf)
299    }
300
301    fn flush(&mut self) -> io::Result<()> {
302        self.writer.flush()
303    }
304}
305
306impl<W: Write> Decoder<'static, W> {
307    /// Creates a new decoder.
308    pub fn new(writer: W) -> io::Result<Self> {
309        Self::with_dictionary(writer, &[])
310    }
311
312    /// Creates a new decoder, using an existing dictionary.
313    ///
314    /// (Provides better compression ratio for small files,
315    /// but requires the dictionary to be present during decompression.)
316    pub fn with_dictionary(writer: W, dictionary: &[u8]) -> io::Result<Self> {
317        let decoder = raw::Decoder::with_dictionary(dictionary)?;
318        let writer = zio::Writer::new(writer, decoder);
319        Ok(Decoder { writer })
320    }
321}
322
323impl<'a, W: Write> Decoder<'a, W> {
324    /// Creates a new decoder, using an existing prepared `DecoderDictionary`.
325    ///
326    /// (Provides better compression ratio for small files,
327    /// but requires the dictionary to be present during decompression.)
328    pub fn with_prepared_dictionary<'b>(
329        writer: W,
330        dictionary: &DecoderDictionary<'b>,
331    ) -> io::Result<Self>
332    where
333        'b: 'a,
334    {
335        let decoder = raw::Decoder::with_prepared_dictionary(dictionary)?;
336        let writer = zio::Writer::new(writer, decoder);
337        Ok(Decoder { writer })
338    }
339
340    /// Acquires a reference to the underlying writer.
341    pub fn get_ref(&self) -> &W {
342        self.writer.writer()
343    }
344
345    /// Acquires a mutable reference to the underlying writer.
346    ///
347    /// Note that mutation of the writer may result in surprising results if
348    /// this decoder is continued to be used.
349    pub fn get_mut(&mut self) -> &mut W {
350        self.writer.writer_mut()
351    }
352
353    /// Returns the inner `Write`.
354    pub fn into_inner(self) -> W {
355        self.writer.into_inner().0
356    }
357
358    /// Return a recommendation for the size of data to write at once.
359    pub fn recommended_input_size() -> usize {
360        zstd_safe::DCtx::in_size()
361    }
362
363    /// Returns a wrapper around `self` that will flush the stream on drop.
364    ///
365    /// # Panic
366    ///
367    /// Panics on drop if an error happens when flushing the stream.
368    pub fn auto_flush(self) -> AutoFlushDecoder<'a, W> {
369        self.on_flush(Box::new(|result| {
370            result.unwrap();
371        }))
372    }
373
374    /// Returns a decoder that will flush the stream on drop.
375    ///
376    /// Calls the given callback with the result from `flush()`.
377    pub fn on_flush<F: FnMut(io::Result<()>)>(
378        self,
379        f: F,
380    ) -> AutoFlushDecoder<'a, W, F> {
381        AutoFlushDecoder::new(self, f)
382    }
383
384    crate::decoder_common!(writer);
385}
386
387impl<W: Write> Write for Decoder<'_, W> {
388    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
389        self.writer.write(buf)
390    }
391
392    fn flush(&mut self) -> io::Result<()> {
393        self.writer.flush()
394    }
395}
396
397fn _assert_traits() {
398    fn _assert_send<T: Send>(_: T) {}
399
400    _assert_send(Decoder::new(Vec::new()));
401    _assert_send(Encoder::new(Vec::new(), 1));
402    _assert_send(Decoder::new(Vec::new()).unwrap().auto_flush());
403    _assert_send(Encoder::new(Vec::new(), 1).unwrap().auto_finish());
404}
zstd/stream/write/mod.rs

zstd/stream/write/
mod.rs
