diff --git a/library/core/src/io/borrowed_buf.rs b/library/core/src/io/borrowed_buf.rs index 6bd9c18e00bdf..b1aa7606b3bd7 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]> for BorrowedBuf<'data> { } } +/// Creates a new `BorrowedBuf` from a cursor. +/// +/// Use `BorrowedCursor::with_unfilled_buf` instead for a safer alternative. +impl<'data> From> 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(&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 fbd3864dcac14..73dbfaf5ee906 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]); +}