public inbox for rust-for-linux@vger.kernel.org
 help / color / mirror / Atom feed
* [PATCH v3] rust: alloc: implement `extend` for `Vec`
@ 2025-04-06 13:01 Alexandre Courbot
  2025-04-07 11:01 ` Danilo Krummrich
  0 siblings, 1 reply; 15+ messages in thread
From: Alexandre Courbot @ 2025-04-06 13:01 UTC (permalink / raw)
  To: Danilo Krummrich, Miguel Ojeda, Alex Gaynor, Boqun Feng, Gary Guo,
	Björn Roy Baron, Benno Lossin, Andreas Hindborg, Alice Ryhl,
	Trevor Gross
  Cc: Joel Fernandes, John Hubbard, rust-for-linux, linux-kernel,
	Alexandre Courbot

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 v3:
- Use repeat_n() instead of repeat() for extend_with() in order to avoid
  an extra clone of the value.
- Be more precise about the cases where extend() will perform optimally
  or not in its documentation, and how the vector might be modified even
  in case of an error.
- Replace into_iter() with iter() and iter_mut() where suggested by the
  kernel test robot.
- Link to v2: https://lore.kernel.org/r/20250405-vec_extend-v2-1-e4a85af43cb3@nvidia.com

Changes in v2:
- Changed the diff algorithm to histogram for a more readable patch.
---
 rust/kernel/alloc/kvec.rs | 111 ++++++++++++++++++++++++++++++++--------------
 1 file changed, 78 insertions(+), 33 deletions(-)

diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index ae9d072741cedbb34bed0be0c20cc75472aa53be..b3c4227e232cca3b5c17930abc63810240b9c4da 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -454,30 +454,86 @@ 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 memory reallocations, but will work even with
+    /// imprecise implementations - albeit in a non-optimal way.
+    ///
+    /// This method returns an error if a memory reallocation required to accommodate the new items
+    /// failed. In this case, callers must assume that some (but not all) elements of `iter` might
+    /// have been added to the vector.
+    ///
+    /// # Note on optimal behavior and correctness
+    ///
+    /// The efficiency of this method depends on how reliable the [`Iterator::size_hint`]
+    /// implementation of the `iter` is.
+    ///
+    /// It performs optimally with at most a single memory reallocation if the lower bound of
+    /// `size_hint` is the exact number of items actually yielded.
+    ///
+    /// If `size_hint` is more vague, there may be as many memory reallocations as necessary to
+    /// cover the whole iterator from the successive lower bounds returned by `size_hint`.
+    ///
+    /// If `size_hint` signals more items than actually yielded by the iterator, some unused memory
+    /// might be reserved.
+    ///
+    /// Finally, whenever `size_hint` returns `(0, Some(0))`, the method assumes that no more items
+    /// are yielded by the iterator and returns. This may result in some items not being added if
+    /// there were still some remaining.
+    ///
+    /// In the kernel most iterators are expected to have a precise and correct `size_hint`
+    /// implementation, so this should nicely optimize out for these cases.
+    pub fn extend<I>(&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 actually added.
+            let added_items = self
+                .spare_capacity_mut()
+                .iter_mut()
+                // 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 lower 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_n(value, n), flags)
     }
 
     /// Pushes clones of the elements of slice into the [`Vec`] instance.
@@ -496,18 +552,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.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>


^ permalink raw reply related	[flat|nested] 15+ messages in thread
* Re: [PATCH v3] rust: alloc: implement `extend` for `Vec`
@ 2025-04-07 16:33 Benno Lossin
  2025-04-08 14:00 ` Alexandre Courbot
  0 siblings, 1 reply; 15+ messages in thread
From: Benno Lossin @ 2025-04-07 16:33 UTC (permalink / raw)
  To: Alexandre Courbot, Danilo Krummrich, Miguel Ojeda, Alex Gaynor,
	Boqun Feng, Gary Guo, Björn Roy Baron, Andreas Hindborg,
	Alice Ryhl, Trevor Gross
  Cc: Joel Fernandes, John Hubbard, rust-for-linux, linux-kernel

