[PATCH v3 6/7] rust: alloc: add Vec::remove

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

This is needed by Rust Binder in the range allocator, and by upcoming
GPU drivers during firmware initialization.

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

diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index 2f894eac02212d15d902fe6702d6155f3128997c..2f28fda793e13841b59e83f34681e71ac815aff2 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -386,6 +386,37 @@ pub fn pop(&mut self) -> Option<T> {
Some(unsafe { removed.read() })
}

+ /// Removes the element at the given index.
+ ///
+ /// # Examples
+ ///
+ /// ```
+ /// let mut v = kernel::kvec![1, 2, 3]?;
+ /// assert_eq!(v.remove(1), 2);
+ /// assert_eq!(v, [1, 3]);
+ /// # Ok::<(), Error>(())
+ /// ```
+ pub fn remove(&mut self, i: usize) -> T {
+ // INVARIANT: This breaks the invariants by invalidating the value at index `i`, but we
+ // restore the invariants below.
+ // SAFETY: Since `&self[i]` did not result in a panic, the value at index `i` is valid.
+ let value = unsafe { ptr::read(&self[i]) };
+
+ // SAFETY: Since the above access did not panic, the length is at least one.
+ unsafe { self.dec_len(1) };
+
+ // SAFETY: We checked that `i` is in-bounds.
+ let p = unsafe { self.as_mut_ptr().add(i) };
+
+ // INVARIANT: This restores the Vec invariants by moving the valid values into the region
+ // that is required to hold valid values.
+ // SAFETY: `p.add(1).add(self.len - i - 1)` is `i+1+len-i-1 == len` elements after the
+ // beginning of the vector, so this is in-bounds of the vector.
+ unsafe { ptr::copy(p.add(1), p, self.len - i - 1) };
+
+ value
+ }
+
/// Creates a new [`Vec`] instance with at least the given capacity.
///
/// # Examples

--
2.49.0.805.g082f7c87e0-goog