Revert #148937 (Remove initialized-bytes tracking from BorrowedBuf and BorrowedCursor)

This caused several performance regressions because of existing code
which uses `Read::read` and therefore requires full buffer
initialization. This is particularly a problem when the same buffer is
re-used for multiple read calls since this means it needs to be fully
re-initialized each time.

There is still some benefit to landing the API changes, but we will have
to add private APIs so that the existing infrastructure can
track and avoid redundant initialization.

Author: Amanieu

library/core/src/io/borrowed_buf.rs

@@ -2,38 +2,42 @@
 
 use crate::fmt::{self, Debug, Formatter};
 use crate::mem::{self, MaybeUninit};
+use crate::{cmp, ptr};
 
-/// A borrowed buffer of initially uninitialized bytes, which is incrementally filled.
+/// A borrowed byte buffer which is incrementally filled and initialized.
 ///
-/// This type makes it safer to work with `MaybeUninit` buffers, such as to read into a buffer
-/// without having to initialize it first. It tracks the region of bytes that have been filled and
-/// the region that remains uninitialized.
+/// This type is a sort of "double cursor". It tracks three regions in the buffer: a region at the beginning of the
+/// buffer that has been logically filled with data, a region that has been initialized at some point but not yet
+/// logically filled, and a region at the end that is fully uninitialized. The filled region is guaranteed to be a
+/// subset of the initialized region.
 ///
-/// The contents of the buffer can be visualized as:
+/// In summary, the contents of the buffer can be visualized as:
 /// ```not_rust
-/// [                         capacity                            ]
-/// [ len: filled and initialized | capacity - len: uninitialized ]
+/// [             capacity              ]
+/// [ filled |         unfilled         ]
+/// [    initialized    | uninitialized ]
 /// ```
 ///
-/// Note that `BorrowedBuf` does not distinguish between uninitialized data and data that was
-/// previously initialized but no longer contains valid data.
-///
-/// A `BorrowedBuf` is created around some existing data (or capacity for data) via a unique
-/// reference (`&mut`). The `BorrowedBuf` can be configured (e.g., using `clear` or `set_len`), but
-/// cannot be directly written. To write into the buffer, use `unfilled` to create a
-/// `BorrowedCursor`. The cursor has write-only access to the unfilled portion of the buffer.
+/// A `BorrowedBuf` is created around some existing data (or capacity for data) via a unique reference
+/// (`&mut`). The `BorrowedBuf` can be configured (e.g., using `clear` or `set_init`), but cannot be
+/// directly written. To write into the buffer, use `unfilled` to create a `BorrowedCursor`. The cursor
+/// has write-only access to the unfilled portion of the buffer (you can think of it as a
+/// write-only iterator).
 ///
 /// The lifetime `'data` is a bound on the lifetime of the underlying data.
 pub struct BorrowedBuf<'data> {
     /// The buffer's underlying data.
     buf: &'data mut [MaybeUninit<u8>],
     /// The length of `self.buf` which is known to be filled.
     filled: usize,
+    /// The length of `self.buf` which is known to be initialized.
+    init: usize,
 }
 
 impl Debug for BorrowedBuf<'_> {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         f.debug_struct("BorrowedBuf")
+            .field("init", &self.init)
             .field("filled", &self.filled)
             .field("capacity", &self.capacity())
             .finish()
@@ -44,22 +48,24 @@ impl Debug for BorrowedBuf<'_> {
 impl<'data> From<&'data mut [u8]> for BorrowedBuf<'data> {
     #[inline]
     fn from(slice: &'data mut [u8]) -> BorrowedBuf<'data> {
+        let len = slice.len();
+
         BorrowedBuf {
-            // SAFETY: Always in bounds. We treat the buffer as uninitialized, even though it's
-            // already initialized.
+            // SAFETY: initialized data never becoming uninitialized is an invariant of BorrowedBuf
             buf: unsafe { (slice as *mut [u8]).as_uninit_slice_mut().unwrap() },
             filled: 0,
+            init: len,
         }
     }
 }
 
 /// Creates a new `BorrowedBuf` from an uninitialized buffer.
 ///