On Sun Apr 6, 2025 at 3:01 PM CEST, Alexandre Courbot wrote:
> diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
> index ae9d072741cedbb34bed0be0c20cc75472aa53be..b3c4227e232cca3b5c17930abc63810240b9c4da 100644
> --- a/rust/kernel/alloc/kvec.rs
> +++ b/rust/kernel/alloc/kvec.rs
> @@ -454,30 +454,86 @@ 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 memory reallocations, but will work even with
> +    /// imprecise implementations - albeit in a non-optimal way.
> +    ///
> +    /// This method returns an error if a memory reallocation required to accommodate the new items
> +    /// failed. In this case, callers must assume that some (but not all) elements of `iter` might
> +    /// have been added to the vector.

I would add that items that haven't been added to the vector remain in
the iterator. Do note that some iterators drops the items when the
iterator is dropped. But if not, one can still access them later:

    let vec1 = vec![...];
    let mut vec2 = KVec::new();
    let mut iter = vec1.iter();
    if let Err(_) = vec2.extend(&mut iter) {
        // can access remaining items:
        for item in iter {
            pr_info!("item: {item:?}\n");
        }
    }

> +    ///
> +    /// # Note on optimal behavior and correctness
> +    ///
> +    /// The efficiency of this method depends on how reliable the [`Iterator::size_hint`]
> +    /// implementation of the `iter` is.
> +    ///
> +    /// It performs optimally with at most a single memory reallocation if the lower bound of
> +    /// `size_hint` is the exact number of items actually yielded.
> +    ///
> +    /// If `size_hint` is more vague, there may be as many memory reallocations as necessary to
> +    /// cover the whole iterator from the successive lower bounds returned by `size_hint`.
> +    ///
> +    /// If `size_hint` signals more items than actually yielded by the iterator, some unused memory
> +    /// might be reserved.
> +    ///
> +    /// Finally, whenever `size_hint` returns `(0, Some(0))`, the method assumes that no more items
> +    /// are yielded by the iterator and returns. This may result in some items not being added if
> +    /// there were still some remaining.

Why do this? Can't we just call `next` and see if that is `None` too?
You could use `Peekable` [1] to call `peek` instead, then the logic
below shouldn't be too complex.

[1]: https://doc.rust-lang.org/core/iter/trait.Iterator.html#method.peekable

> +    ///
> +    /// In the kernel most iterators are expected to have a precise and correct `size_hint`
> +    /// implementation, so this should nicely optimize out for these cases.
> +    pub fn extend<I>(&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 actually added.
> +            let added_items = self
> +                .spare_capacity_mut()
> +                .iter_mut()
> +                // 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,

In my mind this precondition holds, since

    self.spare_capacity_mut().len() + self.len() == self.capacity()

and

    added_items <= self.spare_capacity_mut().len()

But I guess we haven't written the first one down anywhere.

> +            // - 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 lower 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_n(value, n), flags)

`repeat_n` is not stable in 1.78.0, so you need to add `iter_repeat_n`
to the allowed features.

---
Cheers,
Benno

>      }
>  
>      /// Pushes clones of the elements of slice into the [`Vec`] instance.
> @@ -496,18 +552,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.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,



^ permalink raw reply	[flat|nested] 15+ messages in thread
end of thread, other threads:[~2025-04-24 13:36 UTC | newest]

Thread overview: 15+ messages (download: mbox.gz follow: Atom feed
-- links below jump to the message on this page --
2025-04-06 13:01 [PATCH v3] rust: alloc: implement `extend` for `Vec` Alexandre Courbot
2025-04-07 11:01 ` Danilo Krummrich
2025-04-08 13:34   ` Alexandre Courbot
2025-04-21  8:15     ` Alexandre Courbot
2025-04-22 17:03       ` Danilo Krummrich
2025-04-23  1:02         ` Alexandre Courbot
2025-04-23  8:51           ` Alice Ryhl
2025-04-23  9:40             ` Alexandre Courbot
2025-04-23 16:03               ` Boqun Feng
2025-04-24 11:50                 ` Alice Ryhl
2025-04-24 13:36                   ` Boqun Feng
2025-04-23  9:47           ` Danilo Krummrich
2025-04-23 13:15             ` Alexandre Courbot
  -- strict thread matches above, loose matches on Subject: below --
2025-04-07 16:33 Benno Lossin
2025-04-08 14:00 ` Alexandre Courbot

This is a public inbox, see mirroring instructions
for how to clone and mirror all data and code used for this inbox