Add methods to find the next element in an XArray starting from a
given index. The methods return a tuple containing the index where the
element was found and a reference to the element.

The implementation uses the XArray state API via `xas_find` to avoid taking
the xarray lock that is already held by `Guard`.

Signed-off-by: Andreas Hindborg <a.hindborg@kernel.org>
---
 rust/kernel/xarray.rs | 65 +++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 65 insertions(+)

diff --git a/rust/kernel/xarray.rs b/rust/kernel/xarray.rs
index ca97134ba2bd0..9d4589979fd1d 100644
--- a/rust/kernel/xarray.rs
+++ b/rust/kernel/xarray.rs
@@ -255,6 +255,71 @@ pub fn get_mut(&mut self, index: usize) -> Option<T::BorrowedMut<'_>> {
         Some(unsafe { T::borrow_mut(ptr.as_ptr()) })
     }
 
+    fn load_next(&self, index: usize) -> Option<(usize, NonNull<c_void>)> {
+        let mut state = XArrayState::new(self, index);
+        // SAFETY: `state.state` is always valid by the type invariant of
+        // `XArrayState` and the caller holds the lock.
+        let ptr = unsafe { bindings::xas_find(&raw mut state.state, usize::MAX) };
+        NonNull::new(ptr).map(|ptr| (state.state.xa_index, ptr))
+    }
+
+    /// Finds the next element starting from the given index.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use kernel::{prelude::*, xarray::{AllocKind, XArray}};
+    /// let mut xa = KBox::pin_init(XArray::<KBox<u32>>::new(AllocKind::Alloc), GFP_KERNEL)?;
+    /// let mut guard = xa.lock();
+    ///
+    /// guard.store(10, KBox::new(10u32, GFP_KERNEL)?, GFP_KERNEL)?;
+    /// guard.store(20, KBox::new(20u32, GFP_KERNEL)?, GFP_KERNEL)?;
+    ///
+    /// if let Some((found_index, value)) = guard.find_next(11) {
+    ///     assert_eq!(found_index, 20);
+    ///     assert_eq!(*value, 20);
+    /// }
+    ///
+    /// if let Some((found_index, value)) = guard.find_next(5) {
+    ///     assert_eq!(found_index, 10);
+    ///     assert_eq!(*value, 10);
+    /// }
+    ///
+    /// # Ok::<(), kernel::error::Error>(())
+    /// ```
+    pub fn find_next(&self, index: usize) -> Option<(usize, T::Borrowed<'_>)> {
+        self.load_next(index)
+            // SAFETY: `ptr` came from `T::into_foreign`.
+            .map(|(index, ptr)| (index, unsafe { T::borrow(ptr.as_ptr()) }))
+    }
+
+    /// Finds the next element starting from the given index, returning a mutable reference.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use kernel::{prelude::*, xarray::{AllocKind, XArray}};
+    /// let mut xa = KBox::pin_init(XArray::<KBox<u32>>::new(AllocKind::Alloc), GFP_KERNEL)?;
+    /// let mut guard = xa.lock();
+    ///
+    /// guard.store(10, KBox::new(10u32, GFP_KERNEL)?, GFP_KERNEL)?;
+    /// guard.store(20, KBox::new(20u32, GFP_KERNEL)?, GFP_KERNEL)?;
+    ///
+    /// if let Some((found_index, mut_value)) = guard.find_next_mut(5) {
+    ///     assert_eq!(found_index, 10);
+    ///     *mut_value = 0x99;
+    /// }
+    ///
+    /// assert_eq!(guard.get(10).copied(), Some(0x99));
+    ///
+    /// # Ok::<(), kernel::error::Error>(())
+    /// ```
+    pub fn find_next_mut(&mut self, index: usize) -> Option<(usize, T::BorrowedMut<'_>)> {
+        self.load_next(index)
+            // SAFETY: `ptr` came from `T::into_foreign`.
+            .map(move |(index, ptr)| (index, unsafe { T::borrow_mut(ptr.as_ptr()) }))
+    }
+
     /// Removes and returns the element at the given index.
     pub fn remove(&mut self, index: usize) -> Option<T> {
         // SAFETY:

-- 
2.51.2

On Wed, Dec 3, 2025 at 5:27 PM Andreas Hindborg <a.hindborg@kernel.org> wrote:
>
> Add methods to find the next element in an XArray starting from a
> given index. The methods return a tuple containing the index where the
> element was found and a reference to the element.
>
> The implementation uses the XArray state API via `xas_find` to avoid taking
> the xarray lock that is already held by `Guard`.

Similarly to the commit message introducing the use of `xas_load`,
this is not correct because `xa_find` takes and release the RCU lock
only, not the XArray lock.

>
> Signed-off-by: Andreas Hindborg <a.hindborg@kernel.org>
> ---
>  rust/kernel/xarray.rs | 65 +++++++++++++++++++++++++++++++++++++++++++++++++++
>  1 file changed, 65 insertions(+)
>
> diff --git a/rust/kernel/xarray.rs b/rust/kernel/xarray.rs
> index ca97134ba2bd0..9d4589979fd1d 100644
> --- a/rust/kernel/xarray.rs
> +++ b/rust/kernel/xarray.rs
> @@ -255,6 +255,71 @@ pub fn get_mut(&mut self, index: usize) -> Option<T::BorrowedMut<'_>> {
>          Some(unsafe { T::borrow_mut(ptr.as_ptr()) })
>      }
>
> +    fn load_next(&self, index: usize) -> Option<(usize, NonNull<c_void>)> {
> +        let mut state = XArrayState::new(self, index);
> +        // SAFETY: `state.state` is always valid by the type invariant of
> +        // `XArrayState` and the caller holds the lock.
> +        let ptr = unsafe { bindings::xas_find(&raw mut state.state, usize::MAX) };
> +        NonNull::new(ptr).map(|ptr| (state.state.xa_index, ptr))
> +    }

