[PATCH 0/3] rust: xarray: add `insert` and `reserve`

[PATCH 0/3] rust: xarray: add `insert` and `reserve`

[Tamir Duberstein]

Please see individual patches.

Signed-off-by: Tamir Duberstein <tamird@gmail.com>
---
Tamir Duberstein (3):
      rust: xarray: use the prelude
      rust: xarray: implement Default for AllocKind
      rust: xarray: add `insert` and `reserve`

 include/linux/xarray.h |   2 +
 lib/xarray.c           |  28 ++-
 rust/helpers/xarray.c  |   5 +
 rust/kernel/xarray.rs  | 460 +++++++++++++++++++++++++++++++++++++++++++++++--
 4 files changed, 472 insertions(+), 23 deletions(-)
---
base-commit: 769e324b66b0d92d04f315d0c45a0f72737c7494
change-id: 20250701-xarray-insert-reserve-bd811ad46a1d

Best regards,
--  
Tamir Duberstein <tamird@gmail.com>

[PATCH 1/3] rust: xarray: use the prelude

[Tamir Duberstein]

Using the prelude is customary in the kernel crate.

Signed-off-by: Tamir Duberstein <tamird@gmail.com>
---
 rust/kernel/xarray.rs | 34 ++++++++++++++++++++--------------
 1 file changed, 20 insertions(+), 14 deletions(-)

diff --git a/rust/kernel/xarray.rs b/rust/kernel/xarray.rs
index 75719e7bb491..436faad99c89 100644
--- a/rust/kernel/xarray.rs
+++ b/rust/kernel/xarray.rs
@@ -5,16 +5,15 @@
 //! C header: [`include/linux/xarray.h`](srctree/include/linux/xarray.h)
 
 use crate::{
-    alloc, bindings, build_assert,
-    error::{Error, Result},
+    alloc,
+    prelude::*,
     types::{ForeignOwnable, NotThreadSafe, Opaque},
 };
-use core::{iter, marker::PhantomData, mem, pin::Pin, ptr::NonNull};
-use pin_init::{pin_data, pin_init, pinned_drop, PinInit};
+use core::{iter, marker::PhantomData, mem, ptr::NonNull};
 
 /// An array which efficiently maps sparse integer indices to owned objects.
 ///
-/// This is similar to a [`crate::alloc::kvec::Vec<Option<T>>`], but more efficient when there are
+/// This is similar to a [`Vec<Option<T>>`], but more efficient when there are
 /// holes in the index space, and can be efficiently grown.
 ///
 /// # Invariants
@@ -104,16 +103,23 @@ pub fn new(kind: AllocKind) -> impl PinInit<Self> {
     fn iter(&self) -> impl Iterator<Item = NonNull<T::PointedTo>> + '_ {
         let mut index = 0;
 
-        // SAFETY: `self.xa` is always valid by the type invariant.
-        iter::once(unsafe {
-            bindings::xa_find(self.xa.get(), &mut index, usize::MAX, bindings::XA_PRESENT)
-        })
-        .chain(iter::from_fn(move || {
+        core::iter::Iterator::chain(
             // SAFETY: `self.xa` is always valid by the type invariant.
-            Some(unsafe {
-                bindings::xa_find_after(self.xa.get(), &mut index, usize::MAX, bindings::XA_PRESENT)
-            })
-        }))
+            iter::once(unsafe {
+                bindings::xa_find(self.xa.get(), &mut index, usize::MAX, bindings::XA_PRESENT)
+            }),
+            iter::from_fn(move || {
+                // SAFETY: `self.xa` is always valid by the type invariant.
+                Some(unsafe {
+                    bindings::xa_find_after(
+                        self.xa.get(),
+                        &mut index,
+                        usize::MAX,
+                        bindings::XA_PRESENT,
+                    )
+                })
+            }),
+        )
         .map_while(|ptr| NonNull::new(ptr.cast()))
     }
 

-- 
2.50.0

[PATCH 2/3] rust: xarray: implement Default for AllocKind

[Tamir Duberstein]

Most users are likely to want 0-indexed arrays. Clean up the
documentation test accordingly.

Signed-off-by: Tamir Duberstein <tamird@gmail.com>
---
 rust/kernel/xarray.rs | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/rust/kernel/xarray.rs b/rust/kernel/xarray.rs
index 436faad99c89..bbce54ec695c 100644
--- a/rust/kernel/xarray.rs
+++ b/rust/kernel/xarray.rs
@@ -24,10 +24,11 @@
 /// # Examples
 ///
 /// ```rust
-/// use kernel::alloc::KBox;
-/// use kernel::xarray::{AllocKind, XArray};
+/// # use kernel::alloc::KBox;
+/// # use kernel::xarray::XArray;
+/// # use pin_init::stack_pin_init;
 ///
-/// let xa = KBox::pin_init(XArray::new(AllocKind::Alloc1), GFP_KERNEL)?;
+/// stack_pin_init!(let xa = XArray::new(Default::default()));
 ///
 /// let dead = KBox::new(0xdead, GFP_KERNEL)?;
 /// let beef = KBox::new(0xbeef, GFP_KERNEL)?;
@@ -75,8 +76,10 @@ fn drop(self: Pin<&mut Self>) {
 }
 
 /// Flags passed to [`XArray::new`] to configure the array's allocation tracking behavior.