-/// Use `set_filled` if part of the buffer is known to be already filled.
+/// Use `set_init` if part of the buffer is known to be already initialized.
 impl<'data> From<&'data mut [MaybeUninit<u8>]> for BorrowedBuf<'data> {
     #[inline]
     fn from(buf: &'data mut [MaybeUninit<u8>]) -> BorrowedBuf<'data> {
-        BorrowedBuf { buf, filled: 0 }
+        BorrowedBuf { buf, filled: 0, init: 0 }
     }
 }
 
@@ -68,11 +74,14 @@ impl<'data> From<&'data mut [MaybeUninit<u8>]> for BorrowedBuf<'data> {
 /// Use `BorrowedCursor::with_unfilled_buf` instead for a safer alternative.
 impl<'data> From<BorrowedCursor<'data>> for BorrowedBuf<'data> {
     #[inline]
-    fn from(buf: BorrowedCursor<'data>) -> BorrowedBuf<'data> {
+    fn from(mut buf: BorrowedCursor<'data>) -> BorrowedBuf<'data> {
+        let init = buf.init_mut().len();
         BorrowedBuf {
-            // SAFETY: Always in bounds. We treat the buffer as uninitialized.
+            // SAFETY: no initialized byte is ever uninitialized as per
+            // `BorrowedBuf`'s invariant
             buf: unsafe { buf.buf.buf.get_unchecked_mut(buf.buf.filled..) },
             filled: 0,
+            init,
         }
     }
 }
@@ -90,6 +99,12 @@ impl<'data> BorrowedBuf<'data> {
         self.filled
     }
 
+    /// Returns the length of the initialized part of the buffer.
+    #[inline]
+    pub fn init_len(&self) -> usize {
+        self.init
+    }
+
     /// Returns a shared reference to the filled portion of the buffer.
     #[inline]
     pub fn filled(&self) -> &[u8] {
@@ -144,16 +159,33 @@ impl<'data> BorrowedBuf<'data> {
 
     /// Clears the buffer, resetting the filled region to empty.
     ///
-    /// The contents of the buffer are not modified.
+    /// The number of initialized bytes is not changed, and the contents of the buffer are not modified.
     #[inline]
     pub fn clear(&mut self) -> &mut Self {
         self.filled = 0;
         self
     }
+
+    /// Asserts that the first `n` bytes of the buffer are initialized.
+    ///
+    /// `BorrowedBuf` assumes that bytes are never de-initialized, so this method does nothing when called with fewer
+    /// bytes than are already known to be initialized.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that the first `n` unfilled bytes of the buffer have already been initialized.
+    #[inline]
+    pub unsafe fn set_init(&mut self, n: usize) -> &mut Self {
+        self.init = cmp::max(self.init, n);
+        self
+    }
 }
 
 /// A writeable view of the unfilled portion of a [`BorrowedBuf`].
 ///
+/// The unfilled portion consists of an initialized and an uninitialized part; see [`BorrowedBuf`]
+/// for details.
+///
 /// Data can be written directly to the cursor by using [`append`](BorrowedCursor::append) or
 /// indirectly by getting a slice of part or all of the cursor and writing into the slice. In the
 /// indirect case, the caller must call [`advance`](BorrowedCursor::advance) after writing to inform
@@ -206,17 +238,48 @@ impl<'a> BorrowedCursor<'a> {
         self.buf.filled
     }
 
+    /// Returns a mutable reference to the initialized portion of the cursor.
+    #[inline]
+    pub fn init_mut(&mut self) -> &mut [u8] {
+        // SAFETY: We only slice the initialized part of the buffer, which is always valid
+        unsafe {
+            let buf = self.buf.buf.get_unchecked_mut(self.buf.filled..self.buf.init);
+            buf.assume_init_mut()
+        }
+    }
+
     /// Returns a mutable reference to the whole cursor.
     ///
     /// # Safety
     ///
-    /// The caller must not uninitialize any previously initialized bytes.
+    /// The caller must not uninitialize any bytes in the initialized portion of the cursor.
     #[inline]
     pub unsafe fn as_mut(&mut self) -> &mut [MaybeUninit<u8>] {
         // SAFETY: always in bounds
         unsafe { self.buf.buf.get_unchecked_mut(self.buf.filled..) }
     }
 
+    /// Advances the cursor by asserting that `n` bytes have been filled.
+    ///
+    /// After advancing, the `n` bytes are no longer accessible via the cursor and can only be
+    /// accessed via the underlying buffer. I.e., the buffer's filled portion grows by `n` elements
+    /// and its unfilled portion (and the capacity of this cursor) shrinks by `n` elements.
+    ///
+    /// If less than `n` bytes initialized (by the cursor's point of view), `set_init` should be
+    /// called first.
+    ///
+    /// # Panics
+    ///
+    /// Panics if there are less than `n` bytes initialized.
+    #[inline]
+    pub fn advance(&mut self, n: usize) -> &mut Self {
+        // The subtraction cannot underflow by invariant of this type.
+        assert!(n <= self.buf.init - self.buf.filled);
+
+        self.buf.filled += n;
+        self
+    }
+
     /// Advances the cursor by asserting that `n` bytes have been filled.
     ///
     /// After advancing, the `n` bytes are no longer accessible via the cursor and can only be
@@ -225,11 +288,42 @@ impl<'a> BorrowedCursor<'a> {
     ///
     /// # Safety
     ///
-    /// The caller must ensure that the first `n` bytes of the cursor have been initialized. `n`
-    /// must not exceed the remaining capacity of this cursor.
+    /// The caller must ensure that the first `n` bytes of the cursor have been properly
+    /// initialised.
     #[inline]
-    pub unsafe fn advance(&mut self, n: usize) -> &mut Self {
+    pub unsafe fn advance_unchecked(&mut self, n: usize) -> &mut Self {
         self.buf.filled += n;
+        self.buf.init = cmp::max(self.buf.init, self.buf.filled);
+        self
+    }
+
+    /// Initializes all bytes in the cursor.
+    #[inline]
+    pub fn ensure_init(&mut self) -> &mut Self {
+        // SAFETY: always in bounds and we never uninitialize these bytes.
+        let uninit = unsafe { self.buf.buf.get_unchecked_mut(self.buf.init..) };
+
+        // SAFETY: 0 is a valid value for MaybeUninit<u8> and the length matches the allocation
+        // since it is comes from a slice reference.
+        unsafe {
+            ptr::write_bytes(uninit.as_mut_ptr(), 0, uninit.len());
+        }
+        self.buf.init = self.buf.capacity();
+
+        self
+    }
+
+    /// Asserts that the first `n` unfilled bytes of the cursor are initialized.
+    ///
+    /// `BorrowedBuf` assumes that bytes are never de-initialized, so this method does nothing when
+    /// called with fewer bytes than are already known to be initialized.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that the first `n` bytes of the buffer have already been initialized.
+    #[inline]
+    pub unsafe fn set_init(&mut self, n: usize) -> &mut Self {
+        self.buf.init = cmp::max(self.buf.init, self.buf.filled + n);
         self
     }
 
@@ -247,6 +341,10 @@ impl<'a> BorrowedCursor<'a> {
             self.as_mut()[..buf.len()].write_copy_of_slice(buf);
         }
 
+        // SAFETY: We just added the entire contents of buf to the filled section.
+        unsafe {
+            self.set_init(buf.len());
+        }
         self.buf.filled += buf.len();
     }
 
@@ -269,9 +367,17 @@ impl<'a> BorrowedCursor<'a> {
         // there, one could mark some bytes as initialized even though there aren't.
         assert!(core::ptr::addr_eq(prev_ptr, buf.buf));
 
-        // SAFETY: These bytes were filled in the `BorrowedBuf`, so they're filled in the cursor
-        // too, because the buffer wasn't replaced.
-        self.buf.filled += buf.filled;
+        let filled = buf.filled;
+        let init = buf.init;
+
+        // Update `init` and `filled` fields with what was written to the buffer.
+        // `self.buf.filled` was the starting length of the `BorrowedBuf`.
+        //
+        // SAFETY: These amounts of bytes were initialized/filled in the `BorrowedBuf`,
+        // and therefore they are initialized/filled in the cursor too, because the
+        // buffer wasn't replaced.
+        self.buf.init = self.buf.filled + init;
+        self.buf.filled += filled;
 
         res
     }