Can this be a method on XArrayState? It seems odd to document a remote
type's invariant here when we could put that justification on the type
itself.

> +
> +    /// Finds the next element starting from the given index.
> +    ///
> +    /// # Examples
> +    ///
> +    /// ```
> +    /// # use kernel::{prelude::*, xarray::{AllocKind, XArray}};
> +    /// let mut xa = KBox::pin_init(XArray::<KBox<u32>>::new(AllocKind::Alloc), GFP_KERNEL)?;
> +    /// let mut guard = xa.lock();
> +    ///
> +    /// guard.store(10, KBox::new(10u32, GFP_KERNEL)?, GFP_KERNEL)?;
> +    /// guard.store(20, KBox::new(20u32, GFP_KERNEL)?, GFP_KERNEL)?;
> +    ///
> +    /// if let Some((found_index, value)) = guard.find_next(11) {
> +    ///     assert_eq!(found_index, 20);
> +    ///     assert_eq!(*value, 20);
> +    /// }
> +    ///
> +    /// if let Some((found_index, value)) = guard.find_next(5) {
> +    ///     assert_eq!(found_index, 10);
> +    ///     assert_eq!(*value, 10);
> +    /// }
> +    ///
> +    /// # Ok::<(), kernel::error::Error>(())
> +    /// ```
> +    pub fn find_next(&self, index: usize) -> Option<(usize, T::Borrowed<'_>)> {
> +        self.load_next(index)
> +            // SAFETY: `ptr` came from `T::into_foreign`.
> +            .map(|(index, ptr)| (index, unsafe { T::borrow(ptr.as_ptr()) }))
> +    }
> +
> +    /// Finds the next element starting from the given index, returning a mutable reference.
> +    ///
> +    /// # Examples
> +    ///
> +    /// ```
> +    /// # use kernel::{prelude::*, xarray::{AllocKind, XArray}};
> +    /// let mut xa = KBox::pin_init(XArray::<KBox<u32>>::new(AllocKind::Alloc), GFP_KERNEL)?;
> +    /// let mut guard = xa.lock();
> +    ///
> +    /// guard.store(10, KBox::new(10u32, GFP_KERNEL)?, GFP_KERNEL)?;
> +    /// guard.store(20, KBox::new(20u32, GFP_KERNEL)?, GFP_KERNEL)?;
> +    ///
> +    /// if let Some((found_index, mut_value)) = guard.find_next_mut(5) {
> +    ///     assert_eq!(found_index, 10);
> +    ///     *mut_value = 0x99;
> +    /// }
> +    ///
> +    /// assert_eq!(guard.get(10).copied(), Some(0x99));
> +    ///
> +    /// # Ok::<(), kernel::error::Error>(())
> +    /// ```
> +    pub fn find_next_mut(&mut self, index: usize) -> Option<(usize, T::BorrowedMut<'_>)> {
> +        self.load_next(index)
> +            // SAFETY: `ptr` came from `T::into_foreign`.
> +            .map(move |(index, ptr)| (index, unsafe { T::borrow_mut(ptr.as_ptr()) }))
> +    }
> +
>      /// Removes and returns the element at the given index.
>      pub fn remove(&mut self, index: usize) -> Option<T> {
>          // SAFETY:
>
> --
> 2.51.2
>
>

Tamir Duberstein <tamird@gmail.com> writes:

> On Wed, Dec 3, 2025 at 5:27 PM Andreas Hindborg <a.hindborg@kernel.org> wrote:
>>
>> Add methods to find the next element in an XArray starting from a
>> given index. The methods return a tuple containing the index where the
>> element was found and a reference to the element.
>>
>> The implementation uses the XArray state API via `xas_find` to avoid taking
>> the xarray lock that is already held by `Guard`.
>
> Similarly to the commit message introducing the use of `xas_load`,
> this is not correct because `xa_find` takes and release the RCU lock
> only, not the XArray lock.

Right, thanks for pointing that out.

>
>>
>> Signed-off-by: Andreas Hindborg <a.hindborg@kernel.org>
>> ---
>>  rust/kernel/xarray.rs | 65 +++++++++++++++++++++++++++++++++++++++++++++++++++
>>  1 file changed, 65 insertions(+)
>>
>> diff --git a/rust/kernel/xarray.rs b/rust/kernel/xarray.rs
>> index ca97134ba2bd0..9d4589979fd1d 100644
>> --- a/rust/kernel/xarray.rs
>> +++ b/rust/kernel/xarray.rs
>> @@ -255,6 +255,71 @@ pub fn get_mut(&mut self, index: usize) -> Option<T::BorrowedMut<'_>> {
>>          Some(unsafe { T::borrow_mut(ptr.as_ptr()) })
>>      }
>>
>> +    fn load_next(&self, index: usize) -> Option<(usize, NonNull<c_void>)> {
>> +        let mut state = XArrayState::new(self, index);
>> +        // SAFETY: `state.state` is always valid by the type invariant of
>> +        // `XArrayState` and the caller holds the lock.
>> +        let ptr = unsafe { bindings::xas_find(&raw mut state.state, usize::MAX) };
>> +        NonNull::new(ptr).map(|ptr| (state.state.xa_index, ptr))
>> +    }
>
> Can this be a method on XArrayState? It seems odd to document a remote
> type's invariant here when we could put that justification on the type
> itself.

Good idea, I'll move it.


Best regards,
Andreas Hindborg