MIRRORS/linux-imx
1
0
Fork 0
mirror of https://github.com/nxp-imx/linux-imx.git synced 2025-07-17 06:39:37 +02:00
Code Issues Actions Packages Projects Releases Wiki Activity

FROMLIST: rust: rbtree: add red-black tree implementation backed by the C version

Browse Source 
The rust rbtree exposes a map-like interface over keys and values,
backed by the kernel red-black tree implementation. Values can be
inserted, deleted, and retrieved from a `RBTree` by key.

This base abstraction is used by binder to store key/value
pairs and perform lookups, for example the patch
"[PATCH RFC 03/20] rust_binder: add threading support"
in the binder RFC [1].

Link: https://lore.kernel.org/rust-for-linux/20231101-rust-binder-v1-3-08ba9197f637@google.com/ [1]
Change-Id: I1de1370fef24bb6da71becec8c970f6804c75b6a
Signed-off-by: Wedson Almeida Filho <wedsonaf@gmail.com>
Signed-off-by: Matt Gilbride <mattgilbride@google.com>

Bug: 324206405
Link: https://lore.kernel.org/all/20240205-b4-rbtree-v1-2-995e3eee38c0@google.com/
Change-Id: I9a757633b14b9a142106200a44843e65f57369d2
Signed-off-by: Alice Ryhl <aliceryhl@google.com>
This commit is contained in:
Wedson Almeida Filho 2024-02-05 15:50:02 +00:00 committed by Treehugger Robot
parent 4cb6d81d02
commit adfa8383a5
3 changed files with 411 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
rust/helpers.c
View File

