[PATCH v3 0/3] Add Rust abstraction for Maple Trees

[PATCH v3 0/3] Add Rust abstraction for Maple Trees

[Alice Ryhl]

This will be used in the Tyr driver [1] to allocate from the GPU's VA
space that is not owned by userspace, but by the kernel, for kernel GPU
mappings.

Danilo tells me that in nouveau, the maple tree is used for keeping
track of "VM regions" on top of GPUVM, and that he will most likely end
up doing the same in the Rust Nova driver as well.

These abstractions intentionally do not expose any way to make use of
external locking. You are required to use the internal spinlock. For
now, we do not support loads that only utilize rcu for protection.

This contains some parts taken from Andrew Ballance's RFC [2] from
April. However, it has also been reworked significantly compared to that
RFC taking the use-cases in Tyr into account.

[1]: https://lore.kernel.org/r/20250627-tyr-v1-1-cb5f4c6ced46@collabora.com
[2]: https://lore.kernel.org/r/20250405060154.1550858-1-andrewjballance@gmail.com

Signed-off-by: Alice Ryhl <aliceryhl@google.com>
---
Changes in v3:
- Change examples to avoid unwrap(), but not unwrap_err().
- Remove spurious MAINTAINERS change.
- Fix unnnecessary imports.
- Drop lockdep_map_p patch from this series.
- Reword MaState type invariant.
- Various docs improvements.
- Rename mas_find() to find().
- Add to existing MAINTAINERS entry.
- Link to v2: https://lore.kernel.org/r/20250819-maple-tree-v2-0-229b48657bab@google.com

Changes in v2:
- Add MaState abstraction. For now it only has a mas_find method. And
  use it in the destructor.
- Duplicate MA_STATE macro in Rust instead of using a C helper.
- Change maple_tree.h so that cast in ma_lock() is no longer needed.
- Add #[must_use] and #[inline] annotations.
- Rename MapleLock to MapleGuard.
- Change errors to use AllocError.
- Add MAINTAINERS file.
- Link to v1: https://lore.kernel.org/r/20250726-maple-tree-v1-0-27a3da7cb8e5@google.com

---
Alice Ryhl (3):
      rust: maple_tree: add MapleTree
      rust: maple_tree: add lock guard for maple tree
      rust: maple_tree: add MapleTreeAlloc

 MAINTAINERS                |   4 +
 include/linux/maple_tree.h |   3 +
 rust/helpers/helpers.c     |   1 +
 rust/helpers/maple_tree.c  |   8 +
 rust/kernel/lib.rs         |   1 +
 rust/kernel/maple_tree.rs  | 648 +++++++++++++++++++++++++++++++++++++++++++++
 6 files changed, 665 insertions(+)
---
base-commit: 1b237f190eb3d36f52dffe07a40b5eb210280e00
change-id: 20250726-maple-tree-1af0803ac524

Best regards,
-- 
Alice Ryhl <aliceryhl@google.com>

[PATCH v3 1/3] rust: maple_tree: add MapleTree

[Alice Ryhl]

The maple tree will be used in the Tyr driver to allocate and keep track
of GPU allocations created internally (i.e. not by userspace). It will
likely also be used in the Nova driver eventually.

This adds the simplest methods for additional and removal that do not
require any special care with respect to concurrency.

This implementation is based on the RFC by Andrew but with significant
changes to simplify the implementation.

Co-developed-by: Andrew Ballance <andrewjballance@gmail.com>
Signed-off-by: Andrew Ballance <andrewjballance@gmail.com>
Signed-off-by: Alice Ryhl <aliceryhl@google.com>
---
 MAINTAINERS                |   4 +
 include/linux/maple_tree.h |   3 +
 rust/helpers/helpers.c     |   1 +
 rust/helpers/maple_tree.c  |   8 ++
 rust/kernel/lib.rs         |   1 +
 rust/kernel/maple_tree.rs  | 350 +++++++++++++++++++++++++++++++++++++++++++++
 6 files changed, 367 insertions(+)

diff --git a/MAINTAINERS b/MAINTAINERS
index fed6cd812d796a08cebc0c1fd540c8901d1bf448..c076f034562dd3d6b3679e8c2cd390adb312d483 100644
--- a/MAINTAINERS
+++ b/MAINTAINERS
@@ -14672,6 +14672,8 @@ F:	net/mctp/
 
 MAPLE TREE
 M:	Liam R. Howlett <Liam.Howlett@oracle.com>
+R:	Alice Ryhl <aliceryhl@google.com>
+R:	Andrew Ballance <andrewjballance@gmail.com>
 L:	maple-tree@lists.infradead.org
 L:	linux-mm@kvack.org
 S:	Supported
@@ -14680,6 +14682,8 @@ F:	include/linux/maple_tree.h
 F:	include/trace/events/maple_tree.h
 F:	lib/maple_tree.c
 F:	lib/test_maple_tree.c
+F:	rust/helpers/maple_tree.c
+F:	rust/kernel/maple_tree.rs
 F:	tools/testing/radix-tree/maple.c
 F:	tools/testing/shared/linux/maple_tree.h
 
