internal: project slots instead not references

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>

Author: nbdd0121

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},
@@ -229,102 +229,57 @@ 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();
         });

internal/src/pin_data.rs

@@ -341,10 +341,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`.
     fn handle_field(
         Field {
             vis,
@@ -358,51 +357,25 @@ fn generate_the_pin_data(
         let ident = ident
             .as_ref()
             .expect("only structs with named fields are supported");
-        let project_ident = format_ident!("__project_{ident}");
-        let (init_ty, init_fn, pin_marker, pin_safety) = if 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 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 #ident<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.
+            /// - `slot` is a valid, properly aligned and points to uninitialized and exclusively memory.
             #(#attrs)*
-            #vis unsafe fn #project_ident(
+            #vis unsafe fn #ident(
                 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) }
             }
         }
     }

src/__internal.rs

@@ -257,6 +257,82 @@ fn stack_init_reuse() {
 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]
+    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) -> 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]
+    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]
+    pub fn init<E>(self, init: impl PinInit<T, E>) -> Result<DropGuard<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` (but only `Pin<&mut T>`) it has been dropped.
+        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.

src/lib.rs

@@ -867,12 +867,11 @@ pub use pin_init_internal::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());
         };
     };
 

tests/ui/compile-fail/init/colon_instead_of_arrow.stderr

@@ -5,14 +5,23 @@ error[E0308]: mismatched types
    |                 ------------------ the found opaque type
 ...
 21 |         pin_init!(Self { bar: Bar::new() })
-   |                          ---  ^^^^^^^^^^ expected `Bar`, found opaque type
-   |                          |
-   |                          arguments to this function are incorrect
+   |         ----------------------^^^^^^^^^^---
+   |         |                     |
+   |         |                     expected `Bar`, found opaque type
+   |         arguments to this method are incorrect
    |
    = note:   expected struct `Bar`
            found opaque type `impl pin_init::PinInit<Bar>`
-note: function defined here
-  --> $RUST/core/src/ptr/mod.rs
+help: the return type of this call is `impl pin_init::PinInit<Bar>` due to the type of the argument passed
+  --> tests/ui/compile-fail/init/colon_instead_of_arrow.rs:21:9
    |
-   | pub const unsafe fn write<T>(dst: *mut T, src: T) {
-   |                     ^^^^^
+21 |         pin_init!(Self { bar: Bar::new() })
+   |         ^^^^^^^^^^^^^^^^^^^^^^----------^^^
+   |                               |
+   |                               this argument influences the return type of `write`
+note: method defined here
+  --> src/__internal.rs
+   |
+   |     pub fn write(self, value: T) -> DropGuard<P, T>
+   |            ^^^^^
+   = note: this error originates in the macro `pin_init` (in Nightly builds, run with -Z macro-backtrace for more info)

tests/ui/compile-fail/init/field_value_wrong_type.stderr

@@ -2,12 +2,21 @@ error[E0308]: mismatched types
  --> tests/ui/compile-fail/init/field_value_wrong_type.rs:8:28
   |
 8 |     let _ = init!(Foo { a: () });
-  |                         -  ^^ expected `usize`, found `()`
-  |                         |
-  |                         arguments to this function are incorrect
+  |             ---------------^^---
+  |             |              |
+  |             |              expected `usize`, found `()`
+  |             arguments to this method are incorrect
   |
-note: function defined here
- --> $RUST/core/src/ptr/mod.rs
+help: the return type of this call is `()` due to the type of the argument passed
+ --> tests/ui/compile-fail/init/field_value_wrong_type.rs:8:13
   |
-  | pub const unsafe fn write<T>(dst: *mut T, src: T) {
-  |                     ^^^^^
+8 |     let _ = init!(Foo { a: () });
+  |             ^^^^^^^^^^^^^^^--^^^
+  |                            |
+  |                            this argument influences the return type of `write`
+note: method defined here
+ --> src/__internal.rs
+  |
+  |     pub fn write(self, value: T) -> DropGuard<P, T>
+  |            ^^^^^
+  = note: this error originates in the macro `init` (in Nightly builds, run with -Z macro-backtrace for more info)

tests/ui/compile-fail/init/invalid_init.stderr

@@ -1,10 +1,10 @@
-error[E0277]: the trait bound `impl pin_init::PinInit<Bar>: Init<Bar>` is not satisfied
+error[E0277]: the trait bound `impl pin_init::PinInit<Bar>: Init<Bar, _>` is not satisfied
   --> tests/ui/compile-fail/init/invalid_init.rs:19:16
    |
 18 |       let _ = init!(Foo {
    |  _____________-
 19 | |         bar <- Bar::new(),
-   | |                ^^^^^^^^^^ the trait `Init<Bar>` is not implemented for `impl pin_init::PinInit<Bar>`
+   | |                ^^^^^^^^^^ the trait `Init<Bar, _>` is not implemented for `impl pin_init::PinInit<Bar>`
 20 | |     });
    | |______- required by a bound introduced by this call
    |
@@ -19,3 +19,8 @@ help: the following other types implement trait `Init<T, E>`
 ...
    |   unsafe impl<T, E> Init<T, E> for Result<T, E> {
    |   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `Result<T, E>`
+note: required by a bound in `pin_init::__internal::Slot::<pin_init::__internal::Unpinned, T>::init`
+  --> src/__internal.rs
+   |
+   |     pub fn init<E>(self, init: impl Init<T, E>) -> Result<DropGuard<Unpinned, T>, E> {
+   |                                     ^^^^^^^^^^ required by this bound in `Slot::<Unpinned, T>::init`