Rollup merge of #160563 - Ddystopia:improve-buffer-cursor, r=clarfonthey

Make `BorrowedCursor<'a, T>` covariant in `'a` and drop an indirection

This is a solution to #117693 (comment), with some improvements.

Currently `'a` in `BorrowedCursor` is invariant, though people seem to talk about it as if it were covariant, and the feature is in FCP right now. The previous version with `'buf` and `'data` lifetimes had the same flaw: `'data` was invariant.

A later PR landed that merged them and said that `BorrowedCursor` manually ensures that `'data` won't be ever overwritten thus invariance should not be needed. But unfortunately the lifetime is still left invariant. You can see it here, and the error spells it out exactly:

```rust
#![feature(core_io_borrowed_buf)]

use std::io::{BorrowedBuf, BorrowedCursor};

// Accepted.
fn buf_covariant<'short, 'long: 'short>(buf: BorrowedBuf<'long, u8>) -> BorrowedBuf<'short, u8> {
    buf
}

// Rejected.
fn cursor_covariant<'short, 'long: 'short>(
    cursor: BorrowedCursor<'long, u8>,
) -> BorrowedCursor<'short, u8> {
    cursor
}

// Rejected.
fn cursor_contravariant<'short, 'long: 'short>(
    cursor: BorrowedCursor<'short, u8>,
) -> BorrowedCursor<'long, u8> {
    cursor
}

fn main() {}

```

And the errors (also say that `BorrowedCursor` is invariant over `'a`):

```
error: lifetime may not live long enough
  --> src/main.rs:14:5
   |
11 | fn cursor_covariant<'short, 'long: 'short>(
   |                     ------  ----- lifetime `'long` defined here
   |                     |
   |                     lifetime `'short` defined here
...
14 |     cursor
   |     ^^^^^^ function was supposed to return data with lifetime `'long` but it is returning data with lifetime `'short`
   |
   = help: consider adding the following bound: `'short: 'long`
   = note: requirement occurs because of the type `BorrowedCursor<'_, u8>`, which makes the generic argument `'_` invariant
   = note: the struct `BorrowedCursor<'a, T>` is invariant over the parameter `'a`
   = help: see <https://doc.rust-lang.org/nomicon/subtyping.html> for more information about variance

error: lifetime may not live long enough
  --> src/main.rs:21:5
   |
18 | fn cursor_contravariant<'short, 'long: 'short>(
   |                         ------  ----- lifetime `'long` defined here
   |                         |
   |                         lifetime `'short` defined here
...
21 |     cursor
   |     ^^^^^^ function was supposed to return data with lifetime `'long` but it is returning data with lifetime `'short`
   |
   = help: consider adding the following bound: `'short: 'long`
   = note: requirement occurs because of the type `BorrowedCursor<'_, u8>`, which makes the generic argument `'_` invariant
   = note: the struct `BorrowedCursor<'a, T>` is invariant over the parameter `'a`
   = help: see <https://doc.rust-lang.org/nomicon/subtyping.html> for more information about variance

error: could not compile `play` (bin "play") due to 2 previous errors
```

This also removes the two `mem::transmute` calls that `unfilled` and `reborrow` used to shorten `&'this mut BorrowedBuf<'data, T>` into `&'this mut BorrowedBuf<'this, T>`.

---

Additionally I noticed that `BorrowedCursor` is not really as efficient as it could be, for a standard library: it contained a reference to the `BorrowedBuf`, which in turn contains a slice to the data. Without this, the fix is just replacing `&'a mut BorrowedBuf<'a, T>` with `NonNull<BorrowedBuf<'a, T>>`, plus some convenience helpers.

To fix this, I also stored a reborrowed pointer to the first element of the array, with the provenance to access the whole array. `filled` and `init` are still read from the pointer to `BorrowedBuf`, the buffer length is also read from it but carefully, in order to not create a retag which will trigger a foreign access to the pointer stored in `BorrowedCursor`, making it disabled. It increased the size of `BorrowedCursor` from one `usize` to two of them.

It is stored as the pointer rather than `&mut [MaybeUninit<T>]` to save a `usize` from the `BorrowedCursor` size.

Author: jhpratt

library/core/src/io/borrowed_buf.rs

@@ -1,7 +1,9 @@
 #![unstable(feature = "core_io_borrowed_buf", issue = "117693")]
 
 use crate::fmt::{self, Debug, Formatter};
