Date: Tue, 26 Mar 2024 20:48:41 -0700
From: Boqun Feng <boqun.feng@...il.com>
To: Wedson Almeida Filho <wedsonaf@...il.com>
Cc: rust-for-linux@...r.kernel.org, Miguel Ojeda <ojeda@...nel.org>,
	Alex Gaynor <alex.gaynor@...il.com>, Gary Guo <gary@...yguo.net>,
	Björn Roy Baron <bjorn3_gh@...tonmail.com>,
	Benno Lossin <benno.lossin@...ton.me>,
	Andreas Hindborg <a.hindborg@...sung.com>,
	Alice Ryhl <aliceryhl@...gle.com>, linux-kernel@...r.kernel.org,
	Wedson Almeida Filho <walmeida@...rosoft.com>
Subject: Re: [PATCH v2 07/10] rust: alloc: update `VecExt` to take allocation
 flags

On Tue, Mar 26, 2024 at 11:35:28PM -0300, Wedson Almeida Filho wrote:
> From: Wedson Almeida Filho <walmeida@...rosoft.com>
> 
> We also rename the methods by removing the `try_` prefix since the names
> are available due to our usage of the `no_global_oom_handling` config
> when building the `alloc` crate.
> 
> Signed-off-by: Wedson Almeida Filho <walmeida@...rosoft.com>
> ---
>  rust/kernel/alloc/vec_ext.rs | 106 ++++++++++++++++++++++++++++++-----
>  rust/kernel/error.rs         |  11 +---
>  rust/kernel/lib.rs           |   1 -
>  rust/kernel/str.rs           |   6 +-
>  rust/kernel/types.rs         |   4 +-
>  samples/rust/rust_minimal.rs |   6 +-
>  6 files changed, 100 insertions(+), 34 deletions(-)
> 
> diff --git a/rust/kernel/alloc/vec_ext.rs b/rust/kernel/alloc/vec_ext.rs
> index 311e62cc5784..bb6cb1c4ea67 100644
> --- a/rust/kernel/alloc/vec_ext.rs
> +++ b/rust/kernel/alloc/vec_ext.rs
> @@ -2,47 +2,123 @@
>  
>  //! Extensions to [`Vec`] for fallible allocations.
>  
> -use alloc::{collections::TryReserveError, vec::Vec};
> +use super::Flags;
> +use alloc::{alloc::AllocError, vec::Vec};
>  use core::result::Result;
>  
>  /// Extensions to [`Vec`].
>  pub trait VecExt<T>: Sized {
>      /// Creates a new [`Vec`] instance with at least the given capacity.
> -    fn try_with_capacity(capacity: usize) -> Result<Self, TryReserveError>;
> +    fn with_capacity(capacity: usize, flags: Flags) -> Result<Self, AllocError>;
>  
>      /// Appends an element to the back of the [`Vec`] instance.
> -    fn try_push(&mut self, v: T) -> Result<(), TryReserveError>;
> +    fn push(&mut self, v: T, flags: Flags) -> Result<(), AllocError>;
>  
>      /// Pushes clones of the elements of slice into the [`Vec`] instance.
> -    fn try_extend_from_slice(&mut self, other: &[T]) -> Result<(), TryReserveError>
> +    fn extend_from_slice(&mut self, other: &[T], flags: Flags) -> Result<(), AllocError>
>      where
>          T: Clone;
> +
> +    /// Ensures that the capacity exceeds the length by at least `additional` elements.
> +    fn reserve(&mut self, additional: usize, flags: Flags) -> Result<(), AllocError>;
>  }
>  
>  impl<T> VecExt<T> for Vec<T> {
> -    fn try_with_capacity(capacity: usize) -> Result<Self, TryReserveError> {
> +    fn with_capacity(capacity: usize, flags: Flags) -> Result<Self, AllocError> {
>          let mut v = Vec::new();
> -        v.try_reserve(capacity)?;
> +        <Self as VecExt<_>>::reserve(&mut v, capacity, flags)?;
>          Ok(v)
>      }
>  
> -    fn try_push(&mut self, v: T) -> Result<(), TryReserveError> {
> -        if let Err(retry) = self.push_within_capacity(v) {
> -            self.try_reserve(1)?;
> -            let _ = self.push_within_capacity(retry);
> -        }
> +    fn push(&mut self, v: T, flags: Flags) -> Result<(), AllocError> {
> +        <Self as VecExt<_>>::reserve(self, 1, flags)?;
> +        let s = self.spare_capacity_mut();
> +        s[0].write(v);
> +
> +        // SAFETY: We just initialised the first spare entry, so it is safe to increase the length
> +        // by 1. We also know that the new length is <= capacity because of the previous call to
> +        // `reserve` above.
> +        unsafe { self.set_len(self.len() + 1) };
>          Ok(())
>      }
>  
> -    fn try_extend_from_slice(&mut self, other: &[T]) -> Result<(), TryReserveError>
> +    fn extend_from_slice(&mut self, other: &[T], flags: Flags) -> Result<(), AllocError>
>      where
>          T: Clone,
>      {
> -        self.try_reserve(other.len())?;
> -        for item in other {
> -            self.try_push(item.clone())?;
> +        <Self as VecExt<_>>::reserve(self, other.len(), flags)?;
> +        for (slot, item) in core::iter::zip(self.spare_capacity_mut(), other) {
> +            slot.write(item.clone());

/me likes this! ;-) Too bad `write_slice_cloned()` is not stablized:

	https://doc.rust-lang.org/core/mem/union.MaybeUninit.html#method.write_slice_cloned

>          }
>  
> +        // SAFETY: We just initialised the `other.len()` spare entries, so it is safe to increase
> +        // the length by the same amount. We also know that the new length is <= capacity because
> +        // of the previous call to `reserve` above.
> +        unsafe { self.set_len(self.len() + 1) };

should be

	unsafe { self.set_len(self.len() + other.len()) };

, right?

We probably should put some unit tests for these functions.

> +        Ok(())
> +    }
> +
> +    #[cfg(any(test, testlib))]
> +    fn reserve(&mut self, additional: usize, _flags: Flags) -> Result<(), AllocError> {
> +        Vec::reserve(self, additional);
>          Ok(())
>      }
> +
> +    #[cfg(not(any(test, testlib)))]
> +    fn reserve(&mut self, additional: usize, flags: Flags) -> Result<(), AllocError> {
> +        let len = self.len();
> +        let cap = self.capacity();
> +
> +        if cap - len >= additional {
> +            return Ok(());
> +        }
> +
> +        if core::mem::size_of::<T>() == 0 {
> +            // The capacity is already `usize::MAX` for SZTs, we can't go higher.
> +            return Err(AllocError);
> +        }
> +
> +        // We know cap is <= `isize::MAX` because `Layout::array` fails if the resulting byte size
> +        // is greater than `isize::MAX`. So the multiplication by two won't overflow.
> +        let new_cap = core::cmp::max(cap * 2, len.checked_add(additional).ok_or(AllocError)?);
> +        let layout = core::alloc::Layout::array::<T>(new_cap).map_err(|_| AllocError)?;
> +
> +        let (ptr, len, cap) = destructure(self);
> +
> +        // SAFETY: `ptr` is valid because it's either NULL or comes from a previous call to
> +        // `krealloc_aligned`. We also verified that the type is not a ZST.
> +        let new_ptr = unsafe { super::allocator::krealloc_aligned(ptr.cast(), layout, flags.0) };

One optimization we can do later: if the alignment requirement of `T` is
larger than SLAB_MINALIGN, `krealloc_aligned()` will allocate a power
of 2 bytes big enough to hold `layout`, that means we may have enough
room to fit extra `T`s, for example:

	if align of `T` is 32 and size of `T` is 96, `new_cap` is 4,
	then the size of array is 96 * 4 = 384, and `krealloc_aligned()`
	will allocate 512 bytes, and we can fit an extra `T` here
	(because 512 - 384 = 128 > 96)

But of course, I don't think we have such a user yet. I will open an
issue later this week if this sounds reasonable.

The rest looks good to me. (`destructure` and `rebuild` are obviously
very creative to me ;-))

Reviewed-by: Boqun Feng <boqun.feng@...il.com>

Regards,
Boqun

> +        if new_ptr.is_null() {
> +            // SAFETY: We are just rebuilding the existing `Vec` with no changes.
> +            unsafe { rebuild(self, ptr, len, cap) };
> +            Err(AllocError)
> +        } else {
> +            // SAFETY: `ptr` has been reallocated with the layout for `new_cap` elements. New cap
> +            // is greater than `cap`, so it continues to be >= `len`.
> +            unsafe { rebuild(self, new_ptr.cast::<T>(), len, new_cap) };
> +            Ok(())
> +        }
> +    }
> +}
> +
> +#[cfg(not(any(test, testlib)))]
> +fn destructure<T>(v: &mut Vec<T>) -> (*mut T, usize, usize) {
> +    let mut tmp = Vec::new();
> +    core::mem::swap(&mut tmp, v);
> +    let mut tmp = core::mem::ManuallyDrop::new(tmp);
> +    let len = tmp.len();
> +    let cap = tmp.capacity();
> +    (tmp.as_mut_ptr(), len, cap)
> +}
> +
> +/// Rebuilds a `Vec` from a pointer, length, and capacity.
> +///
> +/// # Safety
> +///
> +/// The same as [`Vec::from_raw_parts`].
> +#[cfg(not(any(test, testlib)))]
> +unsafe fn rebuild<T>(v: &mut Vec<T>, ptr: *mut T, len: usize, cap: usize) {
> +    // SAFETY: The safety requirements from this function satisfy those of `from_raw_parts`.
> +    let mut tmp = unsafe { Vec::from_raw_parts(ptr, len, cap) };
> +    core::mem::swap(&mut tmp, v);
>  }
> diff --git a/rust/kernel/error.rs b/rust/kernel/error.rs
> index 4786d3ee1e92..e53466937796 100644
> --- a/rust/kernel/error.rs
> +++ b/rust/kernel/error.rs
> @@ -6,10 +6,7 @@
>  
>  use crate::str::CStr;
>  
> -use alloc::{
> -    alloc::{AllocError, LayoutError},
> -    collections::TryReserveError,
> -};
> +use alloc::alloc::{AllocError, LayoutError};
>  
>  use core::convert::From;
>  use core::fmt;
> @@ -192,12 +189,6 @@ fn from(_: Utf8Error) -> Error {
>      }
>  }
>  
> -impl From<TryReserveError> for Error {
> -    fn from(_: TryReserveError) -> Error {
> -        code::ENOMEM
> -    }
> -}
> -
>  impl From<LayoutError> for Error {
>      fn from(_: LayoutError) -> Error {
>          code::ENOMEM
> diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
> index 7f2841a18d05..51f30e55bd00 100644
> --- a/rust/kernel/lib.rs
> +++ b/rust/kernel/lib.rs
> @@ -19,7 +19,6 @@
>  #![feature(offset_of)]
>  #![feature(receiver_trait)]
>  #![feature(unsize)]
> -#![feature(vec_push_within_capacity)]
>  
>  // Ensure conditional compilation based on the kernel configuration works;
>  // otherwise we may silently break things like initcall handling.
> diff --git a/rust/kernel/str.rs b/rust/kernel/str.rs
> index e1be6c5e72a9..58f2ddafb326 100644
> --- a/rust/kernel/str.rs
> +++ b/rust/kernel/str.rs
> @@ -2,7 +2,7 @@
>  
>  //! String representations.
>  
> -use crate::alloc::vec_ext::VecExt;
> +use crate::alloc::{flags::*, vec_ext::VecExt};
>  use alloc::alloc::AllocError;
>  use alloc::vec::Vec;
>  use core::fmt::{self, Write};
> @@ -730,7 +730,7 @@ pub fn try_from_fmt(args: fmt::Arguments<'_>) -> Result<Self, Error> {
>          let size = f.bytes_written();
>  
>          // Allocate a vector with the required number of bytes, and write to it.
> -        let mut buf = Vec::try_with_capacity(size)?;
> +        let mut buf = <Vec<_> as VecExt<_>>::with_capacity(size, GFP_KERNEL)?;
>          // SAFETY: The buffer stored in `buf` is at least of size `size` and is valid for writes.
>          let mut f = unsafe { Formatter::from_buffer(buf.as_mut_ptr(), size) };
>          f.write_fmt(args)?;
> @@ -771,7 +771,7 @@ impl<'a> TryFrom<&'a CStr> for CString {
>      fn try_from(cstr: &'a CStr) -> Result<CString, AllocError> {
>          let mut buf = Vec::new();
>  
> -        buf.try_extend_from_slice(cstr.as_bytes_with_nul())
> +        <Vec<_> as VecExt<_>>::extend_from_slice(&mut buf, cstr.as_bytes_with_nul(), GFP_KERNEL)
>              .map_err(|_| AllocError)?;
>  
>          // INVARIANT: The `CStr` and `CString` types have the same invariants for
> diff --git a/rust/kernel/types.rs b/rust/kernel/types.rs
> index aa77bad9bce4..8fad61268465 100644
> --- a/rust/kernel/types.rs
> +++ b/rust/kernel/types.rs
> @@ -157,11 +157,11 @@ unsafe fn from_foreign(_: *const core::ffi::c_void) -> Self {}
>  ///     let mut vec =
>  ///         ScopeGuard::new_with_data(Vec::new(), |v| pr_info!("vec had {} elements\n", v.len()));
>  ///
> -///     vec.try_push(10u8)?;
> +///     vec.push(10u8, GFP_KERNEL)?;
>  ///     if arg {
>  ///         return Ok(());
>  ///     }
> -///     vec.try_push(20u8)?;
> +///     vec.push(20u8, GFP_KERNEL)?;
>  ///     Ok(())
>  /// }
>  ///
> diff --git a/samples/rust/rust_minimal.rs b/samples/rust/rust_minimal.rs
> index dc05f4bbe27e..2a9eaab62d1c 100644
> --- a/samples/rust/rust_minimal.rs
> +++ b/samples/rust/rust_minimal.rs
> @@ -22,9 +22,9 @@ fn init(_module: &'static ThisModule) -> Result<Self> {
>          pr_info!("Am I built-in? {}\n", !cfg!(MODULE));
>  
>          let mut numbers = Vec::new();
> -        numbers.try_push(72)?;
> -        numbers.try_push(108)?;
> -        numbers.try_push(200)?;
> +        numbers.push(72, GFP_KERNEL)?;
> +        numbers.push(108, GFP_KERNEL)?;
> +        numbers.push(200, GFP_KERNEL)?;
>  
>          Ok(RustMinimal { numbers })
>      }
> -- 
> 2.34.1
> 

Powered by blists - more mailing lists