[RFC PATCH 01/11] rust: add radix tree abstraction

[Andreas Hindborg]

From: Andreas Hindborg <a.hindborg@xxxxxxxxxxx>

Add abstractions for the C radix_tree. This abstraction allows Rust code to use
the radix_tree as a map from `u64` keys to `ForeignOwnable` values.

Signed-off-by: Andreas Hindborg <a.hindborg@xxxxxxxxxxx>
---
 rust/bindings/bindings_helper.h |   2 +
 rust/bindings/lib.rs            |   1 +
 rust/helpers.c                  |  22 +++++
 rust/kernel/lib.rs              |   1 +
 rust/kernel/radix_tree.rs       | 156 ++++++++++++++++++++++++++++++++
 5 files changed, 182 insertions(+)
 create mode 100644 rust/kernel/radix_tree.rs

diff --git a/rust/bindings/bindings_helper.h b/rust/bindings/bindings_helper.h
index 50e7a76d5455..52834962b94d 100644
--- a/rust/bindings/bindings_helper.h
+++ b/rust/bindings/bindings_helper.h
@@ -10,7 +10,9 @@
 #include <linux/refcount.h>
 #include <linux/wait.h>
 #include <linux/sched.h>
+#include <linux/radix-tree.h>
 
 /* `bindgen` gets confused at certain things. */
 const gfp_t BINDINGS_GFP_KERNEL = GFP_KERNEL;
+const gfp_t BINDINGS_GFP_ATOMIC = GFP_ATOMIC;
 const gfp_t BINDINGS___GFP_ZERO = __GFP_ZERO;
