Skip to main content

compression_codecs/mbrotli/
encoder.rs

1use crate::{mbrotli::params::EncoderParams, EncodeV2};
2use compression_core::util::{PartialBuffer, WriteBuffer};
3use mbrotli::{Compressor, EncoderSessionOwned, EncoderStatus, Operation};
4use std::io;
5
6/// Feeds `input` to the session, writing straight into the uninitialized tail of `output`.
7///
8/// Returns how much input was consumed and what the session needs next.
9fn process(
10    session: &mut EncoderSessionOwned,
11    input: &[u8],
12    output: &mut WriteBuffer<'_>,
13    operation: Operation,
14) -> io::Result<(usize, EncoderStatus)> {
15    // SAFETY: `process_uninit` never de-initializes bytes of `output`.
16    let progress = session.process_uninit(input, unsafe { output.unwritten_mut() }, operation)?;
17    // SAFETY: `process_uninit` initializes exactly `produced` leading bytes of `output`.
18    unsafe { output.assume_init_and_advance(progress.produced) };
19    Ok((progress.consumed, progress.status))
20}
21
22#[derive(Debug)]
23pub struct MbrotliEncoder {
24    session: EncoderSessionOwned,
25}
26
27impl MbrotliEncoder {
28    pub fn new(params: EncoderParams) -> Self {
29        // `EncoderParams` only builds standard windows, the one combination mbrotli rejects at
30        // configuration time is a large window at a low quality.
31        let compressor = Compressor::new(params.config()).unwrap();
32        // A fresh compressor with a zero stream offset has no way to refuse a session other than
33        // failing to allocate its workspace.
34        let session = compressor.into_session(params.stream()).unwrap();
35        Self { session }
36    }
37}
38
39impl EncodeV2 for MbrotliEncoder {
40    fn encode(
41        &mut self,
42        input: &mut PartialBuffer<&[u8]>,
43        output: &mut WriteBuffer<'_>,
44    ) -> io::Result<()> {
45        let (consumed, _) = process(
46            &mut self.session,
47            input.unwritten(),
48            output,
49            Operation::Process,
50        )?;
51        input.advance(consumed);
52
53        Ok(())
54    }
55
56    fn flush(&mut self, output: &mut WriteBuffer<'_>) -> io::Result<bool> {
57        let (_, status) = process(&mut self.session, &[], output, Operation::Flush)?;
58
59        match status {
60            EncoderStatus::NeedsInput | EncoderStatus::Finished => Ok(true),
61            EncoderStatus::NeedsOutput => Ok(false),
62        }
63    }
64
65    fn finish(&mut self, output: &mut WriteBuffer<'_>) -> io::Result<bool> {
66        let (_, status) = process(&mut self.session, &[], output, Operation::Finish)?;
67
68        match status {
69            EncoderStatus::Finished => Ok(true),
70            EncoderStatus::NeedsInput | EncoderStatus::NeedsOutput => Ok(false),
71        }
72    }
73}