Message-ID: <20260130205424.261700-2-shivamklr@cock.li>
Date: Sat, 31 Jan 2026 02:24:21 +0530
From: Shivam Kalra <shivamklr@...k.li>
To: dakr@...nel.org,
	cmllamas@...gle.com,
	gregkh@...uxfoundation.org
Cc: aliceryhl@...gle.com,
	rust-for-linux@...r.kernel.org,
	linux-kernel@...r.kernel.org,
	Shivam Kalra <shivamklr@...k.li>
Subject: [PATCH v1 1/3] rust: alloc: Add shrink_to and shrink_to_fit methods to Vec

Add methods to shrink the capacity of a Vec to free excess memory.
This is useful for drivers that experience variable workloads and want
to reclaim memory after spikes.

- `shrink_to(min_capacity, flags)`: Shrinks the capacity of the vector
  with a lower bound. The capacity will remain at least as large as
  both the length and the supplied value.

- `shrink_to_fit(flags)`: Shrinks the capacity of the vector as much
  as possible.

This implementation guarantees shrinking (unless already optimal),
because the kernel allocators don't support in-place shrinking,
a new allocation is always made.

Suggested-by: Alice Ryhl <aliceryhl@...gle.com>
Signed-off-by: Shivam Kalra <shivamklr@...k.li>
---
 rust/kernel/alloc/kvec.rs | 111 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 111 insertions(+)

diff --git a/rust/kernel/alloc/kvec.rs b/rust/kernel/alloc/kvec.rs
index ac8d6f763ae81..9c02734ced88f 100644
--- a/rust/kernel/alloc/kvec.rs
+++ b/rust/kernel/alloc/kvec.rs
@@ -733,6 +733,117 @@ pub fn retain(&mut self, mut f: impl FnMut(&mut T) -> bool) {
         }
         self.truncate(num_kept);
     }
+
+    /// Shrinks the capacity of the vector with a lower bound.
+    ///
+    /// The capacity will remain at least as large as both the length
+    /// and the supplied value.
+    ///
+    /// If the current capacity is less than the lower limit, this is a no-op.
+    ///
+    /// This requires allocating a new buffer and copying the elements, then freeing
+    /// the old buffer.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// let mut v = KVec::with_capacity(100, GFP_KERNEL)?;
+    /// v.push(1, GFP_KERNEL)?;
+    /// v.push(2, GFP_KERNEL)?;
+    /// assert!(v.capacity() >= 100);
+    ///
+    /// v.shrink_to(50, GFP_KERNEL)?;
+    /// assert!(v.capacity() >= 50);
+    /// assert!(v.capacity() < 100);
+    ///
+    /// v.shrink_to(0, GFP_KERNEL)?;
+    /// assert!(v.capacity() >= 2);
+    /// # Ok::<(), Error>(())
+    /// ```
+    pub fn shrink_to(&mut self, min_capacity: usize, flags: Flags) -> Result<(), AllocError> {
+        // Calculate the target capacity: max(len, min_capacity).
+        let target_cap = core::cmp::max(self.len(), min_capacity);
+
+        // If we're already within limits, nothing to do.
+        if self.capacity() <= target_cap {
+            return Ok(());
+        }
+
+        // ZSTs have no allocation to shrink.
+        if Self::is_zst() {
+            return Ok(());
+        }
+
+        // Handle empty vector or target capacity 0: just free the allocation and reset.
+        if target_cap == 0 {
+            // Only free if we actually have an allocation.
+            if !self.layout.is_empty() {
+                // SAFETY:
+                // - `self.ptr` was previously allocated with `A`.
+                // - `self.layout` matches the `ArrayLayout` of the preceding allocation.
+                unsafe { A::free(self.ptr.cast(), self.layout.into()) };
+            }
+            self.ptr = NonNull::dangling();
+            self.layout = ArrayLayout::empty();
+            return Ok(());
+        }
+
+        // Create a new layout that exactly fits the target capacity.
+        // SAFETY: `target_cap` is guaranteed to be <= `self.capacity()`, and the original
+        // capacity was validated, so `target_cap * size_of::<T>() <= isize::MAX`.
+        let new_layout = unsafe { ArrayLayout::<T>::new_unchecked(target_cap) };
+
+        // Allocate a new, smaller buffer.
+        let new_ptr = A::alloc(new_layout.into(), flags, NumaNode::NO_NODE)?;
+
+        // SAFETY:
+        // - `self.as_ptr()` is valid for reads of `self.len` elements by the type invariant.
+        // - `new_ptr` is valid for writes of `self.len` elements (we just allocated it).
+        // - The regions do not overlap (different allocations).
+        // - Both pointers are properly aligned.
+        unsafe {
+            ptr::copy_nonoverlapping(self.as_ptr(), new_ptr.as_ptr().cast::<T>(), self.len);
+        }
+
+        // Free the old buffer.
+        // SAFETY:
+        // - `self.ptr` was previously allocated with `A`.
+        // - `self.layout` matches the `ArrayLayout` of the preceding allocation.
+        unsafe { A::free(self.ptr.cast(), self.layout.into()) };
+
+        // SAFETY: `new_ptr.as_ptr()` is non-null because `A::alloc` succeeded.
+        self.ptr = unsafe { NonNull::new_unchecked(new_ptr.as_ptr().cast::<T>()) };
+        self.layout = new_layout;
+
+        Ok(())
+    }
+
+    /// Shrinks the capacity of the vector as much as possible.
+    ///
+    /// The capacity will be reduced to match the length, freeing any excess memory.
+    /// This requires allocating a new buffer and copying the elements, then freeing
+    /// the old buffer.
+    ///
+    /// If the allocation of the new buffer fails, the vector is left unchanged and
+    /// an error is returned.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// let mut v = KVec::with_capacity(100, GFP_KERNEL)?;
+    /// v.push(1, GFP_KERNEL)?;
+    /// v.push(2, GFP_KERNEL)?;
+    /// v.push(3, GFP_KERNEL)?;
+    /// assert!(v.capacity() >= 100);
+    ///
+    /// v.shrink_to_fit(GFP_KERNEL)?;
+    /// assert_eq!(v.capacity(), 3);
+    /// assert_eq!(&v, &[1, 2, 3]);
+    /// # Ok::<(), Error>(())
+    /// ```
+    pub fn shrink_to_fit(&mut self, flags: Flags) -> Result<(), AllocError> {
+        self.shrink_to(0, flags)
+    }
 }
 
 impl<T: Clone, A: Allocator> Vec<T, A> {
-- 
2.43.0

Powered by blists - more mailing lists