Re: [PATCH v7 09/26] rust: alloc: implement kernel `Box`

[Gary Guo <gary@garyguo.net>]

On Thu, 12 Sep 2024 00:52:45 +0200
Danilo Krummrich <dakr@kernel.org> wrote:

> `Box` provides the simplest way to allocate memory for a generic type
> with one of the kernel's allocators, e.g. `Kmalloc`, `Vmalloc` or
> `KVmalloc`.
> 
> In contrast to Rust's `Box` type, the kernel `Box` type considers the
> kernel's GFP flags for all appropriate functions, always reports
> allocation failures through `Result<_, AllocError>` and remains
> independent from unstable features.
> 
> Signed-off-by: Danilo Krummrich <dakr@kernel.org>
> ---
>  rust/kernel/alloc.rs      |   6 +
>  rust/kernel/alloc/kbox.rs | 465 ++++++++++++++++++++++++++++++++++++++
>  rust/kernel/prelude.rs    |   2 +-
>  3 files changed, 472 insertions(+), 1 deletion(-)
>  create mode 100644 rust/kernel/alloc/kbox.rs
> 
> diff --git a/rust/kernel/alloc.rs b/rust/kernel/alloc.rs
> index 7a405d6f3034..b1f28334b27a 100644
> --- a/rust/kernel/alloc.rs
> +++ b/rust/kernel/alloc.rs
> @@ -5,6 +5,7 @@
>  #[cfg(not(any(test, testlib)))]
>  pub mod allocator;
>  pub mod box_ext;
> +pub mod kbox;
>  pub mod vec_ext;
>  
>  #[cfg(any(test, testlib))]
> @@ -13,6 +14,11 @@
>  #[cfg(any(test, testlib))]
>  pub use self::allocator_test as allocator;
>  
> +pub use self::kbox::Box;
> +pub use self::kbox::KBox;
> +pub use self::kbox::KVBox;
> +pub use self::kbox::VBox;
> +
>  /// Indicates an allocation error.
>  #[derive(Copy, Clone, PartialEq, Eq, Debug)]
>  pub struct AllocError;
> diff --git a/rust/kernel/alloc/kbox.rs b/rust/kernel/alloc/kbox.rs
> new file mode 100644
> index 000000000000..6188494f040d
> --- /dev/null
> +++ b/rust/kernel/alloc/kbox.rs
> @@ -0,0 +1,465 @@
> +// SPDX-License-Identifier: GPL-2.0
> +
> +//! Implementation of [`Box`].
> +
> +#[allow(unused_imports)] // Used in doc comments.
> +use super::allocator::{KVmalloc, Kmalloc, Vmalloc};
> +use super::{AllocError, Allocator, Flags};
> +use core::fmt;
> +use core::marker::PhantomData;
> +use core::mem::ManuallyDrop;
> +use core::mem::MaybeUninit;
> +use core::ops::{Deref, DerefMut};
> +use core::pin::Pin;
> +use core::ptr::NonNull;
> +use core::result::Result;
> +
> +use crate::init::{InPlaceInit, InPlaceWrite, Init, PinInit};
> +use crate::types::ForeignOwnable;
> +
> +/// The kernel's [`Box`] type -- a heap allocation for a single value of type `T`.
> +///
> +/// This is the kernel's version of the Rust stdlib's `Box`. There are several of differences,
> +/// for example no `noalias` attribute is emitted and partially moving out of a `Box` is not
> +/// supported. There are also several API differences, e.g. `Box` always requires an [`Allocator`]
> +/// implementation to be passed as generic, page [`Flags`] when allocating memory and all functions
> +/// that may allocate memory are fallible.
> +///
> +/// `Box` works with any of the kernel's allocators, e.g. [`Kmalloc`], [`Vmalloc`] or [`KVmalloc`].
> +/// There are aliases for `Box` with these allocators ([`KBox`], [`VBox`], [`KVBox`]).
> +///
> +/// When dropping a [`Box`], the value is also dropped and the heap memory is automatically freed.
> +///
> +/// # Examples
> +///
> +/// ```
> +/// let b = KBox::<u64>::new(24_u64, GFP_KERNEL)?;
> +///
> +/// assert_eq!(*b, 24_u64);
> +/// # Ok::<(), Error>(())
> +/// ```
> +///
> +/// ```
> +/// # use kernel::bindings;
> +/// const SIZE: usize = bindings::KMALLOC_MAX_SIZE as usize + 1;
> +/// struct Huge([u8; SIZE]);
> +///
> +/// assert!(KBox::<Huge>::new_uninit(GFP_KERNEL | __GFP_NOWARN).is_err());
> +/// ```
> +///
> +/// ```
> +/// # use kernel::bindings;
> +/// const SIZE: usize = bindings::KMALLOC_MAX_SIZE as usize + 1;
> +/// struct Huge([u8; SIZE]);
> +///
> +/// assert!(KVBox::<Huge>::new_uninit(GFP_KERNEL).is_ok());
> +/// ```
> +///
> +/// # Invariants
> +///
> +/// `self.0` is always properly aligned and either points to memory allocated with `A` or, for
> +/// zero-sized types, is a dangling, well aligned pointer.
> +#[repr(transparent)]
> +pub struct Box<T: ?Sized, A: Allocator>(NonNull<T>, PhantomData<A>);
> +
> +/// Type alias for [`Box`] with a [`Kmalloc`] allocator.
> +///
> +/// # Examples
> +///
> +/// ```
> +/// let b = KBox::new(24_u64, GFP_KERNEL)?;
> +///
> +/// assert_eq!(*b, 24_u64);
> +/// # Ok::<(), Error>(())
> +/// ```
> +pub type KBox<T> = Box<T, super::allocator::Kmalloc>;
> +
> +/// Type alias for [`Box`] with a [`Vmalloc`] allocator.
> +///
> +/// # Examples
> +///
> +/// ```
> +/// let b = VBox::new(24_u64, GFP_KERNEL)?;
> +///
> +/// assert_eq!(*b, 24_u64);
> +/// # Ok::<(), Error>(())
> +/// ```
> +pub type VBox<T> = Box<T, super::allocator::Vmalloc>;
> +
> +/// Type alias for [`Box`] with a [`KVmalloc`] allocator.
> +///
> +/// # Examples
> +///
> +/// ```
> +/// let b = KVBox::new(24_u64, GFP_KERNEL)?;
> +///
> +/// assert_eq!(*b, 24_u64);
> +/// # Ok::<(), Error>(())
> +/// ```
> +pub type KVBox<T> = Box<T, super::allocator::KVmalloc>;
> +
> +// SAFETY: `Box` is `Send` if `T` is `Send` because the `Box` owns a `T`.
> +unsafe impl<T, A> Send for Box<T, A>
> +where
> +    T: Send + ?Sized,
> +    A: Allocator,
> +{
> +}
> +
> +// SAFETY: `Box` is `Sync` if `T` is `Sync` because the `Box` owns a `T`.
> +unsafe impl<T, A> Sync for Box<T, A>
> +where
> +    T: Sync + ?Sized,
> +    A: Allocator,
> +{
> +}
> +
> +impl<T, A> Box<T, A>
> +where
> +    T: ?Sized,
> +    A: Allocator,
> +{
> +    /// Creates a new `Box<T, A>` from a raw pointer.
> +    ///
> +    /// # Safety
> +    ///
> +    /// For non-ZSTs, `raw` must point at an allocation allocated with `A`that is sufficiently
> +    /// aligned for and holds a valid `T`. The caller passes ownership of the allocation to the
> +    /// `Box`.
> +    ///
> +    /// For ZSTs, `raw` must be a dangling, well aligned pointer.
> +    #[inline]
> +    pub const unsafe fn from_raw(raw: *mut T) -> Self {
> +        // INVARIANT: Validity of `raw` is guaranteed by the safety preconditions of this function.
> +        // SAFETY: By the safety preconditions of this function, `raw` is not a NULL pointer.
> +        Self(unsafe { NonNull::new_unchecked(raw) }, PhantomData::<A>)

nit: the turbofish here is not necessary.

	Self(unsafe { NonNull::new_unchecked(raw) }, PhantomData)

> +    }
> +
> +    /// Consumes the `Box<T, A>` and returns a raw pointer.
> +    ///
> +    /// This will not run the destructor of `T` and for non-ZSTs the allocation will stay alive
> +    /// indefinitely. Use [`Box::from_raw`] to recover the [`Box`], drop the value and free the
> +    /// allocation, if any.
> +    ///
> +    /// # Examples
> +    ///
> +    /// ```
> +    /// let x = KBox::new(24, GFP_KERNEL)?;
> +    /// let ptr = KBox::into_raw(x);
> +    /// let x = unsafe { KBox::from_raw(ptr) };
> +    ///
> +    /// assert_eq!(*x, 24);
> +    /// # Ok::<(), Error>(())
> +    /// ```
> +    #[inline]
> +    pub fn into_raw(b: Self) -> *mut T {
> +        let b = ManuallyDrop::new(b);
> +
> +        b.0.as_ptr()
> +    }
> +
> +    /// Consumes and leaks the `Box<T, A>` and returns a mutable reference.
> +    ///
> +    /// See [Box::into_raw] for more details.
> +    #[inline]
> +    pub fn leak<'a>(b: Self) -> &'a mut T {
> +        // SAFETY: `Box::into_raw` always returns a properly aligned and dereferenceable pointer
> +        // which points to an initialized instance of `T`.
> +        unsafe { &mut *Box::into_raw(b) }
> +    }
> +}
> +
> +impl<T, A> Box<MaybeUninit<T>, A>
> +where
> +    A: Allocator,
> +{
> +    /// Converts a `Box<MaybeUninit<T>, A>` to a `Box<T, A>`.
> +    ///
> +    /// It is undefined behavior to call this function while the value inside of `b` is not yet
> +    /// fully initialized.
> +    ///
> +    /// # Safety
> +    ///
> +    /// Callers must ensure that the value inside of `b` is in an initialized state.
> +    pub unsafe fn assume_init(b: Self) -> Box<T, A> {

I think this can take `self` and shadow the underlying
`MaybeUninit::assume_init` -- the rationale being that the
`MaybeUninit::assume_init` takes `self` so it cannot actually be used
for `Box<MaybeUninit<T>>`. Rust standard library also takes the
approach (with the unstable box assume_init takes `self`).

Note the argument of `Box` not having `self` methods don't apply here,
as it doesn't deref to an arbitrary type, but rather always
`MaybeUninit`.

> +        let raw = Self::into_raw(b);
> +
> +        // SAFETY: `raw` comes from a previous call to `Box::into_raw`. By the safety requirements
> +        // of this function, the value inside the `Box` is in an initialized state. Hence, it is
> +        // safe to reconstruct the `Box` as `Box<T, A>`.
> +        unsafe { Box::from_raw(raw.cast()) }
> +    }
> +
> +    /// Writes the value and converts to `Box<T, A>`.
> +    pub fn write(mut b: Self, value: T) -> Box<T, A> {
> +        (*b).write(value);
> +        // SAFETY: We've just initialized `b`'s value.
> +        unsafe { Self::assume_init(b) }
> +    }

same here.

> +}
> +
> +impl<T, A> Box<T, A>
> +where
> +    A: Allocator,
> +{
> +    fn is_zst() -> bool {
> +        core::mem::size_of::<T>() == 0
> +    }
> +
> +    /// Creates a new `Box<T, A>` and initializes its contents with `x`.
> +    ///
> +    /// New memory is allocated with `A`. The allocation may fail, in which case an error is
> +    /// returned. For ZSTs no memory is allocated.
> +    pub fn new(x: T, flags: Flags) -> Result<Self, AllocError> {
> +        let b = Self::new_uninit(flags)?;
> +        Ok(Box::write(b, x))
> +    }