rust: alloc: add Vec::insert_within_capacity

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>
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-7-06d20ad9366f@google.com
Signed-off-by: Danilo Krummrich <dakr@kernel.org>

2 files changed, 73 insertions, 1 deletions

diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index 3f2617b08753..1a0dd852a468 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -22,7 +22,7 @@ use core::{
 };
 
 mod errors;
-pub use self::errors::{PushError, RemoveError};
+pub use self::errors::{InsertError, PushError, RemoveError};
 
 /// Create a [`KVec`] containing the arguments.
 ///
@@ -358,6 +358,55 @@ where
         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
+    ///
+    /// ```
+    /// use kernel::alloc::kvec::InsertError;
+    ///
+    /// let mut v = KVec::with_capacity(5, GFP_KERNEL)?;
+    /// for i in 0..5 {
+    ///     v.insert_within_capacity(0, i)?;
+    /// }
+    ///
+    /// assert!(matches!(v.insert_within_capacity(0, 5), Err(InsertError::OutOfCapacity(_))));
+    /// assert!(matches!(v.insert_within_capacity(1000, 5), Err(InsertError::IndexOutOfBounds(_))));
+    /// assert_eq!(v, [4, 3, 2, 1, 0]);
+    /// # Ok::<(), Error>(())
+    /// ```
+    pub fn insert_within_capacity(
+        &mut self,
+        index: usize,
+        element: T,
+    ) -> Result<(), InsertError<T>> {
+        let len = self.len();
+        if index > len {
+            return Err(InsertError::IndexOutOfBounds(element));
+        }
+
+        if len >= self.capacity() {
+            return Err(InsertError::OutOfCapacity(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
diff --git a/rust/kernel/alloc/kvec/errors.rs b/rust/kernel/alloc/kvec/errors.rs
index 06fe696e8bc6..348b8d27e102 100644
--- a/rust/kernel/alloc/kvec/errors.rs
+++ b/rust/kernel/alloc/kvec/errors.rs
@@ -36,3 +36,26 @@ impl From<RemoveError> for Error {
         EINVAL
     }
 }
+
+/// Error type for [`Vec::insert_within_capacity`].
+pub enum InsertError<T> {
+    /// The value could not be inserted because the index is out of bounds.
+    IndexOutOfBounds(T),
+    /// The value could not be inserted because the vector is out of capacity.
+    OutOfCapacity(T),
+}
+
+impl<T> Debug for InsertError<T> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        match self {
+            InsertError::IndexOutOfBounds(_) => write!(f, "Index out of bounds"),
+            InsertError::OutOfCapacity(_) => write!(f, "Not enough capacity"),
+        }
+    }
+}
+
+impl<T> From<InsertError<T>> for Error {
+    fn from(_: InsertError<T>) -> Error {
+        EINVAL
+    }
+}