Message-ID: <20240704170738.3621-15-dakr@redhat.com>
Date: Thu,  4 Jul 2024 19:06:42 +0200
From: Danilo Krummrich <dakr@...hat.com>
To: ojeda@...nel.org,
	alex.gaynor@...il.com,
	wedsonaf@...il.com,
	boqun.feng@...il.com,
	gary@...yguo.net,
	bjorn3_gh@...tonmail.com,
	benno.lossin@...ton.me,
	a.hindborg@...sung.com,
	aliceryhl@...gle.com
Cc: daniel.almeida@...labora.com,
	faith.ekstrand@...labora.com,
	boris.brezillon@...labora.com,
	lina@...hilina.net,
	mcanal@...lia.com,
	zhiw@...dia.com,
	acurrid@...dia.com,
	cjia@...dia.com,
	jhubbard@...dia.com,
	airlied@...hat.com,
	ajanulgu@...hat.com,
	lyude@...hat.com,
	linux-kernel@...r.kernel.org,
	rust-for-linux@...r.kernel.org,
	Danilo Krummrich <dakr@...hat.com>
Subject: [PATCH 14/20] rust: alloc: implement `IntoIterator` for `KVec`

Implement `IntoIterator` for `KVec`, `KVec`'s `IntoIter` type, as well
as `Iterator` for `IntoIter`.

`KVec::into_iter` disassembles the `KVec` into its raw parts;
additionally, `IntoIter` keeps track of a separate pointer, which is
incremented correspondingsly as the iterator advances, while the length,
or the count of elements, is decremented.

This also means that `IntoIter` takes the ownership of the backing
buffer and is responsible to drop the remaining elements and free the
backing buffer, if it's dropped.

Signed-off-by: Danilo Krummrich <dakr@...hat.com>
---
 rust/kernel/alloc.rs      |   1 +
 rust/kernel/alloc/kvec.rs | 177 +++++++++++++++++++++++++++++++++++++-
 2 files changed, 177 insertions(+), 1 deletion(-)

diff --git a/rust/kernel/alloc.rs b/rust/kernel/alloc.rs
index 3a0195d23bb4..e6a5d65e0c48 100644
--- a/rust/kernel/alloc.rs
+++ b/rust/kernel/alloc.rs
@@ -15,6 +15,7 @@
 pub use self::allocator_test as allocator;
 
 pub use self::kbox::KBox;
+pub use self::kvec::IntoIter;
 pub use self::kvec::KVec;
 
 /// Indicates an allocation error.
diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index 7ec898657832..ece48930966e 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -11,7 +11,9 @@
     ops::DerefMut,
     ops::Index,
     ops::IndexMut,
-    ptr, slice,
+    ptr,
+    ptr::NonNull,
+    slice,
     slice::SliceIndex,
 };
 
@@ -603,3 +605,176 @@ fn eq(&self, other: &$rhs) -> bool { self[..] == other[..] }
 __impl_slice_eq! { [A: Allocator] [T], KVec<U, A> }
 __impl_slice_eq! { [A: Allocator, const N: usize] KVec<T, A>, [U; N] }
 __impl_slice_eq! { [A: Allocator, const N: usize] KVec<T, A>, &[U; N] }
