rust: pin-init: internal: use marker on drop guard type for pinned fields

Instead of projecting the created reference, simply create drop guards with
different marker types and have the `let_binding()` method of guards of
different marker produce different type instead.

This allows more flexible lifetime as this is now controlled by the guard.
This will be needed when implementing self-referential fields.

Signed-off-by: Gary Guo <gary@garyguo.net>

3 files changed, 68 insertions, 44 deletions

diff --git a/rust/pin-init/internal/src/init.rs b/rust/pin-init/internal/src/init.rs
index 7eda1bcf0f28..a0b3c3790d43 100644
--- a/rust/pin-init/internal/src/init.rs
+++ b/rust/pin-init/internal/src/init.rs
@@ -303,18 +303,31 @@ fn init_fields(
         // `mixed_site` ensures that the guard is not accessible to the user-controlled code.
         let guard = format_ident!("__{ident}_guard", span = Span::mixed_site());
 
-        // NOTE: The reference is derived from the guard so that it only lives as long as the
-        // guard does and cannot escape the scope. If it's created via `&mut (*#slot).#ident`
-        // like the unaligned field guard, it will become effectively `'static`.
-        let accessor = if pinned {
+        let guard_creation = if pinned {
             let project_ident = format_ident!("__project_{ident}");
             quote! {
-                // SAFETY: the initialization is pinned.
-                unsafe { #data.#project_ident(#guard.let_binding()) }
+                // SAFETY:
+                // - `&raw mut (*slot).#ident` points to the `#ident` field of `slot`.
+                // - `&raw mut (*slot).#ident` is valid.
+                // - `make_field_check` checks that `&raw mut (*slot).#ident` is properly aligned.
+                // - `(*slot).#ident` has been initialized above.
+                // - We only need the ownership to the pointee back when initialization has
+                //   succeeded, where we `forget` the guard.
+                unsafe { #data.#project_ident(&raw mut (*slot).#ident) }
             }
         } else {
             quote! {
-                #guard.let_binding()
+                // SAFETY:
+                // - `&raw mut (*slot).#ident` is valid.
+                // - `make_field_check` checks that `&raw mut (*slot).#ident` is properly aligned.
+                // - `(*slot).#ident` has been initialized above.
+                // - We only need the ownership to the pointee back when initialization has
+                //   succeeded, where we `forget` the guard.
+                unsafe {
+                    ::pin_init::__internal::DropGuard::<::pin_init::__internal::Unpinned, _>::new(
+                        &raw mut (*slot).#ident
+                    )
+                }
             }
         };
 
@@ -322,24 +335,16 @@ fn init_fields(
             #init
 
             #(#cfgs)*
-            // Create the drop guard.
-            //
-            // SAFETY:
-            // - `&raw mut (*slot).#ident` is valid.
-            // - `make_field_check` checks that `&raw mut (*slot).#ident` is properly aligned.
-            // - `(*slot).#ident` has been initialized above.
-            // - We only need the ownership to the pointee back when initialization has
-            //   succeeded, where we `forget` the guard.
-            let mut #guard = unsafe {
-                ::pin_init::__internal::DropGuard::new(
-                    &raw mut (*slot).#ident
-                )
-            };
+            let mut #guard = #guard_creation;
 
             #(#cfgs)*
+            // NOTE: The reference is derived from the guard so that it only lives as long as the
+            // guard does and cannot escape the scope. If it's created via `&mut (*#slot).#ident`
+            // like the unaligned field guard, it will become effectively `'static`.
             #[allow(unused_variables)]
-            let #ident = #accessor;
+            let #ident = #guard.let_binding();
         });
+
         guards.push(guard);
         guard_attrs.push(cfgs);
     }
diff --git a/rust/pin-init/internal/src/pin_data.rs b/rust/pin-init/internal/src/pin_data.rs
index 44d0bd18e4ae..90f6b05b957c 100644
--- a/rust/pin-init/internal/src/pin_data.rs
+++ b/rust/pin-init/internal/src/pin_data.rs
@@ -371,29 +371,18 @@ fn generate_the_pin_data(
                 .as_ref()
                 .expect("only structs with named fields are supported");
             let project_ident = format_ident!("__project_{field_name}");
-            let (init_ty, init_fn, project_ty, project_body, pin_safety) = if f.pinned {
+            let (init_ty, init_fn, pin_marker, pin_safety) = if f.pinned {
                 (
                     quote!(PinInit),
                     quote!(__pinned_init),
-                    quote!(::core::pin::Pin<&'__slot mut #ty>),
-                    // SAFETY: this field is structurally pinned.
-                    quote!(unsafe { ::core::pin::Pin::new_unchecked(slot) }),
+                    quote!(Pinned),
                     quote!(
                         /// - `slot` will not move until it is dropped, i.e. it will be pinned.
                     ),
                 )
             } else {
-                (
-                    quote!(Init),
-                    quote!(__init),
-                    quote!(&'__slot mut #ty),
-                    quote!(slot),
-                    quote!(),
-                )
+                (quote!(Init), quote!(__init), quote!(Unpinned), quote!())
             };
-            let slot_safety = format!(
-                " `slot` points at the field `{field_name}` inside of `{struct_name}`, which is pinned.",
-            );
             quote! {
                 /// # Safety
                 ///
@@ -414,13 +403,21 @@ fn generate_the_pin_data(
 
                 /// # Safety
                 ///
-                #[doc = #slot_safety]
+                /// - `slot` points to a `#ident` field of a pinned struct that this `__ThePinData`
+                ///    describes.
+                /// - `slot` is valid and properly aligned.
+                /// - `*slot` is initialized, and the ownership is transferred to the returned
+                ///    guard.
                 #(#attrs)*
-                #vis unsafe fn #project_ident<'__slot>(
+                #vis unsafe fn #project_ident(
                     self,
-                    slot: &'__slot mut #ty,
-                ) -> #project_ty {
-                    #project_body
+                    slot: *mut #ty,
+                ) -> ::pin_init::__internal::DropGuard<::pin_init::__internal::#pin_marker, #ty> {
+                    // SAFETY:
+                    // - If `#pin_marker` is `Pinned`, the corresponding field is structurally
+                    //   pinned.
+                    // - Other safety requirements follows the safety requirement.
+                    unsafe { ::pin_init::__internal::DropGuard::new(slot) }
                 }
             }
         })
diff --git a/rust/pin-init/src/__internal.rs b/rust/pin-init/src/__internal.rs
index e54d90a4742e..010e8bfc6cd3 100644
--- a/rust/pin-init/src/__internal.rs
+++ b/rust/pin-init/src/__internal.rs
@@ -277,6 +277,10 @@ fn stack_init_reuse() {
     println!("{value:?}");
 }
 
+// Marker types that determines type of `DropGuard`'s let bindings.
+pub struct Pinned;
+pub struct Unpinned;
+
 /// When a value of this type is dropped, it drops a `T`.
 ///
 /// Can be forgotten to prevent the drop.
@@ -285,11 +289,13 @@ fn stack_init_reuse() {
 ///
 /// - `ptr` is valid and properly aligned.
 /// - `*ptr` is initialized and owned by this guard.
-pub struct DropGuard<T: ?Sized> {
+/// - if `P` is `Pinned`, `ptr` is pinned.
+pub struct DropGuard<P, T: ?Sized> {
     ptr: *mut T,
+    phantom: PhantomData<P>,
 }
 
-impl<T: ?Sized> DropGuard<T> {
+impl<P, T: ?Sized> DropGuard<P, T> {
     /// Creates a drop guard and transfer the ownership of the pointer content.
     ///
     /// The ownership is only relinguished if the guard is forgotten via [`core::mem::forget`].
@@ -298,12 +304,18 @@ impl<T: ?Sized> DropGuard<T> {
     ///
     /// - `ptr` is valid and properly aligned.
     /// - `*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 }
+        Self {
+            ptr,
+            phantom: PhantomData,
+        }
     }
+}
 
+impl<T: ?Sized> DropGuard<Unpinned, T> {
     /// Create a let binding for accessor use.
     #[inline]
     pub fn let_binding(&mut self) -> &mut T {
@@ -312,7 +324,17 @@ impl<T: ?Sized> DropGuard<T> {
     }
 }
 
-impl<T: ?Sized> Drop for DropGuard<T> {
+impl<T: ?Sized> DropGuard<Pinned, T> {
+    /// Create a let binding for accessor use.
+    #[inline]
+    pub fn let_binding(&mut self) -> Pin<&mut T> {
+        // SAFETY: `self.ptr` is valid, properly aligned, initialized, exclusively accessible and
+        // pinned per type invariant.
+        unsafe { Pin::new_unchecked(&mut *self.ptr) }
+    }
+}
+
+impl<P, T: ?Sized> Drop for DropGuard<P, T> {
     #[inline]
     fn drop(&mut self) {
         // SAFETY: `self.ptr` is valid, properly aligned and `*self.ptr` is owned by this guard.