@ -251,6 +251,13 @@ void rust_helper_init_task_work(struct callback_head *twork,
}
EXPORT_SYMBOL_GPL(rust_helper_init_task_work);
void rust_helper_rb_link_node(struct rb_node *node, struct rb_node *parent,
struct rb_node **rb_link)
{
rb_link_node(node, parent, rb_link);
}
EXPORT_SYMBOL_GPL(rust_helper_rb_link_node);
/*
* `bindgen` binds the C `size_t` type as the Rust `usize` type, so we can
* use it in contexts where Rust expects a `usize` like slice (array) indices.

1
rust/kernel/lib.rs
View File

@ -42,6 +42,7 @@ pub mod ioctl;
pub mod kunit;
pub mod prelude;
pub mod print;
pub mod rbtree;
pub mod security;
mod static_assert;
#[doc(hidden)]

403
rust/kernel/rbtree.rs Normal file
View File

@ -0,0 +1,403 @@
// SPDX-License-Identifier: GPL-2.0
//! Red-black trees.
//!
//! C header: [`include/linux/rbtree.h`](../../../../include/linux/rbtree.h)
//!
//! Reference: <https://www.kernel.org/doc/html/latest/core-api/rbtree.html>
use crate::{bindings, error::Result};
use alloc::boxed::Box;
use core::{
cmp::{Ord, Ordering},
marker::PhantomData,
mem::MaybeUninit,
ptr::{addr_of_mut, NonNull},
};
struct Node<K, V> {
links: bindings::rb_node,
key: K,
value: V,
}
/// A red-black tree with owned nodes.
///
/// It is backed by the kernel C red-black trees.
///
/// # Invariants
///
/// Non-null parent/children pointers stored in instances of the `rb_node` C struct are always
/// valid, and pointing to a field of our internal representation of a node.
///
/// # Examples
///
/// In the example below we do several operations on a tree. We note that insertions may fail if
/// the system is out of memory.
///
/// ```
/// use kernel::rbtree::RBTree;
///
/// // Create a new tree.
/// let mut tree = RBTree::new();
///
/// // Insert three elements.
/// tree.try_create_and_insert(20, 200)?;
/// tree.try_create_and_insert(10, 100)?;
/// tree.try_create_and_insert(30, 300)?;
///
/// // Check the nodes we just inserted.
/// {
/// assert_eq!(tree.get(&10).unwrap(), &100);
/// assert_eq!(tree.get(&20).unwrap(), &200);
/// assert_eq!(tree.get(&30).unwrap(), &300);
/// }
///
/// // Replace one of the elements.
/// tree.try_create_and_insert(10, 1000)?;
///
/// // Check that the tree reflects the replacement.
/// {
/// assert_eq!(tree.get(&10).unwrap(), &1000);
/// assert_eq!(tree.get(&20).unwrap(), &200);
/// assert_eq!(tree.get(&30).unwrap(), &300);
/// }
///
/// // Change the value of one of the elements.
/// *tree.get_mut(&30).unwrap() = 3000;
///
/// // Check that the tree reflects the update.
/// {
/// assert_eq!(tree.get(&10).unwrap(), &1000);
/// assert_eq!(tree.get(&20).unwrap(), &200);
/// assert_eq!(tree.get(&30).unwrap(), &3000);
/// }
///
/// // Remove an element.
/// tree.remove(&10);
///
/// // Check that the tree reflects the removal.
/// {
/// assert_eq!(tree.get(&10), None);
/// assert_eq!(tree.get(&20).unwrap(), &200);
/// assert_eq!(tree.get(&30).unwrap(), &3000);
/// }
///
/// # Ok::<(), Error>(())
/// ```
///
/// In the example below, we first allocate a node, acquire a spinlock, then insert the node into
/// the tree. This is useful when the insertion context does not allow sleeping, for example, when
/// holding a spinlock.
///
/// ```
/// use kernel::{rbtree::RBTree, sync::SpinLock};
///
/// fn insert_test(tree: &SpinLock<RBTree<u32, u32>>) -> Result {
/// // Pre-allocate node. This may fail (as it allocates memory).
/// let node = RBTree::try_allocate_node(10, 100)?;
///
/// // Insert node while holding the lock. It is guaranteed to succeed with no allocation
/// // attempts.
/// let mut guard = tree.lock();
/// guard.insert(node);
/// Ok(())
/// }
/// ```
///
/// In the example below, we reuse an existing node allocation from an element we removed.
///
/// ```
/// use kernel::rbtree::RBTree;
///
/// // Create a new tree.
/// let mut tree = RBTree::new();
///
/// // Insert three elements.
/// tree.try_create_and_insert(20, 200)?;
/// tree.try_create_and_insert(10, 100)?;
/// tree.try_create_and_insert(30, 300)?;
///
/// // Check the nodes we just inserted.
/// {
/// assert_eq!(tree.get(&10).unwrap(), &100);
/// assert_eq!(tree.get(&20).unwrap(), &200);
/// assert_eq!(tree.get(&30).unwrap(), &300);
/// }
///
/// // Remove a node, getting back ownership of it.
/// let existing = tree.remove_node(&30).unwrap();
///
/// // Check that the tree reflects the removal.
/// {
/// assert_eq!(tree.get(&10).unwrap(), &100);
/// assert_eq!(tree.get(&20).unwrap(), &200);
/// assert_eq!(tree.get(&30), None);
/// }
///
/// // Turn the node into a reservation so that we can reuse it with a different key/value.
/// let reservation = existing.into_reservation();
///
/// // Insert a new node into the tree, reusing the previous allocation. This is guaranteed to
/// // succeed (no memory allocations).
/// tree.insert(reservation.into_node(15, 150));
///
/// // Check that the tree reflect the new insertion.
/// {
/// assert_eq!(tree.get(&10).unwrap(), &100);
/// assert_eq!(tree.get(&15).unwrap(), &150);
/// assert_eq!(tree.get(&20).unwrap(), &200);
/// }
///
/// # Ok::<(), Error>(())
/// ```
pub struct RBTree<K, V> {
root: bindings::rb_root,
_p: PhantomData<Node<K, V>>,
}
// SAFETY: An [`RBTree`] allows the same kinds of access to its values that a struct allows to its
// fields, so we use the same Send condition as would be used for a struct with K and V fields.
unsafe impl<K: Send, V: Send> Send for RBTree<K, V> {}
// SAFETY: An [`RBTree`] allows the same kinds of access to its values that a struct allows to its
// fields, so we use the same Sync condition as would be used for a struct with K and V fields.
unsafe impl<K: Sync, V: Sync> Sync for RBTree<K, V> {}
impl<K, V> RBTree<K, V> {
/// Creates a new and empty tree.
pub fn new() -> Self {
Self {
// INVARIANT: There are no nodes in the tree, so the invariant holds vacuously.
root: bindings::rb_root::default(),
_p: PhantomData,
}
}
/// Allocates memory for a node to be eventually initialised and inserted into the tree via a
/// call to [`RBTree::insert`].
pub fn try_reserve_node() -> Result<RBTreeNodeReservation<K, V>> {
Ok(RBTreeNodeReservation {
node: Box::try_new(MaybeUninit::uninit())?,
})
}
/// Allocates and initialises a node that can be inserted into the tree via
/// [`RBTree::insert`].
pub fn try_allocate_node(key: K, value: V) -> Result<RBTreeNode<K, V>> {
Ok(Self::try_reserve_node()?.into_node(key, value))
}
}
impl<K, V> RBTree<K, V>
where
K: Ord,
{
/// Tries to insert a new value into the tree.
///
/// It overwrites a node if one already exists with the same key and returns it (containing the
/// key/value pair). Returns [`None`] if a node with the same key didn't already exist.
///
/// Returns an error if it cannot allocate memory for the new node.
pub fn try_create_and_insert(&mut self, key: K, value: V) -> Result<Option<RBTreeNode<K, V>>> {
Ok(self.insert(Self::try_allocate_node(key, value)?))
}
/// Inserts a new node into the tree.
///
/// It overwrites a node if one already exists with the same key and returns it (containing the
/// key/value pair). Returns [`None`] if a node with the same key didn't already exist.
///
/// This function always succeeds.
pub fn insert(&mut self, node: RBTreeNode<K, V>) -> Option<RBTreeNode<K, V>> {
let RBTreeNode { node } = node;
let node = Box::into_raw(node);
// SAFETY: `node` is valid at least until we call `Box::from_raw`, which only happens when
// the node is removed or replaced.
let node_links = unsafe { addr_of_mut!((*node).links) };
let mut new_link: &mut *mut bindings::rb_node = &mut self.root.rb_node;
let mut parent = core::ptr::null_mut();
while !new_link.is_null() {
// SAFETY: All links fields we create are in a `Node<K, V>`.
let this = unsafe { crate::container_of!(*new_link, Node<K, V>, links) };
parent = *new_link;
// SAFETY: `this` is a non-null node so it is valid by the type invariants. `node` is
// valid until the node is removed.
match unsafe { (*node).key.cmp(&(*this).key) } {
// SAFETY: `parent` is a non-null node so it is valid by the type invariants.
Ordering::Less => new_link = unsafe { &mut (*parent).rb_left },
// SAFETY: `parent` is a non-null node so it is valid by the type invariants.
Ordering::Greater => new_link = unsafe { &mut (*parent).rb_right },
Ordering::Equal => {
// INVARIANT: We are replacing an existing node with a new one, which is valid.
// It remains valid because we "forgot" it with `Box::into_raw`.
// SAFETY: All pointers are non-null and valid (parent, despite the name, really
// is the node we're replacing).
unsafe { bindings::rb_replace_node(parent, node_links, &mut self.root) };
// INVARIANT: The node is being returned and the caller may free it, however,
// it was removed from the tree. So the invariants still hold.
return Some(RBTreeNode {
// SAFETY: `this` was a node in the tree, so it is valid.
node: unsafe { Box::from_raw(this as _) },
});
}
}
}
// INVARIANT: We are linking in a new node, which is valid. It remains valid because we
// "forgot" it with `Box::into_raw`.
// SAFETY: All pointers are non-null and valid (`*new_link` is null, but `new_link` is a
// mutable reference).
unsafe { bindings::rb_link_node(node_links, parent, new_link) };
// SAFETY: All pointers are valid. `node` has just been inserted into the tree.
unsafe { bindings::rb_insert_color(node_links, &mut self.root) };
None
}
/// Returns a node with the given key, if one exists.
fn find(&self, key: &K) -> Option<NonNull<Node<K, V>>> {
let mut node = self.root.rb_node;
while !node.is_null() {
// SAFETY: All links fields we create are in a `Node<K, V>`.
let this = unsafe { crate::container_of!(node, Node<K, V>, links) };
// SAFETY: `this` is a non-null node so it is valid by the type invariants.
node = match key.cmp(unsafe { &(*this).key }) {
// SAFETY: `node` is a non-null node so it is valid by the type invariants.
Ordering::Less => unsafe { (*node).rb_left },
// SAFETY: `node` is a non-null node so it is valid by the type invariants.
Ordering::Greater => unsafe { (*node).rb_right },
Ordering::Equal => return NonNull::new(this as _),
}
}
None
}
/// Returns a reference to the value corresponding to the key.
pub fn get(&self, key: &K) -> Option<&V> {
// SAFETY: The `find` return value is a node in the tree, so it is valid.
self.find(key).map(|node| unsafe { &node.as_ref().value })
}
/// Returns a mutable reference to the value corresponding to the key.
pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
// SAFETY: The `find` return value is a node in the tree, so it is valid.
self.find(key)
.map(|mut node| unsafe { &mut node.as_mut().value })
}
/// Removes the node with the given key from the tree.
///
/// It returns the node that was removed if one exists, or [`None`] otherwise.
fn remove_node(&mut self, key: &K) -> Option<RBTreeNode<K, V>> {
let mut node = self.find(key)?;
// SAFETY: The `find` return value is a node in the tree, so it is valid.
unsafe { bindings::rb_erase(&mut node.as_mut().links, &mut self.root) };
// INVARIANT: The node is being returned and the caller may free it, however, it was
// removed from the tree. So the invariants still hold.
Some(RBTreeNode {
// SAFETY: The `find` return value was a node in the tree, so it is valid.
node: unsafe { Box::from_raw(node.as_ptr()) },
})
}
/// Removes the node with the given key from the tree.
///
/// It returns the value that was removed if one exists, or [`None`] otherwise.
pub fn remove(&mut self, key: &K) -> Option<V> {
let node = self.remove_node(key)?;
let RBTreeNode { node } = node;
let Node {
links: _,
key: _,
value,
} = *node;
Some(value)
}
}
impl<K, V> Default for RBTree<K, V> {
fn default() -> Self {
Self::new()
}
}
impl<K, V> Drop for RBTree<K, V> {
fn drop(&mut self) {
// SAFETY: `root` is valid as it's embedded in `self` and we have a valid `self`.
let mut next = unsafe { bindings::rb_first_postorder(&self.root) };
// INVARIANT: The loop invariant is that all tree nodes from `next` in postorder are valid.
while !next.is_null() {
// SAFETY: All links fields we create are in a `Node<K, V>`.
let this = unsafe { crate::container_of!(next, Node<K, V>, links) };
// Find out what the next node is before disposing of the current one.
// SAFETY: `next` and all nodes in postorder are still valid.
next = unsafe { bindings::rb_next_postorder(next) };
// INVARIANT: This is the destructor, so we break the type invariant during clean-up,
// but it is not observable. The loop invariant is still maintained.
// SAFETY: `this` is valid per the loop invariant.
unsafe { drop(Box::from_raw(this as *mut Node<K, V>)) };
}
}
}
/// A memory reservation for a red-black tree node.
///
/// It contains the memory needed to hold a node that can be inserted into a red-black tree. One
/// can be obtained by directly allocating it ([`RBTree::try_reserve_node`]).
pub struct RBTreeNodeReservation<K, V> {
node: Box<MaybeUninit<Node<K, V>>>,
}
// SAFETY: An [`RBTree`] allows the same kinds of access to its values that a struct allows to its
// fields, so we use the same Send condition as would be used for a struct with K and V fields.
unsafe impl<K: Send, V: Send> Send for RBTreeNodeReservation<K, V> {}
// SAFETY: An [`RBTree`] allows the same kinds of access to its values that a struct allows to its
// fields, so we use the same Sync condition as would be used for a struct with K and V fields.
unsafe impl<K: Sync, V: Sync> Sync for RBTreeNodeReservation<K, V> {}
impl<K, V> RBTreeNodeReservation<K, V> {
/// Initialises a node reservation.
///
/// It then becomes an [`RBTreeNode`] that can be inserted into a tree.
pub fn into_node(mut self, key: K, value: V) -> RBTreeNode<K, V> {
let node_ptr = self.node.as_mut_ptr();
// SAFETY: `node_ptr` is valid, and so are its fields.
unsafe { addr_of_mut!((*node_ptr).links).write(bindings::rb_node::default()) };
// SAFETY: `node_ptr` is valid, and so are its fields.
unsafe { addr_of_mut!((*node_ptr).key).write(key) };
// SAFETY: `node_ptr` is valid, and so are its fields.
unsafe { addr_of_mut!((*node_ptr).value).write(value) };
let raw = Box::into_raw(self.node);
RBTreeNode {
// SAFETY: The pointer came from a `MaybeUninit<Node>` whose fields have all been
// initialised. Additionally, it has the same layout as `Node`.
node: unsafe { Box::from_raw(raw as _) },
}
}
}
/// A red-black tree node.
///
/// The node is fully initialised (with key and value) and can be inserted into a tree without any
/// extra allocations or failure paths.
pub struct RBTreeNode<K, V> {
node: Box<Node<K, V>>,
}
// SAFETY: An [`RBTree`] allows the same kinds of access to its values that a struct allows to its
// fields, so we use the same Send condition as would be used for a struct with K and V fields.
unsafe impl<K: Send, V: Send> Send for RBTreeNode<K, V> {}
// SAFETY: An [`RBTree`] allows the same kinds of access to its values that a struct allows to its
// fields, so we use the same Sync condition as would be used for a struct with K and V fields.
unsafe impl<K: Sync, V: Sync> Sync for RBTreeNode<K, V> {}