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