[PATCH 2/3] rust: maple_tree: add MapleTree::lock() and load()

[Andrew Ballance]

On 7/26/25 8:23 AM, Alice Ryhl 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.
> 
> Co-developed-by: Andrew Ballance <andrewjballance at gmail.com>
> Signed-off-by: Andrew Ballance <andrewjballance at gmail.com>
> Signed-off-by: Alice Ryhl <aliceryhl at google.com>

I have a couple of nits, but overall looks good to me.

> ---
>   rust/kernel/maple_tree.rs | 94 +++++++++++++++++++++++++++++++++++++++++++++++
>   1 file changed, 94 insertions(+)
> 
> diff --git a/rust/kernel/maple_tree.rs b/rust/kernel/maple_tree.rs
> index 0f26c173eedc7c79bb8e2b56fe85e8a266b3ae0c..c7ef504a9c78065b3d5752b4f5337fb6277182d1 100644
> --- a/rust/kernel/maple_tree.rs
> +++ b/rust/kernel/maple_tree.rs
> @@ -206,6 +206,23 @@ pub fn erase(&self, index: usize) -> Option<T> {
>           unsafe { T::try_from_foreign(ret) }
>       }
>   
> +    /// Lock the internal spinlock.

probably should add #[must_use] here.

> +    #[inline]
> +    pub fn lock(&self) -> MapleLock<'_, 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.
> +        MapleLock(self)
> +    }
> +
> +    #[inline]
> +    fn ma_lock(&self) -> *mut bindings::spinlock_t {
> +        // SAFETY: This pointer offset operation stays in-bounds.
> +        let lock = unsafe { &raw mut (*self.tree.get()).__bindgen_anon_1.ma_lock };
> +        lock.cast()

This cast seems unneeded. lock should already be a *mut spinlock_t.

> +    }
> +
>       /// Free all `T` instances in this tree.
>       ///
>       /// # Safety
> @@ -248,6 +265,83 @@ fn drop(mut self: Pin<&mut Self>) {
>       }
>   }
>   
> +/// A reference to a [`MapleTree`] that owns the inner lock.
> +///
> +/// # Invariants
> +///
> +/// This guard owns the inner spinlock.
> +pub struct MapleLock<'tree, T: ForeignOwnable>(&'tree MapleTree<T>);
> +
> +impl<'tree, T: ForeignOwnable> Drop for MapleLock<'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> MapleLock<'tree, T> {
> +    /// 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) })
> +    }
> +}
> +
>   /// Error type for failure to insert a new value.
>   pub struct InsertError<T> {
>       /// The value that could not be inserted.
> 

with or without those fixes, for the entire series,
Reviewed-by: Andrew Ballance <andrewjballance at gmail.com>

Also, if you need one, I would be happy to be a co-maintainer of the
rust maple tree bindings.

Best Regards,
Andrew Ballance