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

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

[Alice Ryhl]

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!");
    }

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

diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index 303198509885f5e24b74da5a92382b518de3e1c0..00dabea8ea6c8a742a7fc95954d8de58be124493 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -588,6 +588,20 @@ pub fn drain_all(&mut self) -> DrainAll<'_, T> {
             elements: self.spare_capacity_mut()[..len].iter_mut(),
         }
     }
+
+    /// Removes all elements that don't match the provided closure.
+    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.rc1.451.g8f38331e32-goog

Re: [PATCH 5/5] rust: alloc: add Vec::retain

[Benno Lossin]

On Thu Mar 20, 2025 at 2:53 PM CET, Alice Ryhl wrote:
> diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
> index 303198509885f5e24b74da5a92382b518de3e1c0..00dabea8ea6c8a742a7fc95954d8de58be124493 100644
> --- a/rust/kernel/alloc/kvec.rs
> +++ b/rust/kernel/alloc/kvec.rs
> @@ -588,6 +588,20 @@ pub fn drain_all(&mut self) -> DrainAll<'_, T> {
>              elements: self.spare_capacity_mut()[..len].iter_mut(),
>          }
>      }
> +
> +    /// Removes all elements that don't match the provided closure.

Can you also add an example here? Otherwise the code looks good.

---
Cheers,
Benno

> +    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> {

Re: [PATCH 5/5] rust: alloc: add Vec::retain

[Tamir Duberstein]

On Thu, Mar 20, 2025 at 9:57 AM Alice Ryhl <aliceryhl@google.com> wrote:
>
> 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!");
>     }

Now that we have kunit in rust-next, should we make this into a test?

Re: [PATCH 5/5] rust: alloc: add Vec::retain

[Alice Ryhl]

On Fri, Mar 21, 2025 at 4:25 PM Tamir Duberstein <tamird@gmail.com> wrote:
> Now that we have kunit in rust-next, should we make this into a test?

Unfortunately, the test uses the Vec::retain from upstream alloc,
which we can't call from a kunit test.

Alice