+
+impl<'a, T, A> IntoIterator for &'a KVec<T, A>
+where
+    A: Allocator,
+{
+    type Item = &'a T;
+    type IntoIter = slice::Iter<'a, T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.iter()
+    }
+}
+
+impl<'a, T, A: Allocator> IntoIterator for &'a mut KVec<T, A>
+where
+    A: Allocator,
+{
+    type Item = &'a mut T;
+    type IntoIter = slice::IterMut<'a, T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.iter_mut()
+    }
+}
+
+/// An iterator that moves out of a vector.
+///
+/// This `struct` is created by the `into_iter` method on [`KVec`] (provided by the [`IntoIterator`]
+/// trait).
+///
+/// # Examples
+///
+/// ```
+/// let v = kernel::kvec![0, 1, 2]?;
+/// let iter: kernel::alloc::IntoIter<_> = v.into_iter();
+///
+/// # Ok::<(), Error>(())
+/// ```
+pub struct IntoIter<T, A: Allocator = Kmalloc> {
+    ptr: *mut T,
+    buf: NonNull<T>,
+    len: usize,
+    cap: usize,
+    alloc: A,
+}
+
+impl<T, A> IntoIter<T, A>
+where
+    A: Allocator,
+{
+    fn as_raw_mut_slice(&mut self) -> *mut [T] {
+        ptr::slice_from_raw_parts_mut(self.ptr, self.len)
+    }
+}
+
+impl<T, A> Iterator for IntoIter<T, A>
+where
+    A: Allocator,
+{
+    type Item = T;
+
+    /// # Examples
+    ///
+    /// ```
+    /// let v = kernel::kvec![1, 2, 3]?;
+    /// let mut it = v.into_iter();
+    ///
+    /// assert_eq!(it.next(), Some(1));
+    /// assert_eq!(it.next(), Some(2));
+    /// assert_eq!(it.next(), Some(3));
+    /// assert_eq!(it.next(), None);
+    ///
+    /// # Ok::<(), Error>(())
+    /// ```
+    fn next(&mut self) -> Option<T> {
+        if self.len == 0 {
+            return None;
+        }
+
+        let ptr = self.ptr;
+        if !KVec::<T>::is_zst() {
+            // SAFETY: We can't overflow; `end` is guaranteed to mark the end of the buffer.
+            unsafe { self.ptr = self.ptr.add(1) };
+        } else {
+            // For ZST `ptr` has to stay where it is to remain aligned, so we just reduce `self.len`
+            // by 1.
+        }
+        self.len -= 1;
+
+        // SAFETY: `ptr` is guaranteed to point at a valid element within the buffer.
+        Some(unsafe { ptr.read() })
+    }
+
+    /// # Examples
+    ///
+    /// ```
+    /// use kernel::alloc::KVec;
+    ///
+    /// let v: KVec<u32, _> = kernel::kvec![1, 2, 3]?;
+    /// let mut iter = v.into_iter();
+    /// let size = iter.size_hint().0;
+    ///
+    /// iter.next();
+    /// assert_eq!(iter.size_hint().0, size - 1);
+    ///
+    /// iter.next();
+    /// assert_eq!(iter.size_hint().0, size - 2);
+    ///
+    /// iter.next();
+    /// assert_eq!(iter.size_hint().0, size - 3);
+    ///
+    /// # Ok::<(), Error>(())
+    /// ```
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (self.len, Some(self.len))
+    }
+}
+
+impl<T, A> Drop for IntoIter<T, A>
+where
+    A: Allocator,
+{
+    fn drop(&mut self) {
+        // SAFETY: Drop the remaining vector's elements in place, before we free the backing
+        // memory.
+        unsafe { ptr::drop_in_place(self.as_raw_mut_slice()) };
+
+        // If `cap == 0` we never allocated any memory in the first place.
+        if self.cap != 0 {
+            // SAFETY: `self.buf` was previously allocated with `self.alloc`.
+            unsafe { self.alloc.free(self.buf.as_ptr().cast()) };
+        }
+    }
+}
+
+impl<T, A> IntoIterator for KVec<T, A>
+where
+    A: Allocator,
+{
+    type Item = T;
+    type IntoIter = IntoIter<T, A>;
+
+    /// Creates a consuming iterator, that is, one that moves each value out of
+    /// the vector (from start to end). The vector cannot be used after calling
+    /// this.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// let v = kernel::kvec![1, 2]?;
+    /// let mut v_iter = v.into_iter();
+    ///
+    /// let first_element: Option<u32> = v_iter.next();
+    ///
+    /// assert_eq!(first_element, Some(1));
+    /// assert_eq!(v_iter.next(), Some(2));
+    /// assert_eq!(v_iter.next(), None);
+    ///
+    /// # Ok::<(), Error>(())
+    /// ```
+    #[inline]
+    fn into_iter(self) -> Self::IntoIter {
+        let (ptr, len, cap, alloc) = self.into_raw_parts_alloc();
+
+        IntoIter {
+            ptr,
+            buf: unsafe { NonNull::new_unchecked(ptr) },
+            len,
+            cap,
+            alloc,
+        }
+    }
+}
-- 
2.45.2

Powered by blists - more mailing lists