diff --git a/include/linux/maple_tree.h b/include/linux/maple_tree.h
index bafe143b1f783202e27b32567fffee4149e8e266..879cd5f24be49548d5fd9d6b19bf069e33c334a1 100644
--- a/include/linux/maple_tree.h
+++ b/include/linux/maple_tree.h
@@ -481,6 +481,9 @@ struct ma_wr_state {
 #define MA_ERROR(err) \
 		((struct maple_enode *)(((unsigned long)err << 2) | 2UL))
 
+/*
+ * When changing MA_STATE, remember to also change rust/kernel/maple_tree.rs
+ */
 #define MA_STATE(name, mt, first, end)					\
 	struct ma_state name = {					\
 		.tree = mt,						\
diff --git a/rust/helpers/helpers.c b/rust/helpers/helpers.c
index 7cf7fe95e41dd51717050648d6160bebebdf4b26..c5d42e0f7ce6786c1e96f8c0f27710959ca3362b 100644
--- a/rust/helpers/helpers.c
+++ b/rust/helpers/helpers.c
@@ -26,6 +26,7 @@
 #include "io.c"
 #include "jump_label.c"
 #include "kunit.c"
+#include "maple_tree.c"
 #include "mm.c"
 #include "mutex.c"
 #include "of.c"
diff --git a/rust/helpers/maple_tree.c b/rust/helpers/maple_tree.c
new file mode 100644
index 0000000000000000000000000000000000000000..1dd9ac84a13feed53c0ed5eec6805517081d0673
--- /dev/null
+++ b/rust/helpers/maple_tree.c
@@ -0,0 +1,8 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include <linux/maple_tree.h>
+
+void rust_helper_mt_init_flags(struct maple_tree *mt, unsigned int flags)
+{
+	mt_init_flags(mt, flags);
+}
diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index ed53169e795c0badf548025a57f946fa18bc73e3..6b0a5689669fa691e366ab3f9d462692c12bd548 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -96,6 +96,7 @@
 #[cfg(CONFIG_KUNIT)]
 pub mod kunit;
 pub mod list;
+pub mod maple_tree;
 pub mod miscdevice;
 pub mod mm;
 #[cfg(CONFIG_NET)]
diff --git a/rust/kernel/maple_tree.rs b/rust/kernel/maple_tree.rs
new file mode 100644
index 0000000000000000000000000000000000000000..92930b408e9683b6615080a53800f1a393c5f31e
--- /dev/null
+++ b/rust/kernel/maple_tree.rs
@@ -0,0 +1,350 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! Maple trees.
+//!
+//! C header: [`include/linux/maple_tree.h`](srctree/include/linux/maple_tree.h)
+//!
+//! Reference: <https://docs.kernel.org/core-api/maple_tree.html>
+
+use core::{
+    marker::PhantomData,
+    ops::{Bound, RangeBounds},
+    ptr,
+};
+
+use kernel::{
+    alloc::Flags,
+    error::to_result,
+    prelude::*,
+    types::{ForeignOwnable, Opaque},
+};
+
+/// A maple tree optimized for storing non-overlapping ranges.
+///
+/// # Invariants
+///
+/// Each range in the maple tree owns an instance of `T`.
+#[pin_data(PinnedDrop)]
+#[repr(transparent)]
+pub struct MapleTree<T: ForeignOwnable> {
+    #[pin]
+    tree: Opaque<bindings::maple_tree>,
+    _p: PhantomData<T>,
+}
+
+#[inline]
+fn to_maple_range(range: impl RangeBounds<usize>) -> Option<(usize, usize)> {
+    let first = match range.start_bound() {
+        Bound::Included(start) => *start,
+        Bound::Excluded(start) => start.checked_add(1)?,
+        Bound::Unbounded => 0,
+    };
+
+    let last = match range.end_bound() {
+        Bound::Included(end) => *end,
+        Bound::Excluded(end) => end.checked_sub(1)?,
+        Bound::Unbounded => usize::MAX,
+    };
+
+    if last < first {
+        return None;
+    }
+
+    Some((first, last))
+}
+
+impl<T: ForeignOwnable> MapleTree<T> {
+    /// Create a new maple tree.
+    ///
+    /// The tree will use the regular implementation with a higher branching factor, rather than
+    /// the allocation tree.
+    #[inline]
+    pub fn new() -> impl PinInit<Self> {
+        pin_init!(MapleTree {
+            // SAFETY: This initializes a maple tree into a pinned slot. The maple tree will be
+            // destroyed in Drop before the memory location becomes invalid.
+            tree <- Opaque::ffi_init(|slot| unsafe { bindings::mt_init_flags(slot, 0) }),
+            _p: PhantomData,
+        })
+    }
+
+    /// Insert the value at the given index.
+    ///
+    /// # Errors
+    ///
+    /// If the maple tree already contains a range using the given index, then this call will
+    /// return an [`InsertError`] with the [`Occupied`] kind. It may also fail if memory
+    /// allocation fails.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use kernel::maple_tree::{InsertErrorKind, MapleTree};
+    ///
+    /// 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)?;
+    /// let the_answer = KBox::new(42, GFP_KERNEL)?;
+    ///
+    /// // These calls will succeed.
+    /// tree.insert(100, ten, GFP_KERNEL)?;
+    /// tree.insert(101, twenty, GFP_KERNEL)?;
+    ///
+    /// // This will fail because the index is already in use.
+    /// assert_eq!(
+    ///     tree.insert(100, the_answer, GFP_KERNEL).unwrap_err().cause,
+    ///     InsertErrorKind::Occupied,
+    /// );
+    /// # Ok::<_, Error>(())
+    /// ```
+    #[inline]
+    pub fn insert(&self, index: usize, value: T, gfp: Flags) -> Result<(), InsertError<T>> {
+        self.insert_range(index..=index, value, gfp)
+    }
+
+    /// Insert a value to the specified range, failing on overlap.
+    ///
+    /// This accepts the usual types of Rust ranges using the `..` and `..=` syntax for exclusive
+    /// and inclusive ranges respectively. The range must not be empty, and must not overlap with
+    /// any existing range.
+    ///
+    /// # Errors
+    ///
+    /// If the maple tree already contains an overlapping range, then this call will return an
+    /// [`InsertError`] with the [`Occupied`] kind. It may also fail if memory allocation fails
+    /// or if the requested range is invalid (e.g. empty).
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use kernel::maple_tree::{InsertErrorKind, MapleTree};
+    ///
+    /// 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)?;
+    /// let the_answer = KBox::new(42, GFP_KERNEL)?;
+    /// let hundred = KBox::new(100, GFP_KERNEL)?;
+    ///
+    /// // Insert the value 10 at the indices 100 to 499.
+    /// tree.insert_range(100..500, ten, GFP_KERNEL)?;
+    ///
+    /// // Insert the value 20 at the indices 500 to 1000.
+    /// tree.insert_range(500..=1000, twenty, GFP_KERNEL)?;
+    ///
+    /// // This will fail due to overlap with the previous range on index 1000.
+    /// assert_eq!(
+    ///     tree.insert_range(1000..1200, the_answer, GFP_KERNEL).unwrap_err().cause,
+    ///     InsertErrorKind::Occupied,
+    /// );
+    ///
+    /// // When using .. to specify the range, you must be careful to ensure that the range is
+    /// // non-empty.
+    /// assert_eq!(
+    ///     tree.insert_range(72..72, hundred, GFP_KERNEL).unwrap_err().cause,
+    ///     InsertErrorKind::InvalidRequest,
+    /// );
+    /// # Ok::<_, Error>(())
+    /// ```
+    pub fn insert_range<R>(&self, range: R, value: T, gfp: Flags) -> Result<(), InsertError<T>>
+    where
+        R: RangeBounds<usize>,
+    {
+        let Some((first, last)) = to_maple_range(range) else {
+            return Err(InsertError {
+                value,
+                cause: InsertErrorKind::InvalidRequest,
+            });
+        };
+
+        let ptr = T::into_foreign(value);
+
+        // SAFETY: The tree is valid, and we are passing a pointer to an owned instance of `T`.
+        let res = to_result(unsafe {
+            bindings::mtree_insert_range(self.tree.get(), first, last, ptr, gfp.as_raw())
+        });
+
+        if let Err(err) = res {
+            // SAFETY: As `mtree_insert_range` failed, it is safe to take back ownership.
+            let value = unsafe { T::from_foreign(ptr) };
+
+            let cause = if err == ENOMEM {
+                InsertErrorKind::AllocError(kernel::alloc::AllocError)
+            } else if err == EEXIST {
+                InsertErrorKind::Occupied
+            } else {
+                InsertErrorKind::InvalidRequest
+            };
+            Err(InsertError { value, cause })
+        } else {
+            Ok(())
+        }
+    }
+
+    /// Erase the range containing the given index.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use kernel::maple_tree::MapleTree;
+    ///
+    /// 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_range(100..500, ten, GFP_KERNEL)?;
+    /// tree.insert(67, twenty, GFP_KERNEL)?;
+    ///
+    /// assert_eq!(tree.erase(67).map(|v| *v), Some(20));
+    /// assert_eq!(tree.erase(275).map(|v| *v), Some(10));
+    ///
+    /// // The previous call erased the entire range, not just index 275.
+    /// assert!(tree.erase(127).is_none());
+    /// # Ok::<_, Error>(())
+    /// ```
+    #[inline]
+    pub fn erase(&self, index: usize) -> Option<T> {
+        // SAFETY: `self.tree` contains a valid maple tree.
+        let ret = unsafe { bindings::mtree_erase(self.tree.get(), index) };
+
+        // SAFETY: If the pointer is not null, then we took ownership of a valid instance of `T`
+        // from the tree.
+        unsafe { T::try_from_foreign(ret) }
+    }
+
+    /// Free all `T` instances in this tree.
+    ///
+    /// # Safety
+    ///
+    /// This frees Rust data referenced by the maple tree without removing it from the maple tree,
+    /// leaving it in an invalid state. The caller must ensure that this invalid state cannot be
+    /// observed by the end-user.
+    unsafe fn free_all_entries(self: Pin<&mut Self>) {
+        // SAFETY: The caller provides exclusive access to the entire maple tree, so we have
+        // exclusive access to the entire maple tree despite not holding the lock.
+        let mut ma_state = unsafe { MaState::new_raw(self.into_ref().get_ref(), 0, usize::MAX) };
+
+        loop {
+            // This uses the raw accessor because we're destroying pointers without removing them
+            // from the maple tree, which is only valid because this is the destructor.
+            let ptr = ma_state.mas_find_raw(usize::MAX);
+            if ptr.is_null() {
+                break;
+            }
+            // SAFETY: By the type invariants, this pointer references a valid value of type `T`.
+            // By the safety requirements, it is okay to free it without removing it from the maple
+            // tree.
+            drop(unsafe { T::from_foreign(ptr) });
+        }
+    }
+}
+
+#[pinned_drop]
+impl<T: ForeignOwnable> PinnedDrop for MapleTree<T> {
+    #[inline]
+    fn drop(mut self: Pin<&mut Self>) {
+        // We only iterate the tree if the Rust value has a destructor.
+        if core::mem::needs_drop::<T>() {
+            // SAFETY: Other than the below `mtree_destroy` call, the tree will not be accessed
+            // after this call.
+            unsafe { self.as_mut().free_all_entries() };
+        }
+
+        // SAFETY: The tree is valid, and will not be accessed after this call.
+        unsafe { bindings::mtree_destroy(self.tree.get()) };
+    }
+}
+
+/// A helper type used for navigating a [`MapleTree`].
+///
+/// # Invariants
+///
+/// For the duration of `'tree`:
+///
+/// * The `ma_state` references a valid `MapleTree<T>`.
+/// * The `ma_state` has read/write access to the tree.
+pub struct MaState<'tree, T: ForeignOwnable> {
+    state: bindings::ma_state,
+    _phantom: PhantomData<&'tree mut MapleTree<T>>,
+}
+
+impl<'tree, T: ForeignOwnable> MaState<'tree, T> {
+    /// Initialize a new `MaState` with the given tree.
+    ///
+    /// # Safety
+    ///
+    /// The caller must ensure that this `MaState` has read/write access to the maple tree.
+    #[inline]
+    unsafe fn new_raw(mt: &'tree MapleTree<T>, first: usize, end: usize) -> Self {
+        // INVARIANT:
+        // * Having a reference ensures that the `MapleTree<T>` is valid for `'tree`.
+        // * The caller ensures that we have read/write access.
+        Self {
+            state: bindings::ma_state {
+                tree: mt.tree.get(),
+                index: first,
+                last: end,
+                node: ptr::null_mut(),
+                status: bindings::maple_status_ma_start,
+                min: 0,
+                max: usize::MAX,
+                alloc: ptr::null_mut(),
+                mas_flags: 0,
+                store_type: bindings::store_type_wr_invalid,
+                ..Default::default()
+            },
+            _phantom: PhantomData,
+        }
+    }
+
+    #[inline]
+    fn as_raw(&mut self) -> *mut bindings::ma_state {
+        &raw mut self.state
+    }
+
+    #[inline]
+    fn mas_find_raw(&mut self, max: usize) -> *mut c_void {
+        // SAFETY: By the type invariants, the `ma_state` is active and we have read/write access
+        // to the tree.
+        unsafe { bindings::mas_find(self.as_raw(), max) }
+    }
+}
+
+/// Error type for failure to insert a new value.
+pub struct InsertError<T> {
+    /// The value that could not be inserted.
+    pub value: T,
+    /// The reason for the failure to insert.
+    pub cause: InsertErrorKind,
+}
+
+/// The reason for the failure to insert.
+#[derive(PartialEq, Eq, Copy, Clone, Debug)]
+pub enum InsertErrorKind {
+    /// There is already a value in the requested range.
+    Occupied,
+    /// Failure to allocate memory.
+    AllocError(kernel::alloc::AllocError),
+    /// The insertion request was invalid.
+    InvalidRequest,
+}
+
+impl From<InsertErrorKind> for Error {
+    #[inline]
+    fn from(kind: InsertErrorKind) -> Error {
+        match kind {
+            InsertErrorKind::Occupied => EEXIST,
+            InsertErrorKind::AllocError(kernel::alloc::AllocError) => ENOMEM,
+            InsertErrorKind::InvalidRequest => EINVAL,
+        }
+    }
+}
+
+impl<T> From<InsertError<T>> for Error {
+    #[inline]
+    fn from(insert_err: InsertError<T>) -> Error {
+        Error::from(insert_err.cause)
+    }
+}

-- 
2.51.0.338.gd7d06c2dae-goog

[PATCH v3 2/3] rust: maple_tree: add lock guard for maple tree

[Alice Ryhl]

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.

There are two ways of using the lock guard: You can call load() directly
to load a value under the lock, or you can create an MaState to iterate
the tree with find().

The find() method does not have the mas_ prefix since it's a method on
MaState, and being a method on that struct serves a similar purpose to
the mas_ prefix in C.

Co-developed-by: Andrew Ballance <andrewjballance@gmail.com>
Signed-off-by: Andrew Ballance <andrewjballance@gmail.com>
Reviewed-by: Andrew Ballance <andrewjballance@gmail.com>
Reviewed-by: Danilo Krummrich <dakr@kernel.org>
Signed-off-by: Alice Ryhl <aliceryhl@google.com>
---
 rust/kernel/maple_tree.rs | 140 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 140 insertions(+)

diff --git a/rust/kernel/maple_tree.rs b/rust/kernel/maple_tree.rs
index 92930b408e9683b6615080a53800f1a393c5f31e..24b674ce07d0481702eccd86a7920f94ca000108 100644
--- a/rust/kernel/maple_tree.rs
+++ b/rust/kernel/maple_tree.rs
@@ -214,6 +214,23 @@ 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.
+        let lock_ptr = unsafe { &raw mut (*self.tree.get()).__bindgen_anon_1.ma_lock };
+        lock_ptr.cast()
+    }
+
     /// Free all `T` instances in this tree.
     ///
     /// # Safety
