[PATCH v5 5/7] rust: alloc: add Vec::retain

Alice Ryhl
2025-05-02T13:19:51Z

This adds a common Vec method called `retain` that removes all elements
that don't match a certain condition. Rust Binder uses it to find all
processes that match a given pid.

The stdlib retain method takes &T rather than &mut T and has a separate
retain_mut for the &mut T case. However, this is considered an API
mistake that can't be fixed now due to backwards compatibility. There's
no reason for us to repeat that mistake.

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

diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index 7f75f3174ccb85daa3a9546d7b110da97f46063f..3298b3b0f32c70f3fe517fcb7af6b9922fea926b 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -610,6 +610,29 @@ pub fn drain_all(&mut self) -> DrainAll<'_, T> {
elements: elems.iter_mut(),
}
}
+
+ /// Removes all elements that don't match the provided closure.
+ ///
+ /// # Examples
+ ///
+ /// ```
+ /// let mut v = kernel::kvec![1, 2, 3, 4]?;
+ /// v.retain(|i| *i % 2 == 0);
+ /// assert_eq!(v, [2, 4]);
+ /// # Ok::<(), Error>(())
+ /// ```
+ pub fn retain(&mut self, mut f: impl FnMut(&mut T) -> bool) {
+ let mut num_kept = 0;
+ let mut next_to_check = 0;
+ while let Some(to_check) = self.get_mut(next_to_check) {
+ if f(to_check) {
+ self.swap(num_kept, next_to_check);
+ num_kept += 1;
+ }
+ next_to_check += 1;
+ }
+ self.truncate(num_kept);
+ }
}

impl<T: Clone, A: Allocator> Vec<T, A> {
@@ -1132,3 +1155,52 @@ fn drop(&mut self) {
}
}
}
+
+#[macros::kunit_tests(rust_kvec_kunit)]
+mod tests {
+ use super::*;
+ use crate::prelude::*;
+
+ #[test]
+ fn test_kvec_retain() {
+ /// Verify correctness for one specific function.
+ #[expect(clippy::needless_range_loop)]
+ fn verify(c: &[bool]) {
+ let mut vec1: KVec<usize> = KVec::with_capacity(c.len(), GFP_KERNEL).unwrap();
+ let mut vec2: KVec<usize> = KVec::with_capacity(c.len(), GFP_KERNEL).unwrap();
+
+ for i in 0..c.len() {
+ vec1.push_within_capacity(i).unwrap();
+ if c[i] {
+ vec2.push_within_capacity(i).unwrap();
+ }
+ }
+
+ vec1.retain(|i| c[*i]);
+
+ assert_eq!(vec1, vec2);
+ }
+
+ /// Add one to a binary integer represented as a boolean array.
+ fn add(value: &mut [bool]) {
+ let mut carry = true;
+ for v in value {
+ let new_v = carry != *v;
+ carry = carry && *v;
+ *v = new_v;
+ }
+ }
+
+ // This boolean array represents a function from index to boolean. We check that `retain`
+ // behaves correctly for all possible boolean arrays of every possible length less than
+ // ten.
+ let mut func = KVec::with_capacity(10, GFP_KERNEL).unwrap();
+ for len in 0..10 {
+ for _ in 0u32..1u32 << len {
+ verify(&func);
+ add(&mut func);
+ }
+ func.push_within_capacity(false).unwrap();
+ }
+ }
+}

--
2.49.0.967.g6a0df3ecc3-goog