[PATCH v3 4/7] rust: alloc: add Vec::drain_all

Alice Ryhl
2025-04-22T09:52:45Z

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 | 59 +++++++++++++++++++++++++++++++++++++++++++++++
1 file changed, 59 insertions(+)

diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index a005a295262cb1e8b7c118125ffa07ae252e257c..4a29ca6e7dedc3e93a58830938f3a51619c270ed 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -583,6 +583,30 @@ pub fn truncate(&mut self, len: usize) {
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();
+ // SAFETY: The length is not greater than the length.
+ let elems = unsafe { self.dec_len(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: elems.iter_mut(),
+ }
+ }
}

impl<T: Clone, A: Allocator> Vec<T, A> {
@@ -1070,3 +1094,38 @@ fn into_iter(self) -> Self::IntoIter {
}
}
}
+
+/// An iterator that owns all items in a vector, but does not own its allocation.
+///
+/// # Invariants
+///
+/// Every `&mut T` returned by the iterator references a `T` that the iterator may take ownership
+/// of.
+pub struct DrainAll<'vec, T> {
+ elements: slice::IterMut<'vec, T>,
+}
+
+impl<'vec, T> Iterator for DrainAll<'vec, T> {
+ type Item = T;
+
+ fn next(&mut self) -> Option<T> {
+ let elem: *mut T = self.elements.next()?;
+ // SAFETY: By the type invariants, we may take ownership of this value.
+ Some(unsafe { elem.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>() {
+ let iter = core::mem::take(&mut self.elements);
+ let ptr: *mut [T] = iter.into_slice();
+ // SAFETY: By the type invariants, we own these values so we may destroy them.
+ unsafe { ptr::drop_in_place(ptr) };
+ }
+ }
+}

--
2.49.0.805.g082f7c87e0-goog