@@ -257,6 +274,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;
+    ///
+    /// 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).map(|v| *v), Some(10));
+    /// assert_eq!(lock.load(200).map(|v| *v), Some(20));
+    /// assert_eq!(lock.load(300).map(|v| *v), None);
+    /// # Ok::<_, Error>(())
+    /// ```
+    ///
+    /// Increment refcount under the lock, to keep value alive afterwards.
+    ///
+    /// ```
+    /// use kernel::maple_tree::MapleTree;
+    /// 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 = tree.lock().load(100).map(Arc::from);
+    ///
+    /// // 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.map(|v| *v), Some(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) })
+    }
+}
+
 /// A helper type used for navigating a [`MapleTree`].
 ///
 /// # Invariants
@@ -310,6 +412,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.
+    ///
+    /// # Examples
+    ///
+    /// Iterate the maple tree.
+    ///
+    /// ```
+    /// use kernel::maple_tree::MapleTree;
+    /// 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.find(usize::MAX).map(|v| *v), Some(10));
+    /// assert_eq!(iter.find(usize::MAX).map(|v| *v), Some(20));
+    /// assert!(iter.find(usize::MAX).is_none());
+    /// # Ok::<_, Error>(())
+    /// ```
+    #[inline]
+    pub fn find(&mut self, max: usize) -> Option<T::BorrowedMut<'_>> {
+        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.338.gd7d06c2dae-goog

