Skip to main content

compression_codecs/mzstd/
decoder.rs

1use crate::{
2    mzstd::params::{decoder_config, DParameter},
3    DecodeV2,
4};
5use compression_core::util::{PartialBuffer, WriteBuffer};
6use mzstd::{
7    dictionary::{DecodeDictionary, DictionaryContent},
8    DecodeOperation, DecoderSessionOwned, DecoderStatus, Decompressor,
9};
10use std::io;
11
12#[derive(Debug)]
13pub struct MzstdDecoder {
14    session: DecoderSessionOwned,
15}
16
17impl Default for MzstdDecoder {
18    fn default() -> Self {
19        Self::new_with_params(&[])
20    }
21}
22
23impl MzstdDecoder {
24    pub fn new() -> Self {
25        Self::default()
26    }
27
28    /// # Panics
29    ///
30    /// Panics if a parameter is out of range, as the `zstd` backend does.
31    pub fn new_with_params(params: &[DParameter]) -> Self {
32        let session = Decompressor::new(decoder_config(params)).into_session();
33        Self { session }
34    }
35
36    pub fn new_with_dict(dictionary: &[u8]) -> io::Result<Self> {
37        let dictionary = DecodeDictionary::new(dictionary, DictionaryContent::Auto)
38            .map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?;
39        let session =
40            Decompressor::new(decoder_config(&[])).into_session_with_dictionary(dictionary);
41        Ok(Self { session })
42    }
43}
44
45/// Feeds `input` to the session, writing straight into the uninitialized tail of `output`.
46///
47/// Returns how much input was consumed and what the session needs next.
48fn process(
49    session: &mut DecoderSessionOwned,
50    input: &[u8],
51    output: &mut WriteBuffer<'_>,
52    operation: DecodeOperation,
53) -> (usize, io::Result<DecoderStatus>) {
54    // SAFETY: `process_uninit` never de-initializes bytes of `output`.
55    let result = session.process_uninit(input, unsafe { output.unwritten_mut() }, operation);
56    // `process_uninit` initializes exactly `produced` leading bytes of `output`, both when it
57    // succeeds and when it fails, so progress made before an error is kept.
58    match result {
59        Ok(progress) => {
60            // SAFETY: see above.
61            unsafe { output.assume_init_and_advance(progress.produced) };
62            (progress.consumed, Ok(progress.status))
63        }
64        Err(failure) => {
65            // SAFETY: see above.
66            unsafe { output.assume_init_and_advance(failure.produced) };
67            (failure.consumed, Err(failure.into_error().into()))
68        }
69    }
70}
71
72impl DecodeV2 for MzstdDecoder {
73    fn reinit(&mut self) -> io::Result<()> {
74        self.session.reinit();
75        Ok(())
76    }
77
78    fn decode(
79        &mut self,
80        input: &mut PartialBuffer<&[u8]>,
81        output: &mut WriteBuffer<'_>,
82    ) -> io::Result<bool> {
83        let (consumed, status) = process(
84            &mut self.session,
85            input.unwritten(),
86            output,
87            DecodeOperation::Process,
88        );
89        input.advance(consumed);
90
91        match status? {
92            DecoderStatus::Finished => Ok(true),
93            DecoderStatus::NeedsInput | DecoderStatus::NeedsOutput => Ok(false),
94        }
95    }
96
97    fn flush(&mut self, output: &mut WriteBuffer<'_>) -> io::Result<bool> {
98        // Empty input with `Process` delivers pending output without declaring EOF.
99        let (_, status) = process(&mut self.session, &[], output, DecodeOperation::Process);
100
101        match status? {
102            DecoderStatus::Finished | DecoderStatus::NeedsInput => Ok(true),
103            DecoderStatus::NeedsOutput => Ok(false),
104        }
105    }
106
107    fn finish(&mut self, output: &mut WriteBuffer<'_>) -> io::Result<bool> {
108        // Declaring EOF inside a frame fails with `UnexpectedEof`.
109        let (_, status) = process(&mut self.session, &[], output, DecodeOperation::Finish);
110
111        match status? {
112            DecoderStatus::Finished => Ok(true),
113            DecoderStatus::NeedsInput | DecoderStatus::NeedsOutput => Ok(false),
114        }
115    }
116}