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

From: Wedson Almeida Filho
Date: Wed Mar 27 2024 - 14:20:18 EST

Next message: Jason Gunthorpe: "Re: [PATCH] selftests: iommu: add config needed for iommufd_fail_nth"
Previous message: Krzysztof Kozlowski: "Re: [PATCH v6 02/16] regulator: dt-bindings: describe the PMU module of the WCN7850 package"
In reply to: Miguel Ojeda: "Re: [PATCH v2 07/10] rust: alloc: update `VecExt` to take allocation flags"
Next in thread: Wedson Almeida Filho: "[PATCH v2 08/10] rust: sync: update `Arc` and `UniqueArc` to take allocation flags"
Messages sorted by: [ date ] [ thread ] [ subject ] [ author ]

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

Next message: Jason Gunthorpe: "Re: [PATCH] selftests: iommu: add config needed for iommufd_fail_nth"
Previous message: Krzysztof Kozlowski: "Re: [PATCH v6 02/16] regulator: dt-bindings: describe the PMU module of the WCN7850 package"
In reply to: Miguel Ojeda: "Re: [PATCH v2 07/10] rust: alloc: update `VecExt` to take allocation flags"
Next in thread: Wedson Almeida Filho: "[PATCH v2 08/10] rust: sync: update `Arc` and `UniqueArc` to take allocation flags"
Messages sorted by: [ date ] [ thread ] [ subject ] [ author ]