[PATCH 0/5] Additional methods for Vec

[PATCH 0/5] Additional methods for Vec

[Alice Ryhl]

This adds various Vec methods. Some of them are needed by Rust Binder,
and others are needed in other places. Each commit explains where it is
needed.

I'm not sure what we concluded on the set_len / dec_len changes, so I
don't depend on that series for now.

This series is based on top of Vec::truncate from
https://lore.kernel.org/rust-for-linux/20250316111644.154602-1-andrewjballance@gmail.com/

Signed-off-by: Alice Ryhl <aliceryhl@google.com>
---
Alice Ryhl (5):
      rust: alloc: add Vec::clear
      rust: alloc: add Vec::pop
      rust: alloc: add Vec::push_within_capacity
      rust: alloc: add Vec::drain_all
      rust: alloc: add Vec::retain

 rust/kernel/alloc/kvec.rs | 147 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 147 insertions(+)
---
base-commit: a337a03281efc4553191b432d757d4c04884bf4c
change-id: 20250320-vec-methods-adfa41e55311

Best regards,
-- 
Alice Ryhl <aliceryhl@google.com>

[PATCH 1/5] rust: alloc: add Vec::clear

[Alice Ryhl]

Our custom Vec type is missing the stdlib method `clear`, thus add it.
It will be used in the miscdevice sample.

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

diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index eb6d40a1bf8ba45126bd47f1dd4a7b7ef86112c4..95e752ed27395fce72d372976b74fb1b0e957194 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -395,6 +395,26 @@ pub fn into_raw_parts(self) -> (*mut T, usize, usize) {
         (ptr, len, capacity)
     }
 
+    /// Clears the vector, removing all values.
+    ///
+    /// Note that this method has no effect on the allocated capacity
+    /// of the vector.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// let mut v = kernel::kvec![1, 2, 3]?;
+    ///
+    /// v.clear();
+    ///
+    /// assert!(v.is_empty());
+    /// # Ok::<(), Error>(())
+    /// ```
+    #[inline]
+    pub fn clear(&mut self) {
+        self.truncate(0);
+    }
+
     /// Ensures that the capacity exceeds the length by at least `additional` elements.
     ///
     /// # Examples

-- 
2.49.0.rc1.451.g8f38331e32-goog

Re: [PATCH 1/5] rust: alloc: add Vec::clear

[Benno Lossin]

On Thu Mar 20, 2025 at 2:52 PM CET, Alice Ryhl wrote:
> Our custom Vec type is missing the stdlib method `clear`, thus add it.
> It will be used in the miscdevice sample.
>
> Signed-off-by: Alice Ryhl <aliceryhl@google.com>

Reviewed-by: Benno Lossin <benno.lossin@proton.me>

---
Cheers,
Benno

> ---
>  rust/kernel/alloc/kvec.rs | 20 ++++++++++++++++++++
>  1 file changed, 20 insertions(+)
>
> diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
> index eb6d40a1bf8ba45126bd47f1dd4a7b7ef86112c4..95e752ed27395fce72d372976b74fb1b0e957194 100644
> --- a/rust/kernel/alloc/kvec.rs
> +++ b/rust/kernel/alloc/kvec.rs
> @@ -395,6 +395,26 @@ pub fn into_raw_parts(self) -> (*mut T, usize, usize) {
>          (ptr, len, capacity)
>      }
>  
> +    /// Clears the vector, removing all values.
> +    ///
> +    /// Note that this method has no effect on the allocated capacity
> +    /// of the vector.
> +    ///
> +    /// # Examples
> +    ///
> +    /// ```
> +    /// let mut v = kernel::kvec![1, 2, 3]?;
> +    ///
> +    /// v.clear();
> +    ///
> +    /// assert!(v.is_empty());
> +    /// # Ok::<(), Error>(())
> +    /// ```
> +    #[inline]
> +    pub fn clear(&mut self) {
> +        self.truncate(0);
> +    }
> +
>      /// Ensures that the capacity exceeds the length by at least `additional` elements.
>      ///
>      /// # Examples

Re: [PATCH 1/5] rust: alloc: add Vec::clear

[Tamir Duberstein]

On Thu, Mar 20, 2025 at 10:06 AM Alice Ryhl <aliceryhl@google.com> wrote:
>
> Our custom Vec type is missing the stdlib method `clear`, thus add it.
> It will be used in the miscdevice sample.
>
> Signed-off-by: Alice Ryhl <aliceryhl@google.com>

