[PATCH v3 7/7] rust: alloc: add Vec::insert_within_capacity

Alice Ryhl
2025-04-22T09:52:53Z

This adds a variant of Vec::insert that does not allocate memory. This
makes it safe to use this function while holding a spinlock. Rust Binder
uses it for the range allocator fast path.

Signed-off-by: Alice Ryhl <aliceryhl@google.com>
---
rust/kernel/alloc/kvec.rs | 39 +++++++++++++++++++++++++++++++++++++++
1 file changed, 39 insertions(+)

diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index 2f28fda793e13841b59e83f34681e71ac815aff2..d96de01a7ab3c808f5aded067ad6f1e3ba8029ce 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -355,6 +355,45 @@ pub unsafe fn push_within_capacity_unchecked(&mut self, v: T) {
unsafe { self.inc_len(1) };
}

+ /// Inserts an element at the given index in the [`Vec`] instance.
+ ///
+ /// Fails if the vector does not have capacity for the new element. Panics if the index is out
+ /// of bounds.
+ ///
+ /// # Examples
+ ///
+ /// ```
+ /// let mut v = KVec::with_capacity(10, GFP_KERNEL);
+ /// for i in 0..10 {
+ /// v.push_within_capacity(i).unwrap();
+ /// }
+ ///
+ /// assert!(v.push_within_capacity(10).is_err());
+ /// # Ok::<(), Error>(())
+ /// ```
+ pub fn insert_within_capacity(&mut self, index: usize, element: T) -> Result<(), T> {
+ let len = self.len();
+ assert!(index <= len);
+
+ if len >= self.capacity() {
+ return Err(element);
+ }
+
+ // SAFETY: This is in bounds since `index <= len < capacity`.
+ let p = unsafe { self.as_mut_ptr().add(index) };
+ // INVARIANT: This breaks the Vec invariants by making `index` contain an invalid element,
+ // but we restore the invariants below.
+ // SAFETY: Both the src and dst ranges end no later than one element after the length.
+ // Since the length is less than the capacity, both ranges are in bounds of the allocation.
+ unsafe { ptr::copy(p, p.add(1), len - index) };
+ // INVARIANT: This restores the Vec invariants.
+ // SAFETY: The pointer is in-bounds of the allocation.
+ unsafe { ptr::write(p, element) };
+ // SAFETY: Index `len` contains a valid element due to the above copy and write.
+ unsafe { self.inc_len(1) };
+ Ok(())
+ }
+
/// Removes the last element from a vector and returns it, or `None` if it is empty.
///
/// # Examples

--
2.49.0.805.g082f7c87e0-goog