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

From: Boqun Feng
Date: Mon Jul 28 2025 - 12:09:28 EST


On Sat, Jul 26, 2025 at 01:23:22PM +0000, 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@xxxxxxxxx>
> Signed-off-by: Andrew Ballance <andrewjballance@xxxxxxxxx>
> Signed-off-by: Alice Ryhl <aliceryhl@xxxxxxxxxx>
> ---
[...]
> + /// Free all `T` instances in this tree.
> + ///
> + /// # Safety
> + ///
> + /// This frees Rust data referenced by the maple tree without removing it from the maple tree.
> + /// The caller must ensure that no reference that remains in the maple tree is used incorrectly
> + /// after this call.
> + unsafe fn free_all_entries(self: Pin<&mut Self>) {
> + // SAFETY: The pointer references a valid maple tree.
> + let ma_state = unsafe { Opaque::new(bindings::MA_STATE(self.tree.get(), 0, usize::MAX)) };
> +

A meta comment here for the future direction: I think it really makes a
lot of sense if we could have the Rust abstraction for struct ma_state,
that'll allow us to have flexible locking strategy and Iterator-like
interface. Maybe it's something Andrew can take a deeper look when
MapleTree binding is in-tree (no word play intented ;-))?

For example, with a ma_state binding, we can do:

let mas = MAState::new(self, 0..);

while let Some(v) = mas.next() {
drop(v)
}

Regards,
Boqun

> + loop {
> + // SAFETY: The maple tree is valid. This call to `free_all_entries` has exclusive
> + // access to the maple tree, so no further synchronization is required.
> + let ptr = unsafe { bindings::mas_find(ma_state.get(), 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.
> + unsafe { drop(T::from_foreign(ptr)) };
> + }
> + }
> +}
> +
[...]