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 e9bf4c97a5a78fc9b54751b57f15a33c716c607b..7f75f3174ccb85daa3a9546d7b110da97f46063f 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -586,6 +586,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);
+ /// # Ok::<(), Error>(())
+ /// ```
+ pub fn drain_all(&mut self) -> DrainAll<'_, T> {
+ // SAFETY: This does not underflow the length.
+ let elems = unsafe { self.dec_len(self.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> {
@@ -1073,3 +1097,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.967.g6a0df3ecc3-goog