[PATCH v4 3/7] rust: alloc: add Vec::push_within_capacity

Alice Ryhl posted 7 patches 9 months, 2 weeks ago
There is a newer version of this series
[PATCH v4 3/7] rust: alloc: add Vec::push_within_capacity
Posted by Alice Ryhl 9 months, 2 weeks ago
This introduces a new method called `push_within_capacity` for appending
to a vector without attempting to allocate if the capacity is full. Rust
Binder will use this in various places to safely push to a vector while
holding a spinlock.

The implementation is moved to a push_within_capacity_unchecked method.
This is preferred over having push() call push_within_capacity()
followed by an unwrap_unchecked() for simpler unsafe.

Reviewed-by: Boqun Feng <boqun.feng@gmail.com>
Signed-off-by: Alice Ryhl <aliceryhl@google.com>
---
 rust/kernel/alloc/kvec.rs | 43 +++++++++++++++++++++++++++++++++++++++----
 1 file changed, 39 insertions(+), 4 deletions(-)

diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index ebca0cfd31c67f3ce13c4825d7039e34bb54f4d4..78a602e0f00494a52df0e0aa5eedc68967a3011e 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -307,17 +307,52 @@ pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<T>] {
     /// ```
     pub fn push(&mut self, v: T, flags: Flags) -> Result<(), AllocError> {
         self.reserve(1, flags)?;
+        // SAFETY: The call to `reserve` was successful, so the capacity is at least one greater
+        // than the length.
+        unsafe { self.push_within_capacity_unchecked(v) };
+        Ok(())
+    }
+
+    /// Appends an element to the back of the [`Vec`] instance without reallocating.
+    ///
+    /// Fails if the vector does not have capacity for the new element.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// let mut v = KVec::with_capacity(10, GFP_KERNEL)?;
+    /// for i in 0..10 {
+    ///     v.push_within_capacity(i).unwrap();
+    /// }
+    ///
+    /// assert!(v.push_within_capacity(10).is_err());
+    /// # Ok::<(), Error>(())
+    /// ```
+    pub fn push_within_capacity(&mut self, v: T) -> Result<(), T> {
+        if self.len() < self.capacity() {
+            // SAFETY: The length is less than the capacity.
+            unsafe { self.push_within_capacity_unchecked(v) };
+            Ok(())
+        } else {
+            Err(v)
+        }
+    }
 
+    /// Appends an element to the back of the [`Vec`] instance without reallocating.
+    ///
+    /// # Safety
+    ///
+    /// The length must be less than the capacity.
+    pub unsafe fn push_within_capacity_unchecked(&mut self, v: T) {
         let spare = self.spare_capacity_mut();
 
-        // SAFETY: The call to `reserve` was successful so the spare capacity is at least 1.
+        // SAFETY: By the safety requirements, `spare` is non-empty.
         unsafe { spare.get_unchecked_mut(0) }.write(v);
 
         // SAFETY: We just initialised the first spare entry, so it is safe to increase the length
-        // by 1. We also know that the new length is <= capacity because of the previous call to
-        // `reserve` above.
+        // by 1. We also know that the new length is <= capacity because the caller guarantees that
+        // the length is less than the capacity at the beginning of this function.
         unsafe { self.inc_len(1) };
-        Ok(())
     }
 
     /// Removes the last element from a vector and returns it, or `None` if it is empty.

-- 
2.49.0.901.g37484f566f-goog
Re: [PATCH v4 3/7] rust: alloc: add Vec::push_within_capacity
Posted by Danilo Krummrich 9 months, 2 weeks ago
On Tue, Apr 29, 2025 at 02:44:23PM +0000, Alice Ryhl wrote:
> 
> +    /// Appends an element to the back of the [`Vec`] instance without reallocating.
> +    ///
> +    /// Fails if the vector does not have capacity for the new element.
> +    ///
> +    /// # Examples
> +    ///
> +    /// ```
> +    /// let mut v = KVec::with_capacity(10, GFP_KERNEL)?;
> +    /// for i in 0..10 {
> +    ///     v.push_within_capacity(i).unwrap();

I'd prefer to make this

	v.push_within_capacity(i).map_err(|_| ENOMEM)?;

instead.

> +    /// }
> +    ///
> +    /// assert!(v.push_within_capacity(10).is_err());
> +    /// # Ok::<(), Error>(())
> +    /// ```
> +    pub fn push_within_capacity(&mut self, v: T) -> Result<(), T> {
> +        if self.len() < self.capacity() {
> +            // SAFETY: The length is less than the capacity.
> +            unsafe { self.push_within_capacity_unchecked(v) };
> +            Ok(())
> +        } else {
> +            Err(v)
> +        }
> +    }
>  
> +    /// Appends an element to the back of the [`Vec`] instance without reallocating.
> +    ///
> +    /// # Safety
> +    ///
> +    /// The length must be less than the capacity.

NIT: Maybe be more specific and say:

"`self.len()` must be less than `self.capacity()`."
Re: [PATCH v4 3/7] rust: alloc: add Vec::push_within_capacity
Posted by Alice Ryhl 9 months, 1 week ago
On Wed, Apr 30, 2025 at 05:34:20PM +0200, Danilo Krummrich wrote:
> On Tue, Apr 29, 2025 at 02:44:23PM +0000, Alice Ryhl wrote:
> > 
> > +    /// Appends an element to the back of the [`Vec`] instance without reallocating.
> > +    ///
> > +    /// Fails if the vector does not have capacity for the new element.
> > +    ///
> > +    /// # Examples
> > +    ///
> > +    /// ```
> > +    /// let mut v = KVec::with_capacity(10, GFP_KERNEL)?;
> > +    /// for i in 0..10 {
> > +    ///     v.push_within_capacity(i).unwrap();
> 
> I'd prefer to make this
> 
> 	v.push_within_capacity(i).map_err(|_| ENOMEM)?;
> 
> instead.

Perhaps we could make a new error type for `push_within_capacity`? That
way, you can use it with question mark directly, and you also get a
proper error message if you unwrap() it.

> > +    /// }
> > +    ///
> > +    /// assert!(v.push_within_capacity(10).is_err());
> > +    /// # Ok::<(), Error>(())
> > +    /// ```
> > +    pub fn push_within_capacity(&mut self, v: T) -> Result<(), T> {
> > +        if self.len() < self.capacity() {
> > +            // SAFETY: The length is less than the capacity.
> > +            unsafe { self.push_within_capacity_unchecked(v) };
> > +            Ok(())
> > +        } else {
> > +            Err(v)
> > +        }
> > +    }
> >  
> > +    /// Appends an element to the back of the [`Vec`] instance without reallocating.
> > +    ///
> > +    /// # Safety
> > +    ///
> > +    /// The length must be less than the capacity.
> 
> NIT: Maybe be more specific and say:
> 
> "`self.len()` must be less than `self.capacity()`."

I try to avoid starting sentences with code, but I can do it if you
prefer that. But saying "the length" and "the capacity" does not seem
ambiguous to me.

Alice
Re: [PATCH v4 3/7] rust: alloc: add Vec::push_within_capacity
Posted by Danilo Krummrich 9 months, 1 week ago
On Thu, May 01, 2025 at 11:03:21AM +0000, Alice Ryhl wrote:
> On Wed, Apr 30, 2025 at 05:34:20PM +0200, Danilo Krummrich wrote:
> > On Tue, Apr 29, 2025 at 02:44:23PM +0000, Alice Ryhl wrote:
> > > 
> > > +    /// Appends an element to the back of the [`Vec`] instance without reallocating.
> > > +    ///
> > > +    /// Fails if the vector does not have capacity for the new element.
> > > +    ///
> > > +    /// # Examples
> > > +    ///
> > > +    /// ```
> > > +    /// let mut v = KVec::with_capacity(10, GFP_KERNEL)?;
> > > +    /// for i in 0..10 {
> > > +    ///     v.push_within_capacity(i).unwrap();
> > 
> > I'd prefer to make this
> > 
> > 	v.push_within_capacity(i).map_err(|_| ENOMEM)?;
> > 
> > instead.
> 
> Perhaps we could make a new error type for `push_within_capacity`? That
> way, you can use it with question mark directly, and you also get a
> proper error message if you unwrap() it.

Generally, that sounds good to me. However, I'd like to avoid unwrap() or
anything that panics from doctests, since they also serve as sample code. Hence,
I think we should showcase how to do things the correct way (as much as
possible).

> > > +    /// }
> > > +    ///
> > > +    /// assert!(v.push_within_capacity(10).is_err());
> > > +    /// # Ok::<(), Error>(())
> > > +    /// ```
> > > +    pub fn push_within_capacity(&mut self, v: T) -> Result<(), T> {
> > > +        if self.len() < self.capacity() {
> > > +            // SAFETY: The length is less than the capacity.
> > > +            unsafe { self.push_within_capacity_unchecked(v) };
> > > +            Ok(())
> > > +        } else {
> > > +            Err(v)
> > > +        }
> > > +    }
> > >  
> > > +    /// Appends an element to the back of the [`Vec`] instance without reallocating.
> > > +    ///
> > > +    /// # Safety
> > > +    ///
> > > +    /// The length must be less than the capacity.
> > 
> > NIT: Maybe be more specific and say:
> > 
> > "`self.len()` must be less than `self.capacity()`."
> 
> I try to avoid starting sentences with code, but I can do it if you
> prefer that. But saying "the length" and "the capacity" does not seem
> ambiguous to me.

I'll leave that up to you. :)