Re: [PATCH v2 07/10] rust: alloc: update `VecExt` to take allocation flags

[Wedson Almeida Filho <wedsonaf@gmail.com>]

On Wed, 27 Mar 2024 at 00:48, Boqun Feng <boqun.feng@gmail.com> wrote:
>
> On Tue, Mar 26, 2024 at 11:35:28PM -0300, Wedson Almeida Filho wrote:
> > From: Wedson Almeida Filho <walmeida@microsoft.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@microsoft.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?

Yes, thanks for catching this!

> We probably should put some unit tests for these functions.

Yeah, I added examples in documentation with asserts that work as
tests in kunit in v3.

>
> > +        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.

It sounds reasonable to me, with one caveat: the information of how
much extra space was allocated should come from `krealloc_aligned`.

IOW, we shouldn't have callers use implementation details from
`krealloc_aligned`. (One possible way to do it is to return the actual
allocated length in addition to the pointer to the allocated block.)

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

They are the functions that make this possible :)

> Reviewed-by: Boqun Feng <boqun.feng@gmail.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
> >