+#[derive(Default)]
 pub enum AllocKind {
     /// Consider the first element to be at index 0.
+    #[default]
     Alloc,
     /// Consider the first element to be at index 1.
     Alloc1,

-- 
2.50.0

[PATCH 3/3] rust: xarray: add `insert` and `reserve`

[Tamir Duberstein]

Add `Guard::{insert,reserve}` and `Guard::{insert,reserve}_limit`, which
are akin to `__xa_{alloc,insert}` in C.

Note that unlike `xa_reserve` which only ensures that memory is
allocated, the semantics of `Reservation` are stricter and require
precise management of the reservation. Indices which have been reserved
can still be overwritten with `Guard::store`, which allows for C-like
semantics if desired.

`__xa_cmpxchg_raw` is exported to facilitate the semantics described
above.

Signed-off-by: Tamir Duberstein <tamird@gmail.com>
---
 include/linux/xarray.h |   2 +
 lib/xarray.c           |  28 +++-
 rust/helpers/xarray.c  |   5 +
 rust/kernel/xarray.rs  | 419 ++++++++++++++++++++++++++++++++++++++++++++++++-
 4 files changed, 447 insertions(+), 7 deletions(-)

diff --git a/include/linux/xarray.h b/include/linux/xarray.h
index be850174e802..64f2a5e06ceb 100644
--- a/include/linux/xarray.h
+++ b/include/linux/xarray.h
@@ -563,6 +563,8 @@ void *__xa_erase(struct xarray *, unsigned long index);
 void *__xa_store(struct xarray *, unsigned long index, void *entry, gfp_t);
 void *__xa_cmpxchg(struct xarray *, unsigned long index, void *old,
 		void *entry, gfp_t);
+void *__xa_cmpxchg_raw(struct xarray *, unsigned long index, void *old,
+		void *entry, gfp_t);
 int __must_check __xa_insert(struct xarray *, unsigned long index,
 		void *entry, gfp_t);
 int __must_check __xa_alloc(struct xarray *, u32 *id, void *entry,
diff --git a/lib/xarray.c b/lib/xarray.c
index 76dde3a1cacf..58202b6fbb59 100644
--- a/lib/xarray.c
+++ b/lib/xarray.c
@@ -1738,9 +1738,6 @@ void *xa_store(struct xarray *xa, unsigned long index, void *entry, gfp_t gfp)
 }
 EXPORT_SYMBOL(xa_store);
 
-static inline void *__xa_cmpxchg_raw(struct xarray *xa, unsigned long index,
-			void *old, void *entry, gfp_t gfp);
-
 /**
  * __xa_cmpxchg() - Conditionally replace an entry in the XArray.
  * @xa: XArray.
@@ -1767,7 +1764,29 @@ void *__xa_cmpxchg(struct xarray *xa, unsigned long index,
 }
 EXPORT_SYMBOL(__xa_cmpxchg);
 
-static inline void *__xa_cmpxchg_raw(struct xarray *xa, unsigned long index,
+/**
+ * __xa_cmpxchg_raw() - Conditionally replace an entry in the XArray.
+ * @xa: XArray.
+ * @index: Index into array.
+ * @old: Old value to test against.
+ * @entry: New value to place in array.
+ * @gfp: Memory allocation flags.
+ *
+ * You must already be holding the xa_lock when calling this function.
+ * It will drop the lock if needed to allocate memory, and then reacquire
+ * it afterwards.
+ *
+ * If the entry at @index is the same as @old, replace it with @entry.
+ * If the return value is equal to @old, then the exchange was successful.
+ *
+ * This function is the same as __xa_cmpxchg() except that it does not coerce
+ * XA_ZERO_ENTRY to NULL on egress.
+ *
+ * Context: Any context.  Expects xa_lock to be held on entry.  May
+ * release and reacquire xa_lock if @gfp flags permit.
+ * Return: The old value at this index or xa_err() if an error happened.
+ */
+void *__xa_cmpxchg_raw(struct xarray *xa, unsigned long index,
 			void *old, void *entry, gfp_t gfp)
 {
 	XA_STATE(xas, xa, index);
@@ -1787,6 +1806,7 @@ static inline void *__xa_cmpxchg_raw(struct xarray *xa, unsigned long index,
 
 	return xas_result(&xas, curr);
 }
+EXPORT_SYMBOL(__xa_cmpxchg_raw);
 
 /**
  * __xa_insert() - Store this entry in the XArray if no entry is present.
diff --git a/rust/helpers/xarray.c b/rust/helpers/xarray.c
index 60b299f11451..b6c078e6a343 100644
--- a/rust/helpers/xarray.c
+++ b/rust/helpers/xarray.c
@@ -2,6 +2,11 @@
 
 #include <linux/xarray.h>
 
+void *rust_helper_xa_zero_entry(void)
+{
+	return XA_ZERO_ENTRY;
+}
+
 int rust_helper_xa_err(void *entry)
 {
 	return xa_err(entry);
diff --git a/rust/kernel/xarray.rs b/rust/kernel/xarray.rs
index bbce54ec695c..87fa3259cdd7 100644
--- a/rust/kernel/xarray.rs
+++ b/rust/kernel/xarray.rs
@@ -9,7 +9,12 @@
     prelude::*,
     types::{ForeignOwnable, NotThreadSafe, Opaque},
 };
-use core::{iter, marker::PhantomData, mem, ptr::NonNull};
+use core::{
+    fmt, iter,
+    marker::PhantomData,
+    mem, ops,
+    ptr::{null_mut, NonNull},
+};
 
 /// An array which efficiently maps sparse integer indices to owned objects.
 ///
@@ -126,6 +131,19 @@ fn iter(&self) -> impl Iterator<Item = NonNull<T::PointedTo>> + '_ {
         .map_while(|ptr| NonNull::new(ptr.cast()))
     }
 
+    fn with_guard<F, U>(&self, guard: Option<&mut Guard<'_, T>>, f: F) -> U
+    where
+        F: FnOnce(&mut Guard<'_, T>) -> U,
+    {
+        match guard {
+            None => f(&mut self.lock()),
+            Some(guard) => {
+                assert_eq!(guard.xa.xa.get(), self.xa.get());
+                f(guard)
+            }
+        }
+    }
+
     /// Attempts to lock the [`XArray`] for exclusive access.
     pub fn try_lock(&self) -> Option<Guard<'_, T>> {
         // SAFETY: `self.xa` is always valid by the type invariant.
@@ -172,6 +190,7 @@ fn drop(&mut self) {
 /// The error returned by [`store`](Guard::store).
 ///
 /// Contains the underlying error and the value that was not stored.
+#[derive(Debug)]
 pub struct StoreError<T> {
     /// The error that occurred.
     pub error: Error,
@@ -185,6 +204,11 @@ fn from(value: StoreError<T>) -> Self {
     }
 }
 
+fn to_usize(i: u32) -> usize {
+    i.try_into()
+        .unwrap_or_else(|_| build_error!("cannot convert u32 to usize"))
+}
+
 impl<'a, T: ForeignOwnable> Guard<'a, T> {
     fn load<F, U>(&self, index: usize, f: F) -> Option<U>
     where
@@ -219,7 +243,7 @@ pub fn remove(&mut self, index: usize) -> Option<T> {
         // - The caller holds the lock.
         let ptr = unsafe { bindings::__xa_erase(self.xa.xa.get(), index) }.cast();
         // SAFETY:
-        // - `ptr` is either NULL or came from `T::into_foreign`.
+        // - `ptr` is either `NULL` or came from `T::into_foreign`.
         // - `&mut self` guarantees that the lifetimes of [`T::Borrowed`] and [`T::BorrowedMut`]
         // borrowed from `self` have ended.
         unsafe { T::try_from_foreign(ptr) }
@@ -267,13 +291,272 @@ pub fn store(
             })
         } else {
             let old = old.cast();
-            // SAFETY: `ptr` is either NULL or came from `T::into_foreign`.
+            // SAFETY: `ptr` is either `NULL` or came from `T::into_foreign`.
             //
             // NB: `XA_ZERO_ENTRY` is never returned by functions belonging to the Normal XArray
             // API; such entries present as `NULL`.
             Ok(unsafe { T::try_from_foreign(old) })
         }
     }
+
+    /// Stores an element at the given index if no entry is present.
+    ///
+    /// May drop the lock if needed to allocate memory, and then reacquire it afterwards.
+    ///
+    /// On failure, returns the element which was attempted to be stored.
+    pub fn insert(
+        &mut self,
+        index: usize,
+        value: T,
+        gfp: alloc::Flags,
+    ) -> Result<(), StoreError<T>> {
+        build_assert!(
+            mem::align_of::<T::PointedTo>() >= 4,
+            "pointers stored in XArray must be 4-byte aligned"
+        );
+        let ptr = value.into_foreign();
+        // SAFETY: `self.xa` is always valid by the type invariant.
+        //
+        // INVARIANT: `ptr` came from `T::into_foreign`.
+        match unsafe { bindings::__xa_insert(self.xa.xa.get(), index, ptr.cast(), gfp.as_raw()) } {
+            0 => Ok(()),
+            errno => {
+                // SAFETY: `ptr` came from `T::into_foreign` and `__xa_insert` does not take
+                // ownership of the value on error.
+                let value = unsafe { T::from_foreign(ptr) };
+                Err(StoreError {
+                    value,
+                    error: Error::from_errno(errno),
+                })
+            }
+        }
+    }
+
+    /// Wrapper around `__xa_alloc`.
+    ///
+    /// On success, takes ownership of pointers passed in `op`.
+    ///
+    /// On failure, ownership returns to the caller.
+    ///
+    /// # Safety
+    ///
+    /// `ptr` must be `NULL` or have come from a previous call to `T::into_foreign`.
+    unsafe fn alloc(
+        &mut self,
+        limit: impl ops::RangeBounds<u32>,
+        ptr: *mut T::PointedTo,
+        gfp: alloc::Flags,
+    ) -> Result<usize> {
+        // NB: `xa_limit::{max,min}` are inclusive.
+        let limit = bindings::xa_limit {
+            max: match limit.end_bound() {
+                ops::Bound::Included(&end) => end,
+                ops::Bound::Excluded(&end) => end - 1,
+                ops::Bound::Unbounded => u32::MAX,
+            },
+            min: match limit.start_bound() {
+                ops::Bound::Included(&start) => start,
+                ops::Bound::Excluded(&start) => start + 1,
+                ops::Bound::Unbounded => 0,
+            },
+        };
+
+        let mut index = u32::MAX;
+
+        // SAFETY:
+        // - `self.xa` is always valid by the type invariant.
+        // - `self.xa` was initialized with `XA_FLAGS_ALLOC` or `XA_FLAGS_ALLOC1`.
+        //
+        // INVARIANT: `ptr` is either `NULL` or came from `T::into_foreign`.
+        match unsafe {
+            bindings::__xa_alloc(
+                self.xa.xa.get(),
+                &mut index,
+                ptr.cast(),
+                limit,
+                gfp.as_raw(),
+            )
+        } {
+            0 => Ok(to_usize(index)),
+            errno => Err(Error::from_errno(errno)),
+        }
+    }
+
+    /// Allocates an entry somewhere in the array.
+    ///
+    /// On success, returns the index at which the entry was stored.
+    ///
+    /// On failure, returns the entry which was attempted to be stored.
+    pub fn insert_limit(
+        &mut self,
+        limit: impl ops::RangeBounds<u32>,
+        value: T,
+        gfp: alloc::Flags,
+    ) -> Result<usize, StoreError<T>> {
+        build_assert!(
+            mem::align_of::<T::PointedTo>() >= 4,
+            "pointers stored in XArray must be 4-byte aligned"
+        );
+        let ptr = value.into_foreign();
+        // SAFETY: `ptr` came from `T::into_foreign`.
+        unsafe { self.alloc(limit, ptr, gfp) }.map_err(|error| {
+            // SAFETY: `ptr` came from `T::into_foreign` and `self.alloc` does not take ownership of
+            // the value on error.
+            let value = unsafe { T::from_foreign(ptr) };
+            StoreError { value, error }
+        })
+    }
+
+    /// Reserves an entry in the array.
+    pub fn reserve(&mut self, index: usize, gfp: alloc::Flags) -> Result<Reservation<'a, T>> {
+        // NB: `__xa_insert` internally coerces `NULL` to `XA_ZERO_ENTRY` on ingress.
+        let ptr = null_mut();
+        // SAFETY: `self.xa` is always valid by the type invariant.
+        //
+        // INVARIANT: `ptr` is `NULL`.
+        match unsafe { bindings::__xa_insert(self.xa.xa.get(), index, ptr, gfp.as_raw()) } {
+            0 => Ok(Reservation { xa: self.xa, index }),
+            errno => Err(Error::from_errno(errno)),
+        }
+    }
+
+    /// Reserves an entry somewhere in the array.
+    pub fn reserve_limit(
+        &mut self,
+        limit: impl ops::RangeBounds<u32>,
+        gfp: alloc::Flags,
+    ) -> Result<Reservation<'a, T>> {
+        // NB: `__xa_alloc` internally coerces `NULL` to `XA_ZERO_ENTRY` on ingress.
+        let ptr = null_mut();
+        // SAFETY: `ptr` is `NULL`.
+        unsafe { self.alloc(limit, ptr, gfp) }.map(|index| Reservation { xa: self.xa, index })
+    }
+}
+
+/// A reserved slot in an array.
+///
+/// The slot is released when the reservation goes out of scope.
+///
+/// Note that the array lock *must not* be held when the reservation is filled or dropped as this
+/// will lead to deadlock. [`Reservation::fill_locked`] and [`Reservation::release_locked`] can be
+/// used in context where the array lock is held.
+#[must_use = "the reservation is released immediately when the reservation is unused"]
+pub struct Reservation<'a, T: ForeignOwnable> {
+    xa: &'a XArray<T>,
+    index: usize,
+}
+
+impl<T: ForeignOwnable> fmt::Debug for Reservation<'_, T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Reservation")
+            .field("index", &self.index())
+            .finish()
+    }
+}
+
+impl<T: ForeignOwnable> Reservation<'_, T> {
+    /// Returns the index of the reservation.
+    pub fn index(&self) -> usize {
+        self.index
+    }
+
+    /// Replaces the reserved entry with the given entry.
+    ///
+    /// # Safety
+    ///
+    /// `ptr` must be `NULL` or have come from a previous call to `T::into_foreign`.
+    unsafe fn replace(guard: &mut Guard<'_, T>, index: usize, ptr: *mut T::PointedTo) -> Result {
+        // SAFETY: `xa_zero_entry` wraps `XA_ZERO_ENTRY` which is always safe to use.
+        let old = unsafe { bindings::xa_zero_entry() };
+
+        // NB: `__xa_cmpxchg_raw` is used over `__xa_cmpxchg` because the latter coerces
+        // `XA_ZERO_ENTRY` to `NULL` on egress, which would prevent us from determining whether a
+        // replacement was made.
+        //
+        // SAFETY: `self.xa` is always valid by the type invariant.
+        //
+        // INVARIANT: `ptr` is either `NULL` or came from `T::into_foreign` and `old` is
+        // `XA_ZERO_ENTRY`.
+        let ret =
+            unsafe { bindings::__xa_cmpxchg_raw(guard.xa.xa.get(), index, old, ptr.cast(), 0) };
+
+        // SAFETY: `__xa_cmpxchg_raw` returns the old entry at this index on success or `xa_err` if
+        // an error happened.
+        match unsafe { bindings::xa_err(ret) } {
+            0 => {
+                if ret == old {
+                    Ok(())
+                } else {
+                    Err(EBUSY)
+                }
+            }
+            errno => Err(Error::from_errno(errno)),
+        }
+    }
+
+    fn fill_inner(&self, guard: Option<&mut Guard<'_, T>>, value: T) -> Result<(), StoreError<T>> {
+        let Self { xa, index } = self;
+        let index = *index;
+
+        let ptr = value.into_foreign();
+        xa.with_guard(guard, |guard| {
+            // SAFETY: `ptr` came from `T::into_foreign`.
+            unsafe { Self::replace(guard, index, ptr) }
+        })
+        .map_err(|error| {
+            // SAFETY: `ptr` came from `T::into_foreign` and `Self::replace` does not take ownership
+            // of the value on error.
+            let value = unsafe { T::from_foreign(ptr) };
+            StoreError { value, error }
+        })
+    }
+
+    /// Fills the reservation.
+    pub fn fill(self, value: T) -> Result<(), StoreError<T>> {
+        let result = self.fill_inner(None, value);
+        mem::forget(self);
+        result
+    }
+
+    /// Fills the reservation without acquiring the array lock.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the passed guard locks a different array.
+    pub fn fill_locked(self, guard: &mut Guard<'_, T>, value: T) -> Result<(), StoreError<T>> {
+        let result = self.fill_inner(Some(guard), value);
+        mem::forget(self);
+        result
+    }
+
+    fn release_inner(&self, guard: Option<&mut Guard<'_, T>>) -> Result {
+        let Self { xa, index } = self;
+        let index = *index;
+
+        xa.with_guard(guard, |guard| {
+            let ptr = null_mut();
+            // SAFETY: `ptr` is `NULL`.
+            unsafe { Self::replace(guard, index, ptr) }
+        })
+    }
+
+    /// Releases the reservation without acquiring the array lock.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the passed guard locks a different array.
+    pub fn release_locked(self, guard: &mut Guard<'_, T>) -> Result {
+        let result = self.release_inner(Some(guard));
+        mem::forget(self);
+        result
+    }
+}
+
+impl<T: ForeignOwnable> Drop for Reservation<'_, T> {
+    fn drop(&mut self) {
+        // NB: Errors here are possible since `Guard::store` does not honor reservations.
+        let _: Result = self.release_inner(None);
+    }
 }
 
 // SAFETY: `XArray<T>` has no shared mutable state so it is `Send` iff `T` is `Send`.
@@ -282,3 +565,133 @@ unsafe impl<T: ForeignOwnable + Send> Send for XArray<T> {}
 // SAFETY: `XArray<T>` serialises the interior mutability it provides so it is `Sync` iff `T` is
 // `Send`.
 unsafe impl<T: ForeignOwnable + Send> Sync for XArray<T> {}
+
+#[macros::kunit_tests(rust_xarray_kunit)]
+mod tests {
+    use super::*;
+    use pin_init::stack_pin_init;
+
+    fn new_kbox<T>(value: T) -> Result<KBox<T>> {
+        KBox::new(value, GFP_KERNEL).map_err(Into::into)
+    }
+
+    #[test]
+    fn test_alloc_kind_alloc() -> Result {
+        test_alloc_kind(AllocKind::Alloc, 0)
+    }
+
+    #[test]
+    fn test_alloc_kind_alloc1() -> Result {
+        test_alloc_kind(AllocKind::Alloc1, 1)
+    }
+
+    fn test_alloc_kind(kind: AllocKind, expected_index: usize) -> Result {
+        stack_pin_init!(let xa = XArray::new(kind));
+        let mut guard = xa.lock();
+
+        let reservation = guard.reserve_limit(.., GFP_KERNEL)?;
+        assert_eq!(reservation.index(), expected_index);
+        reservation.release_locked(&mut guard)?;
+
+        let insertion = guard.insert_limit(.., new_kbox(0x1337)?, GFP_KERNEL);
+        assert!(insertion.is_ok());
+        let insertion_index = insertion.unwrap();
+        assert_eq!(insertion_index, expected_index);
+
+        Ok(())
+    }
+
+    const IDX: usize = 0x1337;
+
+    fn insert<T: ForeignOwnable>(guard: &mut Guard<'_, T>, value: T) -> Result<(), StoreError<T>> {
+        guard.insert(IDX, value, GFP_KERNEL)
+    }
+
+    fn reserve<'a, T: ForeignOwnable>(guard: &mut Guard<'a, T>) -> Result<Reservation<'a, T>> {
+        guard.reserve(IDX, GFP_KERNEL)
+    }
+
+    #[track_caller]
+    fn check_not_vacant<'a>(guard: &mut Guard<'a, KBox<usize>>) -> Result {
+        // Insertion fails.
+        {
+            let beef = new_kbox(0xbeef)?;
+            let ret = insert(guard, beef);
+            assert!(ret.is_err());
+            let StoreError { error, value } = ret.unwrap_err();
+            assert_eq!(error, EBUSY);
+            assert_eq!(*value, 0xbeef);
+        }
+
+        // Reservation fails.
+        {
+            let ret = reserve(guard);
+            assert!(ret.is_err());
+            assert_eq!(ret.unwrap_err(), EBUSY);
+        }
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_insert_and_reserve_interaction() -> Result {
+        stack_pin_init!(let xa = XArray::new(Default::default()));
+        let mut guard = xa.lock();
+
+        // Vacant.
+        assert_eq!(guard.get(IDX), None);
+
+        // Reservation succeeds.
+        let reservation = {
+            let ret = reserve(&mut guard);
+            assert!(ret.is_ok());
+            ret.unwrap()
+        };
+
+        // Reserved presents as vacant.
+        assert_eq!(guard.get(IDX), None);
+
+        check_not_vacant(&mut guard)?;
+
+        // Release reservation.
+        {
+            let ret = reservation.release_locked(&mut guard);
+            assert!(ret.is_ok());
+            let () = ret.unwrap();
+        }
+
+        // Vacant again.
+        assert_eq!(guard.get(IDX), None);
+
+        // Insert succeeds.
+        {
+            let dead = new_kbox(0xdead)?;
+            let ret = insert(&mut guard, dead);
+            assert!(ret.is_ok());
+            let () = ret.unwrap();
+        }
+
+        check_not_vacant(&mut guard)?;
+
+        // Remove.
+        assert_eq!(guard.remove(IDX).as_deref(), Some(&0xdead));
+
+        // Reserve and fill.
+        {
+            let beef = new_kbox(0xbeef)?;
+            let ret = reserve(&mut guard);
+            assert!(ret.is_ok());
+            let reservation = ret.unwrap();
+            let ret = reservation.fill_locked(&mut guard, beef);
+            assert!(ret.is_ok());
+            let () = ret.unwrap();
+        };
+
+        check_not_vacant(&mut guard)?;
+
+        // Remove.
+        assert_eq!(guard.remove(IDX).as_deref(), Some(&0xbeef));
+
+        Ok(())
+    }
+}

-- 
2.50.0

Re: [PATCH 1/3] rust: xarray: use the prelude

[Boqun Feng]

On Tue, Jul 01, 2025 at 12:27:17PM -0400, Tamir Duberstein wrote:
> Using the prelude is customary in the kernel crate.
> 
> Signed-off-by: Tamir Duberstein <tamird@gmail.com>
> ---
>  rust/kernel/xarray.rs | 34 ++++++++++++++++++++--------------
>  1 file changed, 20 insertions(+), 14 deletions(-)
> 
> diff --git a/rust/kernel/xarray.rs b/rust/kernel/xarray.rs
> index 75719e7bb491..436faad99c89 100644
> --- a/rust/kernel/xarray.rs
> +++ b/rust/kernel/xarray.rs
> @@ -5,16 +5,15 @@
>  //! C header: [`include/linux/xarray.h`](srctree/include/linux/xarray.h)
>  
>  use crate::{
> -    alloc, bindings, build_assert,
> -    error::{Error, Result},
> +    alloc,
> +    prelude::*,
>      types::{ForeignOwnable, NotThreadSafe, Opaque},
>  };
> -use core::{iter, marker::PhantomData, mem, pin::Pin, ptr::NonNull};
> -use pin_init::{pin_data, pin_init, pinned_drop, PinInit};
> +use core::{iter, marker::PhantomData, mem, ptr::NonNull};
>  
>  /// An array which efficiently maps sparse integer indices to owned objects.
>  ///
> -/// This is similar to a [`crate::alloc::kvec::Vec<Option<T>>`], but more efficient when there are
> +/// This is similar to a [`Vec<Option<T>>`], but more efficient when there are
>  /// holes in the index space, and can be efficiently grown.
>  ///
>  /// # Invariants
> @@ -104,16 +103,23 @@ pub fn new(kind: AllocKind) -> impl PinInit<Self> {
>      fn iter(&self) -> impl Iterator<Item = NonNull<T::PointedTo>> + '_ {
>          let mut index = 0;
>  
> -        // SAFETY: `self.xa` is always valid by the type invariant.
> -        iter::once(unsafe {
> -            bindings::xa_find(self.xa.get(), &mut index, usize::MAX, bindings::XA_PRESENT)
> -        })
> -        .chain(iter::from_fn(move || {
> +        core::iter::Iterator::chain(

Does this part come from using the prelude? If not, either we need to
split the patch or we need to mention it in the changelog at least.

Also since we `use core::iter` above, we can avoid the `core::` here.

Regards,
Boqun

>              // SAFETY: `self.xa` is always valid by the type invariant.
> -            Some(unsafe {
> -                bindings::xa_find_after(self.xa.get(), &mut index, usize::MAX, bindings::XA_PRESENT)
> -            })
> -        }))
> +            iter::once(unsafe {
> +                bindings::xa_find(self.xa.get(), &mut index, usize::MAX, bindings::XA_PRESENT)
> +            }),
> +            iter::from_fn(move || {
> +                // SAFETY: `self.xa` is always valid by the type invariant.
> +                Some(unsafe {
> +                    bindings::xa_find_after(
> +                        self.xa.get(),
> +                        &mut index,
> +                        usize::MAX,
> +                        bindings::XA_PRESENT,
> +                    )
> +                })
> +            }),
> +        )
>          .map_while(|ptr| NonNull::new(ptr.cast()))
>      }
>  
> 
> -- 
> 2.50.0
> 

Re: [PATCH 1/3] rust: xarray: use the prelude

[Tamir Duberstein]

On Tue, Jul 1, 2025 at 12:35 PM Boqun Feng <boqun.feng@gmail.com> wrote:
>
> On Tue, Jul 01, 2025 at 12:27:17PM -0400, Tamir Duberstein wrote:
> > Using the prelude is customary in the kernel crate.
> >
> > Signed-off-by: Tamir Duberstein <tamird@gmail.com>
> > ---
> >  rust/kernel/xarray.rs | 34 ++++++++++++++++++++--------------
> >  1 file changed, 20 insertions(+), 14 deletions(-)
> >
> > diff --git a/rust/kernel/xarray.rs b/rust/kernel/xarray.rs
> > index 75719e7bb491..436faad99c89 100644
> > --- a/rust/kernel/xarray.rs
> > +++ b/rust/kernel/xarray.rs
> > @@ -5,16 +5,15 @@
> >  //! C header: [`include/linux/xarray.h`](srctree/include/linux/xarray.h)
> >
> >  use crate::{
> > -    alloc, bindings, build_assert,
> > -    error::{Error, Result},
> > +    alloc,
> > +    prelude::*,
> >      types::{ForeignOwnable, NotThreadSafe, Opaque},
> >  };
> > -use core::{iter, marker::PhantomData, mem, pin::Pin, ptr::NonNull};
> > -use pin_init::{pin_data, pin_init, pinned_drop, PinInit};
> > +use core::{iter, marker::PhantomData, mem, ptr::NonNull};
> >
> >  /// An array which efficiently maps sparse integer indices to owned objects.
> >  ///
> > -/// This is similar to a [`crate::alloc::kvec::Vec<Option<T>>`], but more efficient when there are
> > +/// This is similar to a [`Vec<Option<T>>`], but more efficient when there are
> >  /// holes in the index space, and can be efficiently grown.
> >  ///
> >  /// # Invariants
> > @@ -104,16 +103,23 @@ pub fn new(kind: AllocKind) -> impl PinInit<Self> {
> >      fn iter(&self) -> impl Iterator<Item = NonNull<T::PointedTo>> + '_ {
> >          let mut index = 0;
> >
> > -        // SAFETY: `self.xa` is always valid by the type invariant.
> > -        iter::once(unsafe {
> > -            bindings::xa_find(self.xa.get(), &mut index, usize::MAX, bindings::XA_PRESENT)
> > -        })
> > -        .chain(iter::from_fn(move || {
> > +        core::iter::Iterator::chain(
>
> Does this part come from using the prelude? If not, either we need to
> split the patch or we need to mention it in the changelog at least.

Yes, it's from using the prelude - PinInit also has a chain method
that causes ambiguity here.

> Also since we `use core::iter` above, we can avoid the `core::` here.

Good point.

Re: [PATCH 3/3] rust: xarray: add `insert` and `reserve`

[Miguel Ojeda]

On Tue, Jul 1, 2025 at 6:27 PM Tamir Duberstein <tamird@gmail.com> wrote:
>
> Add `Guard::{insert,reserve}` and `Guard::{insert,reserve}_limit`, which
> are akin to `__xa_{alloc,insert}` in C.

Who will be using this? i.e. we need to justify adding code, typically
by mentioning the users.

By the way, it may help splitting the C parts into its own patch, so
that Matthew can focus on that (but maybe you all already discussed
this).

Also, unit tests are great, thanks for adding those. Can we add
examples as well? Maybe on the module-level docs, given the rest are
there.

> +    /// Wrapper around `__xa_alloc`.

We try to mention what something does in the title at least, even if
it is just copied from the C docs -- you can reference those docs too
for more context, e.g. [`__xa_alloc`] to:

    https://docs.kernel.org/core-api/xarray.html#c.__xa_alloc

Cheers,
Miguel

Re: [PATCH 1/3] rust: xarray: use the prelude

[Boqun Feng]

On Tue, Jul 01, 2025 at 12:36:20PM -0400, Tamir Duberstein wrote:
> On Tue, Jul 1, 2025 at 12:35 PM Boqun Feng <boqun.feng@gmail.com> wrote:
> >
> > On Tue, Jul 01, 2025 at 12:27:17PM -0400, Tamir Duberstein wrote:
> > > Using the prelude is customary in the kernel crate.
> > >
> > > Signed-off-by: Tamir Duberstein <tamird@gmail.com>
> > > ---
> > >  rust/kernel/xarray.rs | 34 ++++++++++++++++++++--------------
> > >  1 file changed, 20 insertions(+), 14 deletions(-)
> > >
> > > diff --git a/rust/kernel/xarray.rs b/rust/kernel/xarray.rs
> > > index 75719e7bb491..436faad99c89 100644
> > > --- a/rust/kernel/xarray.rs
> > > +++ b/rust/kernel/xarray.rs
> > > @@ -5,16 +5,15 @@
> > >  //! C header: [`include/linux/xarray.h`](srctree/include/linux/xarray.h)
> > >
> > >  use crate::{
> > > -    alloc, bindings, build_assert,
> > > -    error::{Error, Result},
> > > +    alloc,
> > > +    prelude::*,
> > >      types::{ForeignOwnable, NotThreadSafe, Opaque},
> > >  };
> > > -use core::{iter, marker::PhantomData, mem, pin::Pin, ptr::NonNull};
> > > -use pin_init::{pin_data, pin_init, pinned_drop, PinInit};
> > > +use core::{iter, marker::PhantomData, mem, ptr::NonNull};
> > >
> > >  /// An array which efficiently maps sparse integer indices to owned objects.
> > >  ///
> > > -/// This is similar to a [`crate::alloc::kvec::Vec<Option<T>>`], but more efficient when there are
> > > +/// This is similar to a [`Vec<Option<T>>`], but more efficient when there are
> > >  /// holes in the index space, and can be efficiently grown.
> > >  ///
> > >  /// # Invariants
> > > @@ -104,16 +103,23 @@ pub fn new(kind: AllocKind) -> impl PinInit<Self> {
> > >      fn iter(&self) -> impl Iterator<Item = NonNull<T::PointedTo>> + '_ {
> > >          let mut index = 0;
> > >
> > > -        // SAFETY: `self.xa` is always valid by the type invariant.
> > > -        iter::once(unsafe {
> > > -            bindings::xa_find(self.xa.get(), &mut index, usize::MAX, bindings::XA_PRESENT)
> > > -        })
> > > -        .chain(iter::from_fn(move || {
> > > +        core::iter::Iterator::chain(
> >
> > Does this part come from using the prelude? If not, either we need to
> > split the patch or we need to mention it in the changelog at least.
> 
> Yes, it's from using the prelude - PinInit also has a chain method
> that causes ambiguity here.

Maybe you can mention this in the change log as well. Like "Calling
iter::chain() with the associated function style to disambiguate with
PinInit::chain()". Make it easier to review (and remember) why.

Regards,
Boqun

> 
> > Also since we `use core::iter` above, we can avoid the `core::` here.
> 
> Good point.
> 

Re: [PATCH 3/3] rust: xarray: add `insert` and `reserve`

[Tamir Duberstein]

On Tue, Jul 1, 2025 at 12:56 PM Miguel Ojeda
<miguel.ojeda.sandonis@gmail.com> wrote:
>
> On Tue, Jul 1, 2025 at 6:27 PM Tamir Duberstein <tamird@gmail.com> wrote:
> >
> > Add `Guard::{insert,reserve}` and `Guard::{insert,reserve}_limit`, which
> > are akin to `__xa_{alloc,insert}` in C.
>
> Who will be using this? i.e. we need to justify adding code, typically
> by mentioning the users.

Daniel, could you name your use case?

Alice, could you confirm whether rust binder needs this or not?

Andreas, did you also need this?

> By the way, it may help splitting the C parts into its own patch, so
> that Matthew can focus on that (but maybe you all already discussed
> this).

Happy to do so if it makes Matthew's life easier, but I'd prefer to
keep it together so the motivation is clearer in the git log. Matthew:
any preference?

> Also, unit tests are great, thanks for adding those. Can we add
> examples as well? Maybe on the module-level docs, given the rest are
> there.

Will do. I ported all the examples to unit tests now that we have
kunit, but I'll bring some examples back.

>
> > +    /// Wrapper around `__xa_alloc`.
>
> We try to mention what something does in the title at least, even if
> it is just copied from the C docs -- you can reference those docs too
> for more context, e.g. [`__xa_alloc`] to:
>
>     https://docs.kernel.org/core-api/xarray.html#c.__xa_alloc

👍

Re: [PATCH 1/3] rust: xarray: use the prelude

[Tamir Duberstein]

On Tue, Jul 1, 2025 at 1:02 PM Boqun Feng <boqun.feng@gmail.com> wrote:
>
> On Tue, Jul 01, 2025 at 12:36:20PM -0400, Tamir Duberstein wrote:
> > On Tue, Jul 1, 2025 at 12:35 PM Boqun Feng <boqun.feng@gmail.com> wrote:
> > >
> > > On Tue, Jul 01, 2025 at 12:27:17PM -0400, Tamir Duberstein wrote:
> > > > Using the prelude is customary in the kernel crate.
> > > >
> > > > Signed-off-by: Tamir Duberstein <tamird@gmail.com>
> > > > ---
> > > >  rust/kernel/xarray.rs | 34 ++++++++++++++++++++--------------
> > > >  1 file changed, 20 insertions(+), 14 deletions(-)
> > > >
> > > > diff --git a/rust/kernel/xarray.rs b/rust/kernel/xarray.rs
> > > > index 75719e7bb491..436faad99c89 100644
> > > > --- a/rust/kernel/xarray.rs
> > > > +++ b/rust/kernel/xarray.rs
> > > > @@ -5,16 +5,15 @@
> > > >  //! C header: [`include/linux/xarray.h`](srctree/include/linux/xarray.h)
> > > >
> > > >  use crate::{
> > > > -    alloc, bindings, build_assert,
> > > > -    error::{Error, Result},
> > > > +    alloc,
> > > > +    prelude::*,
> > > >      types::{ForeignOwnable, NotThreadSafe, Opaque},
> > > >  };
> > > > -use core::{iter, marker::PhantomData, mem, pin::Pin, ptr::NonNull};
> > > > -use pin_init::{pin_data, pin_init, pinned_drop, PinInit};
> > > > +use core::{iter, marker::PhantomData, mem, ptr::NonNull};
> > > >
> > > >  /// An array which efficiently maps sparse integer indices to owned objects.
> > > >  ///
> > > > -/// This is similar to a [`crate::alloc::kvec::Vec<Option<T>>`], but more efficient when there are
> > > > +/// This is similar to a [`Vec<Option<T>>`], but more efficient when there are
> > > >  /// holes in the index space, and can be efficiently grown.
> > > >  ///
> > > >  /// # Invariants
> > > > @@ -104,16 +103,23 @@ pub fn new(kind: AllocKind) -> impl PinInit<Self> {
> > > >      fn iter(&self) -> impl Iterator<Item = NonNull<T::PointedTo>> + '_ {
> > > >          let mut index = 0;
> > > >
> > > > -        // SAFETY: `self.xa` is always valid by the type invariant.
> > > > -        iter::once(unsafe {
> > > > -            bindings::xa_find(self.xa.get(), &mut index, usize::MAX, bindings::XA_PRESENT)
> > > > -        })
> > > > -        .chain(iter::from_fn(move || {
> > > > +        core::iter::Iterator::chain(
> > >
> > > Does this part come from using the prelude? If not, either we need to
> > > split the patch or we need to mention it in the changelog at least.
> >
> > Yes, it's from using the prelude - PinInit also has a chain method
> > that causes ambiguity here.
>
> Maybe you can mention this in the change log as well. Like "Calling
> iter::chain() with the associated function style to disambiguate with
> PinInit::chain()". Make it easier to review (and remember) why.

Yep, done for the next spin.

Re: [PATCH 3/3] rust: xarray: add `insert` and `reserve`

[Daniel Almeida]

Hi, this is on my todo list.

Again, thanks Tamir for working on this.

> On 1 Jul 2025, at 14:04, Tamir Duberstein <tamird@gmail.com> wrote:
> 
> On Tue, Jul 1, 2025 at 12:56 PM Miguel Ojeda
> <miguel.ojeda.sandonis@gmail.com> wrote:
>> 
>> On Tue, Jul 1, 2025 at 6:27 PM Tamir Duberstein <tamird@gmail.com> wrote:
>>> 
>>> Add `Guard::{insert,reserve}` and `Guard::{insert,reserve}_limit`, which
>>> are akin to `__xa_{alloc,insert}` in C.
>> 
>> Who will be using this? i.e. we need to justify adding code, typically
>> by mentioning the users.
> 
> Daniel, could you name your use case?

My main use case is for insert_limit.

Tyr uses xarrays to keep track of resources allocated by userspace, and we need
to impose limits on how many resources can be allocated at once.

insert_limits() provides the exact semantics needed, i.e.: insert somewhere in
the array as long as there is vacant space in the given range.

— Daniel

Re: [PATCH 3/3] rust: xarray: add `insert` and `reserve`

[Janne Grunau]

On Tue, Jul 01, 2025 at 06:56:17PM +0200, Miguel Ojeda wrote:
> On Tue, Jul 1, 2025 at 6:27 PM Tamir Duberstein <tamird@gmail.com> wrote:
> >
> > Add `Guard::{insert,reserve}` and `Guard::{insert,reserve}_limit`, which
> > are akin to `__xa_{alloc,insert}` in C.
> 
> Who will be using this? i.e. we need to justify adding code, typically
> by mentioning the users.

xa_alloc() / reserve() is used by asahi. It's still using our own
abstraction but I'm in the progress of rebase onto the upstream xarray
abstractions from v6.16-rc1. Once I'm done I'll reply with "Tested-by:".

Janne

Re: [PATCH 3/3] rust: xarray: add `insert` and `reserve`

[Alice Ryhl]

On Tue, Jul 1, 2025 at 6:56 PM Miguel Ojeda
<miguel.ojeda.sandonis@gmail.com> wrote:
>
> On Tue, Jul 1, 2025 at 6:27 PM Tamir Duberstein <tamird@gmail.com> wrote:
> >
> > Add `Guard::{insert,reserve}` and `Guard::{insert,reserve}_limit`, which
> > are akin to `__xa_{alloc,insert}` in C.
>
> Who will be using this? i.e. we need to justify adding code, typically
> by mentioning the users.

Rust Binder does have a user for `reserve`. As for xa_alloc, Rust
Binder may end up using it, but the current plan is to use Burak's
bitmap work instead.

Alice

Re: [PATCH 3/3] rust: xarray: add `insert` and `reserve`

[Tamir Duberstein]

On Sun, Jul 6, 2025 at 4:29 AM Janne Grunau <j@jannau.net> wrote:
>
> On Tue, Jul 01, 2025 at 06:56:17PM +0200, Miguel Ojeda wrote:
> > On Tue, Jul 1, 2025 at 6:27 PM Tamir Duberstein <tamird@gmail.com> wrote:
> > >
> > > Add `Guard::{insert,reserve}` and `Guard::{insert,reserve}_limit`, which
> > > are akin to `__xa_{alloc,insert}` in C.
> >
> > Who will be using this? i.e. we need to justify adding code, typically
> > by mentioning the users.
>
> xa_alloc() / reserve() is used by asahi. It's still using our own
> abstraction but I'm in the progress of rebase onto the upstream xarray
> abstractions from v6.16-rc1. Once I'm done I'll reply with "Tested-by:".
>
> Janne

Thanks Janne. I'll wait for your Tested-by before sending v2.
Tamir

Re: [PATCH 0/3] rust: xarray: add `insert` and `reserve`

[Janne Grunau]

On Tue, Jul 01, 2025 at 12:27:16PM -0400, Tamir Duberstein wrote:
> Please see individual patches.
> 
> Signed-off-by: Tamir Duberstein <tamird@gmail.com>
> ---
> Tamir Duberstein (3):
>       rust: xarray: use the prelude
>       rust: xarray: implement Default for AllocKind
>       rust: xarray: add `insert` and `reserve`
> 
>  include/linux/xarray.h |   2 +
>  lib/xarray.c           |  28 ++-
>  rust/helpers/xarray.c  |   5 +
>  rust/kernel/xarray.rs  | 460 +++++++++++++++++++++++++++++++++++++++++++++++--
>  4 files changed, 472 insertions(+), 23 deletions(-)

thanks, series is tested with the asahi driver and works as expected.
Usage is limited to ::reserve_limits() and ::fill() of the reservation
so only covering a part of the change.

Whole series
Tested-by: Janne Grunau <j@jannau.net>
Reviewed-by: Janne Grunau <j@jannau.net>

Janne