feature: Add uninitialized output buffer support (#414)

* Add uninitialized buffer API to `WriteBuffer`

* Demono `WriteBuffer::copy_unwritten_from`

 - reduce compile time
 - having one version would make it easier to optimize for compiler

* Use bzip2 0.6.1 uninitialized API

* Use deflate64 0.1.10 uninitialized API

* Use flate2 1.1.4 uninitialized API

* Support uninitialized buffer for lz4 codecs

* Ret early in `Lz4Encoder::write` if output has no spare space

Avoid allocating a new buffer when the output has no spare space,
it makes no sense to allocate an internal buffer
and try compress when output is empty

* Fix `Lz4Encoder`: Use `Vec::spare_capacity_mut()`

This is the safe usage, previous usage is ok in practice,
but is actually considered an out-of-bound access by miri and any santizier.

* Refactor xz2 codecs: Extract new fn `process_stream`

Shared between encoder and decoder

* Use liblzma 0.4.5 uninitialized API

* Use zstd uninitialized API

* Refactor: Extract `tokio::poll_read`

Extract duplicate code, easier to modify

* Support uninitialized buffer for `tokio::bufread`

Author: NobodyXu

crates/async-compression/src/tokio/bufread/generic/decoder.rs

@@ -4,12 +4,11 @@ use crate::{
     generic::bufread::impl_decoder,
 };
 
-use core::{
+use std::{
+    io::{IoSlice, Result},
     pin::Pin,
     task::{Context, Poll},
 };
-use std::io::{IoSlice, Result};
-
 use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, ReadBuf};
 
 impl_decoder!();
@@ -20,19 +19,7 @@ impl<R: AsyncBufRead, D: DecodeV2> AsyncRead for Decoder<R, D> {
         cx: &mut Context<'_>,
         buf: &mut ReadBuf<'_>,
     ) -> Poll<Result<()>> {
-        if buf.remaining() == 0 {
-            return Poll::Ready(Ok(()));
-        }
-
-        let mut output = WriteBuffer::new_initialized(buf.initialize_unfilled());
-        match self.do_poll_read(cx, &mut output)? {
-            Poll::Pending if output.written().is_empty() => Poll::Pending,
-            _ => {
-                let len = output.written_len();
-                buf.advance(len);
-                Poll::Ready(Ok(()))
-            }
-        }
+        super::poll_read(buf, |output| self.do_poll_read(cx, output))
     }
 }
 

crates/async-compression/src/tokio/bufread/generic/encoder.rs

@@ -18,19 +18,7 @@ impl<R: AsyncBufRead, E: EncodeV2> AsyncRead for Encoder<R, E> {
         cx: &mut Context<'_>,
         buf: &mut ReadBuf<'_>,
     ) -> Poll<Result<()>> {
-        if buf.remaining() == 0 {
-            return Poll::Ready(Ok(()));
-        }
-
-        let mut output = WriteBuffer::new_initialized(buf.initialize_unfilled());
-        match self.do_poll_read(cx, &mut output)? {
-            Poll::Pending if output.written().is_empty() => Poll::Pending,
-            _ => {
-                let len = output.written_len();
-                buf.advance(len);
-                Poll::Ready(Ok(()))
-            }
-        }
+        super::poll_read(buf, |output| self.do_poll_read(cx, output))
     }
 }
 

crates/async-compression/src/tokio/bufread/generic/mod.rs

@@ -2,3 +2,36 @@ mod decoder;
 mod encoder;
 
 pub use self::{decoder::Decoder, encoder::Encoder};
