rust: alloc: add Vec::push_within_capacity

This introduces a new method called `push_within_capacity` for appending
to a vector without attempting to allocate if the capacity is full. Rust
Binder will use this in various places to safely push to a vector while
holding a spinlock.

The implementation is moved to a push_within_capacity_unchecked method.
This is preferred over having push() call push_within_capacity()
followed by an unwrap_unchecked() for simpler unsafe.

Panics in the kernel are best avoided when possible, so an error is
returned if the vector does not have sufficient capacity. An error type
is used rather than just returning Result<(),T> to make it more
convenient for callers (i.e. they can use ? or unwrap).

Signed-off-by: Alice Ryhl <aliceryhl@google.com>
Reviewed-by: Greg Kroah-Hartman <gregkh@linuxfoundation.org>
Reviewed-by: Benno Lossin <lossin@kernel.org>
Link: https://lore.kernel.org/r/20250502-vec-methods-v5-3-06d20ad9366f@google.com
[ Remove public visibility from `Vec::push_within_capacity_unchecked()`.
  - Danilo ]
Signed-off-by: Danilo Krummrich <dakr@kernel.org>
This commit is contained in:
Alice Ryhl 2025-05-02 13:19:31 +00:00 committed by Danilo Krummrich
parent f2b4dd7093
commit 9def0d0a2a
2 changed files with 65 additions and 4 deletions

View File

@ -21,6 +21,9 @@
slice::SliceIndex, slice::SliceIndex,
}; };
mod errors;
pub use self::errors::PushError;
/// Create a [`KVec`] containing the arguments. /// Create a [`KVec`] containing the arguments.
/// ///
/// New memory is allocated with `GFP_KERNEL`. /// New memory is allocated with `GFP_KERNEL`.
@ -307,17 +310,52 @@ pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<T>] {
/// ``` /// ```
pub fn push(&mut self, v: T, flags: Flags) -> Result<(), AllocError> { pub fn push(&mut self, v: T, flags: Flags) -> Result<(), AllocError> {
self.reserve(1, flags)?; self.reserve(1, flags)?;
// SAFETY: The call to `reserve` was successful, so the capacity is at least one greater
// than the length.
unsafe { self.push_within_capacity_unchecked(v) };
Ok(())
}
/// Appends an element to the back of the [`Vec`] instance without reallocating.
///
/// Fails if the vector does not have capacity for the new element.
///
/// # Examples
///
/// ```
/// let mut v = KVec::with_capacity(10, GFP_KERNEL)?;
/// for i in 0..10 {
/// v.push_within_capacity(i)?;
/// }
///
/// assert!(v.push_within_capacity(10).is_err());
/// # Ok::<(), Error>(())
/// ```
pub fn push_within_capacity(&mut self, v: T) -> Result<(), PushError<T>> {
if self.len() < self.capacity() {
// SAFETY: The length is less than the capacity.
unsafe { self.push_within_capacity_unchecked(v) };
Ok(())
} else {
Err(PushError(v))
}
}
/// Appends an element to the back of the [`Vec`] instance without reallocating.
///
/// # Safety
///
/// The length must be less than the capacity.
unsafe fn push_within_capacity_unchecked(&mut self, v: T) {
let spare = self.spare_capacity_mut(); let spare = self.spare_capacity_mut();
// SAFETY: The call to `reserve` was successful so the spare capacity is at least 1. // SAFETY: By the safety requirements, `spare` is non-empty.
unsafe { spare.get_unchecked_mut(0) }.write(v); unsafe { spare.get_unchecked_mut(0) }.write(v);
// SAFETY: We just initialised the first spare entry, so it is safe to increase the length // 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 // by 1. We also know that the new length is <= capacity because the caller guarantees that
// `reserve` above. // the length is less than the capacity at the beginning of this function.
unsafe { self.inc_len(1) }; unsafe { self.inc_len(1) };
Ok(())
} }
/// Removes the last element from a vector and returns it, or `None` if it is empty. /// Removes the last element from a vector and returns it, or `None` if it is empty.

View File

@ -0,0 +1,23 @@
// SPDX-License-Identifier: GPL-2.0
//! Errors for the [`Vec`] type.
use core::fmt::{self, Debug, Formatter};
use kernel::prelude::*;
/// Error type for [`Vec::push_within_capacity`].
pub struct PushError<T>(pub T);
impl<T> Debug for PushError<T> {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
write!(f, "Not enough capacity")
}
}
impl<T> From<PushError<T>> for Error {
fn from(_: PushError<T>) -> Error {
// Returning ENOMEM isn't appropriate because the system is not out of memory. The vector
// is just full and we are refusing to resize it.
EINVAL
}
}