diff --git a/rust/kernel/alloc/vec_ext.rs b/rust/kernel/alloc/vec_ext.rs
index 311e62cc578455443160b7ea611fb9555660db58..e24d7c7675cab6eaca58a8ba62746a4b9a0cf324 100644
--- a/rust/kernel/alloc/vec_ext.rs
+++ b/rust/kernel/alloc/vec_ext.rs
@@ -2,47 +2,175 @@
//! 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>;
+ ///
+ /// # Examples
+ ///
+ /// ```
+ /// let v = Vec::<u32>::with_capacity(20, GFP_KERNEL)?;
+ ///
+ /// assert!(v.capacity() >= 20);
+ /// # Ok::<(), Error>(())
+ /// ```
+ 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>;
+ ///
+ /// # Examples
+ ///
+ /// ```
+ /// let mut v = Vec::new();
+ /// v.push(1, GFP_KERNEL)?;
+ /// assert_eq!(&v, &[1]);
+ ///
+ /// v.push(2, GFP_KERNEL)?;
+ /// assert_eq!(&v, &[1, 2]);
+ /// # Ok::<(), Error>(())
+ /// ```
+ 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>
+ ///
+ /// # Examples
+ ///
+ /// ```
+ /// let mut v = Vec::new();
+ /// v.push(1, GFP_KERNEL)?;
+ ///
+ /// v.extend_from_slice(&[20, 30, 40], GFP_KERNEL)?;
+ /// assert_eq!(&v, &[1, 20, 30, 40]);
+ ///
+ /// v.extend_from_slice(&[50, 60], GFP_KERNEL)?;
+ /// assert_eq!(&v, &[1, 20, 30, 40, 50, 60]);
+ /// # Ok::<(), Error>(())
+ /// ```
+ 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.
+ ///
+ /// # Examples
+ ///
+ /// ```
+ /// let mut v = Vec::new();
+ /// v.push(1, GFP_KERNEL)?;
+ ///
+ /// v.reserve(10, GFP_KERNEL)?;
+ /// let cap = v.capacity();
+ /// assert!(cap >= 10);
+ ///
+ /// v.reserve(10, GFP_KERNEL)?;
+ /// let new_cap = v.capacity();
+ /// assert_eq!(new_cap, cap);
+ ///
+ /// # Ok::<(), Error>(())
+ /// ```
+ 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());
}
+ // 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() + other.len()) };
+ 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) };
+ 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 4786d3ee1e92917c4f69f009770063e10dbaca02..e53466937796eda4d0331fd675c7fa272e6c08f2 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 d3d345aed218a993d34ca02c165740f2ab8b8a6e..1e910fe7c2c774af96171bf174ffc9ee6f945c52 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -18,7 +18,6 @@
#![feature(new_uninit)]
#![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 14ef4344cf6e4c630d713e5a3b48321a2138caed..f454252c62158b67210133f9ba3809a2fd077f80 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};
@@ -807,7 +807,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)?;
@@ -856,7 +856,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 aa77bad9bce48655fa1b7ddb700710172cca759a..8fad61268465a717adfaeff119cfdc107f8ccf63 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 dc05f4bbe27ec9228cfb4d79547890e5ab0ef7e6..2a9eaab62d1ca73bfa004b61a18c647e3f15e851 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 })
}