compression_codecs/mzstd/
encoder.rs

1use crate::{
2    mzstd::params::{encoder_config, CParameter},
3    EncodeV2,
4};
5use compression_core::util::{PartialBuffer, WriteBuffer};
6use mzstd::{
7    dictionary::DictionaryBuilder, Compressor, EncoderSessionOwned, EncoderStatus, Operation,
8    StreamConfig,
9};
10use std::io;
11
12/// Feeds `input` to the session, writing straight into the uninitialized tail of `output`.
13///
14/// Returns how much input was consumed and what the session needs next.
15fn process(
16    session: &mut EncoderSessionOwned,
17    input: &[u8],
18    output: &mut WriteBuffer<'_>,
19    operation: Operation,
20) -> io::Result<(usize, EncoderStatus)> {
21    // SAFETY: `process_uninit` never de-initializes bytes of `output`.
22    let progress = session
23        .process_uninit(input, unsafe { output.unwritten_mut() }, operation)
24        .map_err(io::Error::other)?;
25    // SAFETY: `process_uninit` initializes exactly `produced` leading bytes of `output`.
26    unsafe { output.assume_init_and_advance(progress.produced) };
27    Ok((progress.consumed, progress.status))
28}
29
30#[derive(Debug)]
31pub struct MzstdEncoder {
32    session: EncoderSessionOwned,
33}
34
35impl MzstdEncoder {
36    pub fn new(level: i32) -> Self {
37        Self::new_with_params(level, &[])
38    }
39
40    /// # Panics
41    ///
42    /// Panics if a parameter is out of range, as the `zstd` backend does.
43    pub fn new_with_params(level: i32, params: &[CParameter]) -> Self {
44        let compressor = Compressor::new(encoder_config(level, params)).unwrap();
45        Self {
46            session: compressor.into_session(StreamConfig::default()),
47        }
48    }
49
50    pub fn new_with_dict(level: i32, dictionary: &[u8]) -> io::Result<Self> {
51        let config = encoder_config(level, &[]);
52        let dictionary = DictionaryBuilder::new(dictionary)
53            .with_config(config)
54            .build()
55            .map_err(|err| io::Error::new(io::ErrorKind::InvalidInput, err))?;
56        // `encoder_config` clamps the level and sets no other parameters, so it is valid.
57        let session = Compressor::new(config)
58            .unwrap()
59            .into_session_with_dictionary(dictionary, StreamConfig::default())
60            .map_err(io::Error::other)?;
61        Ok(Self { session })
62    }
63}
64
65impl EncodeV2 for MzstdEncoder {
66    fn encode(
67        &mut self,
68        input: &mut PartialBuffer<&[u8]>,
69        output: &mut WriteBuffer<'_>,
70    ) -> io::Result<()> {
71        let (consumed, _) = process(
72            &mut self.session,
73            input.unwritten(),
74            output,
75            Operation::Continue,
76        )?;
77        input.advance(consumed);
78
79        Ok(())
80    }
81
82    fn flush(&mut self, output: &mut WriteBuffer<'_>) -> io::Result<bool> {
83        let (_, status) = process(&mut self.session, &[], output, Operation::Flush)?;
84
85        match status {
86            EncoderStatus::NeedsInput | EncoderStatus::Finished => Ok(true),
87            EncoderStatus::NeedsOutput => Ok(false),
88        }
89    }
90
91    fn finish(&mut self, output: &mut WriteBuffer<'_>) -> io::Result<bool> {
92        let (_, status) = process(&mut self.session, &[], output, Operation::End)?;
93
94        match status {
95            EncoderStatus::Finished => Ok(true),
96            EncoderStatus::NeedsInput | EncoderStatus::NeedsOutput => Ok(false),
97        }
98    }
99}
compression_codecs/mzstd/encoder.rs

compression_codecs/mzstd/
encoder.rs
