[PATCH v5 10/15] rust: init: add `stack_pin_init!` macro

Benno Lossin
2023-04-03T16:06:02Z

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 <y86-dev@protonmail.com>
Cc: Alice Ryhl <aliceryhl@google.com>
Cc: Andreas Hindborg <a.hindborg@samsung.com>
Cc: Gary Guo <gary@garyguo.net>
---
rust/kernel/init.rs | 146 +++++++++++++++++++++++++++++++--
rust/kernel/init/__internal.rs | 50 +++++++++++
2 files changed, 190 insertions(+), 6 deletions(-)

diff --git a/rust/kernel/init.rs b/rust/kernel/init.rs
index 37e8159df24d..1343cafdd760 100644
--- a/rust/kernel/init.rs
+++ b/rust/kernel/init.rs
@@ -12,7 +12,8 @@
//!
//! To initialize a `struct` with an in-place constructor you will need two things:
//! - an in-place constructor,
-//! - a memory location that can hold your `struct`.
+//! - a memory location that can hold your `struct` (this can be the [stack], an [`Arc<T>`],
+//! [`UniqueArc<T>`], [`Box<T>`] or any other smart pointer that implements [`InPlaceInit`]).
//!
//! To get an in-place constructor there are generally three options:
//! - directly creating an in-place constructor using the [`pin_init!`] macro,
@@ -180,6 +181,7 @@
//! [pinning]: https://doc.rust-lang.org/std/pin/index.html
//! [structurally pinned fields]:
//! https://doc.rust-lang.org/std/pin/index.html#pinning-is-structural-for-field
+//! [stack]: crate::stack_pin_init
//! [`Arc<T>`]: crate::sync::Arc
//! [`impl PinInit<Foo>`]: PinInit
//! [`impl PinInit<T, E>`]: PinInit
@@ -202,6 +204,138 @@ pub mod __internal;
#[doc(hidden)]
pub mod macros;

