[PATCH 14/20] rust: pin-init: internal: pin_data: support mutable borrows

[Gary Guo <gary@garyguo.net>]

Allow fields to be mutably referenced by other fields in addition to shared
references. In order for this to be sound, the fields that can be mutably
borrowed are blocked from being accessed via field access syntax or
projection to maintain the aliasing requirements.

Signed-off-by: Gary Guo <gary@garyguo.net>
---
 rust/pin-init/examples/selfref.rs      | 13 +++++
 rust/pin-init/internal/src/pin_data.rs | 94 ++++++++++++++++++++++++++++++----
 rust/pin-init/src/__internal.rs        | 64 ++++++++++++++++++-----
 3 files changed, 148 insertions(+), 23 deletions(-)

diff --git a/rust/pin-init/examples/selfref.rs b/rust/pin-init/examples/selfref.rs
index b5cdf96b6d1c..5e9494dcc17e 100644
--- a/rust/pin-init/examples/selfref.rs
+++ b/rust/pin-init/examples/selfref.rs
@@ -8,12 +8,18 @@
 struct SelfRef {
     part: &'str str,
     str: String,
+
+    mut_part: &'mut_str mut str,
+    #[borrowed(mut)]
+    mut_str: String,
 }
 
 fn use_self_ref() {
     stack_pin_init!(let foo = pin_init!(SelfRef {
         str: "hello world".to_owned(),
         part: &str[..5],
+        mut_str: "hello world".to_owned(),
+        mut_part: &mut mut_str[..5],
     }));
 
     // Access via projection.
@@ -28,6 +34,13 @@ fn use_self_ref() {
     });
 
     println!("{}", foo.part());