-use crate::mem::{self, MaybeUninit};
+use crate::mem::MaybeUninit;
+use crate::ptr::NonNull;
+use crate::slice;
 
 /// A borrowed buffer of initially uninitialized elements, which is incrementally filled.
 ///
@@ -34,11 +36,23 @@ pub struct BorrowedBuf<'data, T> {
 }
 
 impl<T> Debug for BorrowedBuf<'_, T> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        BorrowedBufDebug { init: self.init, filled: self.filled, capacity: self.capacity() }.fmt(f)
+    }
+}
+
+struct BorrowedBufDebug {
+    init: bool,
+    filled: usize,
+    capacity: usize,
+}
+
+impl Debug for BorrowedBufDebug {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         f.debug_struct("BorrowedBuf")
             .field("init", &self.init)
             .field("filled", &self.filled)
-            .field("capacity", &self.capacity())
+            .field("capacity", &self.capacity)
             .finish()
     }
 }
@@ -70,11 +84,15 @@ impl<'data, T: Copy> From<&'data mut [MaybeUninit<T>]> for BorrowedBuf<'data, T>
 impl<'data, T: Copy> From<BorrowedCursor<'data, T>> for BorrowedBuf<'data, T> {
     #[inline]
     fn from(buf: BorrowedCursor<'data, T>) -> BorrowedBuf<'data, T> {
+        let filled = buf.filled();
+        let init = buf.is_buf_init();
+        let len = buf.buf_len();
         BorrowedBuf {
-            // SAFETY: no initialized element is ever uninitialized as per `BorrowedBuf`'s invariant
-            buf: unsafe { buf.buf.buf.get_unchecked_mut(buf.buf.filled..) },
+            // SAFETY: no initialized element is ever uninitialized as per `BorrowedBuf`'s
+            // invariant, and the cursor holds the unique access to those elements for `'data`
+            buf: unsafe { slice::from_raw_parts_mut(buf.buf.as_ptr().add(filled), len - filled) },
             filled: 0,
-            init: buf.buf.init,
+            init,
         }
     }
 }
@@ -144,15 +162,8 @@ impl<'data, T: Copy> BorrowedBuf<'data, T> {
     /// Returns a cursor over the unfilled part of the buffer.
     #[inline]
     pub fn unfilled<'this>(&'this mut self) -> BorrowedCursor<'this, T> {
-        BorrowedCursor {
-            // SAFETY: we never assign into `BorrowedCursor::buf`, so treating its
-            // lifetime covariantly is safe.
-            buf: unsafe {
-                mem::transmute::<&'this mut BorrowedBuf<'data, T>, &'this mut BorrowedBuf<'this, T>>(
-                    self,
-                )
-            },
-        }
+        let borrowed_buf = NonNull::from_mut(self);
+        BorrowedCursor { buf: NonNull::from_mut(self.buf).cast(), borrowed_buf }
     }
 
     /// Clears the buffer, resetting the filled region to empty.
@@ -193,16 +204,98 @@ impl<'data, T: Copy> BorrowedBuf<'data, T> {
 /// The lifetime `'a` is a bound on the lifetime of the underlying buffer (which means it is a bound
 /// on the elements in that buffer by transitivity).
 pub struct BorrowedCursor<'a, T> {
-    /// The underlying buffer.
-    // Safety invariant: we treat the type of buf as covariant in the lifetime of `BorrowedBuf` when
-    // we create a `BorrowedCursor`. This is only safe if we never replace `buf` by assigning into
-    // it, so don't do that!
-    buf: &'a mut BorrowedBuf<'a, T>,
+    /// The start of the elements of the buffer this cursor was created from.
+    /// Safety invariant: this points to the start of the *whole* buffer of `*borrowed_buf` and is
+    /// valid for reads and writes of `(*borrowed_buf).buf.len()` elements, so that
+    /// `(*borrowed_buf).filled` indexes into it.
+    buf: NonNull<MaybeUninit<T>>,
+    /// The buffer this cursor was created from.
+    /// Safety invariants:
+    /// 1. `(*borrowed_buf).buf` is *never* accessed by the owner of the pointee while the `buf`
+    ///    field above is alive, because there is a `&mut` of the pointee while the cursor is alive.
+    /// 2. We promise to only access the `filled` and `init` fields and the metadata of the `buf`
+    ///    field through the `borrowed_buf` pointer, never triggering any retag of `buf`'s pointer,
+    ///    as the `buf` field above holds a reborrow of it and reaching the parent again would be a
+    ///    foreign access for that reborrow. This includes not making a reference to the whole
+    ///    pointee out of `borrowed_buf`, but only accessing those fields directly through pointer
+    ///    manipulation.
+    borrowed_buf: NonNull<BorrowedBuf<'a, T>>,
 }
 
 impl<T> Debug for BorrowedCursor<'_, T> {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