[PATCH v3 3/3] rust: maple_tree: add MapleTreeAlloc

[Alice Ryhl]

To support allocation trees, we introduce a new type MapleTreeAlloc for
the case where the tree is created using MT_FLAGS_ALLOC_RANGE. To ensure
that you can only call mtree_alloc_range on an allocation tree, we
restrict thta method to the new MapleTreeAlloc type. However, all
methods on MapleTree remain accessible to MapleTreeAlloc as allocation
trees can use the other methods without issues.

Reviewed-by: Daniel Almeida <daniel.almeida@collabora.com>
Reviewed-by: Danilo Krummrich <dakr@kernel.org>
Signed-off-by: Alice Ryhl <aliceryhl@google.com>
---
 rust/kernel/maple_tree.rs | 158 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 158 insertions(+)

diff --git a/rust/kernel/maple_tree.rs b/rust/kernel/maple_tree.rs
index 24b674ce07d0481702eccd86a7920f94ca000108..31b6c13d08efab23f7fb20bf97e36c33bf9f6ad9 100644
--- a/rust/kernel/maple_tree.rs
+++ b/rust/kernel/maple_tree.rs
@@ -32,6 +32,26 @@ pub struct MapleTree<T: ForeignOwnable> {
     _p: PhantomData<T>,
 }
 
