Auto merge of #142885 - a1phyr:borrowed_cursor_to_buf, r=Mark-Simulacrum

core: Add `BorrowedCursor::with_unfilled_buf`

Implementation of https://github.com/rust-lang/libs-team/issues/367.

This mainly adds `BorrowedCursor::with_unfilled_buf`, with enables using the unfilled part of a cursor as a `BorrowedBuf`.

Note that unlike the ACP, `BorrowedCursor::unfilled_buf` was moved to a `From` conversion. This is more consistent with other ways of creating a `BorrowedBuf` and hides a bit this conversion that requires unsafe code to be used correctly.

Cc rust-lang/rust#78485 rust-lang/rust#117693

2 files changed, 87 insertions, 0 deletions

diff --git a/library/core/src/io/borrowed_buf.rs b/library/core/src/io/borrowed_buf.rs
index 854e03cf182..b5dfe6dbd71 100644
--- a/library/core/src/io/borrowed_buf.rs
+++ b/library/core/src/io/borrowed_buf.rs
@@ -69,6 +69,23 @@ impl<'data> From<&'data mut [MaybeUninit<u8>]> for BorrowedBuf<'data> {
     }
 }
 
+/// Creates a new `BorrowedBuf` from a cursor.
+///
+/// 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();
+        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,
+        }
+    }
+}
+
 impl<'data> BorrowedBuf<'data> {
     /// Returns the total capacity of the buffer.
     #[inline]
@@ -353,4 +370,38 @@ impl<'a> BorrowedCursor<'a> {
         }
         self.buf.filled += buf.len();
     }
+
+    /// Runs the given closure with a `BorrowedBuf` containing the unfilled part
+    /// of the cursor.
+    ///
+    /// This enables inspecting what was written to the cursor.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the `BorrowedBuf` given to the closure is replaced by another
+    /// one.
+    pub fn with_unfilled_buf<T>(&mut self, f: impl FnOnce(&mut BorrowedBuf<'_>) -> T) -> T {
+        let mut buf = BorrowedBuf::from(self.reborrow());
+        let prev_ptr = buf.buf as *const _;
+        let res = f(&mut buf);
+
+        // 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));
+
+        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
+    }
 }
diff --git a/library/coretests/tests/io/borrowed_buf.rs b/library/coretests/tests/io/borrowed_buf.rs
index fbd3864dcac..73dbfaf5ee9 100644
--- a/library/coretests/tests/io/borrowed_buf.rs
+++ b/library/coretests/tests/io/borrowed_buf.rs
@@ -165,3 +165,39 @@ fn cursor_set_init() {
     assert_eq!(rbuf.unfilled().uninit_mut().len(), 4);
     assert_eq!(unsafe { rbuf.unfilled().as_mut().len() }, 12);
 }
+
+#[test]
+fn cursor_with_unfilled_buf() {
+    let buf: &mut [_] = &mut [MaybeUninit::uninit(); 16];
+    let mut rbuf = BorrowedBuf::from(buf);
+    let mut cursor = rbuf.unfilled();
+
+    cursor.with_unfilled_buf(|buf| {
+        buf.unfilled().append(&[1, 2, 3]);
+        assert_eq!(buf.filled(), &[1, 2, 3]);
+    });
+
+    assert_eq!(cursor.init_mut().len(), 0);
+    assert_eq!(cursor.written(), 3);
+
+    cursor.with_unfilled_buf(|buf| {
+        assert_eq!(buf.capacity(), 13);
+        assert_eq!(buf.init_len(), 0);
+
+        buf.unfilled().ensure_init();
+        buf.unfilled().advance(4);
+    });
+
+    assert_eq!(cursor.init_mut().len(), 9);
+    assert_eq!(cursor.written(), 7);
+
+    cursor.with_unfilled_buf(|buf| {
+        assert_eq!(buf.capacity(), 9);
+        assert_eq!(buf.init_len(), 9);
+    });
+
+    assert_eq!(cursor.init_mut().len(), 9);
+    assert_eq!(cursor.written(), 7);
+
+    assert_eq!(rbuf.filled(), &[1, 2, 3, 0, 0, 0, 0]);
+}