-        f.debug_struct("BorrowedCursor").field("buf", &self.buf).finish()
+        let buf = BorrowedBufDebug {
+            init: self.is_buf_init(),
+            filled: self.filled(),
+            capacity: self.buf_len(),
+        };
+
+        f.debug_struct("BorrowedCursor").field("buf", &buf).finish()
+    }
+}
+
+// Helpers to access underlying buffer state.
+impl<'a, T> BorrowedCursor<'a, T> {
+    #[inline]
+    fn buf_mut(&mut self) -> &mut [MaybeUninit<T>] {
+        let len = self.buf_len();
+        // SAFETY: `buf` points to `len` elements that this cursor borrows exclusively.
+        unsafe { slice::from_raw_parts_mut(self.buf.as_ptr(), len) }
+    }
+
+    #[inline]
+    fn buf_len(&self) -> usize {
+        // SAFETY: We read just the metadata of `buf` and avoid retagging the reference.
+        unsafe {
+            let borrowed_buf = self.borrowed_buf.as_ptr();
+            let buf_ptr: *const &'a mut [MaybeUninit<T>] = &raw const (*borrowed_buf).buf;
+            // Same layout:
+            // https://doc.rust-lang.org/reference/type-layout.html#r-layout.pointer.intro
+            let buf_ptr: *const *const [MaybeUninit<T>] = buf_ptr.cast();
+            let buf: *const [MaybeUninit<T>] = *buf_ptr;
+            buf.len()
+        }
+    }
+
+    #[inline]
+    fn unfilled_slice(&mut self) -> &mut [MaybeUninit<T>] {
+        let filled = self.filled();
+        // SAFETY: always in bounds
+        unsafe { self.buf_mut().get_unchecked_mut(filled..) }
+    }
+
+    #[inline]
+    fn filled(&self) -> usize {
+        // SAFETY: We access just `filled` and avoid foreign read on `buf`.
+        unsafe { (*self.borrowed_buf.as_ptr()).filled }
+    }
+
+    #[inline]
+    fn is_buf_init(&self) -> bool {
+        // SAFETY: We access just `init` and avoid foreign read on `buf`.
+        unsafe { (*self.borrowed_buf.as_ptr()).init }
+    }
+
+    /// # Safety
+    ///
+    /// In case of `true` all the elements of the cursor must be initialized.
+    #[inline]
+    unsafe fn set_buf_init(&mut self, init: bool) {
+        // SAFETY: We access just `init` and avoid foreign read on `buf`.
+        unsafe {
+            (*self.borrowed_buf.as_ptr()).init = init;
+        }
+    }
+
+    /// # Safety
+    ///
+    /// The next `n` elements of the cursor must be initialized.
+    #[inline]
+    unsafe fn add_filled(&mut self, n: usize) {
+        // SAFETY: We access just `filled` and avoid foreign read on `buf`.
+        unsafe {
+            (*self.borrowed_buf.as_ptr()).filled += n;
+        }
     }
 }
 
@@ -213,36 +306,28 @@ impl<'a, T: Copy> BorrowedCursor<'a, T> {
     /// not accessible while the new cursor exists.
     #[inline]
     pub fn reborrow<'this>(&'this mut self) -> BorrowedCursor<'this, T> {
-        BorrowedCursor {
-            // SAFETY: we never assign into `BorrowedCursor::buf`, so treating its
-            // lifetime covariantly is safe.
-            buf: unsafe {
-                mem::transmute::<&'this mut BorrowedBuf<'a, T>, &'this mut BorrowedBuf<'this, T>>(
-                    self.buf,
-                )
-            },
-        }
+        BorrowedCursor { buf: self.buf, borrowed_buf: self.borrowed_buf }
     }
 
     /// Returns the available space in the cursor.
     #[inline]
     pub fn capacity(&self) -> usize {
-        self.buf.capacity() - self.buf.filled
+        self.buf_len() - self.filled()
     }
 
     /// Returns the number of elements written to the `BorrowedBuf` this cursor was created from.
     ///
     /// In particular, the count returned is shared by all reborrows of the cursor.
     #[inline]
     pub fn written(&self) -> usize {
-        self.buf.filled
+        self.filled()
     }
 
     /// Returns `true` if the buffer is initialized.
     #[unstable(feature = "borrowed_buf_init", issue = "160476")]
     #[inline]
     pub fn is_init(&self) -> bool {
-        self.buf.init
+        self.is_buf_init()
     }
 
     /// Set the buffer as fully initialized.
