Add (de)compress_uninit impl for uninit buffer

Author: NobodyXu

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<DirDecompress>,
 }
 
-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<Status, DecompressError> {
         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<Status, DecompressError> {
+        unsafe { self.decompress_inner(input, output.as_mut_ptr(), output.len(), flush) }
+    }
+    fn decompress_uninit(
+        &mut self,
+        input: &[u8],
+        output: &mut [MaybeUninit<u8>],
+        flush: FlushDecompress,
+    ) -> Result<Status, DecompressError> {
+        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<DirCompress>,
 }
 
-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<Status, CompressError> {
         let raw = self.inner.stream_wrapper.inner;
@@ -365,17 +359,17 @@ 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);
 
             // 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();
             (*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<Status, CompressError> {
+        unsafe { self.compress_inner(input, output.as_mut_ptr(), output.len(), flush) }
+    }
+    fn compress_uninit(
+        &mut self,
+        input: &[u8],
+        output: &mut [MaybeUninit<u8>],
+        flush: FlushCompress,
+    ) -> Result<Status, CompressError> {
+        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;

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<u8>]) -> &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<u8>] as *mut [u8])
+    }
+}
 
 /// Traits specifying the interface of the backends.
 ///
@@ -20,6 +29,14 @@ pub trait InflateBackend: Backend {
         output: &mut [u8],
         flush: FlushDecompress,
     ) -> Result<Status, DecompressError>;
+    fn decompress_uninit(
+        &mut self,
+        input: &[u8],
+        output: &mut [MaybeUninit<u8>],
+        flush: FlushDecompress,
+    ) -> Result<Status, DecompressError> {
+        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<Status, CompressError>;
+    fn compress_uninit(
+        &mut self,
+        input: &[u8],
+        output: &mut [MaybeUninit<u8>],
+        flush: FlushCompress,
+    ) -> Result<Status, CompressError> {
+        self.compress(input, initialize_buffer(output), flush)
+    }
     fn reset(&mut self);
 }