+/// A maple tree with `MT_FLAGS_ALLOC_RANGE` set.
+///
+/// All methods on [`MapleTree`] are also accessible on this type.
+#[pin_data]
+#[repr(transparent)]
+pub struct MapleTreeAlloc<T: ForeignOwnable> {
+    #[pin]
+    tree: MapleTree<T>,
+}
+
+// Make MapleTree methods usable on MapleTreeAlloc.
+impl<T: ForeignOwnable> core::ops::Deref for MapleTreeAlloc<T> {
+    type Target = MapleTree<T>;
+
+    #[inline]
+    fn deref(&self) -> &MapleTree<T> {
+        &self.tree
+    }
+}
+
 #[inline]
 fn to_maple_range(range: impl RangeBounds<usize>) -> Option<(usize, usize)> {
     let first = match range.start_bound() {
@@ -359,6 +379,107 @@ pub fn load(&mut self, index: usize) -> Option<T::BorrowedMut<'_>> {
     }
 }
 
+impl<T: ForeignOwnable> MapleTreeAlloc<T> {
+    /// Create a new allocation tree.
+    pub fn new() -> impl PinInit<Self> {
+        let tree = pin_init!(MapleTree {
+            // SAFETY: This initializes a maple tree into a pinned slot. The maple tree will be
+            // destroyed in Drop before the memory location becomes invalid.
+            tree <- Opaque::ffi_init(|slot| unsafe {
+                bindings::mt_init_flags(slot, bindings::MT_FLAGS_ALLOC_RANGE)
+            }),
+            _p: PhantomData,
+        });
+
+        pin_init!(MapleTreeAlloc { tree <- tree })
+    }
+
+    /// Insert an entry with the given size somewhere in the given range.
+    ///
+    /// The maple tree will search for a location in the given range where there is space to insert
+    /// the new range. If there is not enough available space, then an error will be returned.
+    ///
+    /// The index of the new range is returned.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use kernel::maple_tree::{MapleTreeAlloc, AllocErrorKind};
+    ///
+    /// let tree = KBox::pin_init(MapleTreeAlloc::<KBox<i32>>::new(), GFP_KERNEL)?;
+    ///
+    /// let ten = KBox::new(10, GFP_KERNEL)?;
+    /// let twenty = KBox::new(20, GFP_KERNEL)?;
+    /// let thirty = KBox::new(30, GFP_KERNEL)?;
+    /// let hundred = KBox::new(100, GFP_KERNEL)?;
+    ///
+    /// // Allocate three ranges.
+    /// let idx1 = tree.alloc_range(100, ten, ..1000, GFP_KERNEL)?;
+    /// let idx2 = tree.alloc_range(100, twenty, ..1000, GFP_KERNEL)?;
+    /// let idx3 = tree.alloc_range(100, thirty, ..1000, GFP_KERNEL)?;
+    ///
+    /// assert_eq!(idx1, 0);
+    /// assert_eq!(idx2, 100);
+    /// assert_eq!(idx3, 200);
+    ///
+    /// // This will fail because the remaining space is too small.
+    /// assert_eq!(
+    ///     tree.alloc_range(800, hundred, ..1000, GFP_KERNEL).unwrap_err().cause,
+    ///     AllocErrorKind::Busy,
+    /// );
+    /// # Ok::<_, Error>(())
+    /// ```
+    pub fn alloc_range<R>(
+        &self,
+        size: usize,
+        value: T,
+        range: R,
+        gfp: Flags,
+    ) -> Result<usize, AllocError<T>>
+    where
+        R: RangeBounds<usize>,
+    {
+        let Some((min, max)) = to_maple_range(range) else {
+            return Err(AllocError {
+                value,
+                cause: AllocErrorKind::InvalidRequest,
+            });
+        };
+
+        let ptr = T::into_foreign(value);
+        let mut index = 0;
+
+        // SAFETY: The tree is valid, and we are passing a pointer to an owned instance of `T`.
+        let res = to_result(unsafe {
+            bindings::mtree_alloc_range(
+                self.tree.tree.get(),
+                &mut index,
+                ptr,
+                size,
+                min,
+                max,
+                gfp.as_raw(),
+            )
+        });
+
+        if let Err(err) = res {
+            // SAFETY: As `mtree_alloc_range` failed, it is safe to take back ownership.
+            let value = unsafe { T::from_foreign(ptr) };
+
+            let cause = if err == ENOMEM {
+                AllocErrorKind::AllocError(kernel::alloc::AllocError)
+            } else if err == EBUSY {
+                AllocErrorKind::Busy
+            } else {
+                AllocErrorKind::InvalidRequest
+            };
+            Err(AllocError { value, cause })
+        } else {
+            Ok(index)
+        }
+    }
+}
+
 /// A helper type used for navigating a [`MapleTree`].
 ///
 /// # Invariants
