> 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>
Can you sort these alphabetically? Also the constants down below.
>
> /* `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;
I would prefer `use crate::{error::{...}, types::{...}};`.
> +
> +type Key = u64;
Is there a reason to not make this `pub`?
> +
> +/// 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>,
> +}
Design question: why is the tree boxed? Shouldn't the user decide how
they want to manage the memory of the tree? In other words, should
`tree` just be `Opaque<bindings::xarray>` and this type initialized via
`pin-init`?
> +
> +impl<V: ForeignOwnable> RadixTree<V> {
> + /// Create a new radix tree
> + ///
> + /// Note: This function allocates memory with `GFP_ATOMIC`.
There probably is a reason why this uses `GFP_ATOMIC`, but I think if we
decide that the tree allocates memory there should be also a function
with `GFP_KERNEL`. That function should be called `new()` and this one
`new_atomic()` or `new_gfp(gfp: gfp_t)`. Also note that I think the
return type should be `Result<Self, AllocError>`, we should be explicit
where we can be.
> + 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`
Also add that the invariant of only inserting `ForignOwnable` pointers
is upheld.
> + let ret =
> + unsafe { bindings::radix_tree_insert(self.tree.get(), key, value.into_foreign() as _) };
Instead of `as _` prefer to use `.cast::<$type>()` or in this case probably `.cast_mut()`.
> + 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) })?;
You use `NonNull` quiet a bit, consider importing it.
> +
> + // 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)>;
Why is `Target = ScopeGuard`? I think it makes more sense to use
`Target = 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
>
--
Cheers,
Benno