+
+    // Access fields that mutable borrow others are similar to those of shared borrow.
+    println!("{}", foo.as_mut().project().mut_part);
+    println!("{}", foo.mut_part());
+    foo.as_mut().with_project(|proj| {
+        proj.mut_part.make_ascii_uppercase();
+    });
 }
 
 fn main() {
diff --git a/rust/pin-init/internal/src/pin_data.rs b/rust/pin-init/internal/src/pin_data.rs
index 908c0773f791..54f0efef3226 100644
--- a/rust/pin-init/internal/src/pin_data.rs
+++ b/rust/pin-init/internal/src/pin_data.rs
@@ -5,7 +5,7 @@
 use proc_macro2::{Span, TokenStream};
 use quote::{format_ident, quote, quote_spanned, ToTokens};
 use syn::{
-    parse::{End, Nothing, Parse},
+    parse::{End, Nothing, Parse, ParseStream},
     parse_quote, parse_quote_spanned,
     punctuated::Punctuated,
     spanned::Spanned,
@@ -58,6 +58,8 @@ enum BorrowedKind {
     /// `#[borrowed]`, or implicitly inferreed.
     #[default]
     Shared,
+    // `#[borrowed(mut)]`.
+    Mutable,
 }
 
 impl BorrowedKind {
@@ -67,9 +69,17 @@ fn parse(dcx: &mut DiagCtxt, attrs: &mut Vec<Attribute>) -> Option<Self> {
         Some(if let Meta::Path(_) = attr.meta {
             BorrowedKind::Shared
         } else {
-            // Swallow the error and recover by inferring shared.
-            dcx.error(attr.path(), "unexpected `#[borrowed]` attribute");
-            BorrowedKind::Shared
+            match attr.parse_args_with(|input: ParseStream<'_>| {
+                let _: Token![mut] = input.parse()?;
+                Ok(BorrowedKind::Mutable)
+            }) {
+                Ok(v) => v,
+                Err(err) => {
+                    // Swallow the error and recover by inferring shared.
+                    dcx.error(attr.path(), err);
+                    BorrowedKind::Shared
+                }
+            }
         })
     }
 }
@@ -515,8 +525,17 @@ fn generate_struct_def(info: &StructInfo) -> TokenStream {
 
         let mut ty = ty.to_token_stream();
 
-        // Replace lifetime for self-referential fields.
-        if !field.captures.is_empty() {
+        // Replace lifetime for self-referential fields. For mutable fields, this uses `Erase` to
+        // block direct access.
+        if !field.captures.is_empty()
+            || matches!(
+                field.borrowed,
+                Some(BorrowedInfo {
+                    kind: BorrowedKind::Mutable,
+                    ..
+                })
+            )
+        {
             // Build a chain `for<'a> fn(&'a ()) -> ... -> (Ty,)`. Such type will have a `EraseTy`
             // implementation and thus may be used inside `Erased`.
             ty = quote!((#ty,));
@@ -920,9 +939,18 @@ fn generate_projections(info: &StructInfo) -> TokenStream {
                 )
             }
 
-            if !f.captures.iter().all(|b| b.variance == Variance::Covariant) {
+            if !f.captures.iter().all(|b| b.variance == Variance::Covariant)
+                || matches!(
+                    f.borrowed,
+                    Some(BorrowedInfo {
+                        kind: BorrowedKind::Mutable,
+                        ..
+                    })
+                )
+            {
                 // If the type is not covariant, it must omitted, as projection shortens the
                 // lifetime to `'__this`.
+                // Mutable borrow must be omitted for aliasing reason.
                 (
                     quote!(
                         #vis #name ::pin_init::__internal::NotVisible<&'__this #mut_token #ty>,
@@ -990,7 +1018,25 @@ fn generate_projections(info: &StructInfo) -> TokenStream {
                 this_lt.clone()
             };
 
-            if f.pinned {
+            if matches!(
+                f.borrowed,
+                Some(BorrowedInfo {
+                    kind: BorrowedKind::Mutable,
+                    ..
+                })
+            ) {
+                // If the type is not covariant, it must omitted, as projection shortens the
+                // lifetime to `'__this`.
+                // Mutable borrow must be omitted for aliasing reason.
+                (
+                    quote!(
+                        #vis #name ::pin_init::__internal::NotVisible<&#lt #mut_token #ty>,
+                    ),
+                    quote!(
+                        #name ::pin_init::__internal::NotVisible::new(),
+                    ),
+                )
+            } else if f.pinned {
                 (
                     quote!(
                         #vis #name ::core::pin::Pin<&#lt #mut_token #ty>,
@@ -1110,6 +1156,17 @@ fn generate_projections(info: &StructInfo) -> TokenStream {
             continue;
         }
 
+        if matches!(
+            f.borrowed,
+            Some(BorrowedInfo {
+                kind: BorrowedKind::Mutable,
+                ..
+            })
+        ) {
+            // Mutably borrowed fields cannot be accessed directly under any circumstance.
+            continue;
+        }
+
         if f.captures.iter().all(|b| b.variance == Variance::Covariant) {
             let f_doc = format!("Access the `{ident}` field on a shared reference of `Self`.");
             let vis = &f.field.vis;
@@ -1265,7 +1322,26 @@ fn generate_the_pin_data(info: &StructInfo) -> TokenStream {
                     // assumptions on the lifetime except for those implied by the struct's bounds,
                     // and we have validated them in `generate_drop_check`.
                     quote!(SelfRefSlot),
-                    quote!(#lifetime,),
+                    quote!(#lifetime, ::pin_init::__internal::Shared, ),
+                ),
+                Some(BorrowedInfo {
+                    kind: BorrowedKind::Mutable,
+                    lifetime,
+                }) => (
+                    // For borrowed fields, create a `SelfRefSlot`, which after initialization
+                    // turns into a `SelfRefDropGuard` instead of `DropGuard`.
+                    //
+                    // They're mostly the same, except that `SelfRefDropGuard` returns `&'field T`
+                    // instead of `&'guard T` for let bindings; this allows it to be used to be
+                    // used to initialize other fields.
+                    //
+                    // The soundness of doing so relies on fact that `__make_init` requires a
+                    // higher-ranked trait bound on the closure. Within the closure (which is the
+                    // caller of the generated slot projection functions here), it can make no
+                    // assumptions on the lifetime except for those implied by the struct's bounds,
+                    // and we have validated them in `generate_drop_check`.
+                    quote!(SelfRefSlot),
+                    quote!(#lifetime, ::pin_init::__internal::Mutable, ),
                 ),
             };
 
diff --git a/rust/pin-init/src/__internal.rs b/rust/pin-init/src/__internal.rs
index 2410156a0c99..ada07197acd9 100644
--- a/rust/pin-init/src/__internal.rs
+++ b/rust/pin-init/src/__internal.rs
@@ -361,6 +361,9 @@ fn drop(&mut self) {
     }
 }
 
+pub struct Shared;
+pub struct Mutable;
+
 /// Represent an uninitialized field in a pinned struct that will be referenced by other fields.
 ///
 /// # Invariants
@@ -368,12 +371,12 @@ fn drop(&mut self) {
 /// - `ptr` is valid, properly aligned and points to uninitialized and exclusively accessed memory
 ///   and will live longer than `'a`.
 /// - If `P` is `Pinned`, then `ptr` is structurally pinned.
-pub struct SelfRefSlot<'a, P, T: ?Sized> {
+pub struct SelfRefSlot<'a, M, P, T: ?Sized> {
     pub ptr: *mut T,
-    pub _phantom: PhantomData<(P, &'a mut T)>,
+    pub _phantom: PhantomData<(M, P, &'a mut T)>,
 }
 
-impl<'a, P, T: ?Sized> SelfRefSlot<'a, P, T> {
+impl<'a, M, P, T: ?Sized> SelfRefSlot<'a, M, P, T> {
     /// # Safety
     ///
     /// - `ptr` is valid, properly aligned and points to uninitialized and exclusively accessed
@@ -390,7 +393,7 @@ pub unsafe fn new(ptr: *mut T) -> Self {
 
     /// Initialize the field by value.
     #[inline]
-    pub fn write(self, value: T) -> SelfRefDropGuard<'a, P, T>
+    pub fn write(self, value: T) -> SelfRefDropGuard<'a, M, P, T>
     where
         T: Sized,
     {
@@ -404,10 +407,10 @@ pub fn write(self, value: T) -> SelfRefDropGuard<'a, P, T>
     }
 }
 
-impl<'a, T: ?Sized> SelfRefSlot<'a, Unpinned, T> {
+impl<'a, M, T: ?Sized> SelfRefSlot<'a, M, Unpinned, T> {
     /// Initialize the field.
     #[inline]
-    pub fn init<E>(self, init: impl Init<T, E>) -> Result<SelfRefDropGuard<'a, Unpinned, T>, E> {
+    pub fn init<E>(self, init: impl Init<T, E>) -> Result<SelfRefDropGuard<'a, M, Unpinned, T>, E> {
         // SAFETY:
         // - `self.ptr` is valid and properly aligned.
         // - when `Err` is returned, we also propagate the error without touching `slot`;
@@ -421,10 +424,13 @@ pub fn init<E>(self, init: impl Init<T, E>) -> Result<SelfRefDropGuard<'a, Unpin
     }
 }
 
-impl<'a, T: ?Sized> SelfRefSlot<'a, Pinned, T> {
+impl<'a, M, T: ?Sized> SelfRefSlot<'a, M, Pinned, T> {
     /// Initialize the field.
     #[inline]
-    pub fn init<E>(self, init: impl PinInit<T, E>) -> Result<SelfRefDropGuard<'a, Pinned, T>, E> {
+    pub fn init<E>(
+        self,
+        init: impl PinInit<T, E>,
+    ) -> Result<SelfRefDropGuard<'a, M, Pinned, T>, E> {
         // SAFETY:
         // - `ptr` is valid
         // - when `Err` is returned, we also propagate the error without touching `ptr`;
@@ -447,12 +453,12 @@ pub fn init<E>(self, init: impl PinInit<T, E>) -> Result<SelfRefDropGuard<'a, Pi
 /// - `ptr` is valid, properly aligned and live longer than `'a`.
 /// - `*ptr` is initialized and owned by this guard.
 /// - if `P` is `Pinned`, `ptr` is pinned.
-pub struct SelfRefDropGuard<'a, P, T: ?Sized> {
+pub struct SelfRefDropGuard<'a, M, P, T: ?Sized> {
     ptr: *mut T,
-    phantom: PhantomData<(P, &'a mut T)>,
+    phantom: PhantomData<(M, P, &'a mut T)>,
 }
 
-impl<'a, P, T: ?Sized> SelfRefDropGuard<'a, P, T> {
+impl<'a, M, P, T: ?Sized> SelfRefDropGuard<'a, M, P, T> {
     /// Creates a drop guard and transfer the ownership of the pointer content.
     ///
     /// The ownership is only relinquished if the guard is forgotten via [`core::mem::forget`].
@@ -472,7 +478,7 @@ pub unsafe fn new(ptr: *mut T) -> Self {
     }
 }
 
-impl<'a, T: ?Sized> SelfRefDropGuard<'a, Unpinned, T> {
+impl<'a, T: ?Sized> SelfRefDropGuard<'a, Shared, Unpinned, T> {
     /// Create a let binding for accessor use.
     #[inline]
     pub fn let_binding(&mut self) -> &'a T {
@@ -487,7 +493,7 @@ pub fn let_binding_in_dropck(&mut self) -> &T {
     }
 }
 
-impl<'a, T: ?Sized> SelfRefDropGuard<'a, Pinned, T> {
+impl<'a, T: ?Sized> SelfRefDropGuard<'a, Shared, Pinned, T> {
     /// Create a let binding for accessor use.
     #[inline]
     pub fn let_binding(&mut self) -> Pin<&'a T> {
@@ -503,7 +509,37 @@ pub fn let_binding_in_dropck(&mut self) -> Pin<&T> {
     }
 }
 
-impl<P, T: ?Sized> Drop for SelfRefDropGuard<'_, P, T> {
+impl<'a, T: ?Sized> SelfRefDropGuard<'a, Mutable, Unpinned, T> {
+    /// Create a let binding for accessor use.
+    #[inline]
+    pub fn let_binding(&mut self) -> &'a mut T {
+        // SAFETY: Per type invariant.
+        unsafe { &mut *self.ptr }
+    }
+
+    /// Create a let binding for accessor use in dropck.
+    #[inline]
+    pub fn let_binding_in_dropck(&mut self) -> &mut T {
+        self.let_binding()
+    }
+}
+
+impl<'a, T: ?Sized> SelfRefDropGuard<'a, Mutable, Pinned, T> {
+    /// Create a let binding for accessor use.
+    #[inline]
+    pub fn let_binding(&mut self) -> Pin<&'a mut T> {
+        // SAFETY: `self.ptr` is valid, properly aligned, live longer than `'a`, initialized,
+        // exclusively accessible and pinned per type invariant.
+        unsafe { Pin::new_unchecked(&mut *self.ptr) }
+    }
+
+    #[inline]
+    pub fn let_binding_in_dropck(&mut self) -> Pin<&mut T> {
+        self.let_binding()
+    }
+}
+
+impl<M, P, T: ?Sized> Drop for SelfRefDropGuard<'_, M, P, T> {
     #[inline]
     fn drop(&mut self) {
         // SAFETY: `self.ptr` is valid, properly aligned and `*self.ptr` is owned by this guard.

-- 
2.54.0