1use crate::alloc::{Allocator, SliceWrapperMut};
2#[cfg(feature = "std")]
3use std::io;
4#[cfg(feature = "std")]
5use std::io::{Error, ErrorKind, Write};
6
7#[cfg(feature = "std")]
8pub use alloc_stdlib::StandardAlloc;
9use brotli_decompressor::CustomWrite;
10#[cfg(feature = "std")]
11pub use brotli_decompressor::{IntoIoWriter, IoWriterWrapper};
12
13use super::backward_references::BrotliEncoderParams;
14use super::combined_alloc::BrotliAlloc;
15use super::encode::{
16 BrotliEncoderDestroyInstance, BrotliEncoderOperation, BrotliEncoderParameter,
17 BrotliEncoderStateStruct,
18};
19use super::interface;
20use crate::enc::combined_alloc::allocate;
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 Error::new(ErrorKind::WriteZero, "No room in output."),
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 = allocate::<u8, _>(
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 error_if_zero_bytes_written: Option<ErrType>,
132}
133pub fn write_all<ErrType, W: CustomWrite<ErrType>, ErrMaker: FnMut() -> Option<ErrType>>(
134 writer: &mut W,
135 mut buf: &[u8],
136 mut error_to_return_if_zero_bytes_written: ErrMaker,
137) -> Result<(), ErrType> {
138 while !buf.is_empty() {
139 match writer.write(buf) {
140 Ok(bytes_written) => {
141 if bytes_written != 0 {
142 buf = &buf[bytes_written..]
143 } else {
144 match error_to_return_if_zero_bytes_written() {
145 Some(err) => {
146 return Err(err);
147 }
148 _ => {
149 return Ok(());
150 }
151 }
152 }
153 }
154 Err(e) => return Err(e),
155 }
156 }
157 Ok(())
158}
159impl<ErrType, W: CustomWrite<ErrType>, BufferType: SliceWrapperMut<u8>, Alloc: BrotliAlloc>
160 CompressorWriterCustomIo<ErrType, W, BufferType, Alloc>
161{
162 pub fn new(
163 w: W,
164 buffer: BufferType,
165 alloc: Alloc,
166 invalid_data_error_type: ErrType,
167 error_if_zero_bytes_written: ErrType,
168 q: u32,
169 lgwin: u32,
170 ) -> Self {
171 let mut ret = CompressorWriterCustomIo {
172 output_buffer: buffer,
173 total_out: Some(0),
174 output: Some(w),
175 state: BrotliEncoderStateStruct::new(alloc),
176 error_if_invalid_data: Some(invalid_data_error_type),
177 error_if_zero_bytes_written: Some(error_if_zero_bytes_written),
178 };
179 ret.state
180 .set_parameter(BrotliEncoderParameter::BROTLI_PARAM_QUALITY, q);
181 ret.state
182 .set_parameter(BrotliEncoderParameter::BROTLI_PARAM_LGWIN, lgwin);
183
184 ret
185 }
186 fn flush_or_close(&mut self, op: BrotliEncoderOperation) -> Result<(), ErrType> {
187 let mut nop_callback =
188 |_data: &mut interface::PredictionModeContextMap<interface::InputReferenceMut>,
189 _cmds: &mut [interface::StaticCommand],
190 _mb: interface::InputPair,
191 _mfv: &mut Alloc| ();
192
193 loop {
194 let mut avail_in: usize = 0;
195 let mut input_offset: usize = 0;
196 let mut avail_out: usize = self.output_buffer.slice_mut().len();
197 let mut output_offset: usize = 0;
198 let ret = self.state.compress_stream(
199 op,
200 &mut avail_in,
201 &[],
202 &mut input_offset,
203 &mut avail_out,
204 self.output_buffer.slice_mut(),
205 &mut output_offset,
206 &mut self.total_out,
207 &mut nop_callback,
208 );
209 if output_offset > 0 {
210 let zero_err = &mut self.error_if_zero_bytes_written;
211 let fallback = &mut self.error_if_invalid_data;
212 match write_all(
213 self.output.as_mut().unwrap(),
214 &self.output_buffer.slice_mut()[..output_offset],
215 || {
216 if let Some(err) = zero_err.take() {
217 return Some(err);
218 }
219 fallback.take()
220 },
221 ) {
222 Ok(_) => {}
223 Err(e) => return Err(e),
224 }
225 }
226 if !ret {
227 return Err(self.error_if_invalid_data.take().unwrap());
228 }
229 if let BrotliEncoderOperation::BROTLI_OPERATION_FLUSH = op {
230 if self.state.has_more_output() {
231 continue;
232 }
233 return Ok(());
234 }
235 if self.state.is_finished() {
236 return Ok(());
237 }
238 }
239 }
240
241 pub fn get_ref(&self) -> &W {
242 self.output.as_ref().unwrap()
243 }
244 pub fn get_mut(&mut self) -> &mut W {
245 self.output.as_mut().unwrap()
246 }
247 pub fn into_inner(mut self) -> W {
248 match self.flush_or_close(BrotliEncoderOperation::BROTLI_OPERATION_FINISH) {
249 Ok(_) => {}
250 Err(_) => {}
251 }
252 self.output.take().unwrap()
253 }
254}
255
256impl<ErrType, W: CustomWrite<ErrType>, BufferType: SliceWrapperMut<u8>, Alloc: BrotliAlloc> Drop
257 for CompressorWriterCustomIo<ErrType, W, BufferType, Alloc>
258{
259 fn drop(&mut self) {
260 if self.output.is_some() {
261 match self.flush_or_close(BrotliEncoderOperation::BROTLI_OPERATION_FINISH) {
262 Ok(_) => {}
263 Err(_) => {}
264 }
265 }
266 BrotliEncoderDestroyInstance(&mut self.state);
267 }
268}
269impl<ErrType, W: CustomWrite<ErrType>, BufferType: SliceWrapperMut<u8>, Alloc: BrotliAlloc>
270 CustomWrite<ErrType> for CompressorWriterCustomIo<ErrType, W, BufferType, Alloc>
271{
272 fn write(&mut self, buf: &[u8]) -> Result<usize, ErrType> {
273 let mut nop_callback =
274 |_data: &mut interface::PredictionModeContextMap<interface::InputReferenceMut>,
275 _cmds: &mut [interface::StaticCommand],
276 _mb: interface::InputPair,
277 _mfv: &mut Alloc| ();
278 let mut avail_in = buf.len();
279 let mut input_offset: usize = 0;
280 while avail_in != 0 {
281 let mut output_offset = 0;
282 let mut avail_out = self.output_buffer.slice_mut().len();
283 let ret = self.state.compress_stream(
284 BrotliEncoderOperation::BROTLI_OPERATION_PROCESS,
285 &mut avail_in,
286 buf,
287 &mut input_offset,
288 &mut avail_out,
289 self.output_buffer.slice_mut(),
290 &mut output_offset,
291 &mut self.total_out,
292 &mut nop_callback,
293 );
294 if output_offset > 0 {
295 let zero_err = &mut self.error_if_zero_bytes_written;
296 let fallback = &mut self.error_if_invalid_data;
297 match write_all(
298 self.output.as_mut().unwrap(),
299 &self.output_buffer.slice_mut()[..output_offset],
300 || {
301 if let Some(err) = zero_err.take() {
302 return Some(err);
303 }
304 fallback.take()
305 },
306 ) {
307 Ok(_) => {}
308 Err(e) => return Err(e),
309 }
310 }
311 if !ret {
312 return Err(self.error_if_invalid_data.take().unwrap());
313 }
314 }
315 Ok(buf.len())
316 }
317 fn flush(&mut self) -> Result<(), ErrType> {
318 match self.flush_or_close(BrotliEncoderOperation::BROTLI_OPERATION_FLUSH) {
319 Ok(_) => {}
320 Err(e) => return Err(e),
321 }
322 self.output.as_mut().unwrap().flush()
323 }
324}