Reviewed-by: Tamir Duberstein <tamird@gmail.com>

[PATCH 2/5] rust: alloc: add Vec::pop

[Alice Ryhl]

This introduces a basic method that our custom Vec is missing. I expect
that it will be used in many places, but at the time of writing, Rust
Binder has six calls to Vec::pop.

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

diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index 95e752ed27395fce72d372976b74fb1b0e957194..9943358c70aa63f5ad7ed9782cb8879d7a80a8fb 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -302,6 +302,37 @@ pub fn push(&mut self, v: T, flags: Flags) -> Result<(), AllocError> {
         Ok(())
     }
 
+    /// Removes the last element from a vector and returns it, or `None` if it is empty.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// let mut v = KVec::new();
+    /// v.push(1, GFP_KERNEL)?;
+    /// v.push(2, GFP_KERNEL)?;
+    /// assert_eq!(&v, &[1, 2]);
+    ///
+    /// assert_eq!(v.pop(), Some(2));
+    /// assert_eq!(v.pop(), Some(1));
+    /// assert_eq!(v.pop(), None);
+    /// # Ok::<(), Error>(())
+    /// ```
+    pub fn pop(&mut self) -> Option<T> {
+        let Some(len_sub_1) = self.len.checked_sub(1) else {
+            return None;
+        };
+
+        // INVARIANT: If the first `len` elements are valid, then the first `len-1` elements are
+        // valid.
+        self.len = len_sub_1;
+
+        // INVARIANT: This invalidates a value in this vector's allocation, but the Vec invariants
+        // do not require it to be valid because `self.len <= len_sub_1`.
+        // SAFETY: Since `len_sub_1` is less than the value `self.len` had at the beginning of
+        // `pop`, this index holds a valid value.
+        Some(unsafe { self.as_mut_ptr().add(len_sub_1).read() })
+    }
+
     /// Creates a new [`Vec`] instance with at least the given capacity.
     ///
     /// # Examples

-- 
2.49.0.rc1.451.g8f38331e32-goog

[PATCH 3/5] rust: alloc: add Vec::push_within_capacity

[Alice Ryhl]

This introduces a new method called `push_within_capacity` for appending
to a vector without attempting to allocate if the capacity is full. Rust
Binder will use this in various places to safely push to a vector while
holding a spinlock.

The existing Vec::push method is reimplemented in terms of the new
method.

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

diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index 9943358c70aa63f5ad7ed9782cb8879d7a80a8fb..df930ff0d0b85b8b03c9b7932a2b31dfb62612ed 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -284,6 +284,31 @@ pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<T>] {
     /// ```
     pub fn push(&mut self, v: T, flags: Flags) -> Result<(), AllocError> {
         self.reserve(1, flags)?;
+        let err = self.push_within_capacity(v);
+        // SAFETY: The call to `reserve` was successful, so `push_within_capacity` cannot fail.
+        unsafe { err.unwrap_unchecked() };
+        Ok(())
+    }
+
+    /// Appends an element to the back of the [`Vec`] instance.
+    ///
+    /// Fails if the vector does not have capacity for the new element.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// let mut v = KVec::with_capacity(10, GFP_KERNEL);
+    /// for i in 0..10 {
+    ///     v.push_within_capacity(i).unwrap();
+    /// }
+    ///
+    /// assert!(v.push_within_capacity(11).is_err());
+    /// # Ok::<(), Error>(())
+    /// ```
+    pub fn push_within_capacity(&mut self, v: T) -> Result<(), T> {
+        if self.len() >= self.capacity() {
+            return Err(v);
+        }
 
         // SAFETY:
         // - `self.len` is smaller than `self.capacity` and hence, the resulting pointer is

-- 
2.49.0.rc1.451.g8f38331e32-goog

Re: [PATCH 3/5] rust: alloc: add Vec::push_within_capacity

[Benno Lossin]

On Thu Mar 20, 2025 at 2:52 PM CET, Alice Ryhl wrote:
> This introduces a new method called `push_within_capacity` for appending
> to a vector without attempting to allocate if the capacity is full. Rust
> Binder will use this in various places to safely push to a vector while
> holding a spinlock.
>
> The existing Vec::push method is reimplemented in terms of the new
> method.
>
> Signed-off-by: Alice Ryhl <aliceryhl@google.com>

Reviewed-by: Benno Lossin <benno.lossin@proton.me>

---
Cheers,
Benno

> ---
>  rust/kernel/alloc/kvec.rs | 25 +++++++++++++++++++++++++
>  1 file changed, 25 insertions(+)

Re: [PATCH 3/5] rust: alloc: add Vec::push_within_capacity

[Tamir Duberstein]

On Thu, Mar 20, 2025 at 9:57 AM Alice Ryhl <aliceryhl@google.com> wrote:
>
> This introduces a new method called `push_within_capacity` for appending
> to a vector without attempting to allocate if the capacity is full. Rust
> Binder will use this in various places to safely push to a vector while
> holding a spinlock.
>
> The existing Vec::push method is reimplemented in terms of the new
> method.
>
> Signed-off-by: Alice Ryhl <aliceryhl@google.com>
> ---
>  rust/kernel/alloc/kvec.rs | 25 +++++++++++++++++++++++++
>  1 file changed, 25 insertions(+)
>
> diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
> index 9943358c70aa63f5ad7ed9782cb8879d7a80a8fb..df930ff0d0b85b8b03c9b7932a2b31dfb62612ed 100644
> --- a/rust/kernel/alloc/kvec.rs
> +++ b/rust/kernel/alloc/kvec.rs
> @@ -284,6 +284,31 @@ pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<T>] {
>      /// ```
>      pub fn push(&mut self, v: T, flags: Flags) -> Result<(), AllocError> {
>          self.reserve(1, flags)?;
> +        let err = self.push_within_capacity(v);
> +        // SAFETY: The call to `reserve` was successful, so `push_within_capacity` cannot fail.
> +        unsafe { err.unwrap_unchecked() };

I'd prefer an unsafe inner helper to this `unwrap_unchecked` call --
this safety comment is actually describing the semantics of a safe
function which can change without necessarily triggering scrutiny of
its callers.

> +        Ok(())
> +    }
> +
> +    /// Appends an element to the back of the [`Vec`] instance.
> +    ///
> +    /// Fails if the vector does not have capacity for the new element.
> +    ///
> +    /// # Examples
> +    ///
> +    /// ```
> +    /// let mut v = KVec::with_capacity(10, GFP_KERNEL);
> +    /// for i in 0..10 {
> +    ///     v.push_within_capacity(i).unwrap();
> +    /// }
> +    ///
> +    /// assert!(v.push_within_capacity(11).is_err());
> +    /// # Ok::<(), Error>(())
> +    /// ```
> +    pub fn push_within_capacity(&mut self, v: T) -> Result<(), T> {
> +        if self.len() >= self.capacity() {

len() > capacity() should be impossible by the (implied until
https://lore.kernel.org/rust-for-linux/20250318-vec-set-len-v2-1-293d55f82d18@gmail.com/)
type invariant.

> +            return Err(v);
> +        }
>
>          // SAFETY:
>          // - `self.len` is smaller than `self.capacity` and hence, the resulting pointer is
>
> --
> 2.49.0.rc1.451.g8f38331e32-goog
>
>

[PATCH 4/5] rust: alloc: add Vec::drain_all

[Alice Ryhl]

This is like the stdlib method drain, except that it's hard-coded to use
the entire vector's range. Rust Binder uses it in the range allocator to
take ownership of everything in a vector in a case where reusing the
vector is desirable.

Implementing `DrainAll` in terms of `slice::IterMut` lets us reuse some
nice optimizations in core for the case where T is a ZST.

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

diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index df930ff0d0b85b8b03c9b7932a2b31dfb62612ed..303198509885f5e24b74da5a92382b518de3e1c0 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -564,6 +564,30 @@ pub fn truncate(&mut self, len: usize) {
         //   len, therefore we have exclusive access to [`new_len`, `old_len`)
         unsafe { ptr::drop_in_place(ptr) };
     }
+
+    /// Takes ownership of all items in this vector without consuming the allocation.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// let mut v = kernel::kvec![0, 1, 2, 3]?;
+    ///
+    /// for (i, j) in v.drain_all().enumerate() {
+    ///     assert_eq!(i, j);
+    /// }
+    ///
+    /// assert!(v.capacity() >= 4);
+    /// ```
+    pub fn drain_all(&mut self) -> DrainAll<'_, T> {
+        let len = self.len();
+        // INVARIANT: The first 0 elements are valid.
+        self.len = 0;
+        // INVARIANT: The first `len` elements of the spare capacity are valid values, and as we
+        // just set the length to zero, we may transfer ownership to the `DrainAll` object.
+        DrainAll {
+            elements: self.spare_capacity_mut()[..len].iter_mut(),
+        }
+    }
 }
 
 impl<T: Clone, A: Allocator> Vec<T, A> {
@@ -1049,3 +1073,36 @@ fn into_iter(self) -> Self::IntoIter {
         }
     }
 }
+
+/// An iterator that owns all items in a vector, but does not own its allocation.
+///
+/// # Invariants
+///
+/// Every `&mut MaybeUninit<T>` returned by the iterator contains a valid `T` owned by this
+/// `DrainAll`.
+pub struct DrainAll<'vec, T> {
+    elements: slice::IterMut<'vec, MaybeUninit<T>>,
+}
+
+impl<'vec, T> Iterator for DrainAll<'vec, T> {
+    type Item = T;
+
+    fn next(&mut self) -> Option<T> {
+        let elem = self.elements.next()?;
+        // SAFETY: By the type invariants, we may take ownership of the value in this
+        // `MaybeUninit<T>`.
+        Some(unsafe { elem.assume_init_read() })
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.elements.size_hint()
+    }
+}
+
+impl<'vec, T> Drop for DrainAll<'vec, T> {
+    fn drop(&mut self) {
+        if core::mem::needs_drop::<T>() {
+            while self.next().is_some() {}
+        }
+    }
+}

