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

[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

[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

[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

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