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/vecext.rs | 106 ++++++++++++++++++++++++++++-------
rust/kernel/lib.rs | 1 -
rust/kernel/str.rs | 6 +-
rust/kernel/types.rs | 4 +-
samples/rust/rust_minimal.rs | 6 +-
5 files changed, 95 insertions(+), 28 deletions(-)
diff --git a/rust/kernel/alloc/vecext.rs b/rust/kernel/alloc/vecext.rs
index 59e92bab534e..1d4d51b45a49 100644
--- a/rust/kernel/alloc/vecext.rs
+++ b/rust/kernel/alloc/vecext.rs
@@ -2,51 +2,119 @@
//! Extensions to [`Vec`] for fallible allocations.
-use alloc::{collections::TryReserveError, vec::Vec};
-use core::result::Result;
+use super::Flags;
+use alloc::{alloc::AllocError, vec::Vec};
+use core::{mem::ManuallyDrop, 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 (ptr, len, cap) = destructure(self);
+ // SAFETY: ptr is valid for `cap` elements. And `cap` is greater (by at least 1) than
+ // `len` because of the call to `reserve` above. So the pointer after offsetting by `len`
+ // elements is valid for write.
+ unsafe { ptr.wrapping_add(len).write(v) };
+
+ // SAFETY: The only difference from the values returned by `destructure` is that `length`
+ // is incremented by 1, which is fine because we have just initialised the element at
+ // offset `length`.
+ unsafe { rebuild(self, ptr, len + 1, cap) };
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,
{
- let extra_cap = self.capacity() - self.len();
- if extra_cap > 0 {
- self.try_reserve(extra_cap)?;
- }
-
+ <Self as VecExt<_>>::reserve(self, other.len(), flags)?;
for item in other {
- self.try_push(item.clone())?;
+ <Self as VecExt<_>>::push(self, item.clone(), flags)?;
}
+ 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) };
+ 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(())
+ }
+ }
+}
+
+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 = 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`].
+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/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 183748328d43..34dbc85b5220 100644
--- a/rust/kernel/str.rs
+++ b/rust/kernel/str.rs
@@ -2,7 +2,7 @@
//! String representations.
-use crate::alloc::vecext::VecExt;
+use crate::alloc::{flags::*, vecext::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.1On Mon, Mar 25, 2024 at 04:54:15PM -0300, Wedson Almeida Filho wrote:
[...]
> + fn push(&mut self, v: T, flags: Flags) -> Result<(), AllocError> {
> + <Self as VecExt<_>>::reserve(self, 1, flags)?;
> + let (ptr, len, cap) = destructure(self);
> + // SAFETY: ptr is valid for `cap` elements. And `cap` is greater (by at least 1) than
> + // `len` because of the call to `reserve` above. So the pointer after offsetting by `len`
> + // elements is valid for write.
> + unsafe { ptr.wrapping_add(len).write(v) };
> +
> + // SAFETY: The only difference from the values returned by `destructure` is that `length`
> + // is incremented by 1, which is fine because we have just initialised the element at
> + // offset `length`.
> + unsafe { rebuild(self, ptr, len + 1, cap) };
probably use spare_capacity_mut() here to avoid `destructure` and
`rebuild`?
https://doc.rust-lang.org/std/vec/struct.Vec.html#method.spare_capacity_mut
// .. after reserve succeed.
// there must be room for adding one more.
self.spare_capacity_mut()[0].write(v);
// or unsafe { self.spare_capacity_mut().as_mut_ptr().cast().write(v); }
unsafe {
self.set_len(self.len() + 1);
}
Thoughts?
Regards,
Boqun
> Ok(())
> }
>
[...]
On Mon, 25 Mar 2024 at 17:44, Boqun Feng <boqun.feng@gmail.com> wrote:
>
> On Mon, Mar 25, 2024 at 04:54:15PM -0300, Wedson Almeida Filho wrote:
> [...]
> > + fn push(&mut self, v: T, flags: Flags) -> Result<(), AllocError> {
> > + <Self as VecExt<_>>::reserve(self, 1, flags)?;
> > + let (ptr, len, cap) = destructure(self);
> > + // SAFETY: ptr is valid for `cap` elements. And `cap` is greater (by at least 1) than
> > + // `len` because of the call to `reserve` above. So the pointer after offsetting by `len`
> > + // elements is valid for write.
> > + unsafe { ptr.wrapping_add(len).write(v) };
> > +
> > + // SAFETY: The only difference from the values returned by `destructure` is that `length`
> > + // is incremented by 1, which is fine because we have just initialised the element at
> > + // offset `length`.
> > + unsafe { rebuild(self, ptr, len + 1, cap) };
>
> probably use spare_capacity_mut() here to avoid `destructure` and
> `rebuild`?
Ah, yes, this sounds like a better approach. I will use this in v2.
>
> https://doc.rust-lang.org/std/vec/struct.Vec.html#method.spare_capacity_mut
>
> // .. after reserve succeed.
> // there must be room for adding one more.
> self.spare_capacity_mut()[0].write(v);
> // or unsafe { self.spare_capacity_mut().as_mut_ptr().cast().write(v); }
>
> unsafe {
> self.set_len(self.len() + 1);
> }
>
> Thoughts?
>
> Regards,
> Boqun
>
> > Ok(())
> > }
> >
> [...]
On 25.03.24 20:54, 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/vecext.rs | 106 ++++++++++++++++++++++++++++------- > rust/kernel/lib.rs | 1 - > rust/kernel/str.rs | 6 +- > rust/kernel/types.rs | 4 +- > samples/rust/rust_minimal.rs | 6 +- > 5 files changed, 95 insertions(+), 28 deletions(-) With Boqun's change: Reviewed-by: Benno Lossin <benno.lossin@proton.me> One thing that we might consider in the future would be to create our own `Extend` trait to allow extending a Vec with any iterator. -- Cheers, Benno
© 2016 - 2026 Red Hat, Inc.