+/// Initialize and pin a type directly on the stack.
+///
+/// # Examples
+///
+/// ```rust
+/// # #![allow(clippy::disallowed_names, clippy::new_ret_no_self)]
+/// # use kernel::{init, pin_init, stack_pin_init, init::*, sync::Mutex, new_mutex};
+/// # use macros::pin_data;
+/// # use core::pin::Pin;
+/// #[pin_data]
+/// struct Foo {
+/// #[pin]
+/// a: Mutex<usize>,
+/// b: Bar,
+/// }
+///
+/// #[pin_data]
+/// struct Bar {
+/// x: u32,
+/// }
+///
+/// stack_pin_init!(let foo = pin_init!(Foo {
+/// a <- new_mutex!(42),
+/// b: Bar {
+/// x: 64,
+/// },
+/// }));
+/// let foo: Pin<&mut Foo> = foo;
+/// pr_info!("a: {}", &*foo.a.lock());
+/// ```
+///
+/// # Syntax
+///
+/// A normal `let` binding with optional type annotation. The expression is expected to implement
+/// [`PinInit`]/[`Init`] with the error type [`Infallible`]. If you want to use a different error
+/// type, then use [`stack_try_pin_init!`].
+#[macro_export]
+macro_rules! stack_pin_init {
+ (let $var:ident $(: $t:ty)? = $val:expr) => {
+ let mut $var = ::core::pin::pin!($crate::init::__internal::StackInit$(::<$t>)?::uninit());
+ let mut $var = {
+ let val = $val;
+ match $crate::init::__internal::StackInit::init($var, val) {
+ Ok(res) => res,
+ Err(x) => {
+ let x: ::core::convert::Infallible = x;
+ match x {}
+ }
+ }
+ };
+ };
+}
+
+/// Initialize and pin a type directly on the stack.
+///
+/// # Examples
+///
+/// ```rust
+/// # #![allow(clippy::disallowed_names, clippy::new_ret_no_self)]
+/// # use kernel::{init, pin_init, stack_try_pin_init, init::*, sync::Mutex, new_mutex};
+/// # use macros::pin_data;
+/// # use core::{alloc::AllocError, pin::Pin};
+/// #[pin_data]
+/// struct Foo {
+/// #[pin]
+/// a: Mutex<usize>,
+/// b: Box<Bar>,
+/// }
+///
+/// struct Bar {
+/// x: u32,
+/// }
+///
+/// stack_try_pin_init!(let foo: Result<Pin<&mut Foo>, AllocError> = pin_init!(Foo {
+/// a <- new_mutex!(42),
+/// b: Box::try_new(Bar {
+/// x: 64,
+/// })?,
+/// }));
+/// let foo = foo.unwrap();
+/// pr_info!("a: {}", &*foo.a.lock());
+/// ```
+///
+/// ```rust
+/// # #![allow(clippy::disallowed_names, clippy::new_ret_no_self)]
+/// # use kernel::{init, pin_init, stack_try_pin_init, init::*, sync::Mutex, new_mutex};
+/// # use macros::pin_data;
+/// # use core::{alloc::AllocError, pin::Pin};
+/// #[pin_data]
+/// struct Foo {
+/// #[pin]
+/// a: Mutex<usize>,
+/// b: Box<Bar>,
+/// }
+///
+/// struct Bar {
+/// x: u32,
+/// }
+///
+/// stack_try_pin_init!(let foo: Pin<&mut Foo> =? pin_init!(Foo {
+/// a <- new_mutex!(42),
+/// b: Box::try_new(Bar {
+/// x: 64,
+/// })?,
+/// }));
+/// pr_info!("a: {}", &*foo.a.lock());
+/// # Ok::<_, AllocError>(())
+/// ```
+///
+/// # Syntax
+///
+/// A normal `let` binding with optional type annotation. The expression is expected to implement
+/// [`PinInit`]/[`Init`]. This macro assigns a result to the given variable, adding a `?` after the
+/// `=` will propagate this error.
+#[macro_export]
+macro_rules! stack_try_pin_init {
+ (let $var:ident $(: $t:ty)? = $val:expr) => {
+ let mut $var = ::core::pin::pin!($crate::init::__internal::StackInit$(::<$t>)?::uninit());
+ let mut $var = {
+ let val = $val;
+ $crate::init::__internal::StackInit::init($var, val)
+ };
+ };
+ (let $var:ident $(: $t:ty)? =? $val:expr) => {
+ let mut $var = ::core::pin::pin!($crate::init::__internal::StackInit$(::<$t>)?::uninit());
+ let mut $var = {
+ let val = $val;
+ $crate::init::__internal::StackInit::init($var, val)?
+ };
+ };
+}
+
/// Construct an in-place, pinned initializer for `struct`s.
///
/// This macro defaults the error to [`Infallible`]. If you need [`Error`], then use
@@ -913,8 +1047,8 @@ macro_rules! try_init {
/// A pin-initializer for the type `T`.
///
/// To use this initializer, you will need a suitable memory location that can hold a `T`. This can
-/// be [`Box<T>`], [`Arc<T>`], [`UniqueArc<T>`]. Use the [`InPlaceInit::pin_init`] function of a
-/// smart pointer like [`Arc<T>`] on this.
+/// be [`Box<T>`], [`Arc<T>`], [`UniqueArc<T>`] or even the stack (see [`stack_pin_init!`]). Use the
+/// [`InPlaceInit::pin_init`] function of a smart pointer like [`Arc<T>`] on this.
///
/// Also see the [module description](self).
///
@@ -949,9 +1083,9 @@ pub unsafe trait PinInit<T: ?Sized, E = Infallible>: Sized {
/// An initializer for `T`.
///
/// To use this initializer, you will need a suitable memory location that can hold a `T`. This can
-/// be [`Box<T>`], [`Arc<T>`], [`UniqueArc<T>`]. Use the [`InPlaceInit::init`] function of a smart
-/// pointer like [`Arc<T>`] on this. Because [`PinInit<T, E>`] is a super trait, you can
-/// use every function that takes it as well.
+/// be [`Box<T>`], [`Arc<T>`], [`UniqueArc<T>`] or even the stack (see [`stack_pin_init!`]). Use the
+/// [`InPlaceInit::init`] function of a smart pointer like [`Arc<T>`] on this. Because
+/// [`PinInit<T, E>`] is a super trait, you can use every function that takes it as well.
///
/// Also see the [module description](self).
///
diff --git a/rust/kernel/init/__internal.rs b/rust/kernel/init/__internal.rs
index 69be03e17c1f..f4846638cb59 100644
--- a/rust/kernel/init/__internal.rs
+++ b/rust/kernel/init/__internal.rs
@@ -112,6 +112,56 @@ unsafe impl<T: ?Sized> HasInitData for T {
}
}

+/// Stack initializer helper type. Use [`stack_pin_init`] instead of this primitive.
+///
+/// # Invariants
+///
+/// If `self.1` is true, then `self.0` is initialized.
+///
+/// [`stack_pin_init`]: kernel::stack_pin_init
+pub struct StackInit<T>(MaybeUninit<T>, bool);
+
+impl<T> Drop for StackInit<T> {
+ #[inline]
+ fn drop(&mut self) {
+ if self.1 {
+ // SAFETY: As we are being dropped, we only call this once. And since `self.1 == true`,
+ // `self.0` has to be initialized.
+ unsafe { self.0.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(MaybeUninit::uninit(), 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.1 {
+ // SAFETY: `this.1` is true and we set it to false after this.
+ unsafe { this.0.assume_init_drop() };
+ this.1 = false;
+ }
+ // SAFETY: The memory slot is valid and this type ensures that it will stay pinned.
+ unsafe { init.__pinned_init(this.0.as_mut_ptr())? };
+ this.1 = true;
+ // SAFETY: The slot is now pinned, since we will never give access to `&mut T`.
+ Ok(unsafe { Pin::new_unchecked(this.0.assume_init_mut()) })
+ }
+}
+
/// When a value of this type is dropped, it drops a `T`.
///
/// Can be forgotton to prevent the drop.
--
2.39.2