brotli/enc/
writer.rs

1#![cfg_attr(not(feature = "std"), allow(unused_imports))]
2use super::backward_references::BrotliEncoderParams;
3use super::combined_alloc::BrotliAlloc;
4use super::encode::{
5    BrotliEncoderCompressStream, BrotliEncoderCreateInstance, BrotliEncoderDestroyInstance,
6    BrotliEncoderHasMoreOutput, BrotliEncoderIsFinished, BrotliEncoderOperation,
7    BrotliEncoderParameter, BrotliEncoderSetParameter, BrotliEncoderStateStruct,
8};
9use super::interface;
10pub use alloc::{AllocatedStackMemory, Allocator, SliceWrapper, SliceWrapperMut, StackAllocator};
11#[cfg(feature = "std")]
12pub use alloc_stdlib::StandardAlloc;
13use brotli_decompressor::CustomWrite;
14#[cfg(feature = "std")]
15pub use brotli_decompressor::{IntoIoWriter, IoWriterWrapper};
16
17#[cfg(feature = "std")]
18use std::io;
19
20#[cfg(feature = "std")]
21use std::io::{Error, ErrorKind, Write};
22
23#[cfg(feature = "std")]
24pub struct CompressorWriterCustomAlloc<
25    W: Write,
26    BufferType: SliceWrapperMut<u8>,
27    Alloc: BrotliAlloc,
28>(CompressorWriterCustomIo<io::Error, IntoIoWriter<W>, BufferType, Alloc>);
29
30#[cfg(feature = "std")]
31impl<W: Write, BufferType: SliceWrapperMut<u8>, Alloc: BrotliAlloc>
32    CompressorWriterCustomAlloc<W, BufferType, Alloc>
33{
34    pub fn new(w: W, buffer: BufferType, alloc: Alloc, q: u32, lgwin: u32) -> Self {
35        CompressorWriterCustomAlloc::<W, BufferType, Alloc>(CompressorWriterCustomIo::<
36            Error,
37            IntoIoWriter<W>,
38            BufferType,
39            Alloc,
40        >::new(
41            IntoIoWriter::<W>(w),
42            buffer,
43            alloc,
44            Error::new(ErrorKind::InvalidData, "Invalid Data"),
45            q,
46            lgwin,
47        ))
48    }
49
50    pub fn get_ref(&self) -> &W {
51        &self.0.get_ref().0
52    }
53    pub fn get_mut(&mut self) -> &mut W {
54        &mut self.0.get_mut().0
55    }
56    pub fn into_inner(self) -> W {
57        self.0.into_inner().0
58    }
59}
60
61#[cfg(feature = "std")]
62impl<W: Write, BufferType: SliceWrapperMut<u8>, Alloc: BrotliAlloc> Write
63    for CompressorWriterCustomAlloc<W, BufferType, Alloc>
64{
65    fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
66        self.0.write(buf)
67    }
68    fn flush(&mut self) -> Result<(), Error> {
69        self.0.flush()
70    }
71}
72
73#[cfg(feature = "std")]
74pub struct CompressorWriter<W: Write>(
75    CompressorWriterCustomAlloc<
76        W,
77        <StandardAlloc as Allocator<u8>>::AllocatedMemory,
78        StandardAlloc,
79    >,
80);
81
82#[cfg(feature = "std")]
83impl<W: Write> CompressorWriter<W> {
84    pub fn new(w: W, buffer_size: usize, q: u32, lgwin: u32) -> Self {
85        let mut alloc = StandardAlloc::default();
86        let buffer = <StandardAlloc as Allocator<u8>>::alloc_cell(
87            &mut alloc,
88            if buffer_size == 0 { 4096 } else { buffer_size },
89        );
90        CompressorWriter::<W>(CompressorWriterCustomAlloc::new(w, buffer, alloc, q, lgwin))
91    }
92
93    pub fn with_params(w: W, buffer_size: usize, params: &BrotliEncoderParams) -> Self {
94        let mut writer = Self::new(w, buffer_size, params.quality as u32, params.lgwin as u32);
95        (writer.0).0.state.params = params.clone();
96        writer
97    }
98
99    pub fn get_ref(&self) -> &W {
100        self.0.get_ref()
101    }
102    pub fn get_mut(&mut self) -> &mut W {
103        self.0.get_mut()
104    }
105    pub fn into_inner(self) -> W {
106        self.0.into_inner()
107    }
108}
109
110#[cfg(feature = "std")]
111impl<W: Write> Write for CompressorWriter<W> {
112    fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
113        self.0.write(buf)
114    }
115    fn flush(&mut self) -> Result<(), Error> {
116        self.0.flush()
117    }
118}
119
120pub struct CompressorWriterCustomIo<
121    ErrType,
122    W: CustomWrite<ErrType>,
123    BufferType: SliceWrapperMut<u8>,
124    Alloc: BrotliAlloc,
125> {
126    output_buffer: BufferType,
127    total_out: Option<usize>,
128    output: Option<W>,
129    error_if_invalid_data: Option<ErrType>,
130    state: BrotliEncoderStateStruct<Alloc>,
131}
132pub fn write_all<ErrType, W: CustomWrite<ErrType>>(
133    writer: &mut W,
134    mut buf: &[u8],
135) -> Result<(), ErrType> {
136    while !buf.is_empty() {
137        match writer.write(buf) {
138            Ok(bytes_written) => buf = &buf[bytes_written..],
139            Err(e) => return Err(e),
140        }
141    }
142    Ok(())
143}
144impl<ErrType, W: CustomWrite<ErrType>, BufferType: SliceWrapperMut<u8>, Alloc: BrotliAlloc>
145    CompressorWriterCustomIo<ErrType, W, BufferType, Alloc>
146{
147    pub fn new(
148        w: W,
149        buffer: BufferType,
150        alloc: Alloc,
151        invalid_data_error_type: ErrType,
152        q: u32,
153        lgwin: u32,
154    ) -> Self {
155        let mut ret = CompressorWriterCustomIo {
156            output_buffer: buffer,
157            total_out: Some(0),
158            output: Some(w),
159            state: BrotliEncoderCreateInstance(alloc),
160            error_if_invalid_data: Some(invalid_data_error_type),
161        };
162        BrotliEncoderSetParameter(
163            &mut ret.state,
164            BrotliEncoderParameter::BROTLI_PARAM_QUALITY,
165            q,
166        );
167        BrotliEncoderSetParameter(
168            &mut ret.state,
169            BrotliEncoderParameter::BROTLI_PARAM_LGWIN,
170            lgwin,
171        );
172
173        ret
174    }
175    fn flush_or_close(&mut self, op: BrotliEncoderOperation) -> Result<(), ErrType> {
176        let mut nop_callback =
177            |_data: &mut interface::PredictionModeContextMap<interface::InputReferenceMut>,
178             _cmds: &mut [interface::StaticCommand],
179             _mb: interface::InputPair,
180             _mfv: &mut Alloc| ();
181
182        loop {
183            let mut avail_in: usize = 0;
184            let mut input_offset: usize = 0;
185            let mut avail_out: usize = self.output_buffer.slice_mut().len();
186            let mut output_offset: usize = 0;
187            let ret = BrotliEncoderCompressStream(
188                &mut self.state,
189                op,
190                &mut avail_in,
191                &[],
192                &mut input_offset,
193                &mut avail_out,
194                self.output_buffer.slice_mut(),
195                &mut output_offset,
196                &mut self.total_out,
197                &mut nop_callback,
198            );
199            if output_offset > 0 {
200                match write_all(
201                    self.output.as_mut().unwrap(),
202                    &self.output_buffer.slice_mut()[..output_offset],
203                ) {
204                    Ok(_) => {}
205                    Err(e) => return Err(e),
206                }
207            }
208            if ret <= 0 {
209                return Err(self.error_if_invalid_data.take().unwrap());
210            }
211            if let BrotliEncoderOperation::BROTLI_OPERATION_FLUSH = op {
212                if BrotliEncoderHasMoreOutput(&mut self.state) != 0 {
213                    continue;
214                }
215                return Ok(());
216            }
217            if BrotliEncoderIsFinished(&mut self.state) != 0 {
218                return Ok(());
219            }
220        }
221    }
222
223    pub fn get_ref(&self) -> &W {
224        self.output.as_ref().unwrap()
225    }
226    pub fn get_mut(&mut self) -> &mut W {
227        self.output.as_mut().unwrap()
228    }
229    pub fn into_inner(mut self) -> W {
230        match self.flush_or_close(BrotliEncoderOperation::BROTLI_OPERATION_FINISH) {
231            Ok(_) => {}
232            Err(_) => {}
233        }
234        self.output.take().unwrap()
235    }
236}
237
238impl<ErrType, W: CustomWrite<ErrType>, BufferType: SliceWrapperMut<u8>, Alloc: BrotliAlloc> Drop
239    for CompressorWriterCustomIo<ErrType, W, BufferType, Alloc>
240{
241    fn drop(&mut self) {
242        if self.output.is_some() {
243            match self.flush_or_close(BrotliEncoderOperation::BROTLI_OPERATION_FINISH) {
244                Ok(_) => {}
245                Err(_) => {}
246            }
247        }
248        BrotliEncoderDestroyInstance(&mut self.state);
249    }
250}
251impl<ErrType, W: CustomWrite<ErrType>, BufferType: SliceWrapperMut<u8>, Alloc: BrotliAlloc>
252    CustomWrite<ErrType> for CompressorWriterCustomIo<ErrType, W, BufferType, Alloc>
253{
254    fn write(&mut self, buf: &[u8]) -> Result<usize, ErrType> {
255        let mut nop_callback =
256            |_data: &mut interface::PredictionModeContextMap<interface::InputReferenceMut>,
257             _cmds: &mut [interface::StaticCommand],
258             _mb: interface::InputPair,
259             _mfv: &mut Alloc| ();
260        let mut avail_in = buf.len();
261        let mut input_offset: usize = 0;
262        while avail_in != 0 {
263            let mut output_offset = 0;
264            let mut avail_out = self.output_buffer.slice_mut().len();
265            let ret = BrotliEncoderCompressStream(
266                &mut self.state,
267                BrotliEncoderOperation::BROTLI_OPERATION_PROCESS,
268                &mut avail_in,
269                buf,
270                &mut input_offset,
271                &mut avail_out,
272                self.output_buffer.slice_mut(),
273                &mut output_offset,
274                &mut self.total_out,
275                &mut nop_callback,
276            );
277            if output_offset > 0 {
278                match write_all(
279                    self.output.as_mut().unwrap(),
280                    &self.output_buffer.slice_mut()[..output_offset],
281                ) {
282                    Ok(_) => {}
283                    Err(e) => return Err(e),
284                }
285            }
286            if ret <= 0 {
287                return Err(self.error_if_invalid_data.take().unwrap());
288            }
289        }
290        Ok(buf.len())
291    }
292    fn flush(&mut self) -> Result<(), ErrType> {
293        match self.flush_or_close(BrotliEncoderOperation::BROTLI_OPERATION_FLUSH) {
294            Ok(_) => {}
295            Err(e) => return Err(e),
296        }
297        self.output.as_mut().unwrap().flush()
298    }
299}