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

[Benno Lossin]

> 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