[PATCH v2] rust: alloc: implement `extend` for `Vec`

Alexandre Courbot
2025-04-05T13:51:56Z

KVec currently has `extend_with` and `extend_from_slice` methods, but no
way extend a vector from a regular iterator as provided by the `Extend`
trait.

Due to the need to provide the GFP flags, `Extend` cannot be implemented
directly, so simply define a homonymous method that takes an extra
`flags` argument.

The aforementioned `extend_with` and `extend_from_slice` can then be
reimplemented as direct invocations of this new method - maybe they can
eventually be removed.

Signed-off-by: Alexandre Courbot <acourbot@nvidia.com>
---
I was a bit surprised to find no equivalent of the `Extend` trait for
KVec, and while I anticipate to be told the reason for this, I also
didn't hit any hard wall trying to come with my own implementation so
here it is.

I expect the new `extend_with` and `extend_from_slice` to be optimized
into something close to their previous implementations, but am not sure
how I can simply verify that this is the case - any hint would be
appreciated!
---
Changes in v2:
- Changed the diff algorithm to histogram for a more readable patch.
---
rust/kernel/alloc/kvec.rs | 89 +++++++++++++++++++++++++++++------------------
1 file changed, 56 insertions(+), 33 deletions(-)

diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index ae9d072741cedbb34bed0be0c20cc75472aa53be..e78cb5ee575ce01e44283f8b4905689fb1e96165 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -454,30 +454,64 @@ pub fn reserve(&mut self, additional: usize, flags: Flags) -> Result<(), AllocEr
}
}

+impl<T, A: Allocator> Vec<T, A> {
+ /// Extends the vector by the elements of `iter`.
+ ///
+ /// This uses [`Iterator::size_hint`] to optimize reallocation of memory, but will work even
+ /// with imprecise implementations - albeit with potentially more memory reallocations.
+ ///
+ /// In the kernel most iterators are expected to have a precise `size_hint` implementation, so
+ /// this should nicely optimize out in most cases.
+ pub fn extend(&mut self, iter: I, flags: Flags) -> Result<(), AllocError>
+ where
+ I: IntoIterator<Item = T>,
+ {
+ let mut iter = iter.into_iter();
+
+ loop {
+ let low_bound = match iter.size_hint() {
+ // No more items expected, we can return.
+ (0, Some(0)) => break,
+ // Possibly more items but not certain, tentatively add one.
+ (0, _) => 1,
+ // More items pending, reserve space for the lower bound.
+ (low_bound, _) => low_bound,
+ };
+
+ self.reserve(low_bound, flags)?;
+
+ // Number of items we effectively added.
+ let added_items = self
+ .spare_capacity_mut()
+ .into_iter()
+ // Take a mutable reference to the iterator so we can reuse it in the next
+ // iteration of the loop if needed.
+ .zip(&mut iter)
+ .fold(0, |count, (dst, src)| {
+ dst.write(src);
+
+ count + 1
+ });
+
+ // SAFETY:
+ // - `self.len() + added_items <= self.capacity()` due to the call to `reserve` above,
+ // - items `[self.len()..self.len() + added_items - 1]` are initialized.
+ unsafe { self.set_len(self.len() + added_items) };
+
+ // `size_hint` was incorrect and our iterator ended before its advertized low bound.
+ if added_items < low_bound {
+ break;
+ }
+ }
+
+ Ok(())
+ }
+}
+
impl<T: Clone, A: Allocator> Vec<T, A> {
/// Extend the vector by `n` clones of `value`.
pub fn extend_with(&mut self, n: usize, value: T, flags: Flags) -> Result<(), AllocError> {
- if n == 0 {
- return Ok(());
- }
-
- self.reserve(n, flags)?;
-
- let spare = self.spare_capacity_mut();
-
- for item in spare.iter_mut().take(n - 1) {
- item.write(value.clone());
- }
-
- // We can write the last element directly without cloning needlessly.
- spare[n - 1].write(value);
-
- // SAFETY:
- // - `self.len() + n < self.capacity()` due to the call to reserve above,
- // - the loop and the line above initialized the next `n` elements.
- unsafe { self.set_len(self.len() + n) };
-
- Ok(())
+ self.extend(core::iter::repeat(value).take(n), flags)
}

/// Pushes clones of the elements of slice into the [`Vec`] instance.
@@ -496,18 +530,7 @@ pub fn extend_with(&mut self, n: usize, value: T, flags: Flags) -> Result<(), Al
/// # Ok::<(), Error>(())
/// ```
pub fn extend_from_slice(&mut self, other: &[T], flags: Flags) -> Result<(), AllocError> {
- self.reserve(other.len(), flags)?;
- for (slot, item) in core::iter::zip(self.spare_capacity_mut(), other) {
- slot.write(item.clone());
- }
-
- // SAFETY:
- // - `other.len()` spare entries have just been initialized, so it is safe to increase
- // the length by the same number.
- // - `self.len() + other.len() <= self.capacity()` is guaranteed by the preceding `reserve`
- // call.
- unsafe { self.set_len(self.len() + other.len()) };
- Ok(())
+ self.extend(other.into_iter().cloned(), flags)
}

/// Create a new `Vec<T, A>` and extend it by `n` clones of `value`.

---
base-commit: a2cc6ff5ec8f91bc463fd3b0c26b61166a07eb11
change-id: 20250405-vec_extend-4321251acc21

Best regards,
--
Alexandre Courbot <acourbot@nvidia.com>