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

["Liam R. Howlett" <Liam.Howlett@oracle.com>]

* Daniel Almeida <daniel.almeida@collabora.com> [250819 13:08]:
> 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.

The maple tree has two interfaces, the advanced one which starts with
mas_ and the simple on that uses mt_.  This is probably why the mas_ is
here?

> 
> > +        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
> > 
> > 
>