@@ -488,3 +609,40 @@ fn from(insert_err: InsertError<T>) -> Error {
         Error::from(insert_err.cause)
     }
 }
+
+/// Error type for failure to insert a new value.
+pub struct AllocError<T> {
+    /// The value that could not be inserted.
+    pub value: T,
+    /// The reason for the failure to insert.
+    pub cause: AllocErrorKind,
+}
+
+/// The reason for the failure to insert.
+#[derive(PartialEq, Eq, Copy, Clone)]
+pub enum AllocErrorKind {
+    /// There is not enough space for the requested allocation.
+    Busy,
+    /// Failure to allocate memory.
+    AllocError(kernel::alloc::AllocError),
+    /// The insertion request was invalid.
+    InvalidRequest,
+}
+
+impl From<AllocErrorKind> for Error {
+    #[inline]
+    fn from(kind: AllocErrorKind) -> Error {
+        match kind {
+            AllocErrorKind::Busy => EBUSY,
+            AllocErrorKind::AllocError(kernel::alloc::AllocError) => ENOMEM,
+            AllocErrorKind::InvalidRequest => EINVAL,
+        }
+    }
+}
+
+impl<T> From<AllocError<T>> for Error {
+    #[inline]
+    fn from(insert_err: AllocError<T>) -> Error {
+        Error::from(insert_err.cause)
+    }
+}

-- 
2.51.0.338.gd7d06c2dae-goog

Re: [PATCH v3 1/3] rust: maple_tree: add MapleTree

[Danilo Krummrich]

