diff --git a/src/ffi/c.rs b/src/ffi/c.rs index 14e76906..e3fb171e 100644 --- a/src/ffi/c.rs +++ b/src/ffi/c.rs @@ -1,7 +1,7 @@ //! Implementation for C backends. -use std::cmp; use std::fmt; use std::marker; +use std::mem::MaybeUninit; use std::os::raw::{c_int, c_uint}; use std::ptr; @@ -226,34 +226,12 @@ pub struct Inflate { pub inner: Stream, } -impl InflateBackend for Inflate { - fn make(zlib_header: bool, window_bits: u8) -> Self { - unsafe { - let state = StreamWrapper::default(); - let ret = mz_inflateInit2( - state.inner, - if zlib_header { - window_bits as c_int - } else { - -(window_bits as c_int) - }, - ); - assert_eq!(ret, 0); - Inflate { - inner: Stream { - stream_wrapper: state, - total_in: 0, - total_out: 0, - _marker: marker::PhantomData, - }, - } - } - } - - fn decompress( +impl Inflate { + unsafe fn decompress_inner( &mut self, input: &[u8], - output: &mut [u8], + output_ptr: *mut u8, + output_len: usize, flush: FlushDecompress, ) -> Result { let raw = self.inner.stream_wrapper.inner; @@ -263,16 +241,16 @@ impl InflateBackend for Inflate { unsafe { (*raw).msg = ptr::null_mut(); (*raw).next_in = input.as_ptr() as *mut u8; - (*raw).avail_in = cmp::min(input.len(), c_uint::MAX as usize) as c_uint; - (*raw).next_out = output.as_mut_ptr(); - (*raw).avail_out = cmp::min(output.len(), c_uint::MAX as usize) as c_uint; + (*raw).avail_in = input.len().min(c_uint::MAX as usize) as c_uint; + (*raw).next_out = output_ptr; + (*raw).avail_out = output_len.min(c_uint::MAX as usize) as c_uint; let rc = mz_inflate(raw, flush as c_int); // Unfortunately the total counters provided by zlib might be only // 32 bits wide and overflow while processing large amounts of data. self.inner.total_in += ((*raw).next_in as usize - input.as_ptr() as usize) as u64; - self.inner.total_out += ((*raw).next_out as usize - output.as_ptr() as usize) as u64; + self.inner.total_out += ((*raw).next_out as usize - output_ptr as usize) as u64; // reset these pointers so we don't accidentally read them later (*raw).next_in = ptr::null_mut(); @@ -293,6 +271,47 @@ impl InflateBackend for Inflate { } } } +} +impl InflateBackend for Inflate { + fn make(zlib_header: bool, window_bits: u8) -> Self { + unsafe { + let state = StreamWrapper::default(); + let ret = mz_inflateInit2( + state.inner, + if zlib_header { + window_bits as c_int + } else { + -(window_bits as c_int) + }, + ); + assert_eq!(ret, 0); + Inflate { + inner: Stream { + stream_wrapper: state, + total_in: 0, + total_out: 0, + _marker: marker::PhantomData, + }, + } + } + } + + fn decompress( + &mut self, + input: &[u8], + output: &mut [u8], + flush: FlushDecompress, + ) -> Result { + unsafe { self.decompress_inner(input, output.as_mut_ptr(), output.len(), flush) } + } + fn decompress_uninit( + &mut self, + input: &[u8], + output: &mut [MaybeUninit], + flush: FlushDecompress, + ) -> Result { + unsafe { self.decompress_inner(input, output.as_mut_ptr() as *mut _, output.len(), flush) } + } fn reset(&mut self, zlib_header: bool) { let bits = if zlib_header { @@ -325,37 +344,12 @@ pub struct Deflate { pub inner: Stream, } -impl DeflateBackend for Deflate { - fn make(level: Compression, zlib_header: bool, window_bits: u8) -> Self { - unsafe { - let state = StreamWrapper::default(); - let ret = mz_deflateInit2( - state.inner, - level.0 as c_int, - MZ_DEFLATED, - if zlib_header { - window_bits as c_int - } else { - -(window_bits as c_int) - }, - 8, - MZ_DEFAULT_STRATEGY, - ); - assert_eq!(ret, 0); - Deflate { - inner: Stream { - stream_wrapper: state, - total_in: 0, - total_out: 0, - _marker: marker::PhantomData, - }, - } - } - } - fn compress( +impl Deflate { + unsafe fn compress_inner( &mut self, input: &[u8], - output: &mut [u8], + output_ptr: *mut u8, + output_len: usize, flush: FlushCompress, ) -> Result { let raw = self.inner.stream_wrapper.inner; @@ -365,9 +359,9 @@ impl DeflateBackend for Deflate { unsafe { (*raw).msg = ptr::null_mut(); (*raw).next_in = input.as_ptr() as *mut _; - (*raw).avail_in = cmp::min(input.len(), c_uint::MAX as usize) as c_uint; - (*raw).next_out = output.as_mut_ptr(); - (*raw).avail_out = cmp::min(output.len(), c_uint::MAX as usize) as c_uint; + (*raw).avail_in = input.len().min(c_uint::MAX as usize) as c_uint; + (*raw).next_out = output_ptr; + (*raw).avail_out = output_len.min(c_uint::MAX as usize) as c_uint; let rc = mz_deflate(raw, flush as c_int); @@ -375,7 +369,7 @@ impl DeflateBackend for Deflate { // 32 bits wide and overflow while processing large amounts of data. self.inner.total_in += ((*raw).next_in as usize - input.as_ptr() as usize) as u64; - self.inner.total_out += ((*raw).next_out as usize - output.as_ptr() as usize) as u64; + self.inner.total_out += ((*raw).next_out as usize - output_ptr as usize) as u64; // reset these pointers so we don't accidentally read them later (*raw).next_in = ptr::null_mut(); (*raw).avail_in = 0; @@ -391,7 +385,51 @@ impl DeflateBackend for Deflate { } } } +} +impl DeflateBackend for Deflate { + fn make(level: Compression, zlib_header: bool, window_bits: u8) -> Self { + unsafe { + let state = StreamWrapper::default(); + let ret = mz_deflateInit2( + state.inner, + level.0 as c_int, + MZ_DEFLATED, + if zlib_header { + window_bits as c_int + } else { + -(window_bits as c_int) + }, + 8, + MZ_DEFAULT_STRATEGY, + ); + assert_eq!(ret, 0); + Deflate { + inner: Stream { + stream_wrapper: state, + total_in: 0, + total_out: 0, + _marker: marker::PhantomData, + }, + } + } + } + fn compress( + &mut self, + input: &[u8], + output: &mut [u8], + flush: FlushCompress, + ) -> Result { + unsafe { self.compress_inner(input, output.as_mut_ptr(), output.len(), flush) } + } + fn compress_uninit( + &mut self, + input: &[u8], + output: &mut [MaybeUninit], + flush: FlushCompress, + ) -> Result { + unsafe { self.compress_inner(input, output.as_mut_ptr() as *mut _, output.len(), flush) } + } fn reset(&mut self) { self.inner.total_in = 0; self.inner.total_out = 0; diff --git a/src/ffi/mod.rs b/src/ffi/mod.rs index 20b3cae6..856bc651 100644 --- a/src/ffi/mod.rs +++ b/src/ffi/mod.rs @@ -2,6 +2,15 @@ use crate::mem::{CompressError, DecompressError, FlushCompress, FlushDecompress, Status}; use crate::Compression; +use std::mem::MaybeUninit; + +fn initialize_buffer(output: &mut [MaybeUninit]) -> &mut [u8] { + // SAFETY: Here we zero-initialize the output and cast it to [u8] + unsafe { + output.as_mut_ptr().write_bytes(0, output.len()); + &mut *(output as *mut [MaybeUninit] as *mut [u8]) + } +} /// Traits specifying the interface of the backends. /// @@ -20,6 +29,14 @@ pub trait InflateBackend: Backend { output: &mut [u8], flush: FlushDecompress, ) -> Result; + fn decompress_uninit( + &mut self, + input: &[u8], + output: &mut [MaybeUninit], + flush: FlushDecompress, + ) -> Result { + self.decompress(input, initialize_buffer(output), flush) + } fn reset(&mut self, zlib_header: bool); } @@ -31,6 +48,14 @@ pub trait DeflateBackend: Backend { output: &mut [u8], flush: FlushCompress, ) -> Result; + fn compress_uninit( + &mut self, + input: &[u8], + output: &mut [MaybeUninit], + flush: FlushCompress, + ) -> Result { + self.compress(input, initialize_buffer(output), flush) + } fn reset(&mut self); } diff --git a/src/mem.rs b/src/mem.rs index c1edaf5c..8121d98a 100644 --- a/src/mem.rs +++ b/src/mem.rs @@ -1,6 +1,7 @@ use std::error::Error; use std::fmt; use std::io; +use std::mem::MaybeUninit; use crate::ffi::{self, Backend, Deflate, DeflateBackend, ErrorMessage, Inflate, InflateBackend}; use crate::Compression; @@ -337,6 +338,20 @@ impl Compress { self.inner.compress(input, output, flush) } + /// Similar to [`Self::compress`] but accepts uninitialized buffer. + /// + /// If you want to avoid the overhead of zero initializing the + /// buffer and you don't want to use a [`Vec`], then please use + /// this API. + pub fn compress_uninit( + &mut self, + input: &[u8], + output: &mut [MaybeUninit], + flush: FlushCompress, + ) -> Result { + self.inner.compress_uninit(input, output, flush) + } + /// Compresses the input data into the extra space of the output, consuming /// only as much input as needed and writing as much output as possible. /// @@ -351,12 +366,15 @@ impl Compress { output: &mut Vec, flush: FlushCompress, ) -> Result { - write_to_spare_capacity_of_vec(output, |out| { - let before = self.total_out(); - let ret = self.compress(input, out, flush); - let bytes_written = self.total_out() - before; - (bytes_written as usize, ret) - }) + // SAFETY: bytes_written is the number of bytes written into `out` + unsafe { + write_to_spare_capacity_of_vec(output, |out| { + let before = self.total_out(); + let ret = self.compress_uninit(input, out, flush); + let bytes_written = self.total_out() - before; + (bytes_written as usize, ret) + }) + } } } @@ -455,6 +473,20 @@ impl Decompress { self.inner.decompress(input, output, flush) } + /// Similar to [`Self::decompress`] but accepts uninitialized buffer + /// + /// If you want to avoid the overhead of zero initializing the + /// buffer and you don't want to use a [`Vec`], then please use + /// this API. + pub fn decompress_uninit( + &mut self, + input: &[u8], + output: &mut [MaybeUninit], + flush: FlushDecompress, + ) -> Result { + self.inner.decompress_uninit(input, output, flush) + } + /// Decompresses the input data into the extra space in the output vector /// specified by `output`. /// @@ -475,12 +507,15 @@ impl Decompress { output: &mut Vec, flush: FlushDecompress, ) -> Result { - write_to_spare_capacity_of_vec(output, |out| { - let before = self.total_out(); - let ret = self.decompress(input, out, flush); - let bytes_written = self.total_out() - before; - (bytes_written as usize, ret) - }) + // SAFETY: bytes_written is the number of bytes written into `out` + unsafe { + write_to_spare_capacity_of_vec(output, |out| { + let before = self.total_out(); + let ret = self.decompress_uninit(input, out, flush); + let bytes_written = self.total_out() - before; + (bytes_written as usize, ret) + }) + } } /// Specifies the decompression dictionary to use. @@ -580,18 +615,20 @@ impl fmt::Display for CompressError { /// /// `writer` needs to return the number of bytes written (and can also return /// another arbitrary return value). -fn write_to_spare_capacity_of_vec( +/// +/// # Safety: +/// +/// The length returned by the `writer` must be equal to actual number of bytes written +/// to the uninitialized slice passed in and initialized. +unsafe fn write_to_spare_capacity_of_vec( output: &mut Vec, - writer: impl FnOnce(&mut [u8]) -> (usize, T), + writer: impl FnOnce(&mut [MaybeUninit]) -> (usize, T), ) -> T { let cap = output.capacity(); let len = output.len(); - output.resize(output.capacity(), 0); - let (bytes_written, ret) = writer(&mut output[len..]); - - let new_len = core::cmp::min(len + bytes_written, cap); // Sanitizes `bytes_written`. - output.resize(new_len, 0 /* unused */); + let (bytes_written, ret) = writer(output.spare_capacity_mut()); + output.set_len(cap.min(len + bytes_written)); // Sanitizes `bytes_written`. ret }