@@ -253,7 +338,8 @@ impl<'a, T: Copy> BorrowedCursor<'a, T> {
     #[unstable(feature = "borrowed_buf_init", issue = "160476")]
     #[inline]
     pub unsafe fn set_init(&mut self) {
-        self.buf.init = true;
+        // SAFETY: the caller guarantees that all the elements of the cursor are initialized.
+        unsafe { self.set_buf_init(true) }
     }
 
     /// Returns a mutable reference to the whole cursor.
@@ -263,8 +349,7 @@ impl<'a, T: Copy> BorrowedCursor<'a, T> {
     /// The caller must not uninitialize any elements of the cursor if it is initialized.
     #[inline]
     pub unsafe fn as_mut(&mut self) -> &mut [MaybeUninit<T>] {
-        // SAFETY: always in bounds
-        unsafe { self.buf.buf.get_unchecked_mut(self.buf.filled..) }
+        self.unfilled_slice()
     }
 
     /// Advances the cursor by asserting that `n` elements have been filled.
@@ -283,10 +368,11 @@ impl<'a, T: Copy> BorrowedCursor<'a, T> {
     #[inline]
     pub fn advance_checked(&mut self, n: usize) -> &mut Self {
         // The subtraction cannot underflow by invariant of this type.
-        let init_unfilled = if self.buf.init { self.buf.buf.len() - self.buf.filled } else { 0 };
+        let init_unfilled = if self.is_buf_init() { self.buf_len() - self.filled() } else { 0 };
         assert!(n <= init_unfilled);
 
-        self.buf.filled += n;
+        // SAFETY: the next `n` elements are initialized, as asserted above.
+        unsafe { self.advance(n) };
         self
     }
 
@@ -301,7 +387,8 @@ impl<'a, T: Copy> BorrowedCursor<'a, T> {
     /// The caller must ensure that the first `n` elements of the cursor have been initialized.
     #[inline]
     pub unsafe fn advance(&mut self, n: usize) -> &mut Self {
-        self.buf.filled += n;
+        // SAFETY: the caller guarantees that the first `n` elements of the cursor are initialized.
+        unsafe { self.add_filled(n) };
         self
     }
 
@@ -319,7 +406,8 @@ impl<'a, T: Copy> BorrowedCursor<'a, T> {
             self.as_mut()[..buf.len()].write_copy_of_slice(buf);
         }
 
-        self.buf.filled += buf.len();
+        // SAFETY: these elements have just been initialized.
+        unsafe { self.advance(buf.len()) };
     }
 
     /// Runs the given closure with a `BorrowedBuf` containing the unfilled part
@@ -349,8 +437,10 @@ impl<'a, T: Copy> BorrowedCursor<'a, T> {
         //
         // SAFETY: These elements 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 = init;
-        self.buf.filled += filled;
+        unsafe {
+            self.set_buf_init(init);
+            self.advance(filled);
+        }
 
         res
     }
@@ -362,15 +452,13 @@ impl<'a, T: Default + Copy> BorrowedCursor<'a, T> {
     #[unstable(feature = "borrowed_buf_init", issue = "160476")]
     #[inline]
     pub fn ensure_init(&mut self) -> &mut [T] {
-        // SAFETY: always in bounds and we never uninitialize these elements.
-        let unfilled = unsafe { self.buf.buf.get_unchecked_mut(self.buf.filled..) };
-
-        if !self.buf.init {
-            unfilled.write_default();
-            self.buf.init = true;
+        if !self.is_buf_init() {
+            self.unfilled_slice().write_default();
+            // SAFETY: buf is now initialized.
+            unsafe { self.set_buf_init(true) };
         }
 
         // SAFETY: these elements have just been initialized if they weren't before
-        unsafe { unfilled.assume_init_mut() }
+        unsafe { self.unfilled_slice().assume_init_mut() }
     }
 }