Re: [PATCH v17 2/3] rust: xarray: Add an abstraction for XArray

[Danilo Krummrich <dakr@kernel.org>]

On Tue, Feb 18, 2025 at 09:37:44AM -0500, Tamir Duberstein wrote:
> +/// 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
> +/// holes in the index space, and can be efficiently grown.
> +///
> +/// # Invariants
> +///
> +/// `self.xa` is always an initialized and valid [`bindings::xarray`] whose entries are either
> +/// `XA_ZERO_ENTRY` or came from `T::into_foreign`.
> +///
> +/// # Examples
> +///
> +/// ```rust
> +/// use kernel::alloc::KBox;
> +/// use kernel::xarray::{AllocKind, XArray};
> +///
> +/// let xa = KBox::pin_init(XArray::new(AllocKind::Alloc1), GFP_KERNEL)?;
> +///
> +/// let dead = KBox::new(0xdead, GFP_KERNEL)?;
> +/// let beef = KBox::new(0xbeef, GFP_KERNEL)?;
> +///
> +/// let mut guard = xa.lock();
> +///
> +/// assert_eq!(guard.get(0), None);
> +///
> +/// assert_eq!(guard.store(0, dead, GFP_KERNEL)?.as_deref(), None);
> +/// assert_eq!(guard.get(0).copied(), Some(0xdead));
> +///
> +/// *guard.get_mut(0).unwrap() = 0xffff;
> +/// assert_eq!(guard.get(0).copied(), Some(0xffff));
> +///
> +/// assert_eq!(guard.store(0, beef, GFP_KERNEL)?.as_deref().copied(), Some(0xffff));
> +/// assert_eq!(guard.get(0).copied(), Some(0xbeef));
> +///
> +/// guard.remove(0);
> +/// assert_eq!(guard.get(0), None);
> +///
> +/// # Ok::<(), Error>(())
> +/// ```
> +#[pin_data(PinnedDrop)]
> +pub struct XArray<T: ForeignOwnable> {
> +    #[pin]
> +    xa: Opaque<bindings::xarray>,
> +    _p: PhantomData<T>,
> +}
> +

[...]

> +
> +impl<T: ForeignOwnable> XArray<T> {
> +    /// Creates a new [`XArray`].
> +    pub fn new(kind: AllocKind) -> impl PinInit<Self> {
> +        let flags = match kind {
> +            AllocKind::Alloc => bindings::XA_FLAGS_ALLOC,
> +            AllocKind::Alloc1 => bindings::XA_FLAGS_ALLOC1,
> +        };
> +        pin_init!(Self {
> +            // SAFETY: `xa` is valid while the closure is called.
> +            xa <- Opaque::ffi_init(|xa| unsafe {
> +                bindings::xa_init_flags(xa, flags)
> +            }),
> +            _p: PhantomData,
> +        })

I think this needs an `INVARIANT` comment.

[...]

> +/// The error returned by [`store`](Guard::store).
> +///
> +/// Contains the underlying error and the value that was not stored.
> +pub struct StoreError<T> {
> +    /// The error that occurred.
> +    pub error: Error,
> +    /// The value that was not stored.
> +    pub value: T,
> +}
> +
> +impl<T> From<StoreError<T>> for Error {
> +    fn from(value: StoreError<T>) -> Self {
> +        let StoreError { error, value: _ } = value;
> +        error
> +    }

Still think this should just be `value.error`.

If it is important to especially point out that `value` is dropped, maybe a
comment is the better option.

IMHO, adding additionally code here just throws up questions on why that
additional code is needed.

> +}
> +
> +impl<'a, T: ForeignOwnable> Guard<'a, T> {
> +    fn load<F, U>(&self, index: usize, f: F) -> Option<U>
> +    where
> +        F: FnOnce(NonNull<T::PointedTo>) -> U,
> +    {
> +        // SAFETY: `self.xa.xa` is always valid by the type invariant.
> +        let ptr = unsafe { bindings::xa_load(self.xa.xa.get(), index) };
> +        let ptr = NonNull::new(ptr.cast())?;
> +        Some(f(ptr))
> +    }
> +
> +    /// Provides a reference to the element at the given index.
> +    pub fn get(&self, index: usize) -> Option<T::Borrowed<'_>> {
> +        self.load(index, |ptr| {
> +            // SAFETY: `ptr` came from `T::into_foreign`.
> +            unsafe { T::borrow(ptr.as_ptr()) }
> +        })
> +    }
> +
> +    /// Provides a mutable reference to the element at the given index.
> +    pub fn get_mut(&mut self, index: usize) -> Option<T::BorrowedMut<'_>> {
> +        self.load(index, |ptr| {
> +            // SAFETY: `ptr` came from `T::into_foreign`.
> +            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: `self.xa.xa` is always valid by the type invariant.
> +        //
> +        // SAFETY: The caller holds the lock.

I think we only want one `SAFETY` section with an enumeration.

> +        let ptr = unsafe { bindings::__xa_erase(self.xa.xa.get(), index) }.cast();
> +        // SAFETY: `ptr` is either NULL or came from `T::into_foreign`.
> +        //
> +        // SAFETY: `&mut self` guarantees that the lifetimes of [`T::Borrowed`] and
> +        // [`T::BorrowedMut`] borrowed from `self` have ended.

Same here...

> +        unsafe { T::try_from_foreign(ptr) }
> +    }
> +
> +    /// Stores an element at the given index.
> +    ///
> +    /// May drop the lock if needed to allocate memory, and then reacquire it afterwards.
> +    ///
> +    /// On success, returns the element which was previously at the given index.
> +    ///
> +    /// On failure, returns the element which was attempted to be stored.
> +    pub fn store(
> +        &mut self,
> +        index: usize,
> +        value: T,
> +        gfp: alloc::Flags,
> +    ) -> Result<Option<T>, StoreError<T>> {
> +        build_assert!(
> +            mem::align_of::<T::PointedTo>() >= 4,
> +            "pointers stored in XArray must be 4-byte aligned"
> +        );
> +        let new = value.into_foreign();
> +
> +        let old = {
> +            let new = new.cast();
> +            // SAFETY: `self.xa.xa` is always valid by the type invariant.
> +            //
> +            // SAFETY: The caller holds the lock.

...and here.

> +            //
> +            // INVARIANT: `new` came from `T::into_foreign`.
> +            unsafe { bindings::__xa_store(self.xa.xa.get(), index, new, gfp.as_raw()) }
> +        };