rust: init: add `stack_pin_init!` macro

The `stack_pin_init!` macro allows pin-initializing a value on the
stack. It accepts a `impl PinInit<T, E>` to initialize a `T`. It allows
propagating any errors via `?` or handling it normally via `match`.

Signed-off-by: Benno Lossin <benno.lossin@proton.me>
Reviewed-by: Alice Ryhl <aliceryhl@google.com>
Reviewed-by: Andreas Hindborg <a.hindborg@samsung.com>
Reviewed-by: Gary Guo <gary@garyguo.net>
Link: https://lore.kernel.org/r/20230408122429.1103522-11-y86-dev@protonmail.com
Signed-off-by: Miguel Ojeda <ojeda@kernel.org>

1 files changed, 57 insertions, 0 deletions

diff --git a/rust/kernel/init/__internal.rs b/rust/kernel/init/__internal.rs
index 774cb620afa7..44751fb62b51 100644
--- a/rust/kernel/init/__internal.rs
+++ b/rust/kernel/init/__internal.rs
@@ -112,6 +112,63 @@ unsafe impl<T: ?Sized> HasInitData for T {
     }
 }
 
+/// Stack initializer helper type. Use [`stack_pin_init`] instead of this primitive.
+///
+/// # Invariants
+///
+/// If `self.is_init` is true, then `self.value` is initialized.
+///
+/// [`stack_pin_init`]: kernel::stack_pin_init
+pub struct StackInit<T> {
+    value: MaybeUninit<T>,
+    is_init: bool,
+}
+
+impl<T> Drop for StackInit<T> {
+    #[inline]
+    fn drop(&mut self) {
+        if self.is_init {
+            // SAFETY: As we are being dropped, we only call this once. And since `self.is_init` is
+            // true, `self.value` is initialized.
+            unsafe { self.value.assume_init_drop() };
+        }
+    }
+}
+
+impl<T> StackInit<T> {
+    /// Creates a new [`StackInit<T>`] that is uninitialized. Use [`stack_pin_init`] instead of this
+    /// primitive.
+    ///
+    /// [`stack_pin_init`]: kernel::stack_pin_init
+    #[inline]
+    pub fn uninit() -> Self {
+        Self {
+            value: MaybeUninit::uninit(),
+            is_init: false,
+        }
+    }
+
+    /// Initializes the contents and returns the result.
+    #[inline]
+    pub fn init<E>(self: Pin<&mut Self>, init: impl PinInit<T, E>) -> Result<Pin<&mut T>, E> {
+        // SAFETY: We never move out of `this`.
+        let this = unsafe { Pin::into_inner_unchecked(self) };
+        // The value is currently initialized, so it needs to be dropped before we can reuse
+        // the memory (this is a safety guarantee of `Pin`).
+        if this.is_init {
+            this.is_init = false;
+            // SAFETY: `this.is_init` was true and therefore `this.value` is initialized.
+            unsafe { this.value.assume_init_drop() };
+        }
+        // SAFETY: The memory slot is valid and this type ensures that it will stay pinned.
+        unsafe { init.__pinned_init(this.value.as_mut_ptr())? };
+        // INVARIANT: `this.value` is initialized above.
+        this.is_init = true;
+        // SAFETY: The slot is now pinned, since we will never give access to `&mut T`.
+        Ok(unsafe { Pin::new_unchecked(this.value.assume_init_mut()) })
+    }
+}
+
 /// When a value of this type is dropped, it drops a `T`.
 ///
 /// Can be forgotten to prevent the drop.