rust: pin-init: internal: pin_data: implement initialization of borrowed structs

We now have the checks to ensure that lifetime relations are what is
expected, we can generate the slot projections in `generate_pin_data` so
self-referential struct can be implemented.

New slot and guard types are defined (`SelfRefSlot` and `SelfRefDropGuard`)
which gives the generated let bindings longer lifetime than the guard
themselves.

Have `__make_closure` take `data` back as an argument. This gives
`#[pin_data]` an opportunity to change the type to add lifetimes.
Higher-ranked trait bound on `__make_closure` is used to ensure that the
initialization closure cannot make arbitrary assumptions of those
lifetimes.

Acked-by: Benno Lossin <lossin@kernel.org>
Link: https://patch.msgid.link/20261008-dev-selfref-v2-8-e280b3c8fba5@garyguo.net/
Signed-off-by: Gary Guo <gary@garyguo.net>

4 files changed, 305 insertions, 21 deletions

diff --git a/rust/pin-init/internal/src/init.rs b/rust/pin-init/internal/src/init.rs
index 80e4dc068d2e..ded2b5283376 100644
--- a/rust/pin-init/internal/src/init.rs
+++ b/rust/pin-init/internal/src/init.rs
@@ -288,8 +288,10 @@ fn expand(
         },
     };
     // `mixed_site` ensures that the data is not accessible to the user-controlled code.
-    let init_fields = init_fields(&fields, pinned);
+    let init_fields = make_field_init(&fields, pinned, false);
+    let drop_check = make_field_init(&fields, pinned, true);
     let field_check = make_field_check(&fields, init_kind, &path);
