Skip to content

Commit 9f14089

Browse files
DarksonnDanilo Krummrich
authored andcommitted
rust: alloc: add Vec::retain
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> 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-5-06d20ad9366f@google.com Signed-off-by: Danilo Krummrich <dakr@kernel.org>
1 parent 088bf14 commit 9f14089

File tree

1 file changed

+72
-0
lines changed

1 file changed

+72
-0
lines changed

rust/kernel/alloc/kvec.rs

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -610,6 +610,29 @@ where
610610
elements: elems.iter_mut(),
611611
}
612612
}
613+
614+
/// Removes all elements that don't match the provided closure.
615+
///
616+
/// # Examples
617+
///
618+
/// ```
619+
/// let mut v = kernel::kvec![1, 2, 3, 4]?;
620+
/// v.retain(|i| *i % 2 == 0);
621+
/// assert_eq!(v, [2, 4]);
622+
/// # Ok::<(), Error>(())
623+
/// ```
624+
pub fn retain(&mut self, mut f: impl FnMut(&mut T) -> bool) {
625+
let mut num_kept = 0;
626+
let mut next_to_check = 0;
627+
while let Some(to_check) = self.get_mut(next_to_check) {
628+
if f(to_check) {
629+
self.swap(num_kept, next_to_check);
630+
num_kept += 1;
631+
}
632+
next_to_check += 1;
633+
}
634+
self.truncate(num_kept);
635+
}
613636
}
614637

615638
impl<T: Clone, A: Allocator> Vec<T, A> {
@@ -1132,3 +1155,52 @@ impl<'vec, T> Drop for DrainAll<'vec, T> {
11321155
}
11331156
}
11341157
}
1158+
1159+
#[macros::kunit_tests(rust_kvec_kunit)]
1160+
mod tests {
1161+
use super::*;
1162+
use crate::prelude::*;
1163+
1164+
#[test]
1165+
fn test_kvec_retain() {
1166+
/// Verify correctness for one specific function.
1167+
#[expect(clippy::needless_range_loop)]
1168+
fn verify(c: &[bool]) {
1169+
let mut vec1: KVec<usize> = KVec::with_capacity(c.len(), GFP_KERNEL).unwrap();
1170+
let mut vec2: KVec<usize> = KVec::with_capacity(c.len(), GFP_KERNEL).unwrap();
1171+
1172+
for i in 0..c.len() {
1173+
vec1.push_within_capacity(i).unwrap();
1174+
if c[i] {
1175+
vec2.push_within_capacity(i).unwrap();
1176+
}
1177+
}
1178+
1179+
vec1.retain(|i| c[*i]);
1180+
1181+
assert_eq!(vec1, vec2);
1182+
}
1183+
1184+
/// Add one to a binary integer represented as a boolean array.
1185+
fn add(value: &mut [bool]) {
1186+
let mut carry = true;
1187+
for v in value {
1188+
let new_v = carry != *v;
1189+
carry = carry && *v;
1190+
*v = new_v;
1191+
}
1192+
}
1193+
1194+
// This boolean array represents a function from index to boolean. We check that `retain`
1195+
// behaves correctly for all possible boolean arrays of every possible length less than
1196+
// ten.
1197+
let mut func = KVec::with_capacity(10, GFP_KERNEL).unwrap();
1198+
for len in 0..10 {
1199+
for _ in 0u32..1u32 << len {
1200+
verify(&func);
1201+
add(&mut func);
1202+
}
1203+
func.push_within_capacity(false).unwrap();
1204+
}
1205+
}
1206+
}

0 commit comments

Comments
 (0)