-- 
2.49.0.rc1.451.g8f38331e32-goog

Re: [PATCH 4/5] rust: alloc: add Vec::drain_all

[Tamir Duberstein]

On Thu, Mar 20, 2025 at 9:56 AM Alice Ryhl <aliceryhl@google.com> wrote:
>
> This is like the stdlib method drain, except that it's hard-coded to use
> the entire vector's range. Rust Binder uses it in the range allocator to
> take ownership of everything in a vector in a case where reusing the
> vector is desirable.
>
> Implementing `DrainAll` in terms of `slice::IterMut` lets us reuse some
> nice optimizations in core for the case where T is a ZST.
>
> Signed-off-by: Alice Ryhl <aliceryhl@google.com>
> ---
>  rust/kernel/alloc/kvec.rs | 57 +++++++++++++++++++++++++++++++++++++++++++++++
>  1 file changed, 57 insertions(+)
>
> diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
> index df930ff0d0b85b8b03c9b7932a2b31dfb62612ed..303198509885f5e24b74da5a92382b518de3e1c0 100644
> --- a/rust/kernel/alloc/kvec.rs
> +++ b/rust/kernel/alloc/kvec.rs
> @@ -564,6 +564,30 @@ pub fn truncate(&mut self, len: usize) {
>          //   len, therefore we have exclusive access to [`new_len`, `old_len`)
>          unsafe { ptr::drop_in_place(ptr) };
>      }
> +
> +    /// Takes ownership of all items in this vector without consuming the allocation.
> +    ///
> +    /// # Examples
> +    ///
> +    /// ```
> +    /// let mut v = kernel::kvec![0, 1, 2, 3]?;
> +    ///
> +    /// for (i, j) in v.drain_all().enumerate() {
> +    ///     assert_eq!(i, j);
> +    /// }
> +    ///
> +    /// assert!(v.capacity() >= 4);
> +    /// ```
> +    pub fn drain_all(&mut self) -> DrainAll<'_, T> {
> +        let len = self.len();
> +        // INVARIANT: The first 0 elements are valid.
> +        self.len = 0;

Could you use `self.dec_len(self.len)` here? Then you'd have a &mut
[T] rather than `MaybeUninit`. Provided you agree `dec_len` is sound,
of course.

Link: https://lore.kernel.org/rust-for-linux/20250318-vec-set-len-v2-2-293d55f82d18@gmail.com/

Re: [PATCH 4/5] rust: alloc: add Vec::drain_all

[Alice Ryhl]

