From 876346536c1b59a5b1b5e44477b1b3ece77647fd Mon Sep 17 00:00:00 2001 From: Andreas Hindborg Date: Thu, 15 Aug 2024 10:30:26 +0000 Subject: rust: kbuild: split up helpers.c This patch splits up the rust helpers C file. When rebasing patch sets on upstream linux, merge conflicts in helpers.c is common and time consuming [1]. Thus, split the file so that each kernel component can live in a separate file. This patch lists helper files explicitly and thus conflicts in the file list is still likely. However, they should be more simple to resolve than the conflicts usually seen in helpers.c. [ Removed `README.md` and undeleted the original comment since now, in v3 of the series, we have a `helpers.c` again; which also allows us to keep the "Sorted alphabetically" line and makes the diff easier. In addition, updated the Documentation/ mentions of the file, reworded title and removed blank lines at the end of `page.c`. - Miguel ] Link: https://rust-for-linux.zulipchat.com/#narrow/stream/288089-General/topic/Splitting.20up.20helpers.2Ec/near/426694012 [1] Signed-off-by: Andreas Hindborg Reviewed-by: Gary Guo Acked-by: Dirk Behme Reviewed-by: Alice Ryhl Reviewed-by: Benno Lossin Link: https://lore.kernel.org/r/20240815103016.2771842-1-nmi@metaspace.dk Signed-off-by: Miguel Ojeda --- rust/helpers/helpers.c | 38 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 rust/helpers/helpers.c (limited to 'rust/helpers/helpers.c') diff --git a/rust/helpers/helpers.c b/rust/helpers/helpers.c new file mode 100644 index 000000000000..2b54f22e8774 --- /dev/null +++ b/rust/helpers/helpers.c @@ -0,0 +1,38 @@ +// SPDX-License-Identifier: GPL-2.0 +/* + * Non-trivial C macros cannot be used in Rust. Similarly, inlined C functions + * cannot be called either. This file explicitly creates functions ("helpers") + * that wrap those so that they can be called from Rust. + * + * Even though Rust kernel modules should never use the bindings directly, some + * of these helpers need to be exported because Rust generics and inlined + * functions may not get their code generated in the crate where they are + * defined. Other helpers, called from non-inline functions, may not be + * exported, in principle. However, in general, the Rust compiler does not + * guarantee codegen will be performed for a non-inline function either. + * Therefore, this file exports all the helpers. In the future, this may be + * revisited to reduce the number of exports after the compiler is informed + * about the places codegen is required. + * + * All symbols are exported as GPL-only to guarantee no GPL-only feature is + * accidentally exposed. + * + * Sorted alphabetically. + */ + +#include "blk.c" +#include "bug.c" +#include "build_assert.c" +#include "build_bug.c" +#include "err.c" +#include "kunit.c" +#include "mutex.c" +#include "page.c" +#include "refcount.c" +#include "signal.c" +#include "slab.c" +#include "spinlock.c" +#include "task.c" +#include "uaccess.c" +#include "wait.c" +#include "workqueue.c" -- cgit From e26fa546042add70944d018b930530d16b3cf626 Mon Sep 17 00:00:00 2001 From: Gary Guo Date: Sat, 17 Aug 2024 17:51:32 +0100 Subject: rust: kbuild: auto generate helper exports This removes the need to explicitly export all symbols. Generate helper exports similarly to what's currently done for Rust crates. These helpers are exclusively called from within Rust code and therefore can be treated similar as other Rust symbols. Signed-off-by: Gary Guo Reviewed-by: Boqun Feng Tested-by: Boqun Feng Link: https://lore.kernel.org/r/20240817165302.3852499-1-gary@garyguo.net [ Fixed dependency path, reworded slightly, edited comment a bit and rebased on top of the changes made when applying Andreas' patch (e.g. no `README.md` anymore, so moved the edits). - Miguel ] Signed-off-by: Miguel Ojeda --- rust/Makefile | 16 ++++++++++++++-- rust/exports.c | 1 + rust/helpers/blk.c | 2 -- rust/helpers/bug.c | 1 - rust/helpers/build_bug.c | 1 - rust/helpers/err.c | 3 --- rust/helpers/helpers.c | 13 ------------- rust/helpers/kunit.c | 1 - rust/helpers/mutex.c | 1 - rust/helpers/page.c | 3 --- rust/helpers/refcount.c | 3 --- rust/helpers/signal.c | 1 - rust/helpers/slab.c | 1 - rust/helpers/spinlock.c | 3 --- rust/helpers/task.c | 3 --- rust/helpers/uaccess.c | 2 -- rust/helpers/wait.c | 1 - rust/helpers/workqueue.c | 1 - 18 files changed, 15 insertions(+), 42 deletions(-) (limited to 'rust/helpers/helpers.c') diff --git a/rust/Makefile b/rust/Makefile index 99204e33f1dd..7ea5905f544c 100644 --- a/rust/Makefile +++ b/rust/Makefile @@ -16,8 +16,8 @@ no-clean-files += libmacros.so always-$(CONFIG_RUST) += bindings/bindings_generated.rs bindings/bindings_helpers_generated.rs obj-$(CONFIG_RUST) += alloc.o bindings.o kernel.o -always-$(CONFIG_RUST) += exports_alloc_generated.h exports_bindings_generated.h \ - exports_kernel_generated.h +always-$(CONFIG_RUST) += exports_alloc_generated.h exports_helpers_generated.h \ + exports_bindings_generated.h exports_kernel_generated.h always-$(CONFIG_RUST) += uapi/uapi_generated.rs obj-$(CONFIG_RUST) += uapi.o @@ -313,6 +313,18 @@ $(obj)/exports_core_generated.h: $(obj)/core.o FORCE $(obj)/exports_alloc_generated.h: $(obj)/alloc.o FORCE $(call if_changed,exports) +# Even though Rust kernel modules should never use the bindings directly, +# symbols from the `bindings` crate and the C helpers need to be exported +# because Rust generics and inlined functions may not get their code generated +# in the crate where they are defined. Other helpers, called from non-inline +# functions, may not be exported, in principle. However, in general, the Rust +# compiler does not guarantee codegen will be performed for a non-inline +# function either. Therefore, we export all symbols from helpers and bindings. +# In the future, this may be revisited to reduce the number of exports after +# the compiler is informed about the places codegen is required. +$(obj)/exports_helpers_generated.h: $(obj)/helpers/helpers.o FORCE + $(call if_changed,exports) + $(obj)/exports_bindings_generated.h: $(obj)/bindings.o FORCE $(call if_changed,exports) diff --git a/rust/exports.c b/rust/exports.c index 3803c21d1403..e5695f3b45b7 100644 --- a/rust/exports.c +++ b/rust/exports.c @@ -17,6 +17,7 @@ #include "exports_core_generated.h" #include "exports_alloc_generated.h" +#include "exports_helpers_generated.h" #include "exports_bindings_generated.h" #include "exports_kernel_generated.h" diff --git a/rust/helpers/blk.c b/rust/helpers/blk.c index d99c965eb59b..cc9f4e6a2d23 100644 --- a/rust/helpers/blk.c +++ b/rust/helpers/blk.c @@ -7,10 +7,8 @@ void *rust_helper_blk_mq_rq_to_pdu(struct request *rq) { return blk_mq_rq_to_pdu(rq); } -EXPORT_SYMBOL_GPL(rust_helper_blk_mq_rq_to_pdu); struct request *rust_helper_blk_mq_rq_from_pdu(void *pdu) { return blk_mq_rq_from_pdu(pdu); } -EXPORT_SYMBOL_GPL(rust_helper_blk_mq_rq_from_pdu); diff --git a/rust/helpers/bug.c b/rust/helpers/bug.c index e2afbad23dcd..e2d13babc737 100644 --- a/rust/helpers/bug.c +++ b/rust/helpers/bug.c @@ -6,4 +6,3 @@ __noreturn void rust_helper_BUG(void) { BUG(); } -EXPORT_SYMBOL_GPL(rust_helper_BUG); diff --git a/rust/helpers/build_bug.c b/rust/helpers/build_bug.c index f3106f248485..e994f7b5928c 100644 --- a/rust/helpers/build_bug.c +++ b/rust/helpers/build_bug.c @@ -7,4 +7,3 @@ const char *rust_helper_errname(int err) { return errname(err); } -EXPORT_SYMBOL_GPL(rust_helper_errname); diff --git a/rust/helpers/err.c b/rust/helpers/err.c index fba4e0be64f5..be3d45ef78a2 100644 --- a/rust/helpers/err.c +++ b/rust/helpers/err.c @@ -7,16 +7,13 @@ __force void *rust_helper_ERR_PTR(long err) { return ERR_PTR(err); } -EXPORT_SYMBOL_GPL(rust_helper_ERR_PTR); bool rust_helper_IS_ERR(__force const void *ptr) { return IS_ERR(ptr); } -EXPORT_SYMBOL_GPL(rust_helper_IS_ERR); long rust_helper_PTR_ERR(__force const void *ptr) { return PTR_ERR(ptr); } -EXPORT_SYMBOL_GPL(rust_helper_PTR_ERR); diff --git a/rust/helpers/helpers.c b/rust/helpers/helpers.c index 2b54f22e8774..173533616c91 100644 --- a/rust/helpers/helpers.c +++ b/rust/helpers/helpers.c @@ -4,19 +4,6 @@ * cannot be called either. This file explicitly creates functions ("helpers") * that wrap those so that they can be called from Rust. * - * Even though Rust kernel modules should never use the bindings directly, some - * of these helpers need to be exported because Rust generics and inlined - * functions may not get their code generated in the crate where they are - * defined. Other helpers, called from non-inline functions, may not be - * exported, in principle. However, in general, the Rust compiler does not - * guarantee codegen will be performed for a non-inline function either. - * Therefore, this file exports all the helpers. In the future, this may be - * revisited to reduce the number of exports after the compiler is informed - * about the places codegen is required. - * - * All symbols are exported as GPL-only to guarantee no GPL-only feature is - * accidentally exposed. - * * Sorted alphabetically. */ diff --git a/rust/helpers/kunit.c b/rust/helpers/kunit.c index 905e4ff4424a..9d725067eb3b 100644 --- a/rust/helpers/kunit.c +++ b/rust/helpers/kunit.c @@ -7,4 +7,3 @@ struct kunit *rust_helper_kunit_get_current_test(void) { return kunit_get_current_test(); } -EXPORT_SYMBOL_GPL(rust_helper_kunit_get_current_test); diff --git a/rust/helpers/mutex.c b/rust/helpers/mutex.c index 29fd141c387d..200db7e6279f 100644 --- a/rust/helpers/mutex.c +++ b/rust/helpers/mutex.c @@ -7,4 +7,3 @@ void rust_helper_mutex_lock(struct mutex *lock) { mutex_lock(lock); } -EXPORT_SYMBOL_GPL(rust_helper_mutex_lock); diff --git a/rust/helpers/page.c b/rust/helpers/page.c index 7fd333411a88..b3f2b8fbf87f 100644 --- a/rust/helpers/page.c +++ b/rust/helpers/page.c @@ -7,16 +7,13 @@ struct page *rust_helper_alloc_pages(gfp_t gfp_mask, unsigned int order) { return alloc_pages(gfp_mask, order); } -EXPORT_SYMBOL_GPL(rust_helper_alloc_pages); void *rust_helper_kmap_local_page(struct page *page) { return kmap_local_page(page); } -EXPORT_SYMBOL_GPL(rust_helper_kmap_local_page); void rust_helper_kunmap_local(const void *addr) { kunmap_local(addr); } -EXPORT_SYMBOL_GPL(rust_helper_kunmap_local); diff --git a/rust/helpers/refcount.c b/rust/helpers/refcount.c index 13ab64805f77..f47afc148ec3 100644 --- a/rust/helpers/refcount.c +++ b/rust/helpers/refcount.c @@ -7,16 +7,13 @@ refcount_t rust_helper_REFCOUNT_INIT(int n) { return (refcount_t)REFCOUNT_INIT(n); } -EXPORT_SYMBOL_GPL(rust_helper_REFCOUNT_INIT); void rust_helper_refcount_inc(refcount_t *r) { refcount_inc(r); } -EXPORT_SYMBOL_GPL(rust_helper_refcount_inc); bool rust_helper_refcount_dec_and_test(refcount_t *r) { return refcount_dec_and_test(r); } -EXPORT_SYMBOL_GPL(rust_helper_refcount_dec_and_test); diff --git a/rust/helpers/signal.c b/rust/helpers/signal.c index d44e8096b8a9..63c407f80c26 100644 --- a/rust/helpers/signal.c +++ b/rust/helpers/signal.c @@ -7,4 +7,3 @@ int rust_helper_signal_pending(struct task_struct *t) { return signal_pending(t); } -EXPORT_SYMBOL_GPL(rust_helper_signal_pending); diff --git a/rust/helpers/slab.c b/rust/helpers/slab.c index 3e0a1a173d8a..f043e087f9d6 100644 --- a/rust/helpers/slab.c +++ b/rust/helpers/slab.c @@ -7,4 +7,3 @@ rust_helper_krealloc(const void *objp, size_t new_size, gfp_t flags) { return krealloc(objp, new_size, flags); } -EXPORT_SYMBOL_GPL(rust_helper_krealloc); diff --git a/rust/helpers/spinlock.c b/rust/helpers/spinlock.c index 04fd8ddb4986..acc1376b833c 100644 --- a/rust/helpers/spinlock.c +++ b/rust/helpers/spinlock.c @@ -12,16 +12,13 @@ void rust_helper___spin_lock_init(spinlock_t *lock, const char *name, spin_lock_init(lock); #endif } -EXPORT_SYMBOL_GPL(rust_helper___spin_lock_init); void rust_helper_spin_lock(spinlock_t *lock) { spin_lock(lock); } -EXPORT_SYMBOL_GPL(rust_helper_spin_lock); void rust_helper_spin_unlock(spinlock_t *lock) { spin_unlock(lock); } -EXPORT_SYMBOL_GPL(rust_helper_spin_unlock); diff --git a/rust/helpers/task.c b/rust/helpers/task.c index b176c347f0d4..7ac789232d11 100644 --- a/rust/helpers/task.c +++ b/rust/helpers/task.c @@ -7,16 +7,13 @@ struct task_struct *rust_helper_get_current(void) { return current; } -EXPORT_SYMBOL_GPL(rust_helper_get_current); void rust_helper_get_task_struct(struct task_struct *t) { get_task_struct(t); } -EXPORT_SYMBOL_GPL(rust_helper_get_task_struct); void rust_helper_put_task_struct(struct task_struct *t) { put_task_struct(t); } -EXPORT_SYMBOL_GPL(rust_helper_put_task_struct); diff --git a/rust/helpers/uaccess.c b/rust/helpers/uaccess.c index 3d004ac1c180..f49076f813cd 100644 --- a/rust/helpers/uaccess.c +++ b/rust/helpers/uaccess.c @@ -7,11 +7,9 @@ unsigned long rust_helper_copy_from_user(void *to, const void __user *from, { return copy_from_user(to, from, n); } -EXPORT_SYMBOL_GPL(rust_helper_copy_from_user); unsigned long rust_helper_copy_to_user(void __user *to, const void *from, unsigned long n) { return copy_to_user(to, from, n); } -EXPORT_SYMBOL_GPL(rust_helper_copy_to_user); diff --git a/rust/helpers/wait.c b/rust/helpers/wait.c index bf361f40c7cb..c7336bbf2750 100644 --- a/rust/helpers/wait.c +++ b/rust/helpers/wait.c @@ -7,4 +7,3 @@ void rust_helper_init_wait(struct wait_queue_entry *wq_entry) { init_wait(wq_entry); } -EXPORT_SYMBOL_GPL(rust_helper_init_wait); diff --git a/rust/helpers/workqueue.c b/rust/helpers/workqueue.c index 12e2ee66aa4f..f59427acc323 100644 --- a/rust/helpers/workqueue.c +++ b/rust/helpers/workqueue.c @@ -13,4 +13,3 @@ void rust_helper_init_work_with_key(struct work_struct *work, work_func_t func, INIT_LIST_HEAD(&work->entry); work->func = func; } -EXPORT_SYMBOL_GPL(rust_helper_init_work_with_key); -- cgit From a0d13aac7022f95ec161c18d18e3d81172666ed8 Mon Sep 17 00:00:00 2001 From: Wedson Almeida Filho Date: Thu, 22 Aug 2024 16:37:53 +0000 Subject: rust: rbtree: add red-black tree implementation backed by the C version 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] Signed-off-by: Wedson Almeida Filho Reviewed-by: Alice Ryhl Tested-by: Alice Ryhl Reviewed-by: Boqun Feng Reviewed-by: Benno Lossin Signed-off-by: Matt Gilbride Link: https://lore.kernel.org/r/20240822-b4-rbtree-v12-1-014561758a57@google.com [ Updated link to docs.kernel.org. - Miguel ] Signed-off-by: Miguel Ojeda --- rust/helpers/helpers.c | 1 + rust/helpers/rbtree.c | 9 ++ rust/kernel/lib.rs | 1 + rust/kernel/rbtree.rs | 432 +++++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 443 insertions(+) create mode 100644 rust/helpers/rbtree.c create mode 100644 rust/kernel/rbtree.rs (limited to 'rust/helpers/helpers.c') diff --git a/rust/helpers/helpers.c b/rust/helpers/helpers.c index 173533616c91..30f40149f3a9 100644 --- a/rust/helpers/helpers.c +++ b/rust/helpers/helpers.c @@ -15,6 +15,7 @@ #include "kunit.c" #include "mutex.c" #include "page.c" +#include "rbtree.c" #include "refcount.c" #include "signal.c" #include "slab.c" diff --git a/rust/helpers/rbtree.c b/rust/helpers/rbtree.c new file mode 100644 index 000000000000..6d404b84a9b5 --- /dev/null +++ b/rust/helpers/rbtree.c @@ -0,0 +1,9 @@ +// SPDX-License-Identifier: GPL-2.0 + +#include + +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); +} diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs index 9baea9e9ee1a..f10b06a78b9d 100644 --- a/rust/kernel/lib.rs +++ b/rust/kernel/lib.rs @@ -44,6 +44,7 @@ pub mod net; pub mod page; pub mod prelude; pub mod print; +pub mod rbtree; mod static_assert; #[doc(hidden)] pub mod std_vendor; diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs new file mode 100644 index 000000000000..cf25437c795f --- /dev/null +++ b/rust/kernel/rbtree.rs @@ -0,0 +1,432 @@ +// SPDX-License-Identifier: GPL-2.0 + +//! Red-black trees. +//! +//! C header: [`include/linux/rbtree.h`](srctree/include/linux/rbtree.h) +//! +//! Reference: + +use crate::{alloc::Flags, bindings, container_of, error::Result, prelude::*}; +use alloc::boxed::Box; +use core::{ + cmp::{Ord, Ordering}, + marker::PhantomData, + mem::MaybeUninit, + ptr::{addr_of_mut, NonNull}, +}; + +/// A red-black tree with owned nodes. +/// +/// It is backed by the kernel C red-black trees. +/// +/// # 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::{alloc::flags, rbtree::{RBTree, RBTreeNode, RBTreeNodeReservation}}; +/// +/// // Create a new tree. +/// let mut tree = RBTree::new(); +/// +/// // Insert three elements. +/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; +/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; +/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; +/// +/// // 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, flags::GFP_KERNEL)?; +/// +/// // 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::{alloc::flags, rbtree::{RBTree, RBTreeNode}, sync::SpinLock}; +/// +/// fn insert_test(tree: &SpinLock>) -> Result { +/// // Pre-allocate node. This may fail (as it allocates memory). +/// let node = RBTreeNode::new(10, 100, flags::GFP_KERNEL)?; +/// +/// // 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::{alloc::flags, rbtree::{RBTree, RBTreeNodeReservation}}; +/// +/// // Create a new tree. +/// let mut tree = RBTree::new(); +/// +/// // Insert three elements. +/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; +/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; +/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; +/// +/// // 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(&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); +/// } +/// +/// // Create a preallocated reservation that we can re-use later. +/// let reservation = RBTreeNodeReservation::new(flags::GFP_KERNEL)?; +/// +/// // 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>(()) +/// ``` +/// +/// # 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. +pub struct RBTree { + root: bindings::rb_root, + _p: PhantomData>, +} + +// 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 Send for RBTree {} + +// 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 Sync for RBTree {} + +impl RBTree { + /// 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, + } + } +} + +impl RBTree +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, + flags: Flags, + ) -> Result>> { + Ok(self.insert(RBTreeNode::new(key, value, flags)?)) + } + + /// 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, RBTreeNode { node }: RBTreeNode) -> Option> { + 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) }; + + // The parameters of `bindings::rb_link_node` are as follows: + // - `node`: A pointer to an uninitialized node being inserted. + // - `parent`: A pointer to an existing node in the tree. One of its child pointers must be + // null, and `node` will become a child of `parent` by replacing that child pointer + // with a pointer to `node`. + // - `rb_link`: A pointer to either the left-child or right-child field of `parent`. This + // specifies which child of `parent` should hold `node` after this call. The + // value of `*rb_link` must be null before the call to `rb_link_node`. If the + // red/black tree is empty, then it’s also possible for `parent` to be null. In + // this case, `rb_link` is a pointer to the `root` field of the red/black tree. + // + // We will traverse the tree looking for a node that has a null pointer as its child, + // representing an empty subtree where we can insert our new node. We need to make sure + // that we preserve the ordering of the nodes in the tree. In each iteration of the loop + // we store `parent` and `child_field_of_parent`, and the new `node` will go somewhere + // in the subtree of `parent` that `child_field_of_parent` points at. Once + // we find an empty subtree, we can insert the new node using `rb_link_node`. + let mut parent = core::ptr::null_mut(); + let mut child_field_of_parent: &mut *mut bindings::rb_node = &mut self.root.rb_node; + while !child_field_of_parent.is_null() { + parent = *child_field_of_parent; + + // We need to determine whether `node` should be the left or right child of `parent`, + // so we will compare with the `key` field of `parent` a.k.a. `this` below. + // + // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self` + // point to the links field of `Node` objects. + let this = unsafe { container_of!(parent, Node, links) }; + + // 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) } { + // We would like `node` to be the left child of `parent`. Move to this child to check + // whether we can use it, or continue searching, at the next iteration. + // + // SAFETY: `parent` is a non-null node so it is valid by the type invariants. + Ordering::Less => child_field_of_parent = unsafe { &mut (*parent).rb_left }, + // We would like `node` to be the right child of `parent`. Move to this child to check + // whether we can use it, or continue searching, at the next iteration. + // + // SAFETY: `parent` is a non-null node so it is valid by the type invariants. + Ordering::Greater => child_field_of_parent = unsafe { &mut (*parent).rb_right }, + Ordering::Equal => { + // There is an existing node in the tree with this key, and that node is + // `parent`. Thus, we are replacing parent with a new node. + // + // 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. + 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.cast_mut()) }, + }); + } + } + } + + // 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 (`*child_field_of_parent` is null, but `child_field_of_parent` is a + // mutable reference). + unsafe { bindings::rb_link_node(node_links, parent, child_field_of_parent) }; + + // 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>> { + let mut node = self.root.rb_node; + while !node.is_null() { + // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self` + // point to the links field of `Node` objects. + let this = unsafe { container_of!(node, Node, 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.cast_mut()), + } + } + 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> { + 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 { + self.remove_node(key).map(|node| node.node.value) + } +} + +impl Default for RBTree { + fn default() -> Self { + Self::new() + } +} + +impl Drop for RBTree { + 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`. + let this = unsafe { container_of!(next, Node, 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.cast_mut())) }; + } + } +} + +/// 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 ([`RBTreeNodeReservation::new`]). +pub struct RBTreeNodeReservation { + node: Box>>, +} + +impl RBTreeNodeReservation { + /// Allocates memory for a node to be eventually initialised and inserted into the tree via a + /// call to [`RBTree::insert`]. + pub fn new(flags: Flags) -> Result> { + Ok(RBTreeNodeReservation { + node: as BoxExt<_>>::new_uninit(flags)?, + }) + } +} + +// SAFETY: This doesn't actually contain K or V, and is just a memory allocation. Those can always +// be moved across threads. +unsafe impl Send for RBTreeNodeReservation {} + +// SAFETY: This doesn't actually contain K or V, and is just a memory allocation. +unsafe impl Sync for RBTreeNodeReservation {} + +impl RBTreeNodeReservation { + /// Initialises a node reservation. + /// + /// It then becomes an [`RBTreeNode`] that can be inserted into a tree. + pub fn into_node(self, key: K, value: V) -> RBTreeNode { + let node = Box::write( + self.node, + Node { + key, + value, + links: bindings::rb_node::default(), + }, + ); + RBTreeNode { node } + } +} + +/// 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 { + node: Box>, +} + +impl RBTreeNode { + /// Allocates and initialises a node that can be inserted into the tree via + /// [`RBTree::insert`]. + pub fn new(key: K, value: V, flags: Flags) -> Result> { + Ok(RBTreeNodeReservation::new(flags)?.into_node(key, value)) + } +} + +// SAFETY: If K and V can be sent across threads, then it's also okay to send [`RBTreeNode`] across +// threads. +unsafe impl Send for RBTreeNode {} + +// SAFETY: If K and V can be accessed without synchronization, then it's also okay to access +// [`RBTreeNode`] without synchronization. +unsafe impl Sync for RBTreeNode {} + +struct Node { + links: bindings::rb_node, + key: K, + value: V, +} -- cgit