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

[Alice Ryhl <aliceryhl@google.com>]

On Wed, Apr 30, 2025 at 06:26:05PM +0200, Danilo Krummrich wrote:
> On Tue, Apr 29, 2025 at 02:44:25PM +0000, Alice Ryhl 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.
> >
> > 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 72bc743ec88bf7b91a0a1ffd9f830cfe4f983ffd..357f5a37c7b1d15b709a10c162292841eed0e376 100644
> > --- a/rust/kernel/alloc/kvec.rs
> > +++ b/rust/kernel/alloc/kvec.rs
> > @@ -608,6 +608,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);
>
> NIT: What about making this `v.retain(|&mut 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> {
> > @@ -1130,3 +1153,52 @@ fn drop(&mut self) {
> >          }
> >      }
> >  }
> > +
> > +#[macros::kunit_tests(rust_kvec_kunit)]
> > +mod tests {
> > +    use super::*;
> > +    use crate::prelude::*;
> > +
> > +    #[test]
> > +    fn test_kvec_retain() {
>
> Can we have this return a Result, like doctests can do?

I get warning when I try that:

warning: unused `core::result::Result` that must be used
    --> rust/kernel/alloc/kvec.rs:1232:1
     |
1232 | #[macros::kunit_tests(rust_kvec_kunit)]
     | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
     |
     = note: this `Result` may be an `Err` variant, which should be handled
     = note: `#[warn(unused_must_use)]` on by default
     = note: this warning originates in the attribute macro `macros::kunit_tests`
             (in Nightly builds, run with -Z macro-backtrace for more info)

> > +        /// 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);
>
> Don't we have macros around kunit_assert!() and kunit_assert_eq() outside of
> doctests (i.e. dedicated kunit tests)?
>
> I much prefer their output over the kernel panic we get with the "normal"
> asserts, unwraps, etc.
>
> Consistently sticking to the same output on failure makes it easier to catch and
> easier to setup CI environments.

The documentation for those macros says "Public but hidden since it
should only be used from generated tests." so I don't think I'm supposed
to use them.

Alice