On Thu, Mar 20, 2025 at 06:12:50PM -0400, Tamir Duberstein wrote:
> On Thu, Mar 20, 2025 at 9:56 AM Alice Ryhl <aliceryhl@google.com> wrote:
> >
> > This is like the stdlib method drain, except that it's hard-coded to use
> > the entire vector's range. Rust Binder uses it in the range allocator to
> > take ownership of everything in a vector in a case where reusing the
> > vector is desirable.
> >
> > Implementing `DrainAll` in terms of `slice::IterMut` lets us reuse some
> > nice optimizations in core for the case where T is a ZST.
> >
> > Signed-off-by: Alice Ryhl <aliceryhl@google.com>
> > ---
> >  rust/kernel/alloc/kvec.rs | 57 +++++++++++++++++++++++++++++++++++++++++++++++
> >  1 file changed, 57 insertions(+)
> >
> > diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
> > index df930ff0d0b85b8b03c9b7932a2b31dfb62612ed..303198509885f5e24b74da5a92382b518de3e1c0 100644
> > --- a/rust/kernel/alloc/kvec.rs
> > +++ b/rust/kernel/alloc/kvec.rs
> > @@ -564,6 +564,30 @@ pub fn truncate(&mut self, len: usize) {
> >          //   len, therefore we have exclusive access to [`new_len`, `old_len`)
> >          unsafe { ptr::drop_in_place(ptr) };
> >      }
> > +
> > +    /// Takes ownership of all items in this vector without consuming the allocation.
> > +    ///
> > +    /// # Examples
> > +    ///
> > +    /// ```
> > +    /// let mut v = kernel::kvec![0, 1, 2, 3]?;
> > +    ///
> > +    /// for (i, j) in v.drain_all().enumerate() {
> > +    ///     assert_eq!(i, j);
> > +    /// }
> > +    ///
> > +    /// assert!(v.capacity() >= 4);
> > +    /// ```
> > +    pub fn drain_all(&mut self) -> DrainAll<'_, T> {
> > +        let len = self.len();
> > +        // INVARIANT: The first 0 elements are valid.
> > +        self.len = 0;
> 
> Could you use `self.dec_len(self.len)` here? Then you'd have a &mut
> [T] rather than `MaybeUninit`. Provided you agree `dec_len` is sound,
> of course.

I think that `&mut MaybeUninit<T>` is better in this case. Calling
assume_init_read on a `&mut MaybeUninit<T>` does not leave the
MaybeUninit in an invalid state in the same way that calling `ptr::read`
on an `&mut T` does.

Alice

