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

[Benno Lossin <benno.lossin@proton.me>]

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