BorrowedCursor: make init a boolean

library/core/src/io/borrowed_buf.rs

@@ -2,20 +2,18 @@
 
 use crate::fmt::{self, Debug, Formatter};
 use crate::mem::{self, MaybeUninit};
-use crate::{cmp, ptr};
+use crate::ptr;
 
-/// A borrowed byte buffer which is incrementally filled and initialized.
+/// A borrowed byte buffer which is incrementally filled.
 ///
-/// 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.
+/// 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
+/// whether the unfilled region was initialized.
 ///
 /// In summary, the contents of the buffer can be visualized as:
 /// ```not_rust
-/// [             capacity              ]
-/// [ filled |         unfilled         ]
-/// [    initialized    | uninitialized ]
+/// [                capacity                ]
+/// [ filled | unfilled (may be initialized) ]
 /// ```
 ///
 /// A `BorrowedBuf` is created around some existing data (or capacity for data) via a unique reference
@@ -30,8 +28,8 @@ pub struct BorrowedBuf<'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,
+    /// Whether the entire unfilled part of `self.buf` is known to be initialized.
+    init: bool,
 }
 
 impl Debug for BorrowedBuf<'_> {
@@ -48,24 +46,20 @@ 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: initialized data never becoming uninitialized is an invariant of BorrowedBuf
-            buf: unsafe { (slice as *mut [u8]).as_uninit_slice_mut().unwrap() },
+            buf: unsafe { &mut *(slice as *mut [u8] as *mut [MaybeUninit<u8>]) },
             filled: 0,
-            init: len,
+            init: true,
         }
     }
 }
 
 /// Creates a new `BorrowedBuf` from an uninitialized buffer.
-///
-/// 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, init: 0 }
+        BorrowedBuf { buf, filled: 0, init: false }
     }
 }
 
@@ -74,14 +68,13 @@ 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(mut buf: BorrowedCursor<'data>) -> BorrowedBuf<'data> {
-        let init = buf.init_mut().len();
+    fn from(buf: BorrowedCursor<'data>) -> BorrowedBuf<'data> {
         BorrowedBuf {
             // 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,
+            init: buf.buf.init,
         }
     }
 }
@@ -100,8 +93,9 @@ impl<'data> BorrowedBuf<'data> {
     }
 
     /// Returns the length of the initialized part of the buffer.
+    #[unstable(feature = "borrowed_buf_init", issue = "78485")]
     #[inline]
-    pub fn init_len(&self) -> usize {
+    pub fn is_init(&self) -> bool {
         self.init
     }
 
@@ -159,32 +153,29 @@ impl<'data> BorrowedBuf<'data> {
 
     /// Clears the buffer, resetting the filled region to empty.
     ///
-    /// The number of initialized bytes is not changed, and the contents of the buffer are not modified.
+    /// 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.
+    /// Asserts that the unfilled part of the buffer is initialized.
     ///
     /// # Safety
     ///
-    /// The caller must ensure that the first `n` unfilled bytes of the buffer have already been initialized.
+    /// All the bytes of the buffer must be initialized.
+    #[unstable(feature = "borrowed_buf_init", issue = "78485")]
     #[inline]
-    pub unsafe fn set_init(&mut self, n: usize) -> &mut Self {
-        self.init = cmp::max(self.init, n);
+    pub unsafe fn set_init(&mut self) -> &mut Self {
+        self.init = true;
         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.
+/// The unfilled portion may be uninitialized; 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
@@ -238,21 +229,29 @@ impl<'a> BorrowedCursor<'a> {
         self.buf.filled
     }
 
-    /// Returns a mutable reference to the initialized portion of the cursor.
+    /// Returns `true` if the buffer is initialized.
+    #[unstable(feature = "borrowed_buf_init", issue = "78485")]
     #[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()
-        }
+    pub fn is_init(&self) -> bool {
+        self.buf.init
+    }
+
+    /// Set the buffer as fully initialized.
+    ///
+    /// # Safety
+    ///
+    /// All the bytes of the cursor must be initialized.
+    #[unstable(feature = "borrowed_buf_init", issue = "78485")]
+    #[inline]
+    pub unsafe fn set_init(&mut self) {
+        self.buf.init = true;
     }
 
     /// Returns a mutable reference to the whole cursor.
     ///
     /// # Safety
     ///
-    /// The caller must not uninitialize any bytes in the initialized portion of the cursor.
+    /// The caller must not uninitialize any bytes of the cursor if it is initialized.
     #[inline]
     pub unsafe fn as_mut(&mut self) -> &mut [MaybeUninit<u8>] {
         // SAFETY: always in bounds
@@ -271,10 +270,12 @@ impl<'a> BorrowedCursor<'a> {
     /// # Panics
     ///
     /// Panics if there are less than `n` bytes initialized.
+    #[unstable(feature = "borrowed_buf_init", issue = "78485")]
     #[inline]
-    pub fn advance(&mut self, n: usize) -> &mut Self {
+    pub fn advance_checked(&mut self, n: usize) -> &mut Self {
         // The subtraction cannot underflow by invariant of this type.
-        assert!(n <= self.buf.init - self.buf.filled);
+        let init_unfilled = if self.buf.init { self.buf.buf.len() - self.buf.filled } else { 0 };
+        assert!(n <= init_unfilled);
 
         self.buf.filled += n;
         self
@@ -291,40 +292,29 @@ impl<'a> BorrowedCursor<'a> {
     /// The caller must ensure that the first `n` bytes of the cursor have been properly
     /// initialised.
     #[inline]
-    pub unsafe fn advance_unchecked(&mut self, n: usize) -> &mut Self {
+    pub unsafe fn advance(&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.
+    /// Initializes all bytes in the cursor and returns them.
+    #[unstable(feature = "borrowed_buf_init", issue = "78485")]
     #[inline]
-    pub fn ensure_init(&mut self) -> &mut Self {
+    pub fn ensure_init(&mut self) -> &mut [u8] {
         // 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());
+        let unfilled = unsafe { self.buf.buf.get_unchecked_mut(self.buf.filled..) };
+
+        if !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(unfilled.as_mut_ptr(), 0, unfilled.len());
+            }
+            self.buf.init = true;
         }
-        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
+        // SAFETY: these bytes have just been initialized if they weren't before
+        unsafe { unfilled.assume_init_mut() }
     }
 
     /// Appends data to the cursor, advancing position within its buffer.
@@ -341,10 +331,6 @@ 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();
     }
 
@@ -365,7 +351,7 @@ impl<'a> BorrowedCursor<'a> {
         // Check that the caller didn't replace the `BorrowedBuf`.
         // This is necessary for the safety of the code below: if the check wasn't
         // there, one could mark some bytes as initialized even though there aren't.
-        assert!(core::ptr::addr_eq(prev_ptr, buf.buf));
+        assert!(core::ptr::eq(prev_ptr, buf.buf));
 
         let filled = buf.filled;
         let init = buf.init;
@@ -376,7 +362,7 @@ impl<'a> BorrowedCursor<'a> {
         // 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.init = init;
         self.buf.filled += filled;
 
         res