diff --git a/rust/bindings/lib.rs b/rust/bindings/lib.rs
index 7b246454e009..62f36a9eb1f4 100644
--- a/rust/bindings/lib.rs
+++ b/rust/bindings/lib.rs
@@ -51,4 +51,5 @@ mod bindings_helper {
 pub use bindings_raw::*;
 
 pub const GFP_KERNEL: gfp_t = BINDINGS_GFP_KERNEL;
+pub const GFP_ATOMIC: gfp_t = BINDINGS_GFP_ATOMIC;
 pub const __GFP_ZERO: gfp_t = BINDINGS___GFP_ZERO;
diff --git a/rust/helpers.c b/rust/helpers.c
index 81e80261d597..5dd5e325b7cc 100644
--- a/rust/helpers.c
+++ b/rust/helpers.c
@@ -26,6 +26,7 @@
 #include <linux/spinlock.h>
 #include <linux/sched/signal.h>
 #include <linux/wait.h>
+#include <linux/radix-tree.h>
 
 __noreturn void rust_helper_BUG(void)
 {
@@ -128,6 +129,27 @@ void rust_helper_put_task_struct(struct task_struct *t)
 }
 EXPORT_SYMBOL_GPL(rust_helper_put_task_struct);
 
+void rust_helper_init_radix_tree(struct xarray *tree, gfp_t gfp_mask)
+{
+	INIT_RADIX_TREE(tree, gfp_mask);
+}
+EXPORT_SYMBOL_GPL(rust_helper_init_radix_tree);
+
+void **rust_helper_radix_tree_iter_init(struct radix_tree_iter *iter,
+					unsigned long start)
+{
+	return radix_tree_iter_init(iter, start);
+}
+EXPORT_SYMBOL_GPL(rust_helper_radix_tree_iter_init);
+
+void **rust_helper_radix_tree_next_slot(void **slot,
+					struct radix_tree_iter *iter,
+					unsigned flags)
+{
+	return radix_tree_next_slot(slot, iter, flags);
+}
+EXPORT_SYMBOL_GPL(rust_helper_radix_tree_next_slot);
+
 /*
  * We use `bindgen`'s `--size_t-is-usize` option to bind the C `size_t` type
  * as the Rust `usize` type, so we can use it in contexts where Rust
diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index 676995d4e460..a85cb6aae8d6 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -40,6 +40,7 @@ pub mod init;
 pub mod ioctl;
 pub mod prelude;
 pub mod print;
+pub mod radix_tree;
 mod static_assert;
 #[doc(hidden)]
 pub mod std_vendor;
diff --git a/rust/kernel/radix_tree.rs b/rust/kernel/radix_tree.rs
new file mode 100644
index 000000000000..f659ab8b017c
--- /dev/null
+++ b/rust/kernel/radix_tree.rs
@@ -0,0 +1,156 @@
+// SPDX-License-Identifier: GPL-2.0
+
+//! RadixTree abstraction.
+//!
+//! C header: [`include/linux/radix_tree.h`](../../include/linux/radix_tree.h)
+
+use crate::error::to_result;
+use crate::error::Result;
+use crate::types::ForeignOwnable;
+use crate::types::Opaque;
+use crate::types::ScopeGuard;
+use alloc::boxed::Box;
+use core::marker::PhantomData;
+use core::pin::Pin;
+
+type Key = u64;
+
+/// A map of `u64` to `ForeignOwnable`
+///
+/// # Invariants
+///
+/// - `tree` always points to a valid and initialized `struct radix_tree`.
+/// - Pointers stored in the tree are created by a call to `ForignOwnable::into_foreign()`
+pub struct RadixTree<V: ForeignOwnable> {
+    tree: Pin<Box<Opaque<bindings::xarray>>>,
+    _marker: PhantomData<V>,
+}
+
+impl<V: ForeignOwnable> RadixTree<V> {
+    /// Create a new radix tree
+    ///
+    /// Note: This function allocates memory with `GFP_ATOMIC`.
+    pub fn new() -> Result<Self> {
+        let tree = Pin::from(Box::try_new(Opaque::uninit())?);
+
+        // SAFETY: `tree` points to allocated but not initialized memory. This
+        // call will initialize the memory.
+        unsafe { bindings::init_radix_tree(tree.get(), bindings::GFP_ATOMIC) };
+
+        Ok(Self {
+            tree,
+            _marker: PhantomData,
+        })
+    }
+
+    /// Try to insert a value into the tree
+    pub fn try_insert(&mut self, key: Key, value: V) -> Result<()> {
+        // SAFETY: `self.tree` points to a valid and initialized `struct radix_tree`
+        let ret =
+            unsafe { bindings::radix_tree_insert(self.tree.get(), key, value.into_foreign() as _) };
+        to_result(ret)
+    }
+
+    /// Search for `key` in the map. Returns a reference to the associated
+    /// value if found.
+    pub fn get(&self, key: Key) -> Option<V::Borrowed<'_>> {
+        // SAFETY: `self.tree` points to a valid and initialized `struct radix_tree`
+        let item =
+            core::ptr::NonNull::new(unsafe { bindings::radix_tree_lookup(self.tree.get(), key) })?;
+
+        // SAFETY: `item` was created by a call to
+        // `ForeignOwnable::into_foreign()`. As `get_mut()` and `remove()` takes
+        // a `&mut self`, no mutable borrows for `item` can exist and
+        // `ForeignOwnable::from_foreign()` cannot be called until this borrow
+        // is dropped.
+        Some(unsafe { V::borrow(item.as_ptr()) })
+    }
+
+    /// Search for `key` in the map. Return a mutable reference to the
+    /// associated value if found.
+    pub fn get_mut(&mut self, key: Key) -> Option<MutBorrow<'_, V>> {
+        let item =
+            core::ptr::NonNull::new(unsafe { bindings::radix_tree_lookup(self.tree.get(), key) })?;
+
+        // SAFETY: `item` was created by a call to
+        // `ForeignOwnable::into_foreign()`. As `get()` takes a `&self` and
+        // `remove()` takes a `&mut self`, no borrows for `item` can exist and
+        // `ForeignOwnable::from_foreign()` cannot be called until this borrow
+        // is dropped.
+        Some(MutBorrow {
+            guard: unsafe { V::borrow_mut(item.as_ptr()) },
+            _marker: core::marker::PhantomData,
+        })
+    }
+
+    /// Search for `key` in the map. If `key` is found, the key and value is
+    /// removed from the map and the value is returned.
+    pub fn remove(&mut self, key: Key) -> Option<V> {
+        // SAFETY: `self.tree` points to a valid and initialized `struct radix_tree`
+        let item =
+            core::ptr::NonNull::new(unsafe { bindings::radix_tree_delete(self.tree.get(), key) })?;
+
+        // SAFETY: `item` was created by a call to
+        // `ForeignOwnable::into_foreign()` and no borrows to `item` can exist
+        // because this function takes a `&mut self`.
+        Some(unsafe { ForeignOwnable::from_foreign(item.as_ptr()) })
+    }
+}
+
+impl<V: ForeignOwnable> Drop for RadixTree<V> {
+    fn drop(&mut self) {
+        let mut iter = bindings::radix_tree_iter {
+            index: 0,
+            next_index: 0,
+            tags: 0,
+            node: core::ptr::null_mut(),
+        };
+
+        // SAFETY: Iter is valid as we allocated it on the stack above
+        let mut slot = unsafe { bindings::radix_tree_iter_init(&mut iter, 0) };
+        loop {
+            if slot.is_null() {
+                // SAFETY: Both `self.tree` and `iter` are valid
+                slot = unsafe { bindings::radix_tree_next_chunk(self.tree.get(), &mut iter, 0) };
+            }
+
+            if slot.is_null() {
+                break;
+            }
+
+            // SAFETY: `self.tree` is valid and iter is managed by
+            // `radix_tree_next_chunk()` and `radix_tree_next_slot()`
+            let item = unsafe { bindings::radix_tree_delete(self.tree.get(), iter.index) };
+            assert!(!item.is_null());
+
+            // SAFETY: All items in the tree are created by a call to
+            // `ForeignOwnable::into_foreign()`.
+            let _ = unsafe { V::from_foreign(item) };
+
+            // SAFETY: `self.tree` is valid and iter is managed by
+            // `radix_tree_next_chunk()` and `radix_tree_next_slot()`. Slot is
+            // not null.
+            slot = unsafe { bindings::radix_tree_next_slot(slot, &mut iter, 0) };
+        }
+    }
+}
+
+/// A mutable borrow of an object owned by a `RadixTree`
+pub struct MutBorrow<'a, V: ForeignOwnable> {
+    guard: ScopeGuard<V, fn(V)>,
+    _marker: core::marker::PhantomData<&'a mut V>,
+}
+
+impl<'a, V: ForeignOwnable> core::ops::Deref for MutBorrow<'a, V> {
+    type Target = ScopeGuard<V, fn(V)>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.guard
+    }
+}
+
+impl<'a, V: ForeignOwnable> core::ops::DerefMut for MutBorrow<'a, V> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.guard
+    }
+}
-- 
2.40.0