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

Wedson Almeida Filho
2024-03-25T19:55:28Z

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