[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

[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

Re: [PATCH 0/5] Additional methods for Vec

[Alice Ryhl]

On Thu, Mar 20, 2025 at 2:53 PM Alice Ryhl <aliceryhl@google.com> wrote:
>
> This adds various Vec methods. Some of them are needed by Rust Binder,
> and others are needed in other places. Each commit explains where it is
> needed.
>
> I'm not sure what we concluded on the set_len / dec_len changes, so I
> don't depend on that series for now.
>
> This series is based on top of Vec::truncate from
> https://lore.kernel.org/rust-for-linux/20250316111644.154602-1-andrewjballance@gmail.com/
>
> Signed-off-by: Alice Ryhl <aliceryhl@google.com>
> ---
> Alice Ryhl (5):
>       rust: alloc: add Vec::clear
>       rust: alloc: add Vec::pop
>       rust: alloc: add Vec::push_within_capacity
>       rust: alloc: add Vec::drain_all
>       rust: alloc: add Vec::retain
>
>  rust/kernel/alloc/kvec.rs | 147 ++++++++++++++++++++++++++++++++++++++++++++++
>  1 file changed, 147 insertions(+)

I went ahead and sent a new version now. It's a bit quicker than I
would normally do, but I wanted to get it out before I go on vacation.

Alice

Re: [PATCH 4/5] rust: alloc: add Vec::drain_all

[Benno Lossin]

On Thu Mar 20, 2025 at 2:52 PM CET, Alice Ryhl wrote:
> This is like the stdlib method drain, except that it's hard-coded to use
> the entire vector's range. Rust Binder uses it in the range allocator to
> take ownership of everything in a vector in a case where reusing the
> vector is desirable.

Is the reason for not implementing `drain` complexity?

> Implementing `DrainAll` in terms of `slice::IterMut` lets us reuse some
> nice optimizations in core for the case where T is a ZST.
>
> Signed-off-by: Alice Ryhl <aliceryhl@google.com>

The code is good, but I'd like to know the answer to the above question
before giving my RB.

> ---
>  rust/kernel/alloc/kvec.rs | 57 +++++++++++++++++++++++++++++++++++++++++++++++
>  1 file changed, 57 insertions(+)
>
> diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
> index df930ff0d0b85b8b03c9b7932a2b31dfb62612ed..303198509885f5e24b74da5a92382b518de3e1c0 100644
> --- a/rust/kernel/alloc/kvec.rs
> +++ b/rust/kernel/alloc/kvec.rs
> @@ -564,6 +564,30 @@ pub fn truncate(&mut self, len: usize) {
>          //   len, therefore we have exclusive access to [`new_len`, `old_len`)
>          unsafe { ptr::drop_in_place(ptr) };
>      }
> +
> +    /// Takes ownership of all items in this vector without consuming the allocation.
> +    ///
> +    /// # Examples
> +    ///
> +    /// ```
> +    /// let mut v = kernel::kvec![0, 1, 2, 3]?;
> +    ///
> +    /// for (i, j) in v.drain_all().enumerate() {
> +    ///     assert_eq!(i, j);
> +    /// }
> +    ///
> +    /// assert!(v.capacity() >= 4);
> +    /// ```
> +    pub fn drain_all(&mut self) -> DrainAll<'_, T> {
> +        let len = self.len();
> +        // INVARIANT: The first 0 elements are valid.
> +        self.len = 0;

Why not `set_len`?

> +        // INVARIANT: The first `len` elements of the spare capacity are valid values, and as we
> +        // just set the length to zero, we may transfer ownership to the `DrainAll` object.
> +        DrainAll {
> +            elements: self.spare_capacity_mut()[..len].iter_mut(),
> +        }
> +    }
>  }
>  
>  impl<T: Clone, A: Allocator> Vec<T, A> {
> @@ -1049,3 +1073,36 @@ fn into_iter(self) -> Self::IntoIter {
>          }
>      }
>  }
> +
> +/// An iterator that owns all items in a vector, but does not own its allocation.
> +///
> +/// # Invariants
> +///
> +/// Every `&mut MaybeUninit<T>` returned by the iterator contains a valid `T` owned by this
> +/// `DrainAll`.
> +pub struct DrainAll<'vec, T> {
> +    elements: slice::IterMut<'vec, MaybeUninit<T>>,
> +}
> +
> +impl<'vec, T> Iterator for DrainAll<'vec, T> {
> +    type Item = T;
> +
> +    fn next(&mut self) -> Option<T> {
> +        let elem = self.elements.next()?;
> +        // SAFETY: By the type invariants, we may take ownership of the value in this
> +        // `MaybeUninit<T>`.
> +        Some(unsafe { elem.assume_init_read() })
> +    }
> +
> +    fn size_hint(&self) -> (usize, Option<usize>) {
> +        self.elements.size_hint()
> +    }
> +}
> +
> +impl<'vec, T> Drop for DrainAll<'vec, T> {
> +    fn drop(&mut self) {
> +        if core::mem::needs_drop::<T>() {

This is neat!

---
Cheers,
Benno

> +            while self.next().is_some() {}
> +        }
> +    }
> +}

Re: [PATCH 4/5] rust: alloc: add Vec::drain_all

[Alice Ryhl]

On Thu, Mar 20, 2025 at 10:06:18PM +0000, Benno Lossin wrote:
> On Thu Mar 20, 2025 at 2:52 PM CET, Alice Ryhl wrote:
> > This is like the stdlib method drain, except that it's hard-coded to use
> > the entire vector's range. Rust Binder uses it in the range allocator to
> > take ownership of everything in a vector in a case where reusing the
> > vector is desirable.
> 
> Is the reason for not implementing `drain` complexity?

Yes.