On Tue Sep 2, 2025 at 10:35 AM CEST, Alice Ryhl wrote:
> The maple tree will be used in the Tyr driver to allocate and keep track
> of GPU allocations created internally (i.e. not by userspace). It will
> likely also be used in the Nova driver eventually.
>
> This adds the simplest methods for additional and removal that do not
> require any special care with respect to concurrency.
>
> This implementation is based on the RFC by Andrew but with significant
> changes to simplify the implementation.
>
> Co-developed-by: Andrew Ballance <andrewjballance@gmail.com>
> Signed-off-by: Andrew Ballance <andrewjballance@gmail.com>
> Signed-off-by: Alice Ryhl <aliceryhl@google.com>

One nit below, otherwise:

	Reviewed-by: Danilo Krummrich <dakr@kernel.org>

> +    pub fn insert_range<R>(&self, range: R, value: T, gfp: Flags) -> Result<(), InsertError<T>>
> +    where
> +        R: RangeBounds<usize>,
> +    {
> +        let Some((first, last)) = to_maple_range(range) else {
> +            return Err(InsertError {
> +                value,
> +                cause: InsertErrorKind::InvalidRequest,
> +            });
> +        };
> +
> +        let ptr = T::into_foreign(value);
> +
> +        // SAFETY: The tree is valid, and we are passing a pointer to an owned instance of `T`.
> +        let res = to_result(unsafe {
> +            bindings::mtree_insert_range(self.tree.get(), first, last, ptr, gfp.as_raw())
> +        });
> +
> +        if let Err(err) = res {
> +            // SAFETY: As `mtree_insert_range` failed, it is safe to take back ownership.
> +            let value = unsafe { T::from_foreign(ptr) };
> +
> +            let cause = if err == ENOMEM {
> +                InsertErrorKind::AllocError(kernel::alloc::AllocError)
> +            } else if err == EEXIST {
> +                InsertErrorKind::Occupied
> +            } else {
> +                InsertErrorKind::InvalidRequest
> +            };
> +            Err(InsertError { value, cause })
> +        } else {
> +            Ok(())
> +        }
> +    }

	// SAFETY: The tree is valid, and we are passing a pointer to an owned instance of `T`.
	to_result(unsafe {
	    bindings::mtree_insert_range(self.tree.get(), first, last, ptr, gfp.as_raw())
	}).map_err(|err| {
	    // SAFETY: As `mtree_insert_range` failed, it is safe to take back ownership.
	    let value = unsafe { T::from_foreign(ptr) };
	
	    let cause = if err == ENOMEM {
	        InsertErrorKind::AllocError(kernel::alloc::AllocError)
	    } else if err == EEXIST {
	        InsertErrorKind::Occupied
	    } else {
	        InsertErrorKind::InvalidRequest
	    };
	    Err(InsertError { value, cause })
	})

I think that's a bit cleaner than the above (not compile tested).

Re: [PATCH v3 1/3] rust: maple_tree: add MapleTree

[Alice Ryhl]

On Tue, Sep 02, 2025 at 11:01:19AM +0200, Danilo Krummrich wrote:
> On Tue Sep 2, 2025 at 10:35 AM CEST, Alice Ryhl wrote:
> > The maple tree will be used in the Tyr driver to allocate and keep track
> > of GPU allocations created internally (i.e. not by userspace). It will
> > likely also be used in the Nova driver eventually.
> >
> > This adds the simplest methods for additional and removal that do not
> > require any special care with respect to concurrency.
> >
> > This implementation is based on the RFC by Andrew but with significant
> > changes to simplify the implementation.
> >
> > Co-developed-by: Andrew Ballance <andrewjballance@gmail.com>
> > Signed-off-by: Andrew Ballance <andrewjballance@gmail.com>
> > Signed-off-by: Alice Ryhl <aliceryhl@google.com>
> 
> One nit below, otherwise:
> 
> 	Reviewed-by: Danilo Krummrich <dakr@kernel.org>

Thanks!

> > +    pub fn insert_range<R>(&self, range: R, value: T, gfp: Flags) -> Result<(), InsertError<T>>
> > +    where
> > +        R: RangeBounds<usize>,
> > +    {
> > +        let Some((first, last)) = to_maple_range(range) else {
> > +            return Err(InsertError {
> > +                value,
> > +                cause: InsertErrorKind::InvalidRequest,
> > +            });
> > +        };
> > +
> > +        let ptr = T::into_foreign(value);
> > +
> > +        // SAFETY: The tree is valid, and we are passing a pointer to an owned instance of `T`.
> > +        let res = to_result(unsafe {
> > +            bindings::mtree_insert_range(self.tree.get(), first, last, ptr, gfp.as_raw())
> > +        });
> > +
> > +        if let Err(err) = res {
> > +            // SAFETY: As `mtree_insert_range` failed, it is safe to take back ownership.
> > +            let value = unsafe { T::from_foreign(ptr) };
> > +
> > +            let cause = if err == ENOMEM {
> > +                InsertErrorKind::AllocError(kernel::alloc::AllocError)
> > +            } else if err == EEXIST {
> > +                InsertErrorKind::Occupied
> > +            } else {
> > +                InsertErrorKind::InvalidRequest
> > +            };
> > +            Err(InsertError { value, cause })
> > +        } else {
> > +            Ok(())
> > +        }
> > +    }
> 
> 	// SAFETY: The tree is valid, and we are passing a pointer to an owned instance of `T`.
> 	to_result(unsafe {
> 	    bindings::mtree_insert_range(self.tree.get(), first, last, ptr, gfp.as_raw())
> 	}).map_err(|err| {
> 	    // SAFETY: As `mtree_insert_range` failed, it is safe to take back ownership.
> 	    let value = unsafe { T::from_foreign(ptr) };
> 	
> 	    let cause = if err == ENOMEM {
> 	        InsertErrorKind::AllocError(kernel::alloc::AllocError)
> 	    } else if err == EEXIST {
> 	        InsertErrorKind::Occupied
> 	    } else {
> 	        InsertErrorKind::InvalidRequest
> 	    };
> 	    Err(InsertError { value, cause })
> 	})
> 
> I think that's a bit cleaner than the above (not compile tested).

