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

["Alexandre Courbot" <acourbot@nvidia.com>]

On Sun Apr 6, 2025 at 9:59 PM JST, Alexandre Courbot wrote:
> Hi Boqun, thanks for the review!
>
> On Sun Apr 6, 2025 at 4:44 AM JST, Boqun Feng wrote:
>> 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@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<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.
>
> That's a very valid point, and one of the reasons why I would like to
> see how the code is actually optimized in `extend_with` and
> `extend_with_slice`. While the method is designed to handle
> imprecise/incorrect implementations of `size_hint`, the expectation is
> that outside of very unusual uses the code should be able to be
> optimized into a single-allocation, non-loop pass.
>
> We could probably enforce that if we had access to `TrustedLen` or
> defined our own equivalent trait for the kernel. The set of iterators
> that could be passed as arguments would be more limited, but for the
> same reason as above I don't think that would be a big limitation.
>
>> 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`?
>
> I have updated the documentation with more details about the sub-optimal
> and error cases. I am not sure what use a caller would have from the
> remaining items - and after all, the currently existing methods also
> don't return the iterator upon failure. If we want to preserve the
> current behavior, we can always reduce the size of the vector to its
> pre-call value on the error path.

Err of course they don't have any iterator to return so that last point
is mostly moot, except maybe for `extend_with` which doesn't return the
passed value, which the caller hasn't necessarily kept a clone of.