feat: Add output: &mut [MaybeUninit<u8>] support

src/bytes.rs

@@ -1,5 +1,6 @@
 use std::convert::TryInto;
 use std::io;
+use std::mem::MaybeUninit;
 
 /// Read a u16 in little endian format from the beginning of the given slice.
 /// This panics if the slice has length less than 2.
@@ -28,13 +29,11 @@ pub fn io_read_u32_le<R: io::Read>(mut rdr: R) -> io::Result<u32> {
     Ok(u32::from_le_bytes(buf))
 }
 
-/// Write a u16 in little endian format to the beginning of the given slice.
-/// This panics if the slice has length less than 2.
-pub fn write_u16_le(n: u16, slice: &mut [u8]) {
+pub fn write_u16_le(n: u16, slice: &mut [MaybeUninit<u8>]) {
     assert!(slice.len() >= 2);
     let bytes = n.to_le_bytes();
-    slice[0] = bytes[0];
-    slice[1] = bytes[1];
+    slice[0].write(bytes[0]);
+    slice[1].write(bytes[1]);
 }
 
 /// Write a u24 (given as a u32 where the most significant 8 bits are ignored)
@@ -58,14 +57,14 @@ pub fn write_u32_le(n: u32, slice: &mut [u8]) {
 }
 
 /// https://developers.google.com/protocol-buffers/docs/encoding#varints
-pub fn write_varu64(data: &mut [u8], mut n: u64) -> usize {
+pub fn write_varu64(data: &mut [MaybeUninit<u8>], mut n: u64) -> usize {
     let mut i = 0;
     while n >= 0b1000_0000 {
-        data[i] = (n as u8) | 0b1000_0000;
+        data[i].write((n as u8) | 0b1000_0000);
         n >>= 7;
         i += 1;
     }
-    data[i] = n as u8;
+    data[i].write(n as u8);
     i + 1
 }
 

src/compress.rs

@@ -1,6 +1,7 @@
-use std::fmt;
+use std::mem::MaybeUninit;
 use std::ops::{Deref, DerefMut};
 use std::ptr;
+use std::{fmt, slice};
 
 use crate::bytes;
 use crate::error::{Error, Result};
@@ -98,8 +99,46 @@ impl Encoder {
     /// * `output` has length less than `max_compress_len(input.len())`.
     pub fn compress(
         &mut self,
-        mut input: &[u8],
+        input: &[u8],
         output: &mut [u8],
+    ) -> Result<usize> {
+        // SAFETY: `output` is guaranteed to be valid for reads and writes of
+        // `output.len()` bytes. Also compress_uninit does not write any
+        // uninitialised data into the buffer
+        //   (i.e. via array[0] = MaybeUninit::new())
+        let output: &mut [MaybeUninit<u8>] = unsafe {
+            slice::from_raw_parts_mut(
+                output.as_mut_ptr() as *mut MaybeUninit<u8>,
+                output.len(),
+            )
+        };
+        self.compress_uninit(input, output)
+    }
+
+    /// Compresses all bytes in `input` into `output`.
+    ///
+    /// `input` can be any arbitrary sequence of bytes.
+    ///
+    /// `output` must be large enough to hold the maximum possible compressed
+    /// size of `input`, which can be computed using `max_compress_len`.
+    ///
+    /// On success, this returns the number of bytes written to `output`.
+    ///
+    /// The data written to `output` is guaranteed to be initialized to the
+    /// number of bytes written returned by the method. Data outside that
+    /// bound is not initialized and reading from it is UB. It's the callee's
+    /// responsibility to finalize the state of `output`.
+    ///
+    /// # Errors
+    ///
+    /// This method returns an error in the following circumstances:
+    ///
+    /// * The total number of bytes to compress exceeds `2^32 - 1`.
+    /// * `output` has length less than `max_compress_len(input.len())`.
+    pub fn compress_uninit(
+        &mut self,
+        mut input: &[u8],
+        output: &mut [MaybeUninit<u8>],
     ) -> Result<usize> {
         match max_compress_len(input.len()) {
             0 => {
@@ -120,7 +159,7 @@ impl Encoder {
         if input.is_empty() {
             // Encodes a varint of 0, denoting the total size of uncompressed
             // bytes.
-            output[0] = 0;
+            output[0].write(0);
             return Ok(1);
         }
         // Write the Snappy header, which is just the total number of
@@ -162,9 +201,12 @@ impl Encoder {
     /// This method returns an error under the same circumstances that
     /// `compress` does.
     pub fn compress_vec(&mut self, input: &[u8]) -> Result<Vec<u8>> {
-        let mut buf = vec![0; max_compress_len(input.len())];
-        let n = self.compress(input, &mut buf)?;
-        buf.truncate(n);
+        let mut buf = Vec::with_capacity(max_compress_len(input.len()));
+        let n = self.compress_uninit(input, buf.spare_capacity_mut())?;
+
+        // SAFETY: compress_uninit guarantees all the data up to `n` is
+        // initialized
+        unsafe { buf.set_len(n) };
         Ok(buf)
     }
 }
@@ -173,22 +215,15 @@ struct Block<'s, 'd> {
     src: &'s [u8],
     s: usize,
     s_limit: usize,
-    dst: &'d mut [u8],
+    dst: &'d mut [MaybeUninit<u8>],
     d: usize,
     next_emit: usize,
 }
 
 impl<'s, 'd> Block<'s, 'd> {
     #[inline(always)]
-    fn new(src: &'s [u8], dst: &'d mut [u8], d: usize) -> Block<'s, 'd> {
-        Block {
-            src: src,
-            s: 0,
-            s_limit: src.len(),
-            dst: dst,
-            d: d,
-            next_emit: 0,
-        }
+    fn new(src: &'s [u8], dst: &'d mut [MaybeUninit<u8>], d: usize) -> Self {
+        Self { src, s: 0, s_limit: src.len(), dst, d, next_emit: 0 }
     }
 
     #[inline(always)]
@@ -346,10 +381,11 @@ impl<'s, 'd> Block<'s, 'd> {
         }
         // If we can squeeze the last copy into a copy 1 operation, do it.
         if len <= 11 && offset <= 2047 {
-            self.dst[self.d] = (((offset >> 8) as u8) << 5)
+            let tag = (((offset >> 8) as u8) << 5)
                 | (((len - 4) as u8) << 2)
                 | (Tag::Copy1 as u8);
-            self.dst[self.d + 1] = offset as u8;
+            self.dst[self.d].write(tag);
+            self.dst[self.d + 1].write(offset as u8);
             self.d += 2;
         } else {
             self.emit_copy2(offset, len);
@@ -363,7 +399,7 @@ impl<'s, 'd> Block<'s, 'd> {
     fn emit_copy2(&mut self, offset: usize, len: usize) {
         debug_assert!(1 <= offset && offset <= 65535);
         debug_assert!(1 <= len && len <= 64);
-        self.dst[self.d] = (((len - 1) as u8) << 2) | (Tag::Copy2 as u8);
+        self.dst[self.d].write((((len - 1) as u8) << 2) | (Tag::Copy2 as u8));
         bytes::write_u16_le(offset as u16, &mut self.dst[self.d + 1..]);
         self.d += 3;
     }
@@ -435,7 +471,7 @@ impl<'s, 'd> Block<'s, 'd> {
         let len = lit_end - lit_start;
         let n = len.checked_sub(1).unwrap();
         if n <= 59 {
-            self.dst[self.d] = ((n as u8) << 2) | (Tag::Literal as u8);
+            self.dst[self.d].write(((n as u8) << 2) | (Tag::Literal as u8));
             self.d += 1;
             if len <= 16 && lit_start + 16 <= self.src.len() {
                 // SAFETY: lit_start is equivalent to self.next_emit, which is
@@ -448,16 +484,16 @@ impl<'s, 'd> Block<'s, 'd> {
                 // an extra 32 bytes, which exceeds the 16 byte copy here.
                 let srcp = self.src.as_ptr().add(lit_start);
                 let dstp = self.dst.as_mut_ptr().add(self.d);
-                ptr::copy_nonoverlapping(srcp, dstp, 16);
+                ptr::copy_nonoverlapping(srcp, dstp as *mut u8, 16);
                 self.d += len;
                 return;
             }
         } else if n < 256 {
-            self.dst[self.d] = (60 << 2) | (Tag::Literal as u8);
-            self.dst[self.d + 1] = n as u8;
+            self.dst[self.d].write((60 << 2) | (Tag::Literal as u8));
+            self.dst[self.d + 1].write(n as u8);
             self.d += 2;
         } else {
-            self.dst[self.d] = (61 << 2) | (Tag::Literal as u8);
+            self.dst[self.d].write((61 << 2) | (Tag::Literal as u8));
             bytes::write_u16_le(n as u16, &mut self.dst[self.d + 1..]);
             self.d += 3;
         }
@@ -469,7 +505,7 @@ impl<'s, 'd> Block<'s, 'd> {
         // must be guaranteed by the caller and is why this method is unsafe.
         let srcp = self.src.as_ptr().add(lit_start);
         let dstp = self.dst.as_mut_ptr().add(self.d);
-        ptr::copy_nonoverlapping(srcp, dstp, len);
+        ptr::copy_nonoverlapping(srcp, dstp as *mut u8, len);
         self.d += len;
     }
 }

src/decompress.rs

@@ -1,4 +1,5 @@
-use std::ptr;
+use std::mem::MaybeUninit;
+use std::{ptr, slice};
 
 use crate::bytes;
 use crate::error::{Error, Result};
@@ -76,6 +77,47 @@ impl Decoder {
         &mut self,
         input: &[u8],
         output: &mut [u8],
+    ) -> Result<usize> {
+        // SAFETY: `output` is guaranteed to be valid for reads and writes of
+        // `output.len()` bytes. Also compress_uninit does not write any
+        // uninitialised data into the buffer
+        //   (i.e. via array[0] = MaybeUninit::new())
+        let output: &mut [MaybeUninit<u8>] = unsafe {
+            slice::from_raw_parts_mut(
+                output.as_mut_ptr() as *mut MaybeUninit<u8>,
+                output.len(),
+            )
+        };
+        self.decompress_uninit(input, output)
+    }
+
+    /// Decompresses all bytes in `input` into `output`.
+    ///
+    /// `input` must be a sequence of bytes returned by a conforming Snappy
+    /// compressor.
+    ///
+    /// The size of `output` must be large enough to hold all decompressed
+    /// bytes from the `input`. The size required can be queried with the
+    /// `decompress_len` function.
+    ///
+    /// On success, this returns the number of bytes written to `output`.
+    ///
+    /// The data written to `output` is guaranteed to be initialized to the
+    /// number of bytes written returned by the method. Data outside that
+    /// bound is not initialized and reading from it is UB. It's the callee's
+    /// responsibility to finalize the state of `output`.
+    ///
+    /// # Errors
+    ///
+    /// This method returns an error in the following circumstances:
+    ///
+    /// * Invalid compressed Snappy data was seen.
+    /// * The total space required for decompression exceeds `2^32 - 1`.
+    /// * `output` has length less than `decompress_len(input)`.
+    pub fn decompress_uninit(
+        &mut self,
+        input: &[u8],
+        output: &mut [MaybeUninit<u8>],
     ) -> Result<usize> {
         if input.is_empty() {
             return Err(Error::Empty);
@@ -88,8 +130,7 @@ impl Decoder {
             });
         }
         let dst = &mut output[..hdr.decompress_len];
-        let mut dec =
-            Decompress { src: &input[hdr.len..], s: 0, dst: dst, d: 0 };
+        let mut dec = Decompress { src: &input[hdr.len..], s: 0, dst, d: 0 };
         dec.decompress()?;
         Ok(dec.dst.len())
     }
@@ -103,9 +144,15 @@ impl Decoder {
     /// This method returns an error under the same circumstances that
     /// `decompress` does.
     pub fn decompress_vec(&mut self, input: &[u8]) -> Result<Vec<u8>> {
-        let mut buf = vec![0; decompress_len(input)?];
-        let n = self.decompress(input, &mut buf)?;
-        buf.truncate(n);
+        let mut buf = Vec::with_capacity(decompress_len(input)?);
+        let n = self.decompress_uninit(input, buf.spare_capacity_mut())?;
+
+        // SAFETY: compress_uninit guarantees all the data up to `n` is
+        // initialized
+        unsafe {
+            buf.set_len(n);
+        }
+
         Ok(buf)
     }
 }
@@ -117,7 +164,7 @@ struct Decompress<'s, 'd> {
     /// The current position in the compressed bytes.
     s: usize,
     /// The output buffer to write the decompressed bytes.
-    dst: &'d mut [u8],
+    dst: &'d mut [MaybeUninit<u8>],
     /// The current position in the decompressed buffer.
     d: usize,
 }
@@ -178,7 +225,7 @@ impl<'s, 'd> Decompress<'s, 'd> {
                 let srcp = self.src.as_ptr().add(self.s);
                 let dstp = self.dst.as_mut_ptr().add(self.d);
                 // Hopefully uses SIMD registers for 128 bit load/store.
-                ptr::copy_nonoverlapping(srcp, dstp, 16);
+                ptr::copy_nonoverlapping(srcp, dstp as *mut u8, 16);
             }
             self.d += len as usize;
             self.s += len as usize;
@@ -220,7 +267,7 @@ impl<'s, 'd> Decompress<'s, 'd> {
             // is correct.
             let srcp = self.src.as_ptr().add(self.s);
             let dstp = self.dst.as_mut_ptr().add(self.d);
-            ptr::copy_nonoverlapping(srcp, dstp, len as usize);
+            ptr::copy_nonoverlapping(srcp, dstp as *mut u8, len as usize);
         }
         self.s += len as usize;
         self.d += len as usize;