> > Implementing `DrainAll` in terms of `slice::IterMut` lets us reuse some
> > nice optimizations in core for the case where T is a ZST.
> >
> > Signed-off-by: Alice Ryhl <aliceryhl@google.com>
> 
> The code is good, but I'd like to know the answer to the above question
> before giving my RB.
> 
> > ---
> >  rust/kernel/alloc/kvec.rs | 57 +++++++++++++++++++++++++++++++++++++++++++++++
> >  1 file changed, 57 insertions(+)
> >
> > diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
> > index df930ff0d0b85b8b03c9b7932a2b31dfb62612ed..303198509885f5e24b74da5a92382b518de3e1c0 100644
> > --- a/rust/kernel/alloc/kvec.rs
> > +++ b/rust/kernel/alloc/kvec.rs
> > @@ -564,6 +564,30 @@ pub fn truncate(&mut self, len: usize) {
> >          //   len, therefore we have exclusive access to [`new_len`, `old_len`)
> >          unsafe { ptr::drop_in_place(ptr) };
> >      }
> > +
> > +    /// Takes ownership of all items in this vector without consuming the allocation.
> > +    ///
> > +    /// # Examples
> > +    ///
> > +    /// ```
> > +    /// let mut v = kernel::kvec![0, 1, 2, 3]?;
> > +    ///
> > +    /// for (i, j) in v.drain_all().enumerate() {
> > +    ///     assert_eq!(i, j);
> > +    /// }
> > +    ///
> > +    /// assert!(v.capacity() >= 4);
> > +    /// ```
> > +    pub fn drain_all(&mut self) -> DrainAll<'_, T> {
> > +        let len = self.len();
> > +        // INVARIANT: The first 0 elements are valid.
> > +        self.len = 0;
> 
> Why not `set_len`?

I can use set_len.

Alice

Re: [PATCH 4/5] rust: alloc: add Vec::drain_all

[Benno Lossin]

On Fri Mar 21, 2025 at 8:54 AM CET, Alice Ryhl wrote:
> On Thu, Mar 20, 2025 at 10:06:18PM +0000, Benno Lossin wrote:
>> On Thu Mar 20, 2025 at 2:52 PM CET, Alice Ryhl wrote:
>> > This is like the stdlib method drain, except that it's hard-coded to use
>> > the entire vector's range. Rust Binder uses it in the range allocator to
>> > take ownership of everything in a vector in a case where reusing the
>> > vector is desirable.
>> 
>> Is the reason for not implementing `drain` complexity?
>
> Yes.

I thought more about it and as long as the person implementing `drain`,
removes `drain_all`, I have no complaints. (will give my RB in reply to
the patch in hopes that the in-reply-to header is set correctly)

---
Cheers,
Benno

Re: [PATCH 4/5] rust: alloc: add Vec::drain_all

[Benno Lossin]

On Thu Mar 20, 2025 at 2:52 PM CET, Alice Ryhl wrote:
> This is like the stdlib method drain, except that it's hard-coded to use
> the entire vector's range. Rust Binder uses it in the range allocator to
> take ownership of everything in a vector in a case where reusing the
> vector is desirable.
>
> Implementing `DrainAll` in terms of `slice::IterMut` lets us reuse some
> nice optimizations in core for the case where T is a ZST.
>
> Signed-off-by: Alice Ryhl <aliceryhl@google.com>
> ---
>  rust/kernel/alloc/kvec.rs | 57 +++++++++++++++++++++++++++++++++++++++++++++++
>  1 file changed, 57 insertions(+)
>
> diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
> index df930ff0d0b85b8b03c9b7932a2b31dfb62612ed..303198509885f5e24b74da5a92382b518de3e1c0 100644
> --- a/rust/kernel/alloc/kvec.rs
> +++ b/rust/kernel/alloc/kvec.rs
> @@ -564,6 +564,30 @@ pub fn truncate(&mut self, len: usize) {
>          //   len, therefore we have exclusive access to [`new_len`, `old_len`)
>          unsafe { ptr::drop_in_place(ptr) };
>      }
> +
> +    /// Takes ownership of all items in this vector without consuming the allocation.
> +    ///
> +    /// # Examples
> +    ///
> +    /// ```
> +    /// let mut v = kernel::kvec![0, 1, 2, 3]?;
> +    ///
> +    /// for (i, j) in v.drain_all().enumerate() {
> +    ///     assert_eq!(i, j);
> +    /// }
> +    ///
> +    /// assert!(v.capacity() >= 4);
> +    /// ```
> +    pub fn drain_all(&mut self) -> DrainAll<'_, T> {
> +        let len = self.len();
> +        // INVARIANT: The first 0 elements are valid.
> +        self.len = 0;

With this changed to `set_len` (or `dec_len` if only that is available):

Reviewed-by: Benno Lossin <benno.lossin@proton.me>

---
Cheers,
Benno

> +        // INVARIANT: The first `len` elements of the spare capacity are valid values, and as we
> +        // just set the length to zero, we may transfer ownership to the `DrainAll` object.
> +        DrainAll {
> +            elements: self.spare_capacity_mut()[..len].iter_mut(),
> +        }
> +    }
>  }
>  
>  impl<T: Clone, A: Allocator> Vec<T, A> {