Re: [PATCH v2 3/5] rust: maple_tree: add MapleTree::lock() and load()

[Daniel Almeida <daniel.almeida@collabora.com>]

Hi Alice,

Overall LGTM, a few comments below,

> On 19 Aug 2025, at 07:34, Alice Ryhl <aliceryhl@google.com> wrote:
> 
> To load a value, one must be careful to hold the lock while accessing
> it. To enable this, we add a lock() method so that you can perform
> operations on the value before the spinlock is released.
> 
> This adds a MapleGuard type without using the existing SpinLock type.
> This ensures that the MapleGuard type is not unnecessarily large, and
> that it is easy to swap out the type of lock in case the C maple tree is
> changed to use a different kind of lock.
> 
> Co-developed-by: Andrew Ballance <andrewjballance@gmail.com>
> Signed-off-by: Andrew Ballance <andrewjballance@gmail.com>
> Reviewed-by: Andrew Ballance <andrewjballance@gmail.com>
> Signed-off-by: Alice Ryhl <aliceryhl@google.com>
> ---
> rust/kernel/maple_tree.rs | 139 ++++++++++++++++++++++++++++++++++++++++++++++
> 1 file changed, 139 insertions(+)
> 
> diff --git a/rust/kernel/maple_tree.rs b/rust/kernel/maple_tree.rs
> index ea1bd694213b73108732aecc36da95342aeafe04..17e4d8586ebad56aee87a97befdfec5741f147de 100644
> --- a/rust/kernel/maple_tree.rs
> +++ b/rust/kernel/maple_tree.rs
> @@ -220,6 +220,22 @@ pub fn erase(&self, index: usize) -> Option<T> {
>         unsafe { T::try_from_foreign(ret) }
>     }
> 
> +    /// Lock the internal spinlock.
> +    #[inline]
> +    pub fn lock(&self) -> MapleGuard<'_, T> {
> +        // SAFETY: It's safe to lock the spinlock in a maple tree.
> +        unsafe { bindings::spin_lock(self.ma_lock()) };
> +
> +        // INVARIANT: We just took the spinlock.
> +        MapleGuard(self)
> +    }
> +
> +    #[inline]
> +    fn ma_lock(&self) -> *mut bindings::spinlock_t {
> +        // SAFETY: This pointer offset operation stays in-bounds.
> +        unsafe { &raw mut (*self.tree.get()).__bindgen_anon_1.ma_lock }
> +    }
> +
>     /// Free all `T` instances in this tree.
>     ///
>     /// # Safety
> @@ -263,6 +279,91 @@ fn drop(mut self: Pin<&mut Self>) {
>     }
> }
> 
> +/// A reference to a [`MapleTree`] that owns the inner lock.
> +///
> +/// # Invariants
> +///
> +/// This guard owns the inner spinlock.
> +#[must_use = "if unused, the lock will be immediately unlocked"]
> +pub struct MapleGuard<'tree, T: ForeignOwnable>(&'tree MapleTree<T>);
> +
> +impl<'tree, T: ForeignOwnable> Drop for MapleGuard<'tree, T> {
> +    #[inline]
> +    fn drop(&mut self) {
> +        // SAFETY: By the type invariants, we hold this spinlock.
> +        unsafe { bindings::spin_unlock(self.0.ma_lock()) };
> +    }
> +}
> +
> +impl<'tree, T: ForeignOwnable> MapleGuard<'tree, T> {
> +    /// Create a [`MaState`] protected by this lock guard.
> +    pub fn ma_state(&mut self, first: usize, end: usize) -> MaState<'_, T> {
> +        // SAFETY: The `MaState` borrows this `MapleGuard`, so it can also borrow the `MapleGuard`s
> +        // read/write permissions to the maple tree.
> +        unsafe { MaState::new_raw(self.0, first, end) }
> +    }
> +
> +    /// Load the value at the given index.
> +    ///
> +    /// # Examples
> +    ///
> +    /// Read the value while holding the spinlock.
> +    ///
> +    /// ```
> +    /// use kernel::maple_tree::{MapleTree, InsertErrorKind};
> +    ///
> +    /// let tree = KBox::pin_init(MapleTree::<KBox<i32>>::new(), GFP_KERNEL)?;
> +    ///
> +    /// let ten = KBox::new(10, GFP_KERNEL)?;
> +    /// let twenty = KBox::new(20, GFP_KERNEL)?;
> +    /// tree.insert(100, ten, GFP_KERNEL)?;
> +    /// tree.insert(200, twenty, GFP_KERNEL)?;
> +    ///
> +    /// let mut lock = tree.lock();
> +    /// assert_eq!(lock.load(100), Some(&mut 10));
> +    /// assert_eq!(lock.load(200), Some(&mut 20));
> +    /// assert_eq!(lock.load(300), None);
> +    /// # Ok::<_, Error>(())
> +    /// ```
> +    ///
> +    /// Increment refcount while holding spinlock and read afterwards.
> +    ///
> +    /// ```
> +    /// use kernel::maple_tree::{MapleTree, InsertErrorKind};
> +    /// use kernel::sync::Arc;
> +    ///
> +    /// let tree = KBox::pin_init(MapleTree::<Arc<i32>>::new(), GFP_KERNEL)?;
> +    ///
> +    /// let ten = Arc::new(10, GFP_KERNEL)?;
> +    /// let twenty = Arc::new(20, GFP_KERNEL)?;
> +    /// tree.insert(100, ten, GFP_KERNEL)?;
> +    /// tree.insert(200, twenty, GFP_KERNEL)?;
> +    ///
> +    /// // Briefly take the lock to increment the refcount.
> +    /// let value = Arc::from(tree.lock().load(100).unwrap());
> +    ///
> +    /// // At this point, another thread might remove the value.
> +    /// tree.erase(100);
> +    ///
> +    /// // But we can still access it because we took a refcount.
> +    /// assert_eq!(*value, 10);
> +    /// # Ok::<_, Error>(())
> +    /// ```
> +    #[inline]
> +    pub fn load(&mut self, index: usize) -> Option<T::BorrowedMut<'_>> {
> +        // SAFETY: `self.tree` contains a valid maple tree.
> +        let ret = unsafe { bindings::mtree_load(self.0.tree.get(), index) };
> +        if ret.is_null() {
> +            return None;
> +        }
> +
> +        // SAFETY: If the pointer is not null, then it references a valid instance of `T`. It is
> +        // safe to borrow the instance mutably because the signature of this function enforces that
> +        // the mutable borrow is not used after the spinlock is dropped.
> +        Some(unsafe { T::borrow_mut(ret) })
> +    }
> +}
> +
> impl<'tree, T: ForeignOwnable> MaState<'tree, T> {
>     /// Initialize a new `MaState` with the given tree.
>     ///
> @@ -303,6 +404,44 @@ fn mas_find_raw(&mut self, max: usize) -> *mut c_void {
>         // to the tree.
>         unsafe { bindings::mas_find(self.as_raw(), max) }
>     }
> +
> +    /// Find the next entry in the maple tree.

This is not in the commit title.

> +    ///
> +    /// # Examples
> +    ///
> +    /// Iterate the maple tree.
> +    ///
> +    /// ```
> +    /// use kernel::maple_tree::{MapleTree, InsertErrorKind};
> +    /// use kernel::sync::Arc;
> +    ///
> +    /// let tree = KBox::pin_init(MapleTree::<Arc<i32>>::new(), GFP_KERNEL)?;
> +    ///
> +    /// let ten = Arc::new(10, GFP_KERNEL)?;
> +    /// let twenty = Arc::new(20, GFP_KERNEL)?;
> +    /// tree.insert(100, ten, GFP_KERNEL)?;
> +    /// tree.insert(200, twenty, GFP_KERNEL)?;
> +    ///
> +    /// let mut ma_lock = tree.lock();
> +    /// let mut iter = ma_lock.ma_state(0, usize::MAX);
> +    ///
> +    /// assert_eq!(*iter.mas_find(usize::MAX).unwrap(), 10);
> +    /// assert_eq!(*iter.mas_find(usize::MAX).unwrap(), 20);
> +    /// assert!(iter.mas_find(usize::MAX).is_none());
> +    /// # Ok::<_, Error>(())
> +    /// ```
> +    #[inline]
> +    pub fn mas_find(&mut self, max: usize) -> Option<T::BorrowedMut<'_>> {

Should we drop the “mas” prefix here? I think that “find()” is fine.

> +        let ret = self.mas_find_raw(max);
> +        if ret.is_null() {
> +            return None;
> +        }
> +
> +        // SAFETY: If the pointer is not null, then it references a valid instance of `T`. It's
> +        // safe to access it mutably as the returned reference borrows this `MaState`, and the
> +        // `MaState` has read/write access to the maple tree.
> +        Some(unsafe { T::borrow_mut(ret) })
> +    }
> }
> 
> /// Error type for failure to insert a new value.
> 
> -- 
> 2.51.0.rc1.167.g924127e9c0-goog
> 
>