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.

To verify the correctness of this implementation, you may run the
following program in userspace:

    fn retain<T>(vec: &mut Vec<T>, f: impl Fn(&T) -> bool) {
        let mut num_kept = 0;
        let mut next_to_check = 0;
        while let Some(to_check) = vec.get_mut(next_to_check) {
            if f(to_check) {
                vec.swap(num_kept, next_to_check);
                num_kept += 1;
            }
            next_to_check += 1;
        }
        vec.truncate(num_kept);
    }

    fn verify(c: &[bool]) {
        let mut vec1: Vec<usize> = (0..c.len()).collect();
        let mut vec2: Vec<usize> = (0..c.len()).collect();

        vec1.retain(|i| c[*i]);
        retain(&mut vec2, |i| c[*i]);

        assert_eq!(vec1, vec2);
    }

    // Used to loop through all 2^n bit vectors.
    fn add(value: &mut [bool]) -> bool {
        let mut carry = true;
        for v in value {
            let new_v = carry != *v;
            carry = carry && *v;
            *v = new_v;
        }
        carry
    }

    fn main() {
        for len in 0..10 {
            let mut retain = vec![false; len];
            while !add(&mut retain) {
                verify(&retain);
            }
        }
        println!("ok!");
    }

Unfortunately this can't be a kunit test because it uses the upstream
Rust implementation of Vec::retain as a comparison, which we can't call
from a kunit test.

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

diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index 4a29ca6e7dedc3e93a58830938f3a51619c270ed..2f894eac02212d15d902fe6702d6155f3128997c 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -607,6 +607,28 @@ 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]);
+    /// ```
+    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> {

-- 
2.49.0.805.g082f7c87e0-goog

On 4/22/25 11:52 AM, Alice Ryhl wrote:
> Unfortunately this can't be a kunit test because it uses the upstream
> Rust implementation of Vec::retain as a comparison, which we can't call
> from a kunit test.

Would it work from the rusttest target? We can use our kernel Vec type from 
there, but I don't remember if we still can use std Vec from there.

On Wed, Apr 23, 2025 at 02:14:11PM +0200, Danilo Krummrich wrote:
> On 4/22/25 11:52 AM, Alice Ryhl wrote:
> > Unfortunately this can't be a kunit test because it uses the upstream
> > Rust implementation of Vec::retain as a comparison, which we can't call
> > from a kunit test.
> 
> Would it work from the rusttest target? We can use our kernel Vec type from
> there, but I don't remember if we still can use std Vec from there.

My guess is no, but I don't know for sure.

Alice

On Thu, Apr 24, 2025 at 7:46 AM Alice Ryhl <aliceryhl@google.com> wrote:
>
> On Wed, Apr 23, 2025 at 02:14:11PM +0200, Danilo Krummrich wrote:
> > On 4/22/25 11:52 AM, Alice Ryhl wrote:
> > > Unfortunately this can't be a kunit test because it uses the upstream
> > > Rust implementation of Vec::retain as a comparison, which we can't call
> > > from a kunit test.
> >
> > Would it work from the rusttest target? We can use our kernel Vec type from
> > there, but I don't remember if we still can use std Vec from there.
>
> My guess is no, but I don't know for sure.

Isn't it possible to write a test without needing std's implementation?

On Thu, Apr 24, 2025 at 09:49:18AM -0400, Tamir Duberstein wrote:
> On Thu, Apr 24, 2025 at 7:46 AM Alice Ryhl <aliceryhl@google.com> wrote:
> >
> > On Wed, Apr 23, 2025 at 02:14:11PM +0200, Danilo Krummrich wrote:
> > > On 4/22/25 11:52 AM, Alice Ryhl wrote:
> > > > Unfortunately this can't be a kunit test because it uses the upstream
> > > > Rust implementation of Vec::retain as a comparison, which we can't call
> > > > from a kunit test.
> > >
> > > Would it work from the rusttest target? We can use our kernel Vec type from
> > > there, but I don't remember if we still can use std Vec from there.
> >
> > My guess is no, but I don't know for sure.
> 
> Isn't it possible to write a test without needing std's implementation?

I could probably come up with a different test.

Alice