+
     Ok(quote_spanned! { Span::mixed_site() => {
         // Get the data about fields from the supplied type.
         let data = {
@@ -303,17 +305,29 @@ fn expand(
 
         // Ensure that `data` really is of type `data` and help with type inference:
         let init = data.__make_closure::<_, #error>(
-            move |slot| {
+            move |slot, data_lt| {
                 #zeroable_check
                 #this
-                #init_fields
+                // Generate init twice, which is mostly identical except for lifetimes.
+                if true {
+                    // In this path, we use the field lifetime from HRTB to prevent environment
+                    // lifetime from entering the fields, and to ensure that the dependency of
+                    // fields is consistent with the field drop of the struct.
+                    #init_fields
+                } else {
+                    // In this path, we use local lifetime, to make sure that if initialization
+                    // fails, the destructor execution will not cause lifetime issues. This is
+                    // separate as implied bounds between field lifetimes can be inconsistent with
+                    // that of the drop.
+                    #drop_check
+                }
                 #field_check
                 // SAFETY: we are the `init!` macro that is allowed to call this.
                 Ok(unsafe { ::pin_init::__internal::InitOk::new() })
             }
         );
         let init = move |slot| -> ::core::result::Result<(), #error> {
-            init(slot).map(|__InitOk| ())
+            init(slot, data.__with_lt()).map(|__InitOk| ())
         };
         // SAFETY: TODO
         unsafe { ::pin_init::#init_from_closure::<_, #error>(init) }
@@ -360,7 +374,11 @@ fn get_init_kind(rest: Option<(Token![..], Expr)>, dcx: &mut DiagCtxt) -> InitKi
 }
 
 /// Generate the code that initializes the fields of the struct using the initializers in `field`.
-fn init_fields(fields: &Punctuated<InitializerField, Token![,]>, pinned: bool) -> TokenStream {
+fn make_field_init(
+    fields: &Punctuated<InitializerField, Token![,]>,
+    pinned: bool,
+    dropck: bool,
+) -> TokenStream {
     let mut forget_guards = vec![];
     let mut res = TokenStream::new();
     for InitializerField { attrs, kind } in fields {
@@ -388,13 +406,18 @@ fn init_fields(fields: &Punctuated<InitializerField, Token![,]>, pinned: bool) -
         let span = Span::mixed_site().located_at(ident.span());
 
         let slot = if pinned {
+            let data = if !dropck {
+                quote_spanned!(span => data_lt)
+            } else {
+                quote_spanned!(span => data.__with_lt())
+            };
             quote_spanned! { span =>
                 // SAFETY:
                 // - `slot` is valid and properly aligned.
                 // - `make_field_check` checks that `&raw mut (*slot).#member` is properly aligned.
                 // - `make_field_check` prevents `#member` from being used twice, therefore
                 //   `(*slot).#member` is exclusively accessed and has not been initialized.
-                (unsafe { data.#ident(slot) })
+                (unsafe { #data.#ident(slot) })
             }
         } else {
             quote_spanned! { span =>
@@ -455,6 +478,11 @@ fn init_fields(fields: &Punctuated<InitializerField, Token![,]>, pinned: bool) -
 
         // A tuple field has no name that could be bound here (the `_0` identifiers are considered
         // implementation detail and not user-facing).
+        let let_binding_method = if !dropck {
+            format_ident!("let_binding", span = span)
+        } else {
+            format_ident!("let_binding_in_dropck", span = span)
+        };
         let binding = match member {
             Member::Named(ident) => quote_spanned! { span =>
                 #(#cfgs)*
@@ -462,7 +490,7 @@ fn init_fields(fields: &Punctuated<InitializerField, Token![,]>, pinned: bool) -
                 // struct field.
                 #[allow(unused_variables, non_snake_case)]
                 // Include `mut` so that `Pin<&mut T>` bindings can be reborrowed via `.as_mut()`.
-                let mut #ident = #guard.let_binding();
+                let mut #ident = #guard.#let_binding_method();
             },
             Member::Unnamed(_) => quote!(),
         };
diff --git a/rust/pin-init/internal/src/pin_data.rs b/rust/pin-init/internal/src/pin_data.rs
index 84a1fbacac9a..a46c4dcd7da5 100644
--- a/rust/pin-init/internal/src/pin_data.rs
+++ b/rust/pin-init/internal/src/pin_data.rs
@@ -61,7 +61,6 @@ enum BorrowedKind {
 }
 
 /// Information about a borrowed field.
-#[expect(unused)]
 struct BorrowedInfo {
     kind: BorrowedKind,
     /// Field lifetime for this field.
@@ -969,12 +968,36 @@ fn generate_the_pin_data(info: &StructInfo) -> TokenStream {
         generics,
         ..
     } = &info.struct_;
+
+    // Wrap in `CombinedGenerics` because it's ty generics will always output `<>`, so it can be
+    // used with `for`.
+    let field_lts = CombinedGenerics(vec![&info.field_lts]);
+    let generics_with_field_lt = CombinedGenerics(vec![&info.field_lts, generics]);
+
     let (impl_generics, ty_generics, whr) = generics.split_for_impl();
+    let (_, field_lt_ty_generics, _) = field_lts.split_for_impl();
+    let (impl_generics_with_lt, ty_generics_with_field_lt, whr_with_field_lt) =
+        generics_with_field_lt.split_for_impl();
+
+    // Wrap each field in a `PhantomInvariant`. For borrowed fields, additionally
+    // use `&#lt mut #ty` so the `lt` becomes associated with `#ty` which deduces
+    // implied bounds.
+    let phantom_fields = info.fields.iter().map(|f| {
+        let ty = &f.field.ty;
+        let ident = f.member.as_ident();
+
+        if let Some(borrowed) = &f.borrowed {
+            let lt = &borrowed.lifetime;
+            quote!(
+                #ident: ::pin_init::__internal::PhantomInvariant<&#lt mut #ty>,
+            )
+        } else {
+            quote!(
+                #ident: ::pin_init::__internal::PhantomInvariant<#ty>,
+            )
+        }
+    });
 
-    // For every field, we create an initializing projection function according to its projection
-    // type. If a field is structurally pinned, we create a `Slot` with `Pinned` which must be
-    // initialized via `PinInit`; if it is not structurally pinned, then we create a `Slot` with
-    // `Unpinned` which allows initialization via `Init`.
     let field_accessors = info
         .fields
         .iter()
@@ -987,6 +1010,30 @@ fn generate_the_pin_data(info: &StructInfo) -> TokenStream {
             } else {
                 quote!(Unpinned)
             };
+
+            let (slot_ty, slot_arg) = match &f.borrowed {
+                None => (quote!(Slot), quote!()),
+                Some(BorrowedInfo {
+                    kind: BorrowedKind::Shared,
+                    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,),
+                ),
+            };
+
             quote! {
                 /// # Safety
                 ///
@@ -1001,19 +1048,54 @@ fn generate_the_pin_data(info: &StructInfo) -> TokenStream {
                 #vis unsafe fn #field_name(
                     self,
                     slot: *mut #struct_name #ty_generics,
-                ) -> ::pin_init::__internal::Slot<::pin_init::__internal::#pin_marker, #ty> {
+                ) -> ::pin_init::__internal::#slot_ty<
+                    #slot_arg ::pin_init::__internal::#pin_marker, #ty
+                > {
+                    // CAST: `as _` is needed to convert types wrapped inside `SelfRef`.
                     // SAFETY:
                     // - If `#pin_marker` is `Pinned`, the corresponding field is structurally
                     //   pinned.
                     // - Other safety requirements follows the safety requirement.
-                    unsafe { ::pin_init::__internal::Slot::new(&raw mut (*slot).#member) }
+                    // - If `#slot_ty` is `SelfRefSlot`, the lifetime `#lt` represents that of the
+                    //   field.
+                    unsafe { ::pin_init::__internal::#slot_ty::new(&raw mut (*slot).#member as _) }
                 }
             }
         })
         .collect::<TokenStream>();
+
     quote! {
-        // We declare this struct which will host all of the projection function for our type. It
-        // will be invariant over all generic parameters which are inherited from the struct.
+        // We declare this struct which will host all of the projection function for our type.
+        #[doc(hidden)]
+        #[allow(non_snake_case)]
+        #vis struct __PinDataLt #generics_with_field_lt
+            #whr_with_field_lt
+        {
+            #(#phantom_fields)*
+            __pin_phantom: ::core::marker::PhantomData<#struct_name #ty_generics>,
+        }
+
+        impl #impl_generics_with_lt ::core::clone::Clone for __PinDataLt #ty_generics_with_field_lt
+            #whr_with_field_lt
+        {
+            fn clone(&self) -> Self { *self }
+        }
+
+        impl #impl_generics_with_lt ::core::marker::Copy for __PinDataLt #ty_generics_with_field_lt
+            #whr_with_field_lt
+        {}
+
+        #[allow(dead_code)] // Some functions might never be used and private.
+        #[expect(clippy::missing_safety_doc)]
+        impl #impl_generics_with_lt __PinDataLt #ty_generics_with_field_lt
+            #whr_with_field_lt
+        {
+            #field_accessors
+        }
+
+        // Declare a type that serves as the entry point of interaction with the `pin_init!` macro.
+        // We use this type instead of defining methods directly on user's type to avoid possibility
+        // of name conflicts.
         #[doc(hidden)]
         #vis struct __ThePinData #generics
             #whr
@@ -1032,7 +1114,6 @@ fn generate_the_pin_data(info: &StructInfo) -> TokenStream {
             #whr
         {}
 
-        #[allow(dead_code)] // Some functions might never be used and private.
         impl #impl_generics __ThePinData #ty_generics
             #whr
         {
@@ -1040,13 +1121,20 @@ fn generate_the_pin_data(info: &StructInfo) -> TokenStream {
             #[inline(always)]
             #vis fn __make_closure<__F, __E>(self, f: __F) -> __F
             where
-                __F: FnOnce(*mut #struct_name #ty_generics) ->
-                    ::core::result::Result<::pin_init::__internal::InitOk, __E>,
+                __F: for #field_lt_ty_generics ::core::ops::FnOnce(
+                    *mut #struct_name #ty_generics,
+                    __PinDataLt #ty_generics_with_field_lt
+                ) -> ::core::result::Result<::pin_init::__internal::InitOk, __E>,
             {
                 f
             }
 
-            #field_accessors
+            #[inline(always)]
+            #vis fn __with_lt #field_lts(self) -> __PinDataLt #ty_generics_with_field_lt {
+                // Generate a zeroed to avoid naming all fields.
+                // SAFETY: `__PinDataLt` only contains phantom fields.
+                unsafe { ::core::mem::zeroed() }
+            }
         }
 
         // SAFETY: We have added the correct projection functions above to `__ThePinData` and
diff --git a/rust/pin-init/src/__internal.rs b/rust/pin-init/src/__internal.rs
index 32bd6d8c39a1..b639c9204a7d 100644
--- a/rust/pin-init/src/__internal.rs
+++ b/rust/pin-init/src/__internal.rs
@@ -109,10 +109,15 @@ impl<T: ?Sized> InitData<T> {
     #[inline(always)]
     pub fn __make_closure<F, E>(self, f: F) -> F
     where
-        F: FnOnce(*mut T) -> Result<InitOk, E>,
+        F: FnOnce(*mut T, Self) -> Result<InitOk, E>,
     {
         f
     }
+
+    #[inline(always)]
+    pub fn __with_lt(self) -> Self {
+        self
+    }
 }
 
 /// Stack initializer helper type. Use [`stack_pin_init`] instead of this primitive.
@@ -322,6 +327,12 @@ impl<T: ?Sized> DropGuard<Unpinned, 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<T: ?Sized> DropGuard<Pinned, T> {
@@ -332,6 +343,12 @@ impl<T: ?Sized> DropGuard<Pinned, T> {
         // pinned per type invariant.
         unsafe { Pin::new_unchecked(&mut *self.ptr) }
     }
+
+    /// Create a let binding for accessor use in dropck.
+    #[inline]
+    pub fn let_binding_in_dropck(&mut self) -> Pin<&mut T> {
+        self.let_binding()
+    }
 }
 
 impl<P, T: ?Sized> Drop for DropGuard<P, T> {
@@ -342,6 +359,156 @@ impl<P, T: ?Sized> Drop for DropGuard<P, T> {
     }
 }
 
+/// Represent an uninitialized field in a pinned struct that will be referenced by other fields.
+///
+/// # Invariants
+///
+/// - `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 ptr: *mut T,
+    pub _phantom: PhantomData<(P, &'a mut T)>,
+}
+
+impl<'a, P, T: ?Sized> SelfRefSlot<'a, P, T> {
+    /// # Safety
+    ///
+    /// - `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.
+    #[inline]
+    pub unsafe fn new(ptr: *mut T) -> Self {
+        // INVARIANT: Per safety requirement.
+        Self {
+            ptr,
+            _phantom: PhantomData,
+        }
+    }
+
+    /// Initialize the field by value.
+    #[inline]
+    pub fn write(self, value: T) -> SelfRefDropGuard<'a, P, T>
+    where
+        T: Sized,
+    {
+        // SAFETY: `self.ptr` is a valid and aligned pointer for write.
+        unsafe { self.ptr.write(value) }
+        // SAFETY:
+        // - `self.ptr` is valid, properly aligned and live longer than `'a` per type invariant.
+        // - `*self.ptr` is initialized above and the ownership is transferred to the guard.
+        // - If `P` is `Pinned`, `self.ptr` is pinned.
+        unsafe { SelfRefDropGuard::new(self.ptr) }
+    }
+}
+
+impl<'a, T: ?Sized> SelfRefSlot<'a, Unpinned, T> {
+    /// Initialize the field.
+    #[inline]
+    pub fn init<E>(self, init: impl Init<T, E>) -> Result<SelfRefDropGuard<'a, Unpinned, T>, E> {
+        // SAFETY:
+        // - `self.ptr` is valid and properly aligned.
+        // - when `Err` is returned, we also propagate the error without touching `ptr`;
+        //   also `self` is consumed so it cannot be touched further.
+        unsafe { init.__init(self.ptr)? };
+
+        // SAFETY:
+        // - `self.ptr` is valid, properly aligned and live longer than `'a` per type invariant.
+        // - `*self.ptr` is initialized above and the ownership is transferred to the guard.
+        Ok(unsafe { SelfRefDropGuard::new(self.ptr) })
+    }
+}
+
+impl<'a, T: ?Sized> SelfRefSlot<'a, Pinned, T> {
+    /// Initialize the field.
+    #[inline]
+    pub fn init<E>(self, init: impl PinInit<T, E>) -> Result<SelfRefDropGuard<'a, Pinned, T>, E> {
+        // SAFETY:
+        // - `ptr` is valid
+        // - when `Err` is returned, we also propagate the error without touching `ptr`;
+        //   also `self` is consumed so it cannot be touched further.
+        // - the drop guard will not hand out `&mut` (only `Pin<&mut T>`).
+        unsafe { init.__init(self.ptr)? };
+
+        // SAFETY:
+        // - `self.ptr` is valid, properly aligned and live longer than `'a` per type invariant.
+        // - `*self.ptr` is initialized above and the ownership is transferred to the guard.
+        Ok(unsafe { SelfRefDropGuard::new(self.ptr) })
+    }
+}
+/// When a value of this type is dropped, it drops a `T`.
+///
+/// Can be forgotten to prevent the drop.
+///
+/// # Invariants
+///
+/// - `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> {
+    ptr: *mut T,
+    phantom: PhantomData<(P, &'a mut T)>,
+}
+
+impl<'a, P, T: ?Sized> SelfRefDropGuard<'a, 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`].
+    ///
+    /// # Safety
+    ///
+    /// - `ptr` is valid, properly aligned and live longer than `'a`.
+    /// - `*ptr` is initialized, and the ownership is transferred to this guard.
+    /// - if `P` is `Pinned`, `ptr` is pinned.
+    #[inline]
+    pub unsafe fn new(ptr: *mut T) -> Self {
+        // INVARIANT: By safety requirement.
+        Self {
+            ptr,
+            phantom: PhantomData,
+        }
+    }
+}
+
+impl<'a, T: ?Sized> SelfRefDropGuard<'a, Unpinned, T> {
+    /// Create a let binding for accessor use.
+    #[inline]
+    pub fn let_binding(&mut self) -> &'a T {
+        // SAFETY: Per type invariant.
+        unsafe { &*self.ptr }
+    }
+
+    /// Create a let binding for accessor use in dropck.
+    #[inline]
+    pub fn let_binding_in_dropck(&mut self) -> &T {
+        self.let_binding()
+    }
+}
+
+impl<'a, T: ?Sized> SelfRefDropGuard<'a, Pinned, T> {
+    /// Create a let binding for accessor use.
+    #[inline]
+    pub fn let_binding(&mut self) -> Pin<&'a T> {
+        // SAFETY: `self.ptr` is valid, properly aligned, live longer than `'a`, initialized,
+        // exclusively accessible and pinned per type invariant.
+        unsafe { Pin::new_unchecked(&*self.ptr) }
+    }
+
+    /// Create a let binding for accessor use in dropck.
+    #[inline]
+    pub fn let_binding_in_dropck(&mut self) -> Pin<&T> {
+        self.let_binding()
+    }
+}
+
+impl<P, T: ?Sized> Drop for SelfRefDropGuard<'_, P, T> {
+    #[inline]
+    fn drop(&mut self) {
+        // SAFETY: `self.ptr` is valid, properly aligned and `*self.ptr` is owned by this guard.
+        unsafe { ptr::drop_in_place(self.ptr) }
+    }
+}
+
 /// Token used by `PinnedDrop` to prevent calling the function without creating this unsafely
 /// created struct. This is needed, because the `drop` function is safe, but should not be called
 /// manually.
diff --git a/rust/pin-init/src/lib.rs b/rust/pin-init/src/lib.rs
index d1ed0561bb27..9ffa76434310 100644
--- a/rust/pin-init/src/lib.rs
+++ b/rust/pin-init/src/lib.rs
@@ -923,6 +923,7 @@ macro_rules! assert_pinned {
             let data = <$ty as $crate::__internal::HasInitData>::__init_data();
             let data = $crate::__internal::HasPinData::__pin_data(data);
             _ = data
+                .__with_lt()
                 .$field(ptr)
                 .init($crate::__internal::AlwaysFail::<$field_ty>::new());
         };