use alloc::{Allocator, SliceWrapperMut};
#[cfg(feature = "std")]
use std::io;
#[cfg(feature = "std")]
use std::io::{Error, ErrorKind, Write};
#[cfg(feature = "std")]
pub use alloc_stdlib::StandardAlloc;
use brotli_decompressor::CustomWrite;
#[cfg(feature = "std")]
pub use brotli_decompressor::{IntoIoWriter, IoWriterWrapper};
use super::backward_references::BrotliEncoderParams;
use super::combined_alloc::BrotliAlloc;
use super::encode::{
BrotliEncoderDestroyInstance, BrotliEncoderOperation, BrotliEncoderParameter,
BrotliEncoderStateStruct,
};
use super::interface;
use crate::enc::combined_alloc::allocate;
#[cfg(feature = "std")]
pub struct CompressorWriterCustomAlloc<
W: Write,
BufferType: SliceWrapperMut<u8>,
Alloc: BrotliAlloc,
>(CompressorWriterCustomIo<io::Error, IntoIoWriter<W>, BufferType, Alloc>);
#[cfg(feature = "std")]
impl<W: Write, BufferType: SliceWrapperMut<u8>, Alloc: BrotliAlloc>
CompressorWriterCustomAlloc<W, BufferType, Alloc>
{
pub fn new(w: W, buffer: BufferType, alloc: Alloc, q: u32, lgwin: u32) -> Self {
CompressorWriterCustomAlloc::<W, BufferType, Alloc>(CompressorWriterCustomIo::<
Error,
IntoIoWriter<W>,
BufferType,
Alloc,
>::new(
IntoIoWriter::<W>(w),
buffer,
alloc,
Error::new(ErrorKind::InvalidData, "Invalid Data"),
Error::new(ErrorKind::WriteZero, "No room in output."),
q,
lgwin,
))
}
pub fn get_ref(&self) -> &W {
&self.0.get_ref().0
}
pub fn get_mut(&mut self) -> &mut W {
&mut self.0.get_mut().0
}
pub fn into_inner(self) -> W {
self.0.into_inner().0
}
}
#[cfg(feature = "std")]
impl<W: Write, BufferType: SliceWrapperMut<u8>, Alloc: BrotliAlloc> Write
for CompressorWriterCustomAlloc<W, BufferType, Alloc>
{
fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
self.0.write(buf)
}
fn flush(&mut self) -> Result<(), Error> {
self.0.flush()
}
}
#[cfg(feature = "std")]
pub struct CompressorWriter<W: Write>(
CompressorWriterCustomAlloc<
W,
<StandardAlloc as Allocator<u8>>::AllocatedMemory,
StandardAlloc,
>,
);
#[cfg(feature = "std")]
impl<W: Write> CompressorWriter<W> {
pub fn new(w: W, buffer_size: usize, q: u32, lgwin: u32) -> Self {
let mut alloc = StandardAlloc::default();
let buffer = allocate::<u8, _>(
&mut alloc,
if buffer_size == 0 { 4096 } else { buffer_size },
);
CompressorWriter::<W>(CompressorWriterCustomAlloc::new(w, buffer, alloc, q, lgwin))
}
pub fn with_params(w: W, buffer_size: usize, params: &BrotliEncoderParams) -> Self {
let mut writer = Self::new(w, buffer_size, params.quality as u32, params.lgwin as u32);
(writer.0).0.state.params = params.clone();
writer
}
pub fn get_ref(&self) -> &W {
self.0.get_ref()
}
pub fn get_mut(&mut self) -> &mut W {
self.0.get_mut()
}
pub fn into_inner(self) -> W {
self.0.into_inner()
}
}
#[cfg(feature = "std")]
impl<W: Write> Write for CompressorWriter<W> {
fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
self.0.write(buf)
}
fn flush(&mut self) -> Result<(), Error> {
self.0.flush()
}
}
pub struct CompressorWriterCustomIo<
ErrType,
W: CustomWrite<ErrType>,
BufferType: SliceWrapperMut<u8>,
Alloc: BrotliAlloc,
> {
output_buffer: BufferType,
total_out: Option<usize>,
output: Option<W>,
error_if_invalid_data: Option<ErrType>,
state: BrotliEncoderStateStruct<Alloc>,
error_if_zero_bytes_written: Option<ErrType>,
}
pub fn write_all<ErrType, W: CustomWrite<ErrType>, ErrMaker: FnMut() -> Option<ErrType>>(
writer: &mut W,
mut buf: &[u8],
mut error_to_return_if_zero_bytes_written: ErrMaker,
) -> Result<(), ErrType> {
while !buf.is_empty() {
match writer.write(buf) {
Ok(bytes_written) => if bytes_written != 0 {
buf = &buf[bytes_written..]
} else {
if let Some(err) = error_to_return_if_zero_bytes_written() {
return Err(err);
} else {
return Ok(());
}
},
Err(e) => return Err(e),
}
}
Ok(())
}
impl<ErrType, W: CustomWrite<ErrType>, BufferType: SliceWrapperMut<u8>, Alloc: BrotliAlloc>
CompressorWriterCustomIo<ErrType, W, BufferType, Alloc>
{
pub fn new(
w: W,
buffer: BufferType,
alloc: Alloc,
invalid_data_error_type: ErrType,
error_if_zero_bytes_written: ErrType,
q: u32,
lgwin: u32,
) -> Self {
let mut ret = CompressorWriterCustomIo {
output_buffer: buffer,
total_out: Some(0),
output: Some(w),
state: BrotliEncoderStateStruct::new(alloc),
error_if_invalid_data: Some(invalid_data_error_type),
error_if_zero_bytes_written: Some(error_if_zero_bytes_written),
};
ret.state
.set_parameter(BrotliEncoderParameter::BROTLI_PARAM_QUALITY, q);
ret.state
.set_parameter(BrotliEncoderParameter::BROTLI_PARAM_LGWIN, lgwin);
ret
}
fn flush_or_close(&mut self, op: BrotliEncoderOperation) -> Result<(), ErrType> {
let mut nop_callback =
|_data: &mut interface::PredictionModeContextMap<interface::InputReferenceMut>,
_cmds: &mut [interface::StaticCommand],
_mb: interface::InputPair,
_mfv: &mut Alloc| ();
loop {
let mut avail_in: usize = 0;
let mut input_offset: usize = 0;
let mut avail_out: usize = self.output_buffer.slice_mut().len();
let mut output_offset: usize = 0;
let ret = self.state.compress_stream(
op,
&mut avail_in,
&[],
&mut input_offset,
&mut avail_out,
self.output_buffer.slice_mut(),
&mut output_offset,
&mut self.total_out,
&mut nop_callback,
);
if output_offset > 0 {
let zero_err = &mut self.error_if_zero_bytes_written;
let fallback = &mut self.error_if_invalid_data;
match write_all(
self.output.as_mut().unwrap(),
&self.output_buffer.slice_mut()[..output_offset],
|| {
if let Some(err) = zero_err.take() {
return Some(err);
}
fallback.take()
},
) {
Ok(_) => {}
Err(e) => return Err(e),
}
}
if !ret {
return Err(self.error_if_invalid_data.take().unwrap());
}
if let BrotliEncoderOperation::BROTLI_OPERATION_FLUSH = op {
if self.state.has_more_output() {
continue;
}
return Ok(());
}
if self.state.is_finished() {
return Ok(());
}
}
}
pub fn get_ref(&self) -> &W {
self.output.as_ref().unwrap()
}
pub fn get_mut(&mut self) -> &mut W {
self.output.as_mut().unwrap()
}
pub fn into_inner(mut self) -> W {
match self.flush_or_close(BrotliEncoderOperation::BROTLI_OPERATION_FINISH) {
Ok(_) => {}
Err(_) => {}
}
self.output.take().unwrap()
}
}
impl<ErrType, W: CustomWrite<ErrType>, BufferType: SliceWrapperMut<u8>, Alloc: BrotliAlloc> Drop
for CompressorWriterCustomIo<ErrType, W, BufferType, Alloc>
{
fn drop(&mut self) {
if self.output.is_some() {
match self.flush_or_close(BrotliEncoderOperation::BROTLI_OPERATION_FINISH) {
Ok(_) => {}
Err(_) => {}
}
}
BrotliEncoderDestroyInstance(&mut self.state);
}
}
impl<ErrType, W: CustomWrite<ErrType>, BufferType: SliceWrapperMut<u8>, Alloc: BrotliAlloc>
CustomWrite<ErrType> for CompressorWriterCustomIo<ErrType, W, BufferType, Alloc>
{
fn write(&mut self, buf: &[u8]) -> Result<usize, ErrType> {
let mut nop_callback =
|_data: &mut interface::PredictionModeContextMap<interface::InputReferenceMut>,
_cmds: &mut [interface::StaticCommand],
_mb: interface::InputPair,
_mfv: &mut Alloc| ();
let mut avail_in = buf.len();
let mut input_offset: usize = 0;
while avail_in != 0 {
let mut output_offset = 0;
let mut avail_out = self.output_buffer.slice_mut().len();
let ret = self.state.compress_stream(
BrotliEncoderOperation::BROTLI_OPERATION_PROCESS,
&mut avail_in,
buf,
&mut input_offset,
&mut avail_out,
self.output_buffer.slice_mut(),
&mut output_offset,
&mut self.total_out,
&mut nop_callback,
);
if output_offset > 0 {
let zero_err = &mut self.error_if_zero_bytes_written;
let fallback = &mut self.error_if_invalid_data;
match write_all(
self.output.as_mut().unwrap(),
&self.output_buffer.slice_mut()[..output_offset],
|| {
if let Some(err) = zero_err.take() {
return Some(err);
}
fallback.take()
},
) {
Ok(_) => {}
Err(e) => {
return Err(e)
},
}
}
if !ret {
return Err(self.error_if_invalid_data.take().unwrap());
}
}
Ok(buf.len())
}
fn flush(&mut self) -> Result<(), ErrType> {
match self.flush_or_close(BrotliEncoderOperation::BROTLI_OPERATION_FLUSH) {
Ok(_) => {}
Err(e) => return Err(e),
}
self.output.as_mut().unwrap().flush()
}
}