Message-ID: <Z_GIELVB_HI_qeWX@boqun-archlinux>
Date: Sat, 5 Apr 2025 12:44:16 -0700
From: Boqun Feng <boqun.feng@...il.com>
To: Alexandre Courbot <acourbot@...dia.com>
Cc: Danilo Krummrich <dakr@...nel.org>, Miguel Ojeda <ojeda@...nel.org>,
	Alex Gaynor <alex.gaynor@...il.com>, Gary Guo <gary@...yguo.net>,
	Björn Roy Baron <bjorn3_gh@...tonmail.com>,
	Benno Lossin <benno.lossin@...ton.me>,
	Andreas Hindborg <a.hindborg@...nel.org>,
	Alice Ryhl <aliceryhl@...gle.com>, Trevor Gross <tmgross@...ch.edu>,
	Joel Fernandes <joelagnelf@...dia.com>,
	John Hubbard <jhubbard@...dia.com>, rust-for-linux@...r.kernel.org,
	linux-kernel@...r.kernel.org
Subject: Re: [PATCH v2] rust: alloc: implement `extend` for `Vec`

Hi Alexandre,

Thanks for the patch.

On Sat, Apr 05, 2025 at 10:51:41PM +0900, Alexandre Courbot wrote:
> 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@...dia.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<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)?;
> +

I want to point out this might cause a behavior change, previously
extend_with() and extend_with_slice() do a "all-or-nothing" extension
depending on memory allocation, i.e. if there is enough memory for all
the new items, do the extension, otherwise do nothing. Your changes here
make it that extension can fail in-between due to AllocError, that is,
only part of the `iter` is added. Of course, in practice, both
slice::Iter and iter::Take will just return the number of all the items
as the low_bound of .size_hint(), but it's not guaranteed.

I don't see a direct correct-or-wrong answer for what behavior is
desired, but if we are moving to a new behavior, we need to make sure
updating the document of the extend*() function. Plus if failing
in-between, should we return the `iter` so that users can continue do
something about the `iter`?

> +            // 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)

Would this actually call T::clone() n times instead of n - 1 times?

Regards,
Boqun

>      }
>  
>      /// 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@...dia.com>
> 

Powered by blists - more mailing lists