+
+use crate::core::util::WriteBuffer;
+use std::{io::Result, task::Poll};
+use tokio::io::ReadBuf;
+
+fn poll_read(
+    buf: &mut ReadBuf<'_>,
+    do_poll_read: impl FnOnce(&mut WriteBuffer<'_>) -> Poll<Result<()>>,
+) -> Poll<Result<()>> {
+    if buf.remaining() == 0 {
+        return Poll::Ready(Ok(()));
+    }
+
+    let initialized = buf.initialized().len() - buf.filled().len();
+    // Safety: `WriteBuffer` has the same safety invariant as `ReadBuf`
+    let mut output = WriteBuffer::new_uninitialized(unsafe { buf.unfilled_mut() });
+    // Safety: `ReadBuf` ensures that it is initialized
+    unsafe { output.assume_init(initialized) };
+
+    let res = do_poll_read(&mut output);
+
+    let initialized = output.initialized_len();
+    let written = output.written_len();
+
+    // Safety: We trust our implementation to have properly initialized it
+    unsafe { buf.assume_init(initialized) };
+    buf.advance(written);
+
+    match res? {
+        Poll::Pending if written == 0 => Poll::Pending,
+        _ => Poll::Ready(Ok(())),
+    }
+}

crates/compression-codecs/Cargo.toml

@@ -47,12 +47,12 @@ deflate64 = ["dep:deflate64"]
 compression-core.workspace = true
 # features
 brotli = { version = "8", optional = true }
-bzip2 = { version = "0.6", optional = true }
-deflate64 = { version = "0.1.5", optional = true }
-flate2 = { version = "1.0.13", optional = true }
+bzip2 = { version = "0.6.1", optional = true }
+deflate64 = { version = "0.1.10", optional = true }
+flate2 = { version = "1.1.4", optional = true }
 libzstd = { package = "zstd", version = "0.13.1", optional = true, default-features = false }
 lz4 = { version = "1.28.1", optional = true }
-liblzma = { version = "0.4.4", optional = true }
+liblzma = { version = "0.4.5", optional = true }
 memchr = { version = "2", optional = true }
 zstd-safe = { version = "7", optional = true, default-features = false }
 

crates/compression-codecs/src/bzip2/decoder.rs

@@ -43,11 +43,15 @@ impl BzDecoder {
 
         let result = self
             .decompress
-            .decompress(input.unwritten(), output.initialize_unwritten())
+            // Safety: We **trust** bzip2 to only write initialized data to it
+            .decompress_uninit(input.unwritten(), unsafe { output.unwritten_mut() })
             .map_err(io::Error::other);
 
         input.advance((self.decompress.total_in() - prior_in) as usize);
-        output.advance((self.decompress.total_out() - prior_out) as usize);
+        // Safety: We **trust** bzip2 to write bytes properly
+        unsafe {
+            output.assume_init_and_advance((self.decompress.total_out() - prior_out) as usize)
+        };
 
         // Track when stream has properly ended
         if matches!(result, Ok(Status::StreamEnd)) {

crates/compression-codecs/src/bzip2/encoder.rs

@@ -57,11 +57,13 @@ impl BzEncoder {
 
         let result = self
             .compress
-            .compress(input.unwritten(), output.initialize_unwritten(), action)
+            // Safety: We **trust** bzip2 to only write initialized bytes into it
+            .compress_uninit(input.unwritten(), unsafe { output.unwritten_mut() }, action)
             .map_err(io::Error::other);
 
         input.advance((self.compress.total_in() - prior_in) as usize);
-        output.advance((self.compress.total_out() - prior_out) as usize);
+        // Safety: We **trust** bzip2 to properly write bytes into it
+        unsafe { output.assume_init_and_advance((self.compress.total_out() - prior_out) as usize) };
 
         result
     }

crates/compression-codecs/src/deflate64/decoder.rs

@@ -28,10 +28,12 @@ impl Deflate64Decoder {
     ) -> Result<bool> {
         let result = self
             .inflater
-            .inflate(input.unwritten(), output.initialize_unwritten());
+            // Safety: We **trust** deflate64 to not write uninitialized bytes
+            .inflate_uninit(input.unwritten(), unsafe { output.unwritten_mut() });
 
         input.advance(result.bytes_consumed);
-        output.advance(result.bytes_written);
+        // Safety: We **trust** deflate64 to properly write bytes into buffer
+        unsafe { output.assume_init_and_advance(result.bytes_written) };
 
         if result.data_error {
             Err(Error::new(ErrorKind::InvalidData, "invalid data"))

crates/compression-codecs/src/flate/decoder.rs

@@ -26,12 +26,16 @@ impl FlateDecoder {
         let prior_in = self.decompress.total_in();
         let prior_out = self.decompress.total_out();
 
-        let result =
-            self.decompress
-                .decompress(input.unwritten(), output.initialize_unwritten(), flush);
+        let result = self
+            .decompress
+            // Safety: We **trust** flate2 to not write uninitialized bytes into buffer
+            .decompress_uninit(input.unwritten(), unsafe { output.unwritten_mut() }, flush);
 
         input.advance((self.decompress.total_in() - prior_in) as usize);
-        output.advance((self.decompress.total_out() - prior_out) as usize);
+        // Safety: We **trust** flate2 to write bytes into buffer properly
+        unsafe {
+            output.assume_init_and_advance((self.decompress.total_out() - prior_out) as usize)
+        };
 
         Ok(result?)
     }

crates/compression-codecs/src/flate/encoder.rs

@@ -31,12 +31,14 @@ impl FlateEncoder {
         let prior_in = self.compress.total_in();
         let prior_out = self.compress.total_out();
 
-        let result =
-            self.compress
-                .compress(input.unwritten(), output.initialize_unwritten(), flush);
+        let result = self
+            .compress
+            // Safety: We **trust** flate2 to not write uninitialized bytes into buffer
+            .compress_uninit(input.unwritten(), unsafe { output.unwritten_mut() }, flush);
 
         input.advance((self.compress.total_in() - prior_in) as usize);
-        output.advance((self.compress.total_out() - prior_out) as usize);
+        // Safety: We **trust** flate2 to write bytes properly into buffer
+        unsafe { output.assume_init_and_advance((self.compress.total_out() - prior_out) as usize) };
 
         Ok(result?)
     }

crates/compression-codecs/src/lz4/decoder.rs

@@ -61,22 +61,28 @@ impl DecodeV2 for Lz4Decoder {
         input: &mut PartialBuffer<&[u8]>,
         output: &mut WriteBuffer<'_>,
     ) -> Result<bool> {
-        let out_buf = output.initialize_unwritten();
-
-        let mut output_size = out_buf.len();
         let mut input_size = input.unwritten().len();
+
+        // Safety: We **trust** lz4 bytes to properly function as expected,
+        // only write decompressed, initialized data into the buffer properly.
         let result = unsafe {
-            check_error(LZ4F_decompress(
+            let out_buf = output.unwritten_mut();
+
+            let mut output_size = out_buf.len();
+
+            let result = check_error(LZ4F_decompress(
                 self.ctx.get_mut().ctx,
-                out_buf.as_mut_ptr(),
+                out_buf.as_mut_ptr() as *mut _,
                 &mut output_size,
                 input.unwritten().as_ptr(),
                 &mut input_size,
                 core::ptr::null(),
-            ))
+            ));
+            output.assume_init_and_advance(output_size);
+
+            result
         };
         input.advance(input_size);
-        output.advance(output_size);
 
         let finished = result? == 0;
         if finished {