I don't love it. How about a match instead of if/else?

Alice

Re: [PATCH v3 1/3] rust: maple_tree: add MapleTree

[Danilo Krummrich]

On Tue Sep 2, 2025 at 11:28 AM CEST, Alice Ryhl wrote:
> On Tue, Sep 02, 2025 at 11:01:19AM +0200, Danilo Krummrich wrote:
>> On Tue Sep 2, 2025 at 10:35 AM CEST, Alice Ryhl wrote:
>> > The maple tree will be used in the Tyr driver to allocate and keep track
>> > of GPU allocations created internally (i.e. not by userspace). It will
>> > likely also be used in the Nova driver eventually.
>> >
>> > This adds the simplest methods for additional and removal that do not
>> > require any special care with respect to concurrency.
>> >
>> > This implementation is based on the RFC by Andrew but with significant
>> > changes to simplify the implementation.
>> >
>> > Co-developed-by: Andrew Ballance <andrewjballance@gmail.com>
>> > Signed-off-by: Andrew Ballance <andrewjballance@gmail.com>
>> > Signed-off-by: Alice Ryhl <aliceryhl@google.com>
>> 
>> One nit below, otherwise:
>> 
>> 	Reviewed-by: Danilo Krummrich <dakr@kernel.org>
>
> Thanks!
>
>> > +    pub fn insert_range<R>(&self, range: R, value: T, gfp: Flags) -> Result<(), InsertError<T>>
>> > +    where
>> > +        R: RangeBounds<usize>,
>> > +    {
>> > +        let Some((first, last)) = to_maple_range(range) else {
>> > +            return Err(InsertError {
>> > +                value,
>> > +                cause: InsertErrorKind::InvalidRequest,
>> > +            });
>> > +        };
>> > +
>> > +        let ptr = T::into_foreign(value);
>> > +
>> > +        // SAFETY: The tree is valid, and we are passing a pointer to an owned instance of `T`.
>> > +        let res = to_result(unsafe {
>> > +            bindings::mtree_insert_range(self.tree.get(), first, last, ptr, gfp.as_raw())
>> > +        });
>> > +
>> > +        if let Err(err) = res {
>> > +            // SAFETY: As `mtree_insert_range` failed, it is safe to take back ownership.
>> > +            let value = unsafe { T::from_foreign(ptr) };
>> > +
>> > +            let cause = if err == ENOMEM {
>> > +                InsertErrorKind::AllocError(kernel::alloc::AllocError)
>> > +            } else if err == EEXIST {
>> > +                InsertErrorKind::Occupied
>> > +            } else {
>> > +                InsertErrorKind::InvalidRequest
>> > +            };
>> > +            Err(InsertError { value, cause })
>> > +        } else {
>> > +            Ok(())
>> > +        }
>> > +    }
>> 
>> 	// SAFETY: The tree is valid, and we are passing a pointer to an owned instance of `T`.
>> 	to_result(unsafe {
>> 	    bindings::mtree_insert_range(self.tree.get(), first, last, ptr, gfp.as_raw())
>> 	}).map_err(|err| {
>> 	    // SAFETY: As `mtree_insert_range` failed, it is safe to take back ownership.
>> 	    let value = unsafe { T::from_foreign(ptr) };
>> 	
>> 	    let cause = if err == ENOMEM {
>> 	        InsertErrorKind::AllocError(kernel::alloc::AllocError)
>> 	    } else if err == EEXIST {
>> 	        InsertErrorKind::Occupied
>> 	    } else {
>> 	        InsertErrorKind::InvalidRequest
>> 	    };
>> 	    Err(InsertError { value, cause })
>> 	})
>> 
>> I think that's a bit cleaner than the above (not compile tested).
>
> I don't love it. How about a match instead of if/else?

I think if you don't like map_err(), the if/else is fine to keep as is.

Personally, I prefer map_err() over any "manual matching" in this case; it gets
us rid of `ref`, making to_result().map_err() a single statement returning
exactly what we expect. It makes it more self-contained.