[PATCH 7/8] rust: pin-init: internal: project slots instead of references

[Gary Guo <gary@garyguo.net>]

By projecting slots, the `pin_init!` and `init!` code path can be more
unified. This also reduces the amount of macro-generated code and shifts
them to the shared infrastructure.

Signed-off-by: Gary Guo <gary@garyguo.net>
---
 rust/pin-init/internal/src/init.rs     | 113 +++++++++++----------------------
 rust/pin-init/internal/src/pin_data.rs |  52 ++++-----------
 rust/pin-init/src/__internal.rs        |  77 ++++++++++++++++++++++
 rust/pin-init/src/lib.rs               |  12 ++--
 4 files changed, 132 insertions(+), 122 deletions(-)

diff --git a/rust/pin-init/internal/src/init.rs b/rust/pin-init/internal/src/init.rs
index 11affa76d1fc..e6f5ea06f91b 100644
--- a/rust/pin-init/internal/src/init.rs
+++ b/rust/pin-init/internal/src/init.rs
@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0 OR MIT
 
 use proc_macro2::{Span, TokenStream};
-use quote::{format_ident, quote, quote_spanned};
+use quote::{format_ident, quote};
 use syn::{
     braced,
     parse::{End, Parse},
@@ -242,102 +242,61 @@ fn init_fields(
             }
         };
 
