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>

2 files changed, 169 insertions, 1 deletions

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());
         };