-        let init = match kind {
-            InitializerKind::Value { ident, value } => {
-                let mut value_ident = ident.clone();
-                let value_prep = value.as_ref().map(|value| &value.1).map(|value| {
-                    // Setting the span of `value_ident` to `value`'s span improves error messages
-                    // when the type of `value` is wrong.
-                    value_ident.set_span(value.span());
-                    quote!(let #value_ident = #value;)
-                });
-                // Again span for better diagnostics
-                let write = quote_spanned!(ident.span()=> ::core::ptr::write);
-                quote! {
-                    #(#attrs)*
-                    {
-                        #value_prep
-                        // SAFETY: TODO
-                        unsafe { #write(&raw mut (*#slot).#ident, #value_ident) };
-                    }
-                }
-            }
-            InitializerKind::Init { ident, value, .. } => {
-                // Again span for better diagnostics
-                let init = format_ident!("init", span = value.span());
-                let value_init = if pinned {
-                    quote! {
-                        // SAFETY:
-                        // - `slot` is valid, because we are inside of an initializer closure, we
-                        //   return when an error/panic occurs.
-                        // - We also use `#data` to require the correct trait (`Init` or `PinInit`)
-                        //   for `#ident`.
-                        unsafe { #data.#ident(&raw mut (*#slot).#ident, #init)? };
-                    }
-                } else {
-                    quote! {
-                        // SAFETY: `slot` is valid, because we are inside of an initializer
-                        // closure, we return when an error/panic occurs.
-                        unsafe {
-                            ::pin_init::Init::__init(
-                                #init,
-                                &raw mut (*#slot).#ident,
-                            )?
-                        };
-                    }
-                };
-                quote! {
-                    #(#attrs)*
-                    {
-                        let #init = #value;
-                        #value_init
-                    }
-                }
-            }
-            InitializerKind::Code { .. } => unreachable!(),
-        };
-
-        // `mixed_site` ensures that the guard is not accessible to the user-controlled code.
-        let guard = format_ident!("__{ident}_guard", span = Span::mixed_site());
-
-        let guard_creation = if pinned {
-            let project_ident = format_ident!("__project_{ident}");
+        let slot = if pinned {
             quote! {
                 // 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) }
+                // - `make_field_check` prevents `#ident` from being used twice, therefore
+                //   `(*slot).#ident` is exclusively accessed and has not been initialized.
+                (unsafe { #data.#ident(&raw mut (*#slot).#ident) })
             }
         } else {
             quote! {
+                // For `init!()` macro, everything is unpinned.
                 // 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
+                // - `make_field_check` prevents `#ident` from being used twice, therefore
+                //   `(*slot).#ident` is exclusively accessed and has not been initialized.
+                (unsafe {
+                    ::pin_init::__internal::Slot::<::pin_init::__internal::Unpinned, _>::new(
+                        &raw mut (*#slot).#ident
                     )
+                })
+            }
+        };
+
+        // `mixed_site` ensures that the guard is not accessible to the user-controlled code.
+        let guard = format_ident!("__{ident}_guard", span = Span::mixed_site());
+
+        let init = match kind {
+            InitializerKind::Value { ident, value } => {
+                let value = value
+                    .as_ref()
+                    .map(|(_, value)| quote!(#value))
+                    .unwrap_or_else(|| quote!(#ident));
+
+                quote! {
+                    #(#attrs)*
+                    let mut #guard = #slot.write(#value);
+
                 }
             }
+            InitializerKind::Init { value, .. } => {
+                quote! {
+                    #(#attrs)*
+                    let mut #guard = #slot.init(#value)?;
+                }
+            }
+            InitializerKind::Code { .. } => unreachable!(),
         };
 
         res.extend(quote! {
             #init
 
             #(#cfgs)*
-            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 = #guard.let_binding();
         });
diff --git a/rust/pin-init/internal/src/pin_data.rs b/rust/pin-init/internal/src/pin_data.rs
index 713a43c27826..3278a54510e1 100644
--- a/rust/pin-init/internal/src/pin_data.rs
+++ b/rust/pin-init/internal/src/pin_data.rs
@@ -352,10 +352,9 @@ fn generate_the_pin_data(
     let (impl_generics, ty_generics, whr) = generics.split_for_impl();
 
     // For every field, we create an initializing projection function according to its projection
-    // type. If a field is structurally pinned, then it must be initialized via `PinInit`, if it is
-    // not structurally pinned, then it can be initialized via `Init`.
-    //
-    // The functions are `unsafe` to prevent accidentally calling them.
+    // 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 = fields
         .iter()
         .map(|f| {
@@ -370,54 +369,29 @@ fn generate_the_pin_data(
             let field_name = ident
                 .as_ref()
                 .expect("only structs with named fields are supported");
-            let project_ident = format_ident!("__project_{field_name}");
-            let (init_ty, init_fn, pin_marker, pin_safety) = if f.pinned {
-                (
-                    quote!(PinInit),
-                    quote!(__pinned_init),
-                    quote!(Pinned),
-                    quote!(
-                        /// - `slot` will not move until it is dropped, i.e. it will be pinned.
-                    ),
-                )
+            let pin_marker = if f.pinned {
+                quote!(Pinned)
             } else {
-                (quote!(Init), quote!(__init), quote!(Unpinned), quote!())
+                quote!(Unpinned)
             };
             quote! {
-                /// # Safety
-                ///
-                /// - `slot` is a valid pointer to uninitialized memory.
-                /// - the caller does not touch `slot` when `Err` is returned, they are only
-                ///   permitted to deallocate.
-                #pin_safety
-                #(#attrs)*
-                #vis unsafe fn #field_name<E>(
-                    self,
-                    slot: *mut #ty,
-                    init: impl ::pin_init::#init_ty<#ty, E>,
-                ) -> ::core::result::Result<(), E> {
-                    // SAFETY: this function has the same safety requirements as the __init function
-                    // called below.
-                    unsafe { ::pin_init::#init_ty::#init_fn(init, 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.
+                ///   describes.
+                /// - `slot` is a valid, properly aligned and points to uninitialized and
+                ///   exclusively accessed memory.
                 #(#attrs)*
-                #vis unsafe fn #project_ident(
+                #[inline(always)]
+                #vis unsafe fn #field_name(
                     self,
                     slot: *mut #ty,
-                ) -> ::pin_init::__internal::DropGuard<::pin_init::__internal::#pin_marker, #ty> {
+                ) -> ::pin_init::__internal::Slot<::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) }
+                    unsafe { ::pin_init::__internal::Slot::new(slot) }
                 }
             }
         })
diff --git a/rust/pin-init/src/__internal.rs b/rust/pin-init/src/__internal.rs
index d7fdcfef41d2..854fbcaa93f3 100644
--- a/rust/pin-init/src/__internal.rs
+++ b/rust/pin-init/src/__internal.rs
@@ -251,6 +251,83 @@ struct Foo {
 pub struct Pinned;
 pub struct Unpinned;
 
+/// Represent an uninitialized field.
+///
+/// # Invariants
+///
+/// - `ptr` is valid, properly aligned and points to uninitialized and exclusively accessed memory.
+/// - If `P` is `Pinned`, then `ptr` is structurally pinned.
+pub struct Slot<P, T: ?Sized> {
+    ptr: *mut T,
+    _phantom: PhantomData<P>,
+}
+
+impl<P, T: ?Sized> Slot<P, T> {
+    /// # Safety
+    ///
+    /// - `ptr` is valid, properly aligned and points to uninitialized and exclusively accessed
+    ///    memory.
+    /// - If `P` is `Pinned`, then `ptr` is structurally pinned.
+    #[inline(always)]
+    pub unsafe fn new(ptr: *mut T) -> Self {
+        // INVARIANT: Per safety requirement.
+        Self {
+            ptr,
+            _phantom: PhantomData,
+        }
+    }
+
+    /// Initialize the field by value.
+    #[inline(always)]
+    pub fn write(self, value: T) -> DropGuard<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 and properly aligned 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 { DropGuard::new(self.ptr) }
+    }
+}
+
+impl<T: ?Sized> Slot<Unpinned, T> {
+    /// Initialize the field.
+    #[inline(always)]
+    pub fn init<E>(self, init: impl Init<T, E>) -> Result<DropGuard<Unpinned, T>, E> {
+        // SAFETY:
+        // - `self.ptr` is valid and properly aligned.
+        // - when `Err` is returned, we also propagate the error without touching `slot`;
+        //   also `self` is consumed so it cannot be touched further.
+        unsafe { init.__init(self.ptr)? };
+
+        // SAFETY:
+        // - `self.ptr` is valid and properly aligned per type invariant.
+        // - `*self.ptr` is initialized above and the ownership is transferred to the guard.
+        Ok(unsafe { DropGuard::new(self.ptr) })
+    }
+}
+
+impl<T: ?Sized> Slot<Pinned, T> {
+    /// Initialize the field.
+    #[inline(always)]
+    pub fn init<E>(self, init: impl PinInit<T, E>) -> Result<DropGuard<Pinned, 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.
+        // - the drop guard will not hand out `&mut` (only `Pin<&mut T>`).
+        unsafe { init.__pinned_init(self.ptr)? };
+
+        // SAFETY:
+        // - `self.ptr` is valid, properly aligned and pinned per type invariant.
+        // - `*self.ptr` is initialized above and the ownership is transferred to the guard.
+        Ok(unsafe { DropGuard::new(self.ptr) })
+    }
+}
+
 /// When a value of this type is dropped, it drops a `T`.
 ///
 /// Can be forgotten to prevent the drop.
diff --git a/rust/pin-init/src/lib.rs b/rust/pin-init/src/lib.rs
index 4098c65d63c3..e891d65cc469 100644
--- a/rust/pin-init/src/lib.rs
+++ b/rust/pin-init/src/lib.rs
@@ -867,12 +867,12 @@ macro_rules! stack_try_pin_init {
 #[macro_export]
 macro_rules! assert_pinned {
     ($ty:ty, $field:ident, $field_ty:ty, inline) => {
-        let _ = move |ptr: *mut $field_ty| {
-            // SAFETY: This code is unreachable.
-            let data = unsafe { <$ty as $crate::__internal::HasPinData>::__pin_data() };
-            let init = $crate::__internal::AlwaysFail::<$field_ty>::new();
-            // SAFETY: This code is unreachable.
-            unsafe { data.$field(ptr, init) }.ok();
+        // SAFETY: This code is unreachable.
+        let _ = move |ptr: *mut $field_ty| unsafe {
+            let data = <$ty as $crate::__internal::HasPinData>::__pin_data();
+            _ = data
+                .$field(ptr)
+                .init($crate::__internal::AlwaysFail::<$field_ty>::new());
         };
     };
 

-- 
2.51.2