[PATCH v8 0/6] Red-black tree abstraction needed by Rust Binder

[PATCH v8 0/6] Red-black tree abstraction needed by Rust Binder

[Matt Gilbride]

This patchset contains the red-black tree abstractions needed by the Rust
implementation of the Binder driver.

Binder driver benefits from O(log n) search/insertion/deletion of
key/value mappings in various places, including `process.rs` and
`range_alloc.rs`.  In `range_alloc.rs`, the ability to store and
search by a generic key type is also useful.

Please see the Rust Binder RFC for usage examples [1]. Note that
the `container_of` macro is currently used only by `rbtree` itself.

Users of "rust: rbtree: add red-black tree implementation backed by the C version"
    [PATCH RFC 03/20] rust_binder: add threading support
    [PATCH RFC 05/20] rust_binder: add nodes and context managers
    [PATCH RFC 06/20] rust_binder: add oneway transactions

Users of "rust: rbtree: add iterator"
    [PATCH RFC 17/20] rust_binder: add oneway spam detection

Users of "rust: rbtree: add mutable iterator"
    [PATCH RFC 06/20] rust_binder: add oneway transactions

Users of "rust: rbtree: add `RBTreeCursor`"
    [PATCH RFC 06/20] rust_binder: add oneway transactions

Users of "rust: rbtree: add RBTree::entry"
    Not used in the original RFC, but introduced after further
    code review.  See: https://r.android.com/2849906

The Rust Binder RFC addresses the upstream deprecation of red-black
tree. Quoted here for convenience:

"This RFC uses the kernel's red-black tree for key/value mappings, but we
are aware that the red-black tree is deprecated. We did this to make the
performance comparison more fair, since C binder also uses rbtree for
this. We intend to replace these with XArrays instead. That said, we
don't think that XArray is a good fit for the range allocator, and we
propose to continue using the red-black tree for the range allocator."

Link: https://lore.kernel.org/rust-for-linux/20231101-rust-binder-v1-0-08ba9197f637@google.com/ [1]
Signed-off-by: Matt Gilbride <mattgilbride@google.com>
---
Changes in v8:
- Fix small style nit (use ? operator)
- Fix doc comment pointing at a private item
- Link to v7: https://lore.kernel.org/r/20240726-b4-rbtree-v7-0-aee88caaf97c@google.com

Changes in v7:
- make `RawVacantEntry.rbtree` a raw pointer like
  `RawVacantEntry.child_field_of_parent`, since the latter can
  technically point at a field of the former. We prefer that the
  implementation be explicit about the safety guarantees of both because
  of the relationship between them.
- Link to v6: https://lore.kernel.org/r/20240711-b4-rbtree-v6-0-14bef1a8cdba@google.com

Changes in v6:
- Minimize usage of `*mut bindings::rb_node`, replacing with
  `NonNull<bindings::rb_node>`. Specifically, changing
  `RBTreeCursor.current` to be `NonNull<bindings::rb_node>` and updating
  the corresponding functions.
- Update `RBTreeCursor:to_key_value` helpers to have their own lifetime
  (they are not instance methods, using a different lifetime than that
  of the `impl` block they are in makes things more clear.
- Fix misplaced semicolon in `cursor_lower_bound`.
- Link to v5: https://lore.kernel.org/r/20240606-b4-rbtree-v5-0-96fe1a0e97c0@google.com

Changes in v5:
- Used `Box::write` in `RBTreeNodeReservation::into_node`, removing
  unnecessary `unsafe` blocks.
- Updated `RBTreeCursor::remove_current` to return the removed node.
- Link to v4: https://lore.kernel.org/r/20240603-b4-rbtree-v4-0-308e43d6abfc@google.com

Changes in v4:
- rebased onto the tip of rust-for-linux/rust-next (97ab3e8eec0ce79d9e265e6c9e4c480492180409)
- addressed comments from draft PR on GitHub: https://github.com/Rust-for-Linux/linux/pull/1081
- Link to v3: https://lore.kernel.org/r/20240418-b4-rbtree-v3-0-323e134390ce@google.com

Changes in v3:
- Address various feedback re: SAFETY and INVARIANT comments from v2.
- Update variable naming and add detailed comments for the `RBTree::insert` (later moved to
  `RBTree::raw_entry`) implementation.
- Link to v2: https://lore.kernel.org/r/20240219-b4-rbtree-v2-0-0b113aab330d@google.com

Changes in v2:
- Update documentation link to the C header file
- Use `core::convert::Infallible` in try_reserve_node
- Link to v1: https://lore.kernel.org/r/20240205-b4-rbtree-v1-0-995e3eee38c0@google.com

---
Alice Ryhl (1):
      rust: rbtree: add `RBTree::entry`

Benno Lossin (1):
      rust: kernel: add `drop_contents` to `BoxExt`

Matt Gilbride (1):
      rust: rbtree: add `RBTreeCursor`

Wedson Almeida Filho (3):
      rust: rbtree: add red-black tree implementation backed by the C version
      rust: rbtree: add iterator
      rust: rbtree: add mutable iterator

 rust/helpers.c               |    7 +
 rust/kernel/alloc/box_ext.rs |   24 +-
 rust/kernel/lib.rs           |    1 +
 rust/kernel/rbtree.rs        | 1288 ++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 1319 insertions(+), 1 deletion(-)
---
base-commit: 97ab3e8eec0ce79d9e265e6c9e4c480492180409
change-id: 20231205-b4-rbtree-abb1a016f0a0

Best regards,
-- 
Matt Gilbride <mattgilbride@google.com>

[PATCH v8 1/6] rust: kernel: add `drop_contents` to `BoxExt`

[Matt Gilbride]

From: Benno Lossin <benno.lossin@proton.me>

Sometimes (see [1]) it is necessary to drop the value inside of a
`Box<T>`, but retain the allocation. For example to reuse the allocation
in the future.
Introduce a new function `drop_contents` that turns a `Box<T>` into
`Box<MaybeUninit<T>>` by dropping the value.

Signed-off-by: Benno Lossin <benno.lossin@proton.me>
Link: https://lore.kernel.org/rust-for-linux/20240418-b4-rbtree-v3-5-323e134390ce@google.com/ [1]
---
 rust/kernel/alloc/box_ext.rs | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/rust/kernel/alloc/box_ext.rs b/rust/kernel/alloc/box_ext.rs
index 829cb1c1cf9e..557895db4f48 100644
--- a/rust/kernel/alloc/box_ext.rs
+++ b/rust/kernel/alloc/box_ext.rs
@@ -4,7 +4,7 @@
 
 use super::{AllocError, Flags};
 use alloc::boxed::Box;
-use core::mem::MaybeUninit;
+use core::{mem::MaybeUninit, ptr};
 
 /// Extensions to [`Box`].
 pub trait BoxExt<T>: Sized {
@@ -17,6 +17,20 @@ pub trait BoxExt<T>: Sized {
     ///
     /// The allocation may fail, in which case an error is returned.
     fn new_uninit(flags: Flags) -> Result<Box<MaybeUninit<T>>, AllocError>;
+
+    /// Drops the contents, but keeps the allocation.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use kernel::alloc::flags;
+    ///
+    /// let value = Box::new([0; 32], flags::GFP_KERNEL);
+    /// let value = value.unwrap().drop_contents();
+    /// // Now we can re-use `value`:
+    /// Box::write(value, [1; 32]);
+    /// ```
+    fn drop_contents(self) -> Box<MaybeUninit<T>>;
 }
 
 impl<T> BoxExt<T> for Box<T> {
@@ -53,4 +67,12 @@ fn new_uninit(flags: Flags) -> Result<Box<MaybeUninit<T>>, AllocError> {
         // zero-sized types, we use `NonNull::dangling`.
         Ok(unsafe { Box::from_raw(ptr) })
     }
+
+    fn drop_contents(self) -> Box<MaybeUninit<T>> {
+        let ptr = Box::into_raw(self);
+        // SAFETY: `ptr` is valid, because it came from `Box::into_raw`.
+        unsafe { ptr::drop_in_place(ptr) };
+        // SAFETY: `ptr` is valid, because it came from `Box::into_raw`.
+        unsafe { Box::from_raw(ptr.cast()) }
+    }
 }

-- 
2.46.0.rc1.232.g9752f9e123-goog

[PATCH v8 2/6] rust: rbtree: add red-black tree implementation backed by the C version

[Matt Gilbride]

From: Wedson Almeida Filho <wedsonaf@gmail.com>

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 <wedsonaf@gmail.com>
Reviewed-by: Alice Ryhl <aliceryhl@google.com>
Tested-by: Alice Ryhl <aliceryhl@google.com>
Signed-off-by: Matt Gilbride <mattgilbride@google.com>
---
 rust/helpers.c        |   7 +
 rust/kernel/lib.rs    |   1 +
 rust/kernel/rbtree.rs | 432 ++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 440 insertions(+)

diff --git a/rust/helpers.c b/rust/helpers.c
index 4c8b7b92a4f4..608b38c0b3e8 100644
--- a/rust/helpers.c
+++ b/rust/helpers.c
@@ -157,6 +157,13 @@ void rust_helper_init_work_with_key(struct work_struct *work, work_func_t func,
 }
 EXPORT_SYMBOL_GPL(rust_helper_init_work_with_key);
 
+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.
diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
index 9a943d99c71a..dc2678803637 100644
--- a/rust/kernel/lib.rs
+++ b/rust/kernel/lib.rs
@@ -37,6 +37,7 @@
 pub mod net;
 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..74c53e4d5c00
--- /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: <https://www.kernel.org/doc/html/latest/core-api/rbtree.html>
+
+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.
+///
+/// # 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::{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<RBTree<u32, u32>>) -> 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>(())
+/// ```
+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,
+        }
+    }
+}
+
+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,
+        flags: Flags,
+    ) -> Result<Option<RBTreeNode<K, V>>> {
+        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<K, V>) -> Option<RBTreeNode<K, V>> {
+        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 `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<K, V>` objects.
+            let this = unsafe { container_of!(parent, Node<K, V>, 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<NonNull<Node<K, V>>> {
+        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<K, V>` objects.
+            let this = unsafe { 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.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<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> {
+        self.remove_node(key).map(|node| node.node.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 { 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.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<K, V> {
+    node: Box<MaybeUninit<Node<K, V>>>,
+}
+
+impl<K, V> RBTreeNodeReservation<K, V> {
+    /// 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<RBTreeNodeReservation<K, V>> {
+        Ok(RBTreeNodeReservation {
+            node: Box::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<K, V> Send for RBTreeNodeReservation<K, V> {}
+
+// SAFETY: This doesn't actually contain K or V, and is just a memory allocation.
+unsafe impl<K, V> 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(self, key: K, value: V) -> RBTreeNode<K, V> {
+        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<K, V> {
+    node: Box<Node<K, V>>,
+}
+
+impl<K, V> RBTreeNode<K, V> {
+    /// 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<RBTreeNode<K, V>> {
+        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<K: Send, V: Send> Send for RBTreeNode<K, V> {}
+
+// SAFETY: If K and V can be accessed without synchronization, then it's also okay to access
+// [`RBTreeNode`] without synchronization.
+unsafe impl<K: Sync, V: Sync> Sync for RBTreeNode<K, V> {}
+
+struct Node<K, V> {
+    links: bindings::rb_node,
+    key: K,
+    value: V,
+}

-- 
2.46.0.rc1.232.g9752f9e123-goog

[PATCH v8 3/6] rust: rbtree: add iterator

[Matt Gilbride]

From: Wedson Almeida Filho <wedsonaf@gmail.com>

- Add Iterator implementation for `RBTree`, allowing
  iteration over (key, value) pairs in key order.
- Add individual `keys()` and `values()` functions to iterate over keys
  or values alone.
- Update doctests to use iteration instead of explicitly getting items.

Iteration is needed by the binder driver to enumerate all values in a
tree for oneway spam detection [1].

Link: https://lore.kernel.org/rust-for-linux/20231101-rust-binder-v1-17-08ba9197f637@google.com/ [1]
Signed-off-by: Wedson Almeida Filho <wedsonaf@gmail.com>
Reviewed-by: Alice Ryhl <aliceryhl@google.com>
Tested-by: Alice Ryhl <aliceryhl@google.com>
Signed-off-by: Matt Gilbride <mattgilbride@google.com>
---
 rust/kernel/rbtree.rs | 130 +++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 112 insertions(+), 18 deletions(-)

diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs
index 74c53e4d5c00..d10074e4ac58 100644
--- a/rust/kernel/rbtree.rs
+++ b/rust/kernel/rbtree.rs
@@ -47,14 +47,30 @@
 ///     assert_eq!(tree.get(&30).unwrap(), &300);
 /// }
 ///
+/// // Iterate over the nodes we just inserted.
+/// {
+///     let mut iter = tree.iter();
+///     assert_eq!(iter.next().unwrap(), (&10, &100));
+///     assert_eq!(iter.next().unwrap(), (&20, &200));
+///     assert_eq!(iter.next().unwrap(), (&30, &300));
+///     assert!(iter.next().is_none());
+/// }
+///
+/// // Print all elements.
+/// for (key, value) in &tree {
+///     pr_info!("{} = {}\n", key, value);
+/// }
+///
 /// // 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);
+///     let mut iter = tree.iter();
+///     assert_eq!(iter.next().unwrap(), (&10, &1000));
+///     assert_eq!(iter.next().unwrap(), (&20, &200));
+///     assert_eq!(iter.next().unwrap(), (&30, &300));
+///     assert!(iter.next().is_none());
 /// }
 ///
 /// // Change the value of one of the elements.
@@ -62,9 +78,11 @@
 ///
 /// // 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);
+///     let mut iter = tree.iter();
+///     assert_eq!(iter.next().unwrap(), (&10, &1000));
+///     assert_eq!(iter.next().unwrap(), (&20, &200));
+///     assert_eq!(iter.next().unwrap(), (&30, &3000));
+///     assert!(iter.next().is_none());
 /// }
 ///
 /// // Remove an element.
@@ -72,9 +90,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);
+///     let mut iter = tree.iter();
+///     assert_eq!(iter.next().unwrap(), (&20, &200));
+///     assert_eq!(iter.next().unwrap(), (&30, &3000));
+///     assert!(iter.next().is_none());
 /// }
 ///
 /// # Ok::<(), Error>(())
@@ -114,9 +133,11 @@
 ///
 /// // 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);
+///     let mut iter = tree.iter();
+///     assert_eq!(iter.next().unwrap(), (&10, &100));
+///     assert_eq!(iter.next().unwrap(), (&20, &200));
+///     assert_eq!(iter.next().unwrap(), (&30, &300));
+///     assert!(iter.next().is_none());
 /// }
 ///
 /// // Remove a node, getting back ownership of it.
@@ -124,9 +145,10 @@
 ///
 /// // 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);
+///     let mut iter = tree.iter();
+///     assert_eq!(iter.next().unwrap(), (&10, &100));
+///     assert_eq!(iter.next().unwrap(), (&20, &200));
+///     assert!(iter.next().is_none());
 /// }
 ///
 /// // Create a preallocated reservation that we can re-use later.
@@ -138,9 +160,11 @@
 ///
 /// // 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);
+///     let mut iter = tree.iter();
+///     assert_eq!(iter.next().unwrap(), (&10, &100));
+///     assert_eq!(iter.next().unwrap(), (&15, &150));
+///     assert_eq!(iter.next().unwrap(), (&20, &200));
+///     assert!(iter.next().is_none());
 /// }
 ///
 /// # Ok::<(), Error>(())
@@ -167,6 +191,26 @@ pub fn new() -> Self {
             _p: PhantomData,
         }
     }
+
+    /// Returns an iterator over the tree nodes, sorted by key.
+    pub fn iter(&self) -> Iter<'_, K, V> {
+        // INVARIANT: `bindings::rb_first` returns a valid pointer to a tree node given a valid pointer to a tree root.
+        Iter {
+            _tree: PhantomData,
+            // SAFETY: `self.root` is a valid pointer to the tree root.
+            next: unsafe { bindings::rb_first(&self.root) },
+        }
+    }
+
+    /// Returns an iterator over the keys of the nodes in the tree, in sorted order.
+    pub fn keys(&self) -> impl Iterator<Item = &'_ K> {
+        self.iter().map(|(k, _)| k)
+    }
+
+    /// Returns an iterator over the values of the nodes in the tree, sorted by key.
+    pub fn values(&self) -> impl Iterator<Item = &'_ V> {
+        self.iter().map(|(_, v)| v)
+    }
 }
 
 impl<K, V> RBTree<K, V>
@@ -358,6 +402,56 @@ fn drop(&mut self) {
     }
 }
 
+impl<'a, K, V> IntoIterator for &'a RBTree<K, V> {
+    type Item = (&'a K, &'a V);
+    type IntoIter = Iter<'a, K, V>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.iter()
+    }
+}
+
+/// An iterator over the nodes of a [`RBTree`].
+///
+/// Instances are created by calling [`RBTree::iter`].
+///
+/// # Invariants
+/// - `self.next` is a valid pointer.
+/// - `self.next` points to a node stored inside of a valid `RBTree`.
+pub struct Iter<'a, K, V> {
+    _tree: PhantomData<&'a RBTree<K, V>>,
+    next: *mut bindings::rb_node,
+}
+
+// SAFETY: The [`Iter`] gives out immutable references to K and V, so it has the same
+// thread safety requirements as immutable references.
+unsafe impl<'a, K: Sync, V: Sync> Send for Iter<'a, K, V> {}
+
+// SAFETY: The [`Iter`] gives out immutable references to K and V, so it has the same
+// thread safety requirements as immutable references.
+unsafe impl<'a, K: Sync, V: Sync> Sync for Iter<'a, K, V> {}
+
+impl<'a, K, V> Iterator for Iter<'a, K, V> {
+    type Item = (&'a K, &'a V);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.next.is_null() {
+            return None;
+        }
+
+        // SAFETY: By the type invariant of `Iter`, `self.next` is a valid node in an `RBTree`,
+        // and by the type invariant of `RBTree`, all nodes point to the links field of `Node<K, V>` objects.
+        let cur = unsafe { container_of!(self.next, Node<K, V>, links) };
+
+        // SAFETY: `self.next` is a valid tree node by the type invariants.
+        self.next = unsafe { bindings::rb_next(self.next) };
+
+        // SAFETY: By the same reasoning above, it is safe to dereference the node. Additionally,
+        // it is ok to return a reference to members because the iterator must outlive it.
+        Some(unsafe { (&(*cur).key, &(*cur).value) })
+    }
+}
+
 /// A memory reservation for a red-black tree node.
 ///
 ///

-- 
2.46.0.rc1.232.g9752f9e123-goog

[PATCH v8 4/6] rust: rbtree: add mutable iterator

[Matt Gilbride]

From: Wedson Almeida Filho <wedsonaf@gmail.com>

Add mutable Iterator implementation for `RBTree`,
allowing iteration over (key, value) pairs in key order. Only values are
mutable, as mutating keys implies modifying a node's position in the tree.

Mutable iteration is used by the binder driver during shutdown to
clean up the tree maintained by the "range allocator" [1].

Link: https://lore.kernel.org/rust-for-linux/20231101-rust-binder-v1-6-08ba9197f637@google.com/ [1]
Signed-off-by: Wedson Almeida Filho <wedsonaf@gmail.com>
Signed-off-by: Matt Gilbride <mattgilbride@google.com>
Reviewed-by: Alice Ryhl <aliceryhl@google.com>
Tested-by: Alice Ryhl <aliceryhl@google.com>
---
 rust/kernel/rbtree.rs | 98 ++++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 86 insertions(+), 12 deletions(-)

diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs
index d10074e4ac58..d7514ebadfa8 100644
--- a/rust/kernel/rbtree.rs
+++ b/rust/kernel/rbtree.rs
@@ -197,8 +197,26 @@ pub fn iter(&self) -> Iter<'_, K, V> {
         // INVARIANT: `bindings::rb_first` returns a valid pointer to a tree node given a valid pointer to a tree root.
         Iter {
             _tree: PhantomData,
-            // SAFETY: `self.root` is a valid pointer to the tree root.
-            next: unsafe { bindings::rb_first(&self.root) },
+            iter_raw: IterRaw {
+                // SAFETY: by the invariants, all pointers are valid.
+                next: unsafe { bindings::rb_first(&self.root) },
+                _phantom: PhantomData,
+            },
+        }
+    }
+
+    /// Returns a mutable iterator over the tree nodes, sorted by key.
+    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
+        IterMut {
+            _tree: PhantomData,
+            // INVARIANT:
+            //   - `self.root` is a valid pointer to a tree root.
+            //   - `bindings::rb_first` produces a valid pointer to a node given `root` is valid.
+            iter_raw: IterRaw {
+                // SAFETY: by the invariants, all pointers are valid.
+                next: unsafe { bindings::rb_first(&self.root) },
+                _phantom: PhantomData,
+            },
         }
     }
 
@@ -211,6 +229,11 @@ pub fn keys(&self) -> impl Iterator<Item = &'_ K> {
     pub fn values(&self) -> impl Iterator<Item = &'_ V> {
         self.iter().map(|(_, v)| v)
     }
+
+    /// Returns a mutable iterator over the values of the nodes in the tree, sorted by key.
+    pub fn values_mut(&mut self) -> impl Iterator<Item = &'_ mut V> {
+        self.iter_mut().map(|(_, v)| v)
+    }
 }
 
 impl<K, V> RBTree<K, V>
@@ -414,13 +437,9 @@ fn into_iter(self) -> Self::IntoIter {
 /// An iterator over the nodes of a [`RBTree`].
 ///
 /// Instances are created by calling [`RBTree::iter`].
-///
-/// # Invariants
-/// - `self.next` is a valid pointer.
-/// - `self.next` points to a node stored inside of a valid `RBTree`.
 pub struct Iter<'a, K, V> {
     _tree: PhantomData<&'a RBTree<K, V>>,
-    next: *mut bindings::rb_node,
+    iter_raw: IterRaw<K, V>,
 }
 
 // SAFETY: The [`Iter`] gives out immutable references to K and V, so it has the same
@@ -434,21 +453,76 @@ unsafe impl<'a, K: Sync, V: Sync> Sync for Iter<'a, K, V> {}
 impl<'a, K, V> Iterator for Iter<'a, K, V> {
     type Item = (&'a K, &'a V);
 
+    fn next(&mut self) -> Option<Self::Item> {
+        self.iter_raw.next().map(|(k, v)|
+            // SAFETY: Due to `self._tree`, `k` and `v` are valid for the lifetime of `'a`.
+            unsafe { (&*k, &*v) })
+    }
+}
+
+impl<'a, K, V> IntoIterator for &'a mut RBTree<K, V> {
+    type Item = (&'a K, &'a mut V);
+    type IntoIter = IterMut<'a, K, V>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.iter_mut()
+    }
+}
+
+/// A mutable iterator over the nodes of a [`RBTree`].
+///
+/// Instances are created by calling [`RBTree::iter_mut`].
+pub struct IterMut<'a, K, V> {
+    _tree: PhantomData<&'a mut RBTree<K, V>>,
+    iter_raw: IterRaw<K, V>,
+}
+
+// SAFETY: The [`IterMut`] gives out immutable references to K and mutable references to V, so it has the same
+// thread safety requirements as mutable references.
+unsafe impl<'a, K: Send, V: Send> Send for IterMut<'a, K, V> {}
+
+// SAFETY: The [`IterMut`] gives out immutable references to K and mutable references to V, so it has the same
+// thread safety requirements as mutable references.
+unsafe impl<'a, K: Sync, V: Sync> Sync for IterMut<'a, K, V> {}
+
+impl<'a, K, V> Iterator for IterMut<'a, K, V> {
+    type Item = (&'a K, &'a mut V);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.iter_raw.next().map(|(k, v)|
+            // SAFETY: Due to `&mut self`, we have exclusive access to `k` and `v`, for the lifetime of `'a`.
+            unsafe { (&*k, &mut *v) })
+    }
+}
+
+/// A raw iterator over the nodes of a [`RBTree`].
+///
+/// # Invariants
+/// - `self.next` is a valid pointer.
+/// - `self.next` points to a node stored inside of a valid `RBTree`.
+struct IterRaw<K, V> {
+    next: *mut bindings::rb_node,
+    _phantom: PhantomData<fn() -> (K, V)>,
+}
+
+impl<K, V> Iterator for IterRaw<K, V> {
+    type Item = (*mut K, *mut V);
+
     fn next(&mut self) -> Option<Self::Item> {
         if self.next.is_null() {
             return None;
         }
 
-        // SAFETY: By the type invariant of `Iter`, `self.next` is a valid node in an `RBTree`,
+        // SAFETY: By the type invariant of `IterRaw`, `self.next` is a valid node in an `RBTree`,
         // and by the type invariant of `RBTree`, all nodes point to the links field of `Node<K, V>` objects.
-        let cur = unsafe { container_of!(self.next, Node<K, V>, links) };
+        let cur: *mut Node<K, V> =
+            unsafe { container_of!(self.next, Node<K, V>, links) }.cast_mut();
 
         // SAFETY: `self.next` is a valid tree node by the type invariants.
         self.next = unsafe { bindings::rb_next(self.next) };
 
-        // SAFETY: By the same reasoning above, it is safe to dereference the node. Additionally,
-        // it is ok to return a reference to members because the iterator must outlive it.
-        Some(unsafe { (&(*cur).key, &(*cur).value) })
+        // SAFETY: By the same reasoning above, it is safe to dereference the node.
+        Some(unsafe { (addr_of_mut!((*cur).key), addr_of_mut!((*cur).value)) })
     }
 }
 

-- 
2.46.0.rc1.232.g9752f9e123-goog

[PATCH v8 5/6] rust: rbtree: add `RBTreeCursor`

[Matt Gilbride]

Add a cursor interface to `RBTree`, supporting the following use cases:
- Inspect the current node pointed to by the cursor, inspect/move to
  it's neighbors in sort order (bidirectionally).
- Mutate the tree itself by removing the current node pointed to by the
  cursor, or one of its neighbors.

Add functions to obtain a cursor to the tree by key:
- The node with the smallest key
- The node with the largest key
- The node matching the given key, or the one with the next larger key

The cursor abstraction is needed by the binder driver to efficiently
search for nodes and (conditionally) modify them, as well as their
neighbors [1].

Link: https://lore.kernel.org/rust-for-linux/20231101-rust-binder-v1-6-08ba9197f637@google.com/ [1]
Co-developed-by: Alice Ryhl <aliceryhl@google.com>
Signed-off-by: Alice Ryhl <aliceryhl@google.com>
Reviewed-by: Alice Ryhl <aliceryhl@google.com>
Tested-by: Alice Ryhl <aliceryhl@google.com>
Signed-off-by: Matt Gilbride <mattgilbride@google.com>
---
 rust/kernel/rbtree.rs | 536 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 536 insertions(+)

diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs
index d7514ebadfa8..5d37aa373685 100644
--- a/rust/kernel/rbtree.rs
+++ b/rust/kernel/rbtree.rs
@@ -234,6 +234,40 @@ pub fn values(&self) -> impl Iterator<Item = &'_ V> {
     pub fn values_mut(&mut self) -> impl Iterator<Item = &'_ mut V> {
         self.iter_mut().map(|(_, v)| v)
     }
+
+    /// Returns a cursor over the tree nodes, starting with the smallest key.
+    pub fn cursor_front(&mut self) -> Option<RBTreeCursor<'_, K, V>> {
+        let root = addr_of_mut!(self.root);
+        // SAFETY: `self.root` is always a valid root node
+        let current = unsafe { bindings::rb_first(root) };
+        NonNull::new(current).map(|current| {
+            // INVARIANT:
+            // - `current` is a valid node in the [`RBTree`] pointed to by `self`.
+            // - Due to the type signature of this function, the returned [`RBTreeCursor`]
+            //   borrows mutably from `self`.
+            RBTreeCursor {
+                current,
+                tree: self,
+            }
+        })
+    }
+
+    /// Returns a cursor over the tree nodes, starting with the largest key.
+    pub fn cursor_back(&mut self) -> Option<RBTreeCursor<'_, K, V>> {
+        let root = addr_of_mut!(self.root);
+        // SAFETY: `self.root` is always a valid root node
+        let current = unsafe { bindings::rb_last(root) };
+        NonNull::new(current).map(|current| {
+            // INVARIANT:
+            // - `current` is a valid node in the [`RBTree`] pointed to by `self`.
+            // - Due to the type signature of this function, the returned [`RBTreeCursor`]
+            //   borrows mutably from `self`.
+            RBTreeCursor {
+                current,
+                tree: self,
+            }
+        })
+    }
 }
 
 impl<K, V> RBTree<K, V>
@@ -394,6 +428,75 @@ fn remove_node(&mut self, key: &K) -> Option<RBTreeNode<K, V>> {
     pub fn remove(&mut self, key: &K) -> Option<V> {
         self.remove_node(key).map(|node| node.node.value)
     }
+
+    /// Returns a cursor over the tree nodes based on the given key.
+    ///
+    /// If the given key exists, the cursor starts there.
+    /// Otherwise it starts with the first larger key in sort order.
+    /// If there is no larger key, it returns [`None`].
+    pub fn cursor_lower_bound(&mut self, key: &K) -> Option<RBTreeCursor<'_, K, V>>
+    where
+        K: Ord,
+    {
+        let mut node = self.root.rb_node;
+        let mut best_match: Option<NonNull<Node<K, V>>> = None;
+        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<K, V>` objects.
+            let this = unsafe { container_of!(node, Node<K, V>, links) }.cast_mut();
+            // SAFETY: `this` is a non-null node so it is valid by the type invariants.
+            let this_key = unsafe { &(*this).key };
+            // SAFETY: `node` is a non-null node so it is valid by the type invariants.
+            let left_child = unsafe { (*node).rb_left };
+            // SAFETY: `node` is a non-null node so it is valid by the type invariants.
+            let right_child = unsafe { (*node).rb_right };
+            if key == this_key {
+                return NonNull::new(node).map(|current| {
+                    // INVARIANT:
+                    // - `node` is a valid node in the [`RBTree`] pointed to by `self`.
+                    // - Due to the type signature of this function, the returned [`RBTreeCursor`]
+                    //   borrows mutably from `self`.
+                    RBTreeCursor {
+                        current,
+                        tree: self,
+                    }
+                });
+            } else {
+                node = if key > this_key {
+                    right_child
+                } else {
+                    let is_better_match = match best_match {
+                        None => true,
+                        Some(best) => {
+                            // SAFETY: `best` is a non-null node so it is valid by the type invariants.
+                            let best_key = unsafe { &(*best.as_ptr()).key };
+                            best_key > this_key
+                        }
+                    };
+                    if is_better_match {
+                        best_match = NonNull::new(this);
+                    }
+                    left_child
+                };
+            }
+        }
+
+        let best = best_match?;
+
+        // SAFETY: `best` is a non-null node so it is valid by the type invariants.
+        let links = unsafe { addr_of_mut!((*best.as_ptr()).links) };
+
+        NonNull::new(links).map(|current| {
+            // INVARIANT:
+            // - `current` is a valid node in the [`RBTree`] pointed to by `self`.
+            // - Due to the type signature of this function, the returned [`RBTreeCursor`]
+            //   borrows mutably from `self`.
+            RBTreeCursor {
+                current,
+                tree: self,
+            }
+        })
+    }
 }
 
 impl<K, V> Default for RBTree<K, V> {
@@ -425,6 +528,434 @@ fn drop(&mut self) {
     }
 }
 
+/// A bidirectional cursor over the tree nodes, sorted by key.
+///
+/// # Examples
+///
+/// In the following example, we obtain a cursor to the first element in the tree.
+/// The cursor allows us to iterate bidirectionally over key/value pairs in the tree.
+///
+/// ```
+/// use kernel::{alloc::flags, rbtree::RBTree};
+///
+/// // Create a new tree.
+/// let mut tree = RBTree::new();
+///
+/// // Insert three elements.
+/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
+/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
+/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
+///
+/// // Get a cursor to the first element.
+/// let mut cursor = tree.cursor_front().unwrap();
+/// let mut current = cursor.current();
+/// assert_eq!(current, (&10, &100));
+///
+/// // Move the cursor, updating it to the 2nd element.
+/// cursor = cursor.move_next().unwrap();
+/// current = cursor.current();
+/// assert_eq!(current, (&20, &200));
+///
+/// // Peek at the next element without impacting the cursor.
+/// let next = cursor.peek_next().unwrap();
+/// assert_eq!(next, (&30, &300));
+/// current = cursor.current();
+/// assert_eq!(current, (&20, &200));
+///
+/// // Moving past the last element causes the cursor to return [`None`].
+/// cursor = cursor.move_next().unwrap();
+/// current = cursor.current();
+/// assert_eq!(current, (&30, &300));
+/// let cursor = cursor.move_next();
+/// assert!(cursor.is_none());
+///
+/// # Ok::<(), Error>(())
+/// ```
+///
+/// A cursor can also be obtained at the last element in the tree.
+///
+/// ```
+/// use kernel::{alloc::flags, rbtree::RBTree};
+///
+/// // Create a new tree.
+/// let mut tree = RBTree::new();
+///
+/// // Insert three elements.
+/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
+/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
+/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
+///
+/// let mut cursor = tree.cursor_back().unwrap();
+/// let current = cursor.current();
+/// assert_eq!(current, (&30, &300));
+///
+/// # Ok::<(), Error>(())
+/// ```
+///
+/// Obtaining a cursor returns [`None`] if the tree is empty.
+///
+/// ```
+/// use kernel::rbtree::RBTree;
+///
+/// let mut tree: RBTree<u16, u16> = RBTree::new();
+/// assert!(tree.cursor_front().is_none());
+///
+/// # Ok::<(), Error>(())
+/// ```
+///
+/// [`RBTree::cursor_lower_bound`] can be used to start at an arbitrary node in the tree.
+///
+/// ```
+/// use kernel::{alloc::flags, rbtree::RBTree};
+///
+/// // Create a new tree.
+/// let mut tree = RBTree::new();
+///
+/// // Insert five elements.
+/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
+/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
+/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
+/// tree.try_create_and_insert(40, 400, flags::GFP_KERNEL)?;
+/// tree.try_create_and_insert(50, 500, flags::GFP_KERNEL)?;
+///
+/// // If the provided key exists, a cursor to that key is returned.
+/// let cursor = tree.cursor_lower_bound(&20).unwrap();
+/// let current = cursor.current();
+/// assert_eq!(current, (&20, &200));
+///
+/// // If the provided key doesn't exist, a cursor to the first larger element in sort order is returned.
+/// let cursor = tree.cursor_lower_bound(&25).unwrap();
+/// let current = cursor.current();
+/// assert_eq!(current, (&30, &300));
+///
+/// // If there is no larger key, [`None`] is returned.
+/// let cursor = tree.cursor_lower_bound(&55);
+/// assert!(cursor.is_none());
+///
+/// # Ok::<(), Error>(())
+/// ```
+///
+/// The cursor allows mutation of values in the tree.
+///
+/// ```
+/// use kernel::{alloc::flags, rbtree::RBTree};
+///
+/// // Create a new tree.
+/// let mut tree = RBTree::new();
+///
+/// // Insert three elements.
+/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
+/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
+/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
+///
+/// // Retrieve a cursor.
+/// let mut cursor = tree.cursor_front().unwrap();
+///
+/// // Get a mutable reference to the current value.
+/// let (k, v) = cursor.current_mut();
+/// *v = 1000;
+///
+/// // The updated value is reflected in the tree.
+/// let updated = tree.get(&10).unwrap();
+/// assert_eq!(updated, &1000);
+///
+/// # Ok::<(), Error>(())
+/// ```
+///
+/// It also allows node removal. The following examples demonstrate the behavior of removing the current node.
+///
+/// ```
+/// use kernel::{alloc::flags, rbtree::RBTree};
+///
+/// // Create a new tree.
+/// let mut tree = RBTree::new();
+///
+/// // Insert three elements.
+/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
+/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
+/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
+///
+/// // Remove the first element.
+/// let mut cursor = tree.cursor_front().unwrap();
+/// let mut current = cursor.current();
+/// assert_eq!(current, (&10, &100));
+/// cursor = cursor.remove_current().0.unwrap();
+///
+/// // If a node exists after the current element, it is returned.
+/// current = cursor.current();
+/// assert_eq!(current, (&20, &200));
+///
+/// // Get a cursor to the last element, and remove it.
+/// cursor = tree.cursor_back().unwrap();
+/// current = cursor.current();
+/// assert_eq!(current, (&30, &300));
+///
+/// // Since there is no next node, the previous node is returned.
+/// cursor = cursor.remove_current().0.unwrap();
+/// current = cursor.current();
+/// assert_eq!(current, (&20, &200));
+///
+/// // Removing the last element in the tree returns [`None`].
+/// assert!(cursor.remove_current().0.is_none());
+///
+/// # Ok::<(), Error>(())
+/// ```
+///
+/// Nodes adjacent to the current node can also be removed.
+///
+/// ```
+/// use kernel::{alloc::flags, rbtree::RBTree};
+///
+/// // Create a new tree.
+/// let mut tree = RBTree::new();
+///
+/// // Insert three elements.
+/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
+/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
+/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
+///
+/// // Get a cursor to the first element.
+/// let mut cursor = tree.cursor_front().unwrap();
+/// let mut current = cursor.current();
+/// assert_eq!(current, (&10, &100));
+///
+/// // Calling `remove_prev` from the first element returns [`None`].
+/// assert!(cursor.remove_prev().is_none());
+///
+/// // Get a cursor to the last element.
+/// cursor = tree.cursor_back().unwrap();
+/// current = cursor.current();
+/// assert_eq!(current, (&30, &300));
+///
+/// // Calling `remove_prev` removes and returns the middle element.
+/// assert_eq!(cursor.remove_prev().unwrap().to_key_value(), (20, 200));
+///
+/// // Calling `remove_next` from the last element returns [`None`].
+/// assert!(cursor.remove_next().is_none());
+///
+/// // Move to the first element
+/// cursor = cursor.move_prev().unwrap();
+/// current = cursor.current();
+/// assert_eq!(current, (&10, &100));
+///
+/// // Calling `remove_next` removes and returns the last element.
+/// assert_eq!(cursor.remove_next().unwrap().to_key_value(), (30, 300));
+///
+/// # Ok::<(), Error>(())
+/// ```
+/// # Invariants
+/// - `current` points to a node that is in the same [`RBTree`] as `tree`.
+pub struct RBTreeCursor<'a, K, V> {
+    tree: &'a mut RBTree<K, V>,
+    current: NonNull<bindings::rb_node>,
+}
+
+// SAFETY: The [`RBTreeCursor`] gives out immutable references to K and mutable references to V,
+// so it has the same thread safety requirements as mutable references.
+unsafe impl<'a, K: Send, V: Send> Send for RBTreeCursor<'a, K, V> {}
+
+// SAFETY: The [`RBTreeCursor`] gives out immutable references to K and mutable references to V,
+// so it has the same thread safety requirements as mutable references.
+unsafe impl<'a, K: Sync, V: Sync> Sync for RBTreeCursor<'a, K, V> {}
+
+impl<'a, K, V> RBTreeCursor<'a, K, V> {
+    /// The current node
+    pub fn current(&self) -> (&K, &V) {
+        // SAFETY:
+        // - `self.current` is a valid node by the type invariants.
+        // - We have an immutable reference by the function signature.
+        unsafe { Self::to_key_value(self.current) }
+    }
+
+    /// The current node, with a mutable value
+    pub fn current_mut(&mut self) -> (&K, &mut V) {
+        // SAFETY:
+        // - `self.current` is a valid node by the type invariants.
+        // - We have an mutable reference by the function signature.
+        unsafe { Self::to_key_value_mut(self.current) }
+    }
+
+    /// Remove the current node from the tree.
+    ///
+    /// Returns a tuple where the first element is a cursor to the next node, if it exists,
+    /// else the previous node, else [`None`] (if the tree becomes empty). The second element
+    /// is the removed node.
+    pub fn remove_current(self) -> (Option<Self>, RBTreeNode<K, V>) {
+        let prev = self.get_neighbor_raw(Direction::Prev);
+        let next = self.get_neighbor_raw(Direction::Next);
+        // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`
+        // point to the links field of `Node<K, V>` objects.
+        let this = unsafe { container_of!(self.current.as_ptr(), Node<K, V>, links) }.cast_mut();
+        // SAFETY: `this` is valid by the type invariants as described above.
+        let node = unsafe { Box::from_raw(this) };
+        let node = RBTreeNode { node };
+        // SAFETY: The reference to the tree used to create the cursor outlives the cursor, so
+        // the tree cannot change. By the tree invariant, all nodes are valid.
+        unsafe { bindings::rb_erase(&mut (*this).links, addr_of_mut!(self.tree.root)) };
+
+        let current = match (prev, next) {
+            (_, Some(next)) => next,
+            (Some(prev), None) => prev,
+            (None, None) => {
+                return (None, node);
+            }
+        };
+
+        (
+            // INVARIANT:
+            // - `current` is a valid node in the [`RBTree`] pointed to by `self.tree`.
+            // - Due to the function signature, `self` is an owned [`RBTreeCursor`],
+            //   and [`RBTreeCursor`]s are only created via functions with a mutable reference
+            //   to an [`RBTree`].
+            Some(Self {
+                current,
+                tree: self.tree,
+            }),
+            node,
+        )
+    }
+
+    /// Remove the previous node, returning it if it exists.
+    pub fn remove_prev(&mut self) -> Option<RBTreeNode<K, V>> {
+        self.remove_neighbor(Direction::Prev)
+    }
+
+    /// Remove the next node, returning it if it exists.
+    pub fn remove_next(&mut self) -> Option<RBTreeNode<K, V>> {
+        self.remove_neighbor(Direction::Next)
+    }
+
+    fn remove_neighbor(&mut self, direction: Direction) -> Option<RBTreeNode<K, V>> {
+        if let Some(neighbor) = self.get_neighbor_raw(direction) {
+            let neighbor = neighbor.as_ptr();
+            // SAFETY: The reference to the tree used to create the cursor outlives the cursor, so
+            // the tree cannot change. By the tree invariant, all nodes are valid.
+            unsafe { bindings::rb_erase(neighbor, addr_of_mut!(self.tree.root)) };
+            // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`
+            // point to the links field of `Node<K, V>` objects.
+            let this = unsafe { container_of!(neighbor, Node<K, V>, links) }.cast_mut();
+            // SAFETY: `this` is valid by the type invariants as described above.
+            let node = unsafe { Box::from_raw(this) };
+            return Some(RBTreeNode { node });
+        }
+        None
+    }
+
+    /// Move the cursor to the previous node, returning [`None`] if it doesn't exist.
+    pub fn move_prev(self) -> Option<Self> {
+        self.mv(Direction::Prev)
+    }
+
+    /// Move the cursor to the next node, returning [`None`] if it doesn't exist.
+    pub fn move_next(self) -> Option<Self> {
+        self.mv(Direction::Next)
+    }
+
+    fn mv(self, direction: Direction) -> Option<Self> {
+        // INVARIANT:
+        // - `neighbor` is a valid node in the [`RBTree`] pointed to by `self.tree`.
+        // - Due to the function signature, `self` is an owned [`RBTreeCursor`],
+        //   and [`RBTreeCursor`]s are only created via functions with a mutable reference
+        //   to an [`RBTree`].
+        self.get_neighbor_raw(direction).map(|neighbor| Self {
+            tree: self.tree,
+            current: neighbor,
+        })
+    }
+
+    /// Access the previous node without moving the cursor.
+    pub fn peek_prev(&self) -> Option<(&K, &V)> {
+        self.peek(Direction::Prev)
+    }
+
+    /// Access the previous node without moving the cursor.
+    pub fn peek_next(&self) -> Option<(&K, &V)> {
+        self.peek(Direction::Next)
+    }
+
+    fn peek(&self, direction: Direction) -> Option<(&K, &V)> {
+        self.get_neighbor_raw(direction)
+            // SAFETY:
+            // - `neighbor` is a valid tree node.
+            // - By the function signature, we have an immutable reference to `self`.
+            .map(|neighbor| unsafe { Self::to_key_value(neighbor) })
+    }
+
+    /// Access the previous node mutably without moving the cursor.
+    pub fn peek_prev_mut(&mut self) -> Option<(&K, &mut V)> {
+        self.peek_mut(Direction::Prev)
+    }
+
+    /// Access the next node mutably without moving the cursor.
+    pub fn peek_next_mut(&mut self) -> Option<(&K, &mut V)> {
+        self.peek_mut(Direction::Next)
+    }
+
+    fn peek_mut(&mut self, direction: Direction) -> Option<(&K, &mut V)> {
+        self.get_neighbor_raw(direction)
+            // SAFETY:
+            // - `neighbor` is a valid tree node.
+            // - By the function signature, we have a mutable reference to `self`.
+            .map(|neighbor| unsafe { Self::to_key_value_mut(neighbor) })
+    }
+
+    fn get_neighbor_raw(&self, direction: Direction) -> Option<NonNull<bindings::rb_node>> {
+        // SAFETY: `self.current` is valid by the type invariants.
+        let neighbor = unsafe {
+            match direction {
+                Direction::Prev => bindings::rb_prev(self.current.as_ptr()),
+                Direction::Next => bindings::rb_next(self.current.as_ptr()),
+            }
+        };
+
+        NonNull::new(neighbor)
+    }
+
+    /// SAFETY:
+    /// - `node` must be a valid pointer to a node in an [`RBTree`].
+    /// - The caller has immutable access to `node` for the duration of 'b.
+    unsafe fn to_key_value<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, &'b V) {
+        // SAFETY: the caller guarantees that `node` is a valid pointer in an `RBTree`.
+        let (k, v) = unsafe { Self::to_key_value_raw(node) };
+        // SAFETY: the caller guarantees immutable access to `node`.
+        (k, unsafe { &*v })
+    }
+
+    /// SAFETY:
+    /// - `node` must be a valid pointer to a node in an [`RBTree`].
+    /// - The caller has mutable access to `node` for the duration of 'b.
+    unsafe fn to_key_value_mut<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, &'b mut V) {
+        // SAFETY: the caller guarantees that `node` is a valid pointer in an `RBTree`.
+        let (k, v) = unsafe { Self::to_key_value_raw(node) };
+        // SAFETY: the caller guarantees mutable access to `node`.
+        (k, unsafe { &mut *v })
+    }
+
+    /// SAFETY:
+    /// - `node` must be a valid pointer to a node in an [`RBTree`].
+    /// - The caller has immutable access to the key for the duration of 'b.
+    unsafe fn to_key_value_raw<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, *mut V) {
+        // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`
+        // point to the links field of `Node<K, V>` objects.
+        let this = unsafe { container_of!(node.as_ptr(), Node<K, V>, links) }.cast_mut();
+        // SAFETY: The passed `node` is the current node or a non-null neighbor,
+        // thus `this` is valid by the type invariants.
+        let k = unsafe { &(*this).key };
+        // SAFETY: The passed `node` is the current node or a non-null neighbor,
+        // thus `this` is valid by the type invariants.
+        let v = unsafe { addr_of_mut!((*this).value) };
+        (k, v)
+    }
+}
+
+/// Direction for [`RBTreeCursor`] operations.
+enum Direction {
+    /// the node immediately before, in sort order
+    Prev,
+    /// the node immediately after, in sort order
+    Next,
+}
+
 impl<'a, K, V> IntoIterator for &'a RBTree<K, V> {
     type Item = (&'a K, &'a V);
     type IntoIter = Iter<'a, K, V>;
@@ -583,6 +1114,11 @@ impl<K, V> RBTreeNode<K, V> {
     pub fn new(key: K, value: V, flags: Flags) -> Result<RBTreeNode<K, V>> {
         Ok(RBTreeNodeReservation::new(flags)?.into_node(key, value))
     }
+
+    /// Get the key and value from inside the node.
+    pub fn to_key_value(self) -> (K, V) {
+        (self.node.key, self.node.value)
+    }
 }
 
 // SAFETY: If K and V can be sent across threads, then it's also okay to send [`RBTreeNode`] across

-- 
2.46.0.rc1.232.g9752f9e123-goog

[PATCH v8 6/6] rust: rbtree: add `RBTree::entry`

[Matt Gilbride]

From: Alice Ryhl <aliceryhl@google.com>

This mirrors the entry API [1] from the Rust standard library on
`RBTree`. This API can be used to access the entry at a specific key and
make modifications depending on whether the key is vacant or occupied.
This API is useful because it can often be used to avoid traversing the
tree multiple times.

This is used by binder to look up and conditionally access or insert a
value, depending on whether it is there or not [2].

Link: https://doc.rust-lang.org/stable/std/collections/btree_map/enum.Entry.html [1]
Link: https://android-review.googlesource.com/c/kernel/common/+/2849906 [2]
Signed-off-by: Alice Ryhl <aliceryhl@google.com>
Tested-by: Alice Ryhl <aliceryhl@google.com>
Signed-off-by: Matt Gilbride <mattgilbride@google.com>
---
 rust/kernel/rbtree.rs | 302 +++++++++++++++++++++++++++++++++++++-------------
 1 file changed, 227 insertions(+), 75 deletions(-)

diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs
index 5d37aa373685..428f8be8f3a2 100644
--- a/rust/kernel/rbtree.rs
+++ b/rust/kernel/rbtree.rs
@@ -295,12 +295,19 @@ pub fn try_create_and_insert(
     /// 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<K, V>) -> Option<RBTreeNode<K, V>> {
-        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) };
+    pub fn insert(&mut self, node: RBTreeNode<K, V>) -> Option<RBTreeNode<K, V>> {
+        match self.raw_entry(&node.node.key) {
+            RawEntry::Occupied(entry) => Some(entry.replace(node)),
+            RawEntry::Vacant(entry) => {
+                entry.insert(node);
+                None
+            }
+        }
+    }
 
+    fn raw_entry(&mut self, key: &K) -> RawEntry<'_, K, V> {
+        let raw_self: *mut RBTree<K, V> = self;
+        // The returned `RawEntry` is used to call either `rb_link_node` or `rb_replace_node`.
         // The parameters of `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
@@ -319,62 +326,56 @@ pub fn insert(&mut self, RBTreeNode { node }: RBTreeNode<K, V>) -> Option<RBTree
         // 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;
+        let mut child_field_of_parent: &mut *mut bindings::rb_node =
+            // SAFETY: `raw_self` is a valid pointer to the `RBTree` (created from `self` above).
+            unsafe { &mut (*raw_self).root.rb_node };
+        while !(*child_field_of_parent).is_null() {
+            let curr = *child_field_of_parent;
+            // SAFETY: All links fields we create are in a `Node<K, V>`.
+            let node = unsafe { container_of!(curr, Node<K, V>, links) };
 
-            // 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<K, V>` objects.
-            let this = unsafe { container_of!(parent, Node<K, V>, 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 },
+            // SAFETY: `node` is a non-null node so it is valid by the type invariants.
+            match key.cmp(unsafe { &(*node).key }) {
+                // SAFETY: `curr` is a non-null node so it is valid by the type invariants.
+                Ordering::Less => child_field_of_parent = unsafe { &mut (*curr).rb_left },
+                // SAFETY: `curr` is a non-null node so it is valid by the type invariants.
+                Ordering::Greater => child_field_of_parent = unsafe { &mut (*curr).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()) },
-                    });
+                    return RawEntry::Occupied(OccupiedEntry {
+                        rbtree: self,
+                        node_links: curr,
+                    })
                 }
             }
+            parent = curr;
         }
 
-        // 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) };
+        RawEntry::Vacant(RawVacantEntry {
+            rbtree: raw_self,
+            parent,
+            child_field_of_parent,
+            _phantom: PhantomData,
+        })
+    }
 
-        // SAFETY: All pointers are valid. `node` has just been inserted into the tree.
-        unsafe { bindings::rb_insert_color(node_links, &mut self.root) };
-        None
+    /// Gets the given key's corresponding entry in the map for in-place manipulation.
+    pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
+        match self.raw_entry(&key) {
+            RawEntry::Occupied(entry) => Entry::Occupied(entry),
+            RawEntry::Vacant(entry) => Entry::Vacant(VacantEntry { raw: entry, key }),
+        }
+    }
+
+    /// Used for accessing the given node, if it exists.
+    pub fn find_mut(&mut self, key: &K) -> Option<OccupiedEntry<'_, K, V>> {
+        match self.raw_entry(key) {
+            RawEntry::Occupied(entry) => Some(entry),
+            RawEntry::Vacant(_entry) => None,
+        }
     }
 
-    /// Returns a node with the given key, if one exists.
-    fn find(&self, key: &K) -> Option<NonNull<Node<K, V>>> {
+    /// Returns a reference to the value corresponding to the key.
+    pub fn get(&self, key: &K) -> Option<&V> {
         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`
@@ -386,47 +387,30 @@ fn find(&self, key: &K) -> Option<NonNull<Node<K, V>>> {
                 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()),
+                // SAFETY: `node` is a non-null node so it is valid by the type invariants.
+                Ordering::Equal => return Some(unsafe { &(*this).value }),
             }
         }
         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 })
+        self.find_mut(key).map(|node| node.into_mut())
     }
 
     /// 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()) },
-        })
+    pub fn remove_node(&mut self, key: &K) -> Option<RBTreeNode<K, V>> {
+        self.find_mut(key).map(OccupiedEntry::remove_node)
     }
 
     /// 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> {
-        self.remove_node(key).map(|node| node.node.value)
+        self.find_mut(key).map(OccupiedEntry::remove)
     }
 
     /// Returns a cursor over the tree nodes based on the given key.
@@ -1129,6 +1113,174 @@ unsafe impl<K: Send, V: Send> Send for RBTreeNode<K, V> {}
 // [`RBTreeNode`] without synchronization.
 unsafe impl<K: Sync, V: Sync> Sync for RBTreeNode<K, V> {}
 
+impl<K, V> RBTreeNode<K, V> {
+    /// Drop the key and value, but keep the allocation.
+    ///
+    /// It then becomes a reservation that can be re-initialised into a different node (i.e., with
+    /// a different key and/or value).
+    ///
+    /// The existing key and value are dropped in-place as part of this operation, that is, memory
+    /// may be freed (but only for the key/value; memory for the node itself is kept for reuse).
+    pub fn into_reservation(self) -> RBTreeNodeReservation<K, V> {
+        RBTreeNodeReservation {
+            node: Box::drop_contents(self.node),
+        }
+    }
+}
+
+/// A view into a single entry in a map, which may either be vacant or occupied.
+///
+/// This enum is constructed from the [`RBTree::entry`].
+///
+/// [`entry`]: fn@RBTree::entry
+pub enum Entry<'a, K, V> {
+    /// This [`RBTree`] does not have a node with this key.
+    Vacant(VacantEntry<'a, K, V>),
+    /// This [`RBTree`] already has a node with this key.
+    Occupied(OccupiedEntry<'a, K, V>),
+}
+
+/// Like [`Entry`], except that it doesn't have ownership of the key.
+enum RawEntry<'a, K, V> {
+    Vacant(RawVacantEntry<'a, K, V>),
+    Occupied(OccupiedEntry<'a, K, V>),
+}
+
+/// A view into a vacant entry in a [`RBTree`]. It is part of the [`Entry`] enum.
+pub struct VacantEntry<'a, K, V> {
+    key: K,
+    raw: RawVacantEntry<'a, K, V>,
+}
+
+/// Like [`VacantEntry`], but doesn't hold on to the key.a
+///
+/// # Invariants
+/// - `parent` may be null if the new node becomes the root.
+/// - `child_field_of_parent` is a valid pointer to the left-child or right-child of `parent`. If `parent` is
+///     null, it is a pointer to the root of the [`RBTree`].
+struct RawVacantEntry<'a, K, V> {
+    rbtree: *mut RBTree<K, V>,
+    /// The node that will become the parent of the new node if we insert one.
+    parent: *mut bindings::rb_node,
+    /// This points to the left-child or right-child field of `parent`, or `root` if `parent` is
+    /// null.
+    child_field_of_parent: *mut *mut bindings::rb_node,
+    _phantom: PhantomData<&'a mut RBTree<K, V>>,
+}
+
+impl<'a, K, V> RawVacantEntry<'a, K, V> {
+    /// Inserts the given node into the [`RBTree`] at this entry.
+    ///
+    /// The `node` must have a key such that inserting it here does not break the ordering of this
+    /// [`RBTree`].
+    fn insert(self, node: RBTreeNode<K, V>) -> &'a mut V {
+        let node = Box::into_raw(node.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) };
+
+        // INVARIANT: We are linking in a new node, which is valid. It remains valid because we
+        // "forgot" it with `Box::into_raw`.
+        // SAFETY: The type invariants of `RawVacantEntry` are exactly the safety requirements of `rb_link_node`.
+        unsafe { bindings::rb_link_node(node_links, self.parent, self.child_field_of_parent) };
+
+        // SAFETY: All pointers are valid. `node` has just been inserted into the tree.
+        unsafe { bindings::rb_insert_color(node_links, addr_of_mut!((*self.rbtree).root)) };
+
+        // SAFETY: The node is valid until we remove it from the tree.
+        unsafe { &mut (*node).value }
+    }
+}
+
+impl<'a, K, V> VacantEntry<'a, K, V> {
+    /// Inserts the given node into the [`RBTree`] at this entry.
+    pub fn insert(self, value: V, reservation: RBTreeNodeReservation<K, V>) -> &'a mut V {
+        self.raw.insert(reservation.into_node(self.key, value))
+    }
+}
+
+/// A view into an occupied entry in a [`RBTree`]. It is part of the [`Entry`] enum.
+///
+/// # Invariants
+/// - `node_links` is a valid, non-null pointer to a tree node in `self.rbtree`
+pub struct OccupiedEntry<'a, K, V> {
+    rbtree: &'a mut RBTree<K, V>,
+    /// The node that this entry corresponds to.
+    node_links: *mut bindings::rb_node,
+}
+
+impl<'a, K, V> OccupiedEntry<'a, K, V> {
+    fn node_ptr(&self) -> *mut Node<K, V> {
+        // SAFETY: By the type invariant of `Self`, all `node_links` pointers stored in `self`
+        // point to the links field of `Node<K, V>` objects.
+        unsafe { container_of!(self.node_links, Node<K, V>, links) }.cast_mut()
+    }
+
+    /// Gets a reference to the value in the entry.
+    pub fn get(&self) -> &V {
+        // SAFETY: `self.node_ptr` produces a valid pointer to a node in the tree.
+        unsafe { &(*self.node_ptr()).value }
+    }
+
+    /// Gets a mutable reference to the value in the entry.
+    pub fn get_mut(&mut self) -> &mut V {
+        // SAFETY: `self.node_ptr` produces a valid pointer to a node in the tree.
+        unsafe { &mut (*self.node_ptr()).value }
+    }
+
+    /// Converts the entry into a mutable reference to its value.
+    ///
+    /// If you need multiple references to the `OccupiedEntry`, see [`self#get_mut`].
+    pub fn into_mut(self) -> &'a mut V {
+        // SAFETY: `self.node_ptr` produces a valid pointer to a node in the tree.
+        unsafe { &mut (*self.node_ptr()).value }
+    }
+
+    /// Remove this entry from the [`RBTree`].
+    pub fn remove_node(self) -> RBTreeNode<K, V> {
+        // SAFETY: The node is a node in the tree, so it is valid.
+        unsafe { bindings::rb_erase(self.node_links, &mut self.rbtree.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.
+        RBTreeNode {
+            // SAFETY: The node was a node in the tree, but we removed it, so we can convert it
+            // back into a box.
+            node: unsafe { Box::from_raw(self.node_ptr()) },
+        }
+    }
+
+    /// Takes the value of the entry out of the map, and returns it.
+    pub fn remove(self) -> V {
+        self.remove_node().node.value
+    }
+
+    /// Swap the current node for the provided node.
+    ///
+    /// The key of both nodes must be equal.
+    fn replace(self, node: RBTreeNode<K, V>) -> RBTreeNode<K, V> {
+        let node = Box::into_raw(node.node);
+
+        // SAFETY: `node` is valid at least until we call `Box::from_raw`, which only happens when
+        // the node is removed or replaced.
+        let new_node_links = unsafe { addr_of_mut!((*node).links) };
+
+        // SAFETY: This updates the pointers so that `new_node_links` is in the tree where
+        // `self.node_links` used to be.
+        unsafe {
+            bindings::rb_replace_node(self.node_links, new_node_links, &mut self.rbtree.root)
+        };
+
+        // SAFETY:
+        // - `self.node_ptr` produces a valid pointer to a node in the tree.
+        // - Now that we removed this entry from the tree, we can convert the node to a box.
+        let old_node = unsafe { Box::from_raw(self.node_ptr()) };
+
+        RBTreeNode { node: old_node }
+    }
+}
+
 struct Node<K, V> {
     links: bindings::rb_node,
     key: K,

-- 
2.46.0.rc1.232.g9752f9e123-goog

Re: [PATCH v8 2/6] rust: rbtree: add red-black tree implementation backed by the C version

[Boqun Feng]

On Sat, Jul 27, 2024 at 08:30:47PM +0000, Matt Gilbride wrote:
[...]
> +/// 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<K, V> {
> +    node: Box<MaybeUninit<Node<K, V>>>,
> +}
> +
> +impl<K, V> RBTreeNodeReservation<K, V> {
> +    /// 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<RBTreeNodeReservation<K, V>> {
> +        Ok(RBTreeNodeReservation {
> +            node: Box::new_uninit(flags)?,

This will cause a rusttest error, although I'm not sure whether we want
to keep doing <Box<_> as BoxExt<_>>::new_uninit() to shutdown this kind
of errors since we are going to have our own alloc mod.

Regards,
Boqun

> +        })
> +    }
> +}
> +
[...]

Re: [PATCH v8 2/6] rust: rbtree: add red-black tree implementation backed by the C version

[Boqun Feng]

On Sat, Jul 27, 2024 at 08:30:47PM +0000, Matt Gilbride wrote:
[...]
> diff --git a/rust/helpers.c b/rust/helpers.c
> index 4c8b7b92a4f4..608b38c0b3e8 100644
> --- a/rust/helpers.c
> +++ b/rust/helpers.c
> @@ -157,6 +157,13 @@ void rust_helper_init_work_with_key(struct work_struct *work, work_func_t func,
>  }
>  EXPORT_SYMBOL_GPL(rust_helper_init_work_with_key);
>  
> +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);
> +

Just FYI, this will trigger a conflict because of:

	53ed0af49642 ("rust: add a rust helper for krealloc()")

Regards,
Boqun

>  /*
>   * `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.
> diff --git a/rust/kernel/lib.rs b/rust/kernel/lib.rs
> index 9a943d99c71a..dc2678803637 100644
[...]

Re: [PATCH v8 0/6] Red-black tree abstraction needed by Rust Binder

[Boqun Feng]

On Sat, Jul 27, 2024 at 08:30:45PM +0000, Matt Gilbride wrote:
> This patchset contains the red-black tree abstractions needed by the Rust
> implementation of the Binder driver.
> 
> Binder driver benefits from O(log n) search/insertion/deletion of
> key/value mappings in various places, including `process.rs` and
> `range_alloc.rs`.  In `range_alloc.rs`, the ability to store and
> search by a generic key type is also useful.
> 
> Please see the Rust Binder RFC for usage examples [1]. Note that
> the `container_of` macro is currently used only by `rbtree` itself.
> 
> Users of "rust: rbtree: add red-black tree implementation backed by the C version"
>     [PATCH RFC 03/20] rust_binder: add threading support
>     [PATCH RFC 05/20] rust_binder: add nodes and context managers
>     [PATCH RFC 06/20] rust_binder: add oneway transactions
> 
> Users of "rust: rbtree: add iterator"
>     [PATCH RFC 17/20] rust_binder: add oneway spam detection
> 
> Users of "rust: rbtree: add mutable iterator"
>     [PATCH RFC 06/20] rust_binder: add oneway transactions
> 
> Users of "rust: rbtree: add `RBTreeCursor`"
>     [PATCH RFC 06/20] rust_binder: add oneway transactions
> 
> Users of "rust: rbtree: add RBTree::entry"
>     Not used in the original RFC, but introduced after further
>     code review.  See: https://r.android.com/2849906
> 
> The Rust Binder RFC addresses the upstream deprecation of red-black
> tree. Quoted here for convenience:
> 
> "This RFC uses the kernel's red-black tree for key/value mappings, but we
> are aware that the red-black tree is deprecated. We did this to make the
> performance comparison more fair, since C binder also uses rbtree for
> this. We intend to replace these with XArrays instead. That said, we
> don't think that XArray is a good fit for the range allocator, and we
> propose to continue using the red-black tree for the range allocator."
> 
> Link: https://lore.kernel.org/rust-for-linux/20231101-rust-binder-v1-0-08ba9197f637@google.com/ [1]
> Signed-off-by: Matt Gilbride <mattgilbride@google.com>

Reviewed-by: Boqun Feng <boqun.feng@gmail.com>

Again, I drop patch #1 and put it in rust-dev with conflicts fixed.

Regards,
Boqun

> ---
> Changes in v8:
> - Fix small style nit (use ? operator)
> - Fix doc comment pointing at a private item
> - Link to v7: https://lore.kernel.org/r/20240726-b4-rbtree-v7-0-aee88caaf97c@google.com
> 
> Changes in v7:
> - make `RawVacantEntry.rbtree` a raw pointer like
>   `RawVacantEntry.child_field_of_parent`, since the latter can
>   technically point at a field of the former. We prefer that the
>   implementation be explicit about the safety guarantees of both because
>   of the relationship between them.
> - Link to v6: https://lore.kernel.org/r/20240711-b4-rbtree-v6-0-14bef1a8cdba@google.com
> 
> Changes in v6:
> - Minimize usage of `*mut bindings::rb_node`, replacing with
>   `NonNull<bindings::rb_node>`. Specifically, changing
>   `RBTreeCursor.current` to be `NonNull<bindings::rb_node>` and updating
>   the corresponding functions.
> - Update `RBTreeCursor:to_key_value` helpers to have their own lifetime
>   (they are not instance methods, using a different lifetime than that
>   of the `impl` block they are in makes things more clear.
> - Fix misplaced semicolon in `cursor_lower_bound`.
> - Link to v5: https://lore.kernel.org/r/20240606-b4-rbtree-v5-0-96fe1a0e97c0@google.com
> 
> Changes in v5:
> - Used `Box::write` in `RBTreeNodeReservation::into_node`, removing
>   unnecessary `unsafe` blocks.
> - Updated `RBTreeCursor::remove_current` to return the removed node.
> - Link to v4: https://lore.kernel.org/r/20240603-b4-rbtree-v4-0-308e43d6abfc@google.com
> 
> Changes in v4:
> - rebased onto the tip of rust-for-linux/rust-next (97ab3e8eec0ce79d9e265e6c9e4c480492180409)
> - addressed comments from draft PR on GitHub: https://github.com/Rust-for-Linux/linux/pull/1081
> - Link to v3: https://lore.kernel.org/r/20240418-b4-rbtree-v3-0-323e134390ce@google.com
> 
> Changes in v3:
> - Address various feedback re: SAFETY and INVARIANT comments from v2.
> - Update variable naming and add detailed comments for the `RBTree::insert` (later moved to
>   `RBTree::raw_entry`) implementation.
> - Link to v2: https://lore.kernel.org/r/20240219-b4-rbtree-v2-0-0b113aab330d@google.com
> 
> Changes in v2:
> - Update documentation link to the C header file
> - Use `core::convert::Infallible` in try_reserve_node
> - Link to v1: https://lore.kernel.org/r/20240205-b4-rbtree-v1-0-995e3eee38c0@google.com
> 
> ---
> Alice Ryhl (1):
>       rust: rbtree: add `RBTree::entry`
> 
> Benno Lossin (1):
>       rust: kernel: add `drop_contents` to `BoxExt`
> 
> Matt Gilbride (1):
>       rust: rbtree: add `RBTreeCursor`
> 
> Wedson Almeida Filho (3):
>       rust: rbtree: add red-black tree implementation backed by the C version
>       rust: rbtree: add iterator
>       rust: rbtree: add mutable iterator
> 
>  rust/helpers.c               |    7 +
>  rust/kernel/alloc/box_ext.rs |   24 +-
>  rust/kernel/lib.rs           |    1 +
>  rust/kernel/rbtree.rs        | 1288 ++++++++++++++++++++++++++++++++++++++++++
>  4 files changed, 1319 insertions(+), 1 deletion(-)
> ---
> base-commit: 97ab3e8eec0ce79d9e265e6c9e4c480492180409
> change-id: 20231205-b4-rbtree-abb1a016f0a0
> 
> Best regards,
> -- 
> Matt Gilbride <mattgilbride@google.com>
> 

Re: [PATCH v8 1/6] rust: kernel: add `drop_contents` to `BoxExt`

[Alice Ryhl]

On Sat, Jul 27, 2024 at 10:30 PM Matt Gilbride <mattgilbride@google.com> wrote:
>
> From: Benno Lossin <benno.lossin@proton.me>
>
> Sometimes (see [1]) it is necessary to drop the value inside of a
> `Box<T>`, but retain the allocation. For example to reuse the allocation
> in the future.
> Introduce a new function `drop_contents` that turns a `Box<T>` into
> `Box<MaybeUninit<T>>` by dropping the value.
>
> Signed-off-by: Benno Lossin <benno.lossin@proton.me>
> Link: https://lore.kernel.org/rust-for-linux/20240418-b4-rbtree-v3-5-323e134390ce@google.com/ [1]

Reviewed-by: Alice Ryhl <aliceryhl@google.com>

Re: [PATCH v8 1/6] rust: kernel: add `drop_contents` to `BoxExt`

[Alice Ryhl]

On Thu, Aug 1, 2024 at 11:00 AM Alice Ryhl <aliceryhl@google.com> wrote:
>
> On Sat, Jul 27, 2024 at 10:30 PM Matt Gilbride <mattgilbride@google.com> wrote:
> >
> > From: Benno Lossin <benno.lossin@proton.me>
> >
> > Sometimes (see [1]) it is necessary to drop the value inside of a
> > `Box<T>`, but retain the allocation. For example to reuse the allocation
> > in the future.
> > Introduce a new function `drop_contents` that turns a `Box<T>` into
> > `Box<MaybeUninit<T>>` by dropping the value.
> >
> > Signed-off-by: Benno Lossin <benno.lossin@proton.me>
> > Link: https://lore.kernel.org/rust-for-linux/20240418-b4-rbtree-v3-5-323e134390ce@google.com/ [1]
>
> Reviewed-by: Alice Ryhl <aliceryhl@google.com>

(We may prefer to pick this from Benno's series:
https://lore.kernel.org/rust-for-linux/20240708205325.1275473-1-benno.lossin@proton.me/
)

Re: [PATCH v8 2/6] rust: rbtree: add red-black tree implementation backed by the C version

[Benno Lossin]

On 27.07.24 22:30, Matt Gilbride wrote:
> diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs
> new file mode 100644
> index 000000000000..74c53e4d5c00
> --- /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: <https://www.kernel.org/doc/html/latest/core-api/rbtree.html>
> +
> +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.
> +///
> +/// # 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.

I think we should standardize that `Invariants` and `Safety` sections
are put last. This is because people reading the HTML version are
probably not interested in the inner workings. This also makes it
possible to have the invariants and the struct definition on the same
screen.

> +///
> +/// # Examples

[...]

> +/// ```
> +pub struct RBTree<K, V> {
> +    root: bindings::rb_root,

It has been a while, so I might have already asked this, but do we need
`Opaque` here? We would need it, if C changes anything inside `root`
while Rust holds a `&RBTree` or `&mut RBTree`.
I don't think that this is the case (ie we don't need `Opaque`), but I
am not sure.

> +    _p: PhantomData<Node<K, V>>,
> +}
> +

[...]

> +    /// 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<K, V>) -> Option<RBTreeNode<K, V>> {
> +        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 `rb_link_node` are as follows:

Can you write `bindings::rb_link_node`?

> +        // - `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<K, V>` objects.
> +            let this = unsafe { container_of!(parent, Node<K, V>, 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.

Missing `` around "parent", double space after '.'.

> +                    //
> +                    // 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()) },
> +                    });
> +                }
> +            }
> +        }

[...]

> +struct Node<K, V> {
> +    links: bindings::rb_node,

Same question as with `rb_root`, do we need `Opaque`?


Otherwise the patch looks good.

---
Cheers,
Benno

> +    key: K,
> +    value: V,
> +}
> 
> --
> 2.46.0.rc1.232.g9752f9e123-goog
> 

Re: [PATCH v8 3/6] rust: rbtree: add iterator

[Benno Lossin]

On 27.07.24 22:30, Matt Gilbride wrote:
> From: Wedson Almeida Filho <wedsonaf@gmail.com>
> 
> - Add Iterator implementation for `RBTree`, allowing
>   iteration over (key, value) pairs in key order.
> - Add individual `keys()` and `values()` functions to iterate over keys
>   or values alone.
> - Update doctests to use iteration instead of explicitly getting items.
> 
> Iteration is needed by the binder driver to enumerate all values in a
> tree for oneway spam detection [1].
> 
> Link: https://lore.kernel.org/rust-for-linux/20231101-rust-binder-v1-17-08ba9197f637@google.com/ [1]
> Signed-off-by: Wedson Almeida Filho <wedsonaf@gmail.com>
> Reviewed-by: Alice Ryhl <aliceryhl@google.com>
> Tested-by: Alice Ryhl <aliceryhl@google.com>
> Signed-off-by: Matt Gilbride <mattgilbride@google.com>
> ---
>  rust/kernel/rbtree.rs | 130 +++++++++++++++++++++++++++++++++++++++++++-------
>  1 file changed, 112 insertions(+), 18 deletions(-)

Reviewed-by: Benno Lossin <benno.lossin@proton.me>

---
Cheers,
Benno

Re: [PATCH v8 4/6] rust: rbtree: add mutable iterator

[Benno Lossin]

On 27.07.24 22:30, Matt Gilbride wrote:
> From: Wedson Almeida Filho <wedsonaf@gmail.com>
> 
> Add mutable Iterator implementation for `RBTree`,
> allowing iteration over (key, value) pairs in key order. Only values are
> mutable, as mutating keys implies modifying a node's position in the tree.
> 
> Mutable iteration is used by the binder driver during shutdown to
> clean up the tree maintained by the "range allocator" [1].
> 
> Link: https://lore.kernel.org/rust-for-linux/20231101-rust-binder-v1-6-08ba9197f637@google.com/ [1]
> Signed-off-by: Wedson Almeida Filho <wedsonaf@gmail.com>
> Signed-off-by: Matt Gilbride <mattgilbride@google.com>
> Reviewed-by: Alice Ryhl <aliceryhl@google.com>
> Tested-by: Alice Ryhl <aliceryhl@google.com>
> ---
>  rust/kernel/rbtree.rs | 98 ++++++++++++++++++++++++++++++++++++++++++++-------
>  1 file changed, 86 insertions(+), 12 deletions(-)
> 
> diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs
> index d10074e4ac58..d7514ebadfa8 100644
> --- a/rust/kernel/rbtree.rs
> +++ b/rust/kernel/rbtree.rs
> @@ -197,8 +197,26 @@ pub fn iter(&self) -> Iter<'_, K, V> {
>          // INVARIANT: `bindings::rb_first` returns a valid pointer to a tree node given a valid pointer to a tree root.

This INVARIANT is out of place, `Iter` doesn't have any INVARIANT any
more.

>          Iter {
>              _tree: PhantomData,
> -            // SAFETY: `self.root` is a valid pointer to the tree root.
> -            next: unsafe { bindings::rb_first(&self.root) },
> +            iter_raw: IterRaw {

This `IterRaw` construction is missing an INVARIANT comment. I think you
can copy paste from below.

> +                // SAFETY: by the invariants, all pointers are valid.
> +                next: unsafe { bindings::rb_first(&self.root) },
> +                _phantom: PhantomData,
> +            },
> +        }
> +    }
> +
> +    /// Returns a mutable iterator over the tree nodes, sorted by key.
> +    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
> +        IterMut {
> +            _tree: PhantomData,
> +            // INVARIANT:
> +            //   - `self.root` is a valid pointer to a tree root.
> +            //   - `bindings::rb_first` produces a valid pointer to a node given `root` is valid.
> +            iter_raw: IterRaw {
> +                // SAFETY: by the invariants, all pointers are valid.
> +                next: unsafe { bindings::rb_first(&self.root) },

Does this really derive a mutable reference? Ie shouldn't this be:?

    next: unsafe { bindings::rb_first(&mut self.root) },

> +                _phantom: PhantomData,
> +            },
>          }
>      }
> 
> @@ -211,6 +229,11 @@ pub fn keys(&self) -> impl Iterator<Item = &'_ K> {
>      pub fn values(&self) -> impl Iterator<Item = &'_ V> {
>          self.iter().map(|(_, v)| v)
>      }
> +
> +    /// Returns a mutable iterator over the values of the nodes in the tree, sorted by key.
> +    pub fn values_mut(&mut self) -> impl Iterator<Item = &'_ mut V> {
> +        self.iter_mut().map(|(_, v)| v)
> +    }
>  }
> 
>  impl<K, V> RBTree<K, V>
> @@ -414,13 +437,9 @@ fn into_iter(self) -> Self::IntoIter {
>  /// An iterator over the nodes of a [`RBTree`].
>  ///
>  /// Instances are created by calling [`RBTree::iter`].
> -///
> -/// # Invariants
> -/// - `self.next` is a valid pointer.
> -/// - `self.next` points to a node stored inside of a valid `RBTree`.
>  pub struct Iter<'a, K, V> {
>      _tree: PhantomData<&'a RBTree<K, V>>,
> -    next: *mut bindings::rb_node,
> +    iter_raw: IterRaw<K, V>,
>  }
> 
>  // SAFETY: The [`Iter`] gives out immutable references to K and V, so it has the same
> @@ -434,21 +453,76 @@ unsafe impl<'a, K: Sync, V: Sync> Sync for Iter<'a, K, V> {}
>  impl<'a, K, V> Iterator for Iter<'a, K, V> {
>      type Item = (&'a K, &'a V);
> 
> +    fn next(&mut self) -> Option<Self::Item> {
> +        self.iter_raw.next().map(|(k, v)|
> +            // SAFETY: Due to `self._tree`, `k` and `v` are valid for the lifetime of `'a`.
> +            unsafe { (&*k, &*v) })

I don't really like the formatting here, can you move the SAFETY one
line upwards? It should format nicely then.

> +    }
> +}
> +
> +impl<'a, K, V> IntoIterator for &'a mut RBTree<K, V> {
> +    type Item = (&'a K, &'a mut V);
> +    type IntoIter = IterMut<'a, K, V>;
> +
> +    fn into_iter(self) -> Self::IntoIter {
> +        self.iter_mut()
> +    }
> +}
> +
> +/// A mutable iterator over the nodes of a [`RBTree`].
> +///
> +/// Instances are created by calling [`RBTree::iter_mut`].
> +pub struct IterMut<'a, K, V> {
> +    _tree: PhantomData<&'a mut RBTree<K, V>>,
> +    iter_raw: IterRaw<K, V>,
> +}
> +
> +// SAFETY: The [`IterMut`] gives out immutable references to K and mutable references to V, so it has the same
> +// thread safety requirements as mutable references.
> +unsafe impl<'a, K: Send, V: Send> Send for IterMut<'a, K, V> {}

Since we only borrow `K` immutably, would it make sense to have `K:
Sync`?

---
Cheers,
Benno

> +
> +// SAFETY: The [`IterMut`] gives out immutable references to K and mutable references to V, so it has the same
> +// thread safety requirements as mutable references.
> +unsafe impl<'a, K: Sync, V: Sync> Sync for IterMut<'a, K, V> {}
> +
> +impl<'a, K, V> Iterator for IterMut<'a, K, V> {
> +    type Item = (&'a K, &'a mut V);
> +
> +    fn next(&mut self) -> Option<Self::Item> {
> +        self.iter_raw.next().map(|(k, v)|
> +            // SAFETY: Due to `&mut self`, we have exclusive access to `k` and `v`, for the lifetime of `'a`.
> +            unsafe { (&*k, &mut *v) })
> +    }
> +}
> +
> +/// A raw iterator over the nodes of a [`RBTree`].
> +///
> +/// # Invariants
> +/// - `self.next` is a valid pointer.
> +/// - `self.next` points to a node stored inside of a valid `RBTree`.
> +struct IterRaw<K, V> {
> +    next: *mut bindings::rb_node,
> +    _phantom: PhantomData<fn() -> (K, V)>,
> +}
> +
> +impl<K, V> Iterator for IterRaw<K, V> {
> +    type Item = (*mut K, *mut V);
> +
>      fn next(&mut self) -> Option<Self::Item> {
>          if self.next.is_null() {
>              return None;
>          }
> 
> -        // SAFETY: By the type invariant of `Iter`, `self.next` is a valid node in an `RBTree`,
> +        // SAFETY: By the type invariant of `IterRaw`, `self.next` is a valid node in an `RBTree`,
>          // and by the type invariant of `RBTree`, all nodes point to the links field of `Node<K, V>` objects.
> -        let cur = unsafe { container_of!(self.next, Node<K, V>, links) };
> +        let cur: *mut Node<K, V> =
> +            unsafe { container_of!(self.next, Node<K, V>, links) }.cast_mut();
> 
>          // SAFETY: `self.next` is a valid tree node by the type invariants.
>          self.next = unsafe { bindings::rb_next(self.next) };
> 
> -        // SAFETY: By the same reasoning above, it is safe to dereference the node. Additionally,
> -        // it is ok to return a reference to members because the iterator must outlive it.
> -        Some(unsafe { (&(*cur).key, &(*cur).value) })
> +        // SAFETY: By the same reasoning above, it is safe to dereference the node.
> +        Some(unsafe { (addr_of_mut!((*cur).key), addr_of_mut!((*cur).value)) })
>      }
>  }
> 
> 
> --
> 2.46.0.rc1.232.g9752f9e123-goog
> 

Re: [PATCH v8 5/6] rust: rbtree: add `RBTreeCursor`

[Benno Lossin]

On 27.07.24 22:30, Matt Gilbride wrote:
> Add a cursor interface to `RBTree`, supporting the following use cases:
> - Inspect the current node pointed to by the cursor, inspect/move to
>   it's neighbors in sort order (bidirectionally).
> - Mutate the tree itself by removing the current node pointed to by the
>   cursor, or one of its neighbors.
> 
> Add functions to obtain a cursor to the tree by key:
> - The node with the smallest key
> - The node with the largest key
> - The node matching the given key, or the one with the next larger key
> 
> The cursor abstraction is needed by the binder driver to efficiently
> search for nodes and (conditionally) modify them, as well as their
> neighbors [1].
> 
> Link: https://lore.kernel.org/rust-for-linux/20231101-rust-binder-v1-6-08ba9197f637@google.com/ [1]
> Co-developed-by: Alice Ryhl <aliceryhl@google.com>
> Signed-off-by: Alice Ryhl <aliceryhl@google.com>
> Reviewed-by: Alice Ryhl <aliceryhl@google.com>
> Tested-by: Alice Ryhl <aliceryhl@google.com>
> Signed-off-by: Matt Gilbride <mattgilbride@google.com>
> ---
>  rust/kernel/rbtree.rs | 536 ++++++++++++++++++++++++++++++++++++++++++++++++++
>  1 file changed, 536 insertions(+)
> 
> diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs
> index d7514ebadfa8..5d37aa373685 100644
> --- a/rust/kernel/rbtree.rs
> +++ b/rust/kernel/rbtree.rs
> @@ -234,6 +234,40 @@ pub fn values(&self) -> impl Iterator<Item = &'_ V> {
>      pub fn values_mut(&mut self) -> impl Iterator<Item = &'_ mut V> {
>          self.iter_mut().map(|(_, v)| v)
>      }
> +
> +    /// Returns a cursor over the tree nodes, starting with the smallest key.
> +    pub fn cursor_front(&mut self) -> Option<RBTreeCursor<'_, K, V>> {
> +        let root = addr_of_mut!(self.root);
> +        // SAFETY: `self.root` is always a valid root node
> +        let current = unsafe { bindings::rb_first(root) };
> +        NonNull::new(current).map(|current| {
> +            // INVARIANT:
> +            // - `current` is a valid node in the [`RBTree`] pointed to by `self`.
> +            // - Due to the type signature of this function, the returned [`RBTreeCursor`]
> +            //   borrows mutably from `self`.
> +            RBTreeCursor {
> +                current,
> +                tree: self,
> +            }
> +        })
> +    }
> +
> +    /// Returns a cursor over the tree nodes, starting with the largest key.
> +    pub fn cursor_back(&mut self) -> Option<RBTreeCursor<'_, K, V>> {
> +        let root = addr_of_mut!(self.root);
> +        // SAFETY: `self.root` is always a valid root node
> +        let current = unsafe { bindings::rb_last(root) };
> +        NonNull::new(current).map(|current| {
> +            // INVARIANT:
> +            // - `current` is a valid node in the [`RBTree`] pointed to by `self`.
> +            // - Due to the type signature of this function, the returned [`RBTreeCursor`]
> +            //   borrows mutably from `self`.
> +            RBTreeCursor {
> +                current,
> +                tree: self,
> +            }
> +        })
> +    }
>  }
> 
>  impl<K, V> RBTree<K, V>
> @@ -394,6 +428,75 @@ fn remove_node(&mut self, key: &K) -> Option<RBTreeNode<K, V>> {
>      pub fn remove(&mut self, key: &K) -> Option<V> {
>          self.remove_node(key).map(|node| node.node.value)
>      }
> +
> +    /// Returns a cursor over the tree nodes based on the given key.
> +    ///
> +    /// If the given key exists, the cursor starts there.
> +    /// Otherwise it starts with the first larger key in sort order.
> +    /// If there is no larger key, it returns [`None`].
> +    pub fn cursor_lower_bound(&mut self, key: &K) -> Option<RBTreeCursor<'_, K, V>>
> +    where
> +        K: Ord,
> +    {
> +        let mut node = self.root.rb_node;
> +        let mut best_match: Option<NonNull<Node<K, V>>> = None;
> +        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<K, V>` objects.
> +            let this = unsafe { container_of!(node, Node<K, V>, links) }.cast_mut();
> +            // SAFETY: `this` is a non-null node so it is valid by the type invariants.
> +            let this_key = unsafe { &(*this).key };
> +            // SAFETY: `node` is a non-null node so it is valid by the type invariants.
> +            let left_child = unsafe { (*node).rb_left };
> +            // SAFETY: `node` is a non-null node so it is valid by the type invariants.
> +            let right_child = unsafe { (*node).rb_right };
> +            if key == this_key {
> +                return NonNull::new(node).map(|current| {
> +                    // INVARIANT:
> +                    // - `node` is a valid node in the [`RBTree`] pointed to by `self`.
> +                    // - Due to the type signature of this function, the returned [`RBTreeCursor`]
> +                    //   borrows mutably from `self`.
> +                    RBTreeCursor {
> +                        current,
> +                        tree: self,
> +                    }
> +                });
> +            } else {
> +                node = if key > this_key {
> +                    right_child
> +                } else {
> +                    let is_better_match = match best_match {
> +                        None => true,
> +                        Some(best) => {
> +                            // SAFETY: `best` is a non-null node so it is valid by the type invariants.
> +                            let best_key = unsafe { &(*best.as_ptr()).key };
> +                            best_key > this_key
> +                        }
> +                    };
> +                    if is_better_match {
> +                        best_match = NonNull::new(this);
> +                    }
> +                    left_child
> +                };
> +            }
> +        }
> +
> +        let best = best_match?;
> +
> +        // SAFETY: `best` is a non-null node so it is valid by the type invariants.
> +        let links = unsafe { addr_of_mut!((*best.as_ptr()).links) };
> +
> +        NonNull::new(links).map(|current| {

Why would `links` be a null pointer? AFAIK it just came from `best`
which is non-null. (I don't know if we want to use `new_unchecked`
instead, but wanted to mention it)

> +            // INVARIANT:
> +            // - `current` is a valid node in the [`RBTree`] pointed to by `self`.
> +            // - Due to the type signature of this function, the returned [`RBTreeCursor`]
> +            //   borrows mutably from `self`.
> +            RBTreeCursor {
> +                current,
> +                tree: self,
> +            }
> +        })
> +    }

[...]

> +/// // Calling `remove_next` removes and returns the last element.
> +/// assert_eq!(cursor.remove_next().unwrap().to_key_value(), (30, 300));
> +///
> +/// # Ok::<(), Error>(())
> +/// ```

I would put a newline here.

> +/// # Invariants
> +/// - `current` points to a node that is in the same [`RBTree`] as `tree`.
> +pub struct RBTreeCursor<'a, K, V> {

I think we can name it just `Cursor`, since one can refer to it as
`rbtree::Cursor` and then it also follows the naming scheme for `Iter`
etc.

> +    tree: &'a mut RBTree<K, V>,
> +    current: NonNull<bindings::rb_node>,
> +}
> +
> +// SAFETY: The [`RBTreeCursor`] gives out immutable references to K and mutable references to V,
> +// so it has the same thread safety requirements as mutable references.
> +unsafe impl<'a, K: Send, V: Send> Send for RBTreeCursor<'a, K, V> {}

Again, do we want to use `K: Sync` here instead?

> +
> +// SAFETY: The [`RBTreeCursor`] gives out immutable references to K and mutable references to V,
> +// so it has the same thread safety requirements as mutable references.
> +unsafe impl<'a, K: Sync, V: Sync> Sync for RBTreeCursor<'a, K, V> {}
> +
> +impl<'a, K, V> RBTreeCursor<'a, K, V> {
> +    /// The current node
> +    pub fn current(&self) -> (&K, &V) {
> +        // SAFETY:
> +        // - `self.current` is a valid node by the type invariants.
> +        // - We have an immutable reference by the function signature.
> +        unsafe { Self::to_key_value(self.current) }
> +    }
> +
> +    /// The current node, with a mutable value
> +    pub fn current_mut(&mut self) -> (&K, &mut V) {
> +        // SAFETY:
> +        // - `self.current` is a valid node by the type invariants.
> +        // - We have an mutable reference by the function signature.
> +        unsafe { Self::to_key_value_mut(self.current) }
> +    }
> +
> +    /// Remove the current node from the tree.
> +    ///
> +    /// Returns a tuple where the first element is a cursor to the next node, if it exists,
> +    /// else the previous node, else [`None`] (if the tree becomes empty). The second element
> +    /// is the removed node.
> +    pub fn remove_current(self) -> (Option<Self>, RBTreeNode<K, V>) {
> +        let prev = self.get_neighbor_raw(Direction::Prev);
> +        let next = self.get_neighbor_raw(Direction::Next);
> +        // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`
> +        // point to the links field of `Node<K, V>` objects.
> +        let this = unsafe { container_of!(self.current.as_ptr(), Node<K, V>, links) }.cast_mut();
> +        // SAFETY: `this` is valid by the type invariants as described above.
> +        let node = unsafe { Box::from_raw(this) };
> +        let node = RBTreeNode { node };
> +        // SAFETY: The reference to the tree used to create the cursor outlives the cursor, so
> +        // the tree cannot change. By the tree invariant, all nodes are valid.
> +        unsafe { bindings::rb_erase(&mut (*this).links, addr_of_mut!(self.tree.root)) };
> +
> +        let current = match (prev, next) {
> +            (_, Some(next)) => next,
> +            (Some(prev), None) => prev,
> +            (None, None) => {
> +                return (None, node);
> +            }
> +        };
> +
> +        (
> +            // INVARIANT:
> +            // - `current` is a valid node in the [`RBTree`] pointed to by `self.tree`.
> +            // - Due to the function signature, `self` is an owned [`RBTreeCursor`],
> +            //   and [`RBTreeCursor`]s are only created via functions with a mutable reference
> +            //   to an [`RBTree`].
> +            Some(Self {
> +                current,
> +                tree: self.tree,
> +            }),
> +            node,
> +        )
> +    }
> +
> +    /// Remove the previous node, returning it if it exists.
> +    pub fn remove_prev(&mut self) -> Option<RBTreeNode<K, V>> {
> +        self.remove_neighbor(Direction::Prev)
> +    }
> +
> +    /// Remove the next node, returning it if it exists.
> +    pub fn remove_next(&mut self) -> Option<RBTreeNode<K, V>> {
> +        self.remove_neighbor(Direction::Next)
> +    }
> +
> +    fn remove_neighbor(&mut self, direction: Direction) -> Option<RBTreeNode<K, V>> {
> +        if let Some(neighbor) = self.get_neighbor_raw(direction) {
> +            let neighbor = neighbor.as_ptr();
> +            // SAFETY: The reference to the tree used to create the cursor outlives the cursor, so
> +            // the tree cannot change. By the tree invariant, all nodes are valid.
> +            unsafe { bindings::rb_erase(neighbor, addr_of_mut!(self.tree.root)) };
> +            // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`
> +            // point to the links field of `Node<K, V>` objects.
> +            let this = unsafe { container_of!(neighbor, Node<K, V>, links) }.cast_mut();
> +            // SAFETY: `this` is valid by the type invariants as described above.
> +            let node = unsafe { Box::from_raw(this) };
> +            return Some(RBTreeNode { node });
> +        }
> +        None
> +    }
> +
> +    /// Move the cursor to the previous node, returning [`None`] if it doesn't exist.
> +    pub fn move_prev(self) -> Option<Self> {
> +        self.mv(Direction::Prev)
> +    }
> +
> +    /// Move the cursor to the next node, returning [`None`] if it doesn't exist.
> +    pub fn move_next(self) -> Option<Self> {
> +        self.mv(Direction::Next)
> +    }
> +
> +    fn mv(self, direction: Direction) -> Option<Self> {
> +        // INVARIANT:
> +        // - `neighbor` is a valid node in the [`RBTree`] pointed to by `self.tree`.
> +        // - Due to the function signature, `self` is an owned [`RBTreeCursor`],
> +        //   and [`RBTreeCursor`]s are only created via functions with a mutable reference
> +        //   to an [`RBTree`].
> +        self.get_neighbor_raw(direction).map(|neighbor| Self {
> +            tree: self.tree,
> +            current: neighbor,
> +        })
> +    }
> +
> +    /// Access the previous node without moving the cursor.
> +    pub fn peek_prev(&self) -> Option<(&K, &V)> {
> +        self.peek(Direction::Prev)
> +    }
> +
> +    /// Access the previous node without moving the cursor.
> +    pub fn peek_next(&self) -> Option<(&K, &V)> {
> +        self.peek(Direction::Next)
> +    }
> +
> +    fn peek(&self, direction: Direction) -> Option<(&K, &V)> {
> +        self.get_neighbor_raw(direction)
> +            // SAFETY:
> +            // - `neighbor` is a valid tree node.
> +            // - By the function signature, we have an immutable reference to `self`.
> +            .map(|neighbor| unsafe { Self::to_key_value(neighbor) })

Alternative way of formatting this:

        self.get_neighbor_raw(direction).map(|neighbor| {
            // SAFETY:
            // - `neighbor` is a valid tree node.
            // - By the function signature, we have an immutable reference to `self`.
            unsafe { Self::to_key_value(neighbor) }
        })

I think it looks nicer, but we should probably have a written
preference.

> +    }
> +
> +    /// Access the previous node mutably without moving the cursor.
> +    pub fn peek_prev_mut(&mut self) -> Option<(&K, &mut V)> {
> +        self.peek_mut(Direction::Prev)
> +    }
> +
> +    /// Access the next node mutably without moving the cursor.
> +    pub fn peek_next_mut(&mut self) -> Option<(&K, &mut V)> {
> +        self.peek_mut(Direction::Next)
> +    }
> +
> +    fn peek_mut(&mut self, direction: Direction) -> Option<(&K, &mut V)> {
> +        self.get_neighbor_raw(direction)
> +            // SAFETY:
> +            // - `neighbor` is a valid tree node.
> +            // - By the function signature, we have a mutable reference to `self`.
> +            .map(|neighbor| unsafe { Self::to_key_value_mut(neighbor) })

Ditto.

---
Cheers,
Benno

> +    }
> +
> +    fn get_neighbor_raw(&self, direction: Direction) -> Option<NonNull<bindings::rb_node>> {
> +        // SAFETY: `self.current` is valid by the type invariants.
> +        let neighbor = unsafe {
> +            match direction {
> +                Direction::Prev => bindings::rb_prev(self.current.as_ptr()),
> +                Direction::Next => bindings::rb_next(self.current.as_ptr()),
> +            }
> +        };
> +
> +        NonNull::new(neighbor)
> +    }
> +
> +    /// SAFETY:
> +    /// - `node` must be a valid pointer to a node in an [`RBTree`].
> +    /// - The caller has immutable access to `node` for the duration of 'b.
> +    unsafe fn to_key_value<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, &'b V) {
> +        // SAFETY: the caller guarantees that `node` is a valid pointer in an `RBTree`.
> +        let (k, v) = unsafe { Self::to_key_value_raw(node) };
> +        // SAFETY: the caller guarantees immutable access to `node`.
> +        (k, unsafe { &*v })
> +    }
> +
> +    /// SAFETY:
> +    /// - `node` must be a valid pointer to a node in an [`RBTree`].
> +    /// - The caller has mutable access to `node` for the duration of 'b.
> +    unsafe fn to_key_value_mut<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, &'b mut V) {
> +        // SAFETY: the caller guarantees that `node` is a valid pointer in an `RBTree`.
> +        let (k, v) = unsafe { Self::to_key_value_raw(node) };
> +        // SAFETY: the caller guarantees mutable access to `node`.
> +        (k, unsafe { &mut *v })
> +    }
> +
> +    /// SAFETY:
> +    /// - `node` must be a valid pointer to a node in an [`RBTree`].
> +    /// - The caller has immutable access to the key for the duration of 'b.
> +    unsafe fn to_key_value_raw<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, *mut V) {
> +        // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`
> +        // point to the links field of `Node<K, V>` objects.
> +        let this = unsafe { container_of!(node.as_ptr(), Node<K, V>, links) }.cast_mut();
> +        // SAFETY: The passed `node` is the current node or a non-null neighbor,
> +        // thus `this` is valid by the type invariants.
> +        let k = unsafe { &(*this).key };
> +        // SAFETY: The passed `node` is the current node or a non-null neighbor,
> +        // thus `this` is valid by the type invariants.
> +        let v = unsafe { addr_of_mut!((*this).value) };
> +        (k, v)
> +    }
> +}
> +
> +/// Direction for [`RBTreeCursor`] operations.
> +enum Direction {
> +    /// the node immediately before, in sort order
> +    Prev,
> +    /// the node immediately after, in sort order
> +    Next,
> +}
> +
>  impl<'a, K, V> IntoIterator for &'a RBTree<K, V> {
>      type Item = (&'a K, &'a V);
>      type IntoIter = Iter<'a, K, V>;
> @@ -583,6 +1114,11 @@ impl<K, V> RBTreeNode<K, V> {
>      pub fn new(key: K, value: V, flags: Flags) -> Result<RBTreeNode<K, V>> {
>          Ok(RBTreeNodeReservation::new(flags)?.into_node(key, value))
>      }
> +
> +    /// Get the key and value from inside the node.
> +    pub fn to_key_value(self) -> (K, V) {
> +        (self.node.key, self.node.value)
> +    }
>  }
> 
>  // SAFETY: If K and V can be sent across threads, then it's also okay to send [`RBTreeNode`] across
> 
> --
> 2.46.0.rc1.232.g9752f9e123-goog
> 

Re: [PATCH v8 6/6] rust: rbtree: add `RBTree::entry`

[Benno Lossin]

On 27.07.24 22:30, Matt Gilbride wrote:
> From: Alice Ryhl <aliceryhl@google.com>
> 
> This mirrors the entry API [1] from the Rust standard library on
> `RBTree`. This API can be used to access the entry at a specific key and
> make modifications depending on whether the key is vacant or occupied.
> This API is useful because it can often be used to avoid traversing the
> tree multiple times.
> 
> This is used by binder to look up and conditionally access or insert a
> value, depending on whether it is there or not [2].
> 
> Link: https://doc.rust-lang.org/stable/std/collections/btree_map/enum.Entry.html [1]
> Link: https://android-review.googlesource.com/c/kernel/common/+/2849906 [2]
> Signed-off-by: Alice Ryhl <aliceryhl@google.com>
> Tested-by: Alice Ryhl <aliceryhl@google.com>
> Signed-off-by: Matt Gilbride <mattgilbride@google.com>
> ---
>  rust/kernel/rbtree.rs | 302 +++++++++++++++++++++++++++++++++++++-------------
>  1 file changed, 227 insertions(+), 75 deletions(-)
> 
> diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs
> index 5d37aa373685..428f8be8f3a2 100644
> --- a/rust/kernel/rbtree.rs
> +++ b/rust/kernel/rbtree.rs
> @@ -295,12 +295,19 @@ pub fn try_create_and_insert(
>      /// 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<K, V>) -> Option<RBTreeNode<K, V>> {
> -        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) };
> +    pub fn insert(&mut self, node: RBTreeNode<K, V>) -> Option<RBTreeNode<K, V>> {
> +        match self.raw_entry(&node.node.key) {
> +            RawEntry::Occupied(entry) => Some(entry.replace(node)),
> +            RawEntry::Vacant(entry) => {
> +                entry.insert(node);
> +                None
> +            }
> +        }
> +    }
> 
> +    fn raw_entry(&mut self, key: &K) -> RawEntry<'_, K, V> {
> +        let raw_self: *mut RBTree<K, V> = self;
> +        // The returned `RawEntry` is used to call either `rb_link_node` or `rb_replace_node`.
>          // The parameters of `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
> @@ -319,62 +326,56 @@ pub fn insert(&mut self, RBTreeNode { node }: RBTreeNode<K, V>) -> Option<RBTree
>          // 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;
> +        let mut child_field_of_parent: &mut *mut bindings::rb_node =
> +            // SAFETY: `raw_self` is a valid pointer to the `RBTree` (created from `self` above).
> +            unsafe { &mut (*raw_self).root.rb_node };
> +        while !(*child_field_of_parent).is_null() {

Why do you manually dereference `child_field_of_parent` here?

> +            let curr = *child_field_of_parent;
> +            // SAFETY: All links fields we create are in a `Node<K, V>`.

I think the SAFETY comment from below that argues via the type invariant
of `Self` is better.

> +            let node = unsafe { container_of!(curr, Node<K, V>, links) };
> 
> -            // 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<K, V>` objects.
> -            let this = unsafe { container_of!(parent, Node<K, V>, 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 },
> +            // SAFETY: `node` is a non-null node so it is valid by the type invariants.
> +            match key.cmp(unsafe { &(*node).key }) {
> +                // SAFETY: `curr` is a non-null node so it is valid by the type invariants.
> +                Ordering::Less => child_field_of_parent = unsafe { &mut (*curr).rb_left },
> +                // SAFETY: `curr` is a non-null node so it is valid by the type invariants.
> +                Ordering::Greater => child_field_of_parent = unsafe { &mut (*curr).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()) },
> -                    });
> +                    return RawEntry::Occupied(OccupiedEntry {
> +                        rbtree: self,
> +                        node_links: curr,
> +                    })
>                  }
>              }
> +            parent = curr;
>          }
> 
> -        // 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) };
> +        RawEntry::Vacant(RawVacantEntry {

RawVacantEntry has Invariants, so missing INVARIANT comment.

> +            rbtree: raw_self,
> +            parent,
> +            child_field_of_parent,
> +            _phantom: PhantomData,
> +        })
> +    }

[...]

> +/// A view into a vacant entry in a [`RBTree`]. It is part of the [`Entry`] enum.
> +pub struct VacantEntry<'a, K, V> {
> +    key: K,
> +    raw: RawVacantEntry<'a, K, V>,
> +}
> +
> +/// Like [`VacantEntry`], but doesn't hold on to the key.a

Typo: trailing 'a'.

> +///
> +/// # Invariants
> +/// - `parent` may be null if the new node becomes the root.
> +/// - `child_field_of_parent` is a valid pointer to the left-child or right-child of `parent`. If `parent` is
> +///     null, it is a pointer to the root of the [`RBTree`].
> +struct RawVacantEntry<'a, K, V> {
> +    rbtree: *mut RBTree<K, V>,
> +    /// The node that will become the parent of the new node if we insert one.
> +    parent: *mut bindings::rb_node,
> +    /// This points to the left-child or right-child field of `parent`, or `root` if `parent` is
> +    /// null.
> +    child_field_of_parent: *mut *mut bindings::rb_node,
> +    _phantom: PhantomData<&'a mut RBTree<K, V>>,
> +}
> +
> +impl<'a, K, V> RawVacantEntry<'a, K, V> {
> +    /// Inserts the given node into the [`RBTree`] at this entry.
> +    ///
> +    /// The `node` must have a key such that inserting it here does not break the ordering of this
> +    /// [`RBTree`].
> +    fn insert(self, node: RBTreeNode<K, V>) -> &'a mut V {
> +        let node = Box::into_raw(node.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) };
> +
> +        // INVARIANT: We are linking in a new node, which is valid. It remains valid because we
> +        // "forgot" it with `Box::into_raw`.
> +        // SAFETY: The type invariants of `RawVacantEntry` are exactly the safety requirements of `rb_link_node`.
> +        unsafe { bindings::rb_link_node(node_links, self.parent, self.child_field_of_parent) };
> +
> +        // SAFETY: All pointers are valid. `node` has just been inserted into the tree.
> +        unsafe { bindings::rb_insert_color(node_links, addr_of_mut!((*self.rbtree).root)) };
> +
> +        // SAFETY: The node is valid until we remove it from the tree.
> +        unsafe { &mut (*node).value }
> +    }
> +}
> +
> +impl<'a, K, V> VacantEntry<'a, K, V> {
> +    /// Inserts the given node into the [`RBTree`] at this entry.
> +    pub fn insert(self, value: V, reservation: RBTreeNodeReservation<K, V>) -> &'a mut V {
> +        self.raw.insert(reservation.into_node(self.key, value))
> +    }
> +}
> +
> +/// A view into an occupied entry in a [`RBTree`]. It is part of the [`Entry`] enum.
> +///
> +/// # Invariants
> +/// - `node_links` is a valid, non-null pointer to a tree node in `self.rbtree`
> +pub struct OccupiedEntry<'a, K, V> {
> +    rbtree: &'a mut RBTree<K, V>,
> +    /// The node that this entry corresponds to.
> +    node_links: *mut bindings::rb_node,
> +}
> +
> +impl<'a, K, V> OccupiedEntry<'a, K, V> {
> +    fn node_ptr(&self) -> *mut Node<K, V> {
> +        // SAFETY: By the type invariant of `Self`, all `node_links` pointers stored in `self`
> +        // point to the links field of `Node<K, V>` objects.
> +        unsafe { container_of!(self.node_links, Node<K, V>, links) }.cast_mut()

You should not call `cast_mut` here, see below

> +    }
> +
> +    /// Gets a reference to the value in the entry.
> +    pub fn get(&self) -> &V {
> +        // SAFETY: `self.node_ptr` produces a valid pointer to a node in the tree.

Can you add a `# Guarantees` section to `node_ptr` that states exactly
this?

> +        unsafe { &(*self.node_ptr()).value }
> +    }
> +
> +    /// Gets a mutable reference to the value in the entry.
> +    pub fn get_mut(&mut self) -> &mut V {
> +        // SAFETY: `self.node_ptr` produces a valid pointer to a node in the tree.
> +        unsafe { &mut (*self.node_ptr()).value }

This is sadly UB, you are creating a `&mut` reference from a pointer
that was derived from a `&` reference:
- `node_ptr` takes `&self`, thus it converts the `&mut self` to that.
- `container_of!` inside of `node_ptr` is used to create a read-only
  pointer to the `links` field (it is casted to `*mut`, but that doesn't
  change the fact that you are only allowed to use it for reads)
- `get_mut` turns it again into a `&mut` reference.

One solution is to make `note_ptr` take a `*mut Self`/`*const Self`.

> +    }
> +
> +    /// Converts the entry into a mutable reference to its value.
> +    ///
> +    /// If you need multiple references to the `OccupiedEntry`, see [`self#get_mut`].
> +    pub fn into_mut(self) -> &'a mut V {
> +        // SAFETY: `self.node_ptr` produces a valid pointer to a node in the tree.
> +        unsafe { &mut (*self.node_ptr()).value }
> +    }
> +
> +    /// Remove this entry from the [`RBTree`].
> +    pub fn remove_node(self) -> RBTreeNode<K, V> {
> +        // SAFETY: The node is a node in the tree, so it is valid.
> +        unsafe { bindings::rb_erase(self.node_links, &mut self.rbtree.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.
> +        RBTreeNode {
> +            // SAFETY: The node was a node in the tree, but we removed it, so we can convert it
> +            // back into a box.
> +            node: unsafe { Box::from_raw(self.node_ptr()) },
> +        }
> +    }
> +
> +    /// Takes the value of the entry out of the map, and returns it.
> +    pub fn remove(self) -> V {
> +        self.remove_node().node.value
> +    }
> +
> +    /// Swap the current node for the provided node.
> +    ///
> +    /// The key of both nodes must be equal.

Is this a safety requirement? Ie if it is violated, can memory bugs
occur, or is it only going to lead to logic bugs?

---
Cheers,
Benno

> +    fn replace(self, node: RBTreeNode<K, V>) -> RBTreeNode<K, V> {
> +        let node = Box::into_raw(node.node);
> +
> +        // SAFETY: `node` is valid at least until we call `Box::from_raw`, which only happens when
> +        // the node is removed or replaced.
> +        let new_node_links = unsafe { addr_of_mut!((*node).links) };
> +
> +        // SAFETY: This updates the pointers so that `new_node_links` is in the tree where
> +        // `self.node_links` used to be.
> +        unsafe {
> +            bindings::rb_replace_node(self.node_links, new_node_links, &mut self.rbtree.root)
> +        };
> +
> +        // SAFETY:
> +        // - `self.node_ptr` produces a valid pointer to a node in the tree.
> +        // - Now that we removed this entry from the tree, we can convert the node to a box.
> +        let old_node = unsafe { Box::from_raw(self.node_ptr()) };
> +
> +        RBTreeNode { node: old_node }
> +    }
> +}
> +
>  struct Node<K, V> {
>      links: bindings::rb_node,
>      key: K,
> 
> --
> 2.46.0.rc1.232.g9752f9e123-goog
> 

Re: [PATCH v8 5/6] rust: rbtree: add `RBTreeCursor`

[Alice Ryhl]

On Mon, Aug 5, 2024 at 9:35 PM Benno Lossin <benno.lossin@proton.me> wrote:
>
> On 27.07.24 22:30, Matt Gilbride wrote:
> > +    /// Returns a cursor over the tree nodes based on the given key.
> > +    ///
> > +    /// If the given key exists, the cursor starts there.
> > +    /// Otherwise it starts with the first larger key in sort order.
> > +    /// If there is no larger key, it returns [`None`].
> > +    pub fn cursor_lower_bound(&mut self, key: &K) -> Option<RBTreeCursor<'_, K, V>>
> > +    where
> > +        K: Ord,
> > +    {
> > +        let mut node = self.root.rb_node;
> > +        let mut best_match: Option<NonNull<Node<K, V>>> = None;
> > +        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<K, V>` objects.
> > +            let this = unsafe { container_of!(node, Node<K, V>, links) }.cast_mut();
> > +            // SAFETY: `this` is a non-null node so it is valid by the type invariants.
> > +            let this_key = unsafe { &(*this).key };
> > +            // SAFETY: `node` is a non-null node so it is valid by the type invariants.
> > +            let left_child = unsafe { (*node).rb_left };
> > +            // SAFETY: `node` is a non-null node so it is valid by the type invariants.
> > +            let right_child = unsafe { (*node).rb_right };
> > +            if key == this_key {
> > +                return NonNull::new(node).map(|current| {
> > +                    // INVARIANT:
> > +                    // - `node` is a valid node in the [`RBTree`] pointed to by `self`.
> > +                    // - Due to the type signature of this function, the returned [`RBTreeCursor`]
> > +                    //   borrows mutably from `self`.
> > +                    RBTreeCursor {
> > +                        current,
> > +                        tree: self,
> > +                    }
> > +                });
> > +            } else {
> > +                node = if key > this_key {
> > +                    right_child
> > +                } else {
> > +                    let is_better_match = match best_match {
> > +                        None => true,
> > +                        Some(best) => {
> > +                            // SAFETY: `best` is a non-null node so it is valid by the type invariants.
> > +                            let best_key = unsafe { &(*best.as_ptr()).key };
> > +                            best_key > this_key
> > +                        }
> > +                    };
> > +                    if is_better_match {
> > +                        best_match = NonNull::new(this);
> > +                    }
> > +                    left_child
> > +                };
> > +            }
> > +        }
> > +
> > +        let best = best_match?;
> > +
> > +        // SAFETY: `best` is a non-null node so it is valid by the type invariants.
> > +        let links = unsafe { addr_of_mut!((*best.as_ptr()).links) };
> > +
> > +        NonNull::new(links).map(|current| {
>
> Why would `links` be a null pointer? AFAIK it just came from `best`
> which is non-null. (I don't know if we want to use `new_unchecked`
> instead, but wanted to mention it)

It's never a null pointer in this branch. Do you prefer an extra
unsafe block to call new_unchecked?

> > +            // INVARIANT:
> > +            // - `current` is a valid node in the [`RBTree`] pointed to by `self`.
> > +            // - Due to the type signature of this function, the returned [`RBTreeCursor`]
> > +            //   borrows mutably from `self`.
> > +            RBTreeCursor {
> > +                current,
> > +                tree: self,
> > +            }
> > +        })
> > +    }
>
> [...]
>
> > +/// // Calling `remove_next` removes and returns the last element.
> > +/// assert_eq!(cursor.remove_next().unwrap().to_key_value(), (30, 300));
> > +///
> > +/// # Ok::<(), Error>(())
> > +/// ```
>
> I would put a newline here.

Ok.

> > +/// # Invariants
> > +/// - `current` points to a node that is in the same [`RBTree`] as `tree`.
> > +pub struct RBTreeCursor<'a, K, V> {
>
> I think we can name it just `Cursor`, since one can refer to it as
> `rbtree::Cursor` and then it also follows the naming scheme for `Iter`
> etc.

You are welcome to submit that as a follow-up change.

> > +    tree: &'a mut RBTree<K, V>,
> > +    current: NonNull<bindings::rb_node>,
> > +}
> > +
> > +// SAFETY: The [`RBTreeCursor`] gives out immutable references to K and mutable references to V,
> > +// so it has the same thread safety requirements as mutable references.
> > +unsafe impl<'a, K: Send, V: Send> Send for RBTreeCursor<'a, K, V> {}
>
> Again, do we want to use `K: Sync` here instead?

In this case, `K: Send` and `K: Sync` are both sufficient conditions,
but `K: Send` will generally be less restrictive for the user.

> > +    fn peek(&self, direction: Direction) -> Option<(&K, &V)> {
> > +        self.get_neighbor_raw(direction)
> > +            // SAFETY:
> > +            // - `neighbor` is a valid tree node.
> > +            // - By the function signature, we have an immutable reference to `self`.
> > +            .map(|neighbor| unsafe { Self::to_key_value(neighbor) })
>
> Alternative way of formatting this:
>
>         self.get_neighbor_raw(direction).map(|neighbor| {
>             // SAFETY:
>             // - `neighbor` is a valid tree node.
>             // - By the function signature, we have an immutable reference to `self`.
>             unsafe { Self::to_key_value(neighbor) }
>         })
>
> I think it looks nicer, but we should probably have a written
> preference.

We can reformat since we need another version anyway, but otherwise I
would have asked you to make this a follow-up change.

> > +    }
> > +
> > +    /// Access the previous node mutably without moving the cursor.
> > +    pub fn peek_prev_mut(&mut self) -> Option<(&K, &mut V)> {
> > +        self.peek_mut(Direction::Prev)
> > +    }
> > +
> > +    /// Access the next node mutably without moving the cursor.
> > +    pub fn peek_next_mut(&mut self) -> Option<(&K, &mut V)> {
> > +        self.peek_mut(Direction::Next)
> > +    }
> > +
> > +    fn peek_mut(&mut self, direction: Direction) -> Option<(&K, &mut V)> {
> > +        self.get_neighbor_raw(direction)
> > +            // SAFETY:
> > +            // - `neighbor` is a valid tree node.
> > +            // - By the function signature, we have a mutable reference to `self`.
> > +            .map(|neighbor| unsafe { Self::to_key_value_mut(neighbor) })
>
> Ditto.

Ditto.

Alice

Re: [PATCH v8 4/6] rust: rbtree: add mutable iterator

[Alice Ryhl]

On Mon, Aug 5, 2024 at 9:22 PM Benno Lossin <benno.lossin@proton.me> wrote:
>
> On 27.07.24 22:30, Matt Gilbride wrote:
> > From: Wedson Almeida Filho <wedsonaf@gmail.com>
> >
> > Add mutable Iterator implementation for `RBTree`,
> > allowing iteration over (key, value) pairs in key order. Only values are
> > mutable, as mutating keys implies modifying a node's position in the tree.
> >
> > Mutable iteration is used by the binder driver during shutdown to
> > clean up the tree maintained by the "range allocator" [1].
> >
> > Link: https://lore.kernel.org/rust-for-linux/20231101-rust-binder-v1-6-08ba9197f637@google.com/ [1]
> > Signed-off-by: Wedson Almeida Filho <wedsonaf@gmail.com>
> > Signed-off-by: Matt Gilbride <mattgilbride@google.com>
> > Reviewed-by: Alice Ryhl <aliceryhl@google.com>
> > Tested-by: Alice Ryhl <aliceryhl@google.com>
> > ---
> >  rust/kernel/rbtree.rs | 98 ++++++++++++++++++++++++++++++++++++++++++++-------
> >  1 file changed, 86 insertions(+), 12 deletions(-)
> >
> > diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs
> > index d10074e4ac58..d7514ebadfa8 100644
> > --- a/rust/kernel/rbtree.rs
> > +++ b/rust/kernel/rbtree.rs
> > @@ -197,8 +197,26 @@ pub fn iter(&self) -> Iter<'_, K, V> {
> >          // INVARIANT: `bindings::rb_first` returns a valid pointer to a tree node given a valid pointer to a tree root.
>
> This INVARIANT is out of place, `Iter` doesn't have any INVARIANT any
> more.

We can delete it.

> >          Iter {
> >              _tree: PhantomData,
> > -            // SAFETY: `self.root` is a valid pointer to the tree root.
> > -            next: unsafe { bindings::rb_first(&self.root) },
> > +            iter_raw: IterRaw {
>
> This `IterRaw` construction is missing an INVARIANT comment. I think you
> can copy paste from below.

We can copy from below.

> > +                // SAFETY: by the invariants, all pointers are valid.
> > +                next: unsafe { bindings::rb_first(&self.root) },
> > +                _phantom: PhantomData,
> > +            },
> > +        }
> > +    }
> > +
> > +    /// Returns a mutable iterator over the tree nodes, sorted by key.
> > +    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
> > +        IterMut {
> > +            _tree: PhantomData,
> > +            // INVARIANT:
> > +            //   - `self.root` is a valid pointer to a tree root.
> > +            //   - `bindings::rb_first` produces a valid pointer to a node given `root` is valid.
> > +            iter_raw: IterRaw {
> > +                // SAFETY: by the invariants, all pointers are valid.
> > +                next: unsafe { bindings::rb_first(&self.root) },
>
> Does this really derive a mutable reference? Ie shouldn't this be:?
>
>     next: unsafe { bindings::rb_first(&mut self.root) },

Let's change this to:

next: unsafe { bindings::rb_first(ptr::from_mut(&mut self.root)) }

This way, the pointer will be derived from a mutable reference even if
it becomes a `*const` through intermediate operations.


> > +                _phantom: PhantomData,
> > +            },
> >          }
> >      }
> >
> > @@ -211,6 +229,11 @@ pub fn keys(&self) -> impl Iterator<Item = &'_ K> {
> >      pub fn values(&self) -> impl Iterator<Item = &'_ V> {
> >          self.iter().map(|(_, v)| v)
> >      }
> > +
> > +    /// Returns a mutable iterator over the values of the nodes in the tree, sorted by key.
> > +    pub fn values_mut(&mut self) -> impl Iterator<Item = &'_ mut V> {
> > +        self.iter_mut().map(|(_, v)| v)
> > +    }
> >  }
> >
> >  impl<K, V> RBTree<K, V>
> > @@ -414,13 +437,9 @@ fn into_iter(self) -> Self::IntoIter {
> >  /// An iterator over the nodes of a [`RBTree`].
> >  ///
> >  /// Instances are created by calling [`RBTree::iter`].
> > -///
> > -/// # Invariants
> > -/// - `self.next` is a valid pointer.
> > -/// - `self.next` points to a node stored inside of a valid `RBTree`.
> >  pub struct Iter<'a, K, V> {
> >      _tree: PhantomData<&'a RBTree<K, V>>,
> > -    next: *mut bindings::rb_node,
> > +    iter_raw: IterRaw<K, V>,
> >  }
> >
> >  // SAFETY: The [`Iter`] gives out immutable references to K and V, so it has the same
> > @@ -434,21 +453,76 @@ unsafe impl<'a, K: Sync, V: Sync> Sync for Iter<'a, K, V> {}
> >  impl<'a, K, V> Iterator for Iter<'a, K, V> {
> >      type Item = (&'a K, &'a V);
> >
> > +    fn next(&mut self) -> Option<Self::Item> {
> > +        self.iter_raw.next().map(|(k, v)|
> > +            // SAFETY: Due to `self._tree`, `k` and `v` are valid for the lifetime of `'a`.
> > +            unsafe { (&*k, &*v) })
>
> I don't really like the formatting here, can you move the SAFETY one
> line upwards? It should format nicely then.

You suggested exactly the reverse formatting change on RBTreeCursor?

> > +    }
> > +}
> > +
> > +impl<'a, K, V> IntoIterator for &'a mut RBTree<K, V> {
> > +    type Item = (&'a K, &'a mut V);
> > +    type IntoIter = IterMut<'a, K, V>;
> > +
> > +    fn into_iter(self) -> Self::IntoIter {
> > +        self.iter_mut()
> > +    }
> > +}
> > +
> > +/// A mutable iterator over the nodes of a [`RBTree`].
> > +///
> > +/// Instances are created by calling [`RBTree::iter_mut`].
> > +pub struct IterMut<'a, K, V> {
> > +    _tree: PhantomData<&'a mut RBTree<K, V>>,
> > +    iter_raw: IterRaw<K, V>,
> > +}
> > +
> > +// SAFETY: The [`IterMut`] gives out immutable references to K and mutable references to V, so it has the same
> > +// thread safety requirements as mutable references.
> > +unsafe impl<'a, K: Send, V: Send> Send for IterMut<'a, K, V> {}
>
> Since we only borrow `K` immutably, would it make sense to have `K:
> Sync`?

No, `K: Send` is better because it's less restrictive in practice.

Alice

Re: [PATCH v8 6/6] rust: rbtree: add `RBTree::entry`

[Alice Ryhl]

On Mon, Aug 5, 2024 at 10:03 PM Benno Lossin <benno.lossin@proton.me> wrote:
>
> On 27.07.24 22:30, Matt Gilbride wrote:
> > From: Alice Ryhl <aliceryhl@google.com>
> >
> > This mirrors the entry API [1] from the Rust standard library on
> > `RBTree`. This API can be used to access the entry at a specific key and
> > make modifications depending on whether the key is vacant or occupied.
> > This API is useful because it can often be used to avoid traversing the
> > tree multiple times.
> >
> > This is used by binder to look up and conditionally access or insert a
> > value, depending on whether it is there or not [2].
> >
> > Link: https://doc.rust-lang.org/stable/std/collections/btree_map/enum.Entry.html [1]
> > Link: https://android-review.googlesource.com/c/kernel/common/+/2849906 [2]
> > Signed-off-by: Alice Ryhl <aliceryhl@google.com>
> > Tested-by: Alice Ryhl <aliceryhl@google.com>
> > Signed-off-by: Matt Gilbride <mattgilbride@google.com>
> > ---
> >  rust/kernel/rbtree.rs | 302 +++++++++++++++++++++++++++++++++++++-------------
> >  1 file changed, 227 insertions(+), 75 deletions(-)
> >
> > diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs
> > index 5d37aa373685..428f8be8f3a2 100644
> > --- a/rust/kernel/rbtree.rs
> > +++ b/rust/kernel/rbtree.rs
> > @@ -295,12 +295,19 @@ pub fn try_create_and_insert(
> >      /// 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<K, V>) -> Option<RBTreeNode<K, V>> {
> > -        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) };
> > +    pub fn insert(&mut self, node: RBTreeNode<K, V>) -> Option<RBTreeNode<K, V>> {
> > +        match self.raw_entry(&node.node.key) {
> > +            RawEntry::Occupied(entry) => Some(entry.replace(node)),
> > +            RawEntry::Vacant(entry) => {
> > +                entry.insert(node);
> > +                None
> > +            }
> > +        }
> > +    }
> >
> > +    fn raw_entry(&mut self, key: &K) -> RawEntry<'_, K, V> {
> > +        let raw_self: *mut RBTree<K, V> = self;
> > +        // The returned `RawEntry` is used to call either `rb_link_node` or `rb_replace_node`.
> >          // The parameters of `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
> > @@ -319,62 +326,56 @@ pub fn insert(&mut self, RBTreeNode { node }: RBTreeNode<K, V>) -> Option<RBTree
> >          // 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;
> > +        let mut child_field_of_parent: &mut *mut bindings::rb_node =
> > +            // SAFETY: `raw_self` is a valid pointer to the `RBTree` (created from `self` above).
> > +            unsafe { &mut (*raw_self).root.rb_node };
> > +        while !(*child_field_of_parent).is_null() {
>
> Why do you manually dereference `child_field_of_parent` here?

I think it helps for clarity. It makes it clear to the reader that
`child_field_of_parent` isn't null, but that it's pointing at a null
pointer.

> > +impl<'a, K, V> OccupiedEntry<'a, K, V> {
> > +    fn node_ptr(&self) -> *mut Node<K, V> {
> > +        // SAFETY: By the type invariant of `Self`, all `node_links` pointers stored in `self`
> > +        // point to the links field of `Node<K, V>` objects.
> > +        unsafe { container_of!(self.node_links, Node<K, V>, links) }.cast_mut()
>
> You should not call `cast_mut` here, see below
>
> > +    }
> > +
> > +    /// Gets a reference to the value in the entry.
> > +    pub fn get(&self) -> &V {
> > +        // SAFETY: `self.node_ptr` produces a valid pointer to a node in the tree.
>
> Can you add a `# Guarantees` section to `node_ptr` that states exactly
> this?
>
> > +        unsafe { &(*self.node_ptr()).value }
> > +    }
> > +
> > +    /// Gets a mutable reference to the value in the entry.
> > +    pub fn get_mut(&mut self) -> &mut V {
> > +        // SAFETY: `self.node_ptr` produces a valid pointer to a node in the tree.
> > +        unsafe { &mut (*self.node_ptr()).value }
>
> This is sadly UB, you are creating a `&mut` reference from a pointer
> that was derived from a `&` reference:
> - `node_ptr` takes `&self`, thus it converts the `&mut self` to that.
> - `container_of!` inside of `node_ptr` is used to create a read-only
>   pointer to the `links` field (it is casted to `*mut`, but that doesn't
>   change the fact that you are only allowed to use it for reads)
> - `get_mut` turns it again into a `&mut` reference.
>
> One solution is to make `note_ptr` take a `*mut Self`/`*const Self`.

We will get rid of node_ptr and duplicate its contents into get and get_mut.

> > +    }
> > +
> > +    /// Converts the entry into a mutable reference to its value.
> > +    ///
> > +    /// If you need multiple references to the `OccupiedEntry`, see [`self#get_mut`].
> > +    pub fn into_mut(self) -> &'a mut V {
> > +        // SAFETY: `self.node_ptr` produces a valid pointer to a node in the tree.
> > +        unsafe { &mut (*self.node_ptr()).value }
> > +    }
> > +
> > +    /// Remove this entry from the [`RBTree`].
> > +    pub fn remove_node(self) -> RBTreeNode<K, V> {
> > +        // SAFETY: The node is a node in the tree, so it is valid.
> > +        unsafe { bindings::rb_erase(self.node_links, &mut self.rbtree.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.
> > +        RBTreeNode {
> > +            // SAFETY: The node was a node in the tree, but we removed it, so we can convert it
> > +            // back into a box.
> > +            node: unsafe { Box::from_raw(self.node_ptr()) },
> > +        }
> > +    }
> > +
> > +    /// Takes the value of the entry out of the map, and returns it.
> > +    pub fn remove(self) -> V {
> > +        self.remove_node().node.value
> > +    }
> > +
> > +    /// Swap the current node for the provided node.
> > +    ///
> > +    /// The key of both nodes must be equal.
>
> Is this a safety requirement? Ie if it is violated, can memory bugs
> occur, or is it only going to lead to logic bugs?

No, it's not a safety requirement.

Alice

Re: [PATCH v8 2/6] rust: rbtree: add red-black tree implementation backed by the C version

[Alice Ryhl]

On Mon, Aug 5, 2024 at 9:02 PM Benno Lossin <benno.lossin@proton.me> wrote:
>
> On 27.07.24 22:30, Matt Gilbride wrote:
> > diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs
> > new file mode 100644
> > index 000000000000..74c53e4d5c00
> > --- /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: <https://www.kernel.org/doc/html/latest/core-api/rbtree.html>
> > +
> > +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.
> > +///
> > +/// # 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.
>
> I think we should standardize that `Invariants` and `Safety` sections
> are put last. This is because people reading the HTML version are
> probably not interested in the inner workings. This also makes it
> possible to have the invariants and the struct definition on the same
> screen.

We can reorder.

> > +/// ```
> > +pub struct RBTree<K, V> {
> > +    root: bindings::rb_root,
>
> It has been a while, so I might have already asked this, but do we need
> `Opaque` here? We would need it, if C changes anything inside `root`
> while Rust holds a `&RBTree` or `&mut RBTree`.
> I don't think that this is the case (ie we don't need `Opaque`), but I
> am not sure.

It's not needed.

> > +    _p: PhantomData<Node<K, V>>,
> > +}
> > +
>
> [...]
>
> > +    /// 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<K, V>) -> Option<RBTreeNode<K, V>> {
> > +        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 `rb_link_node` are as follows:
>
> Can you write `bindings::rb_link_node`?

Ok.

> > +        // - `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<K, V>` objects.
> > +            let this = unsafe { container_of!(parent, Node<K, V>, 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.
>
> Missing `` around "parent", double space after '.'.

Ok.

> > +                    //
> > +                    // 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()) },
> > +                    });
> > +                }
> > +            }
> > +        }
>
> [...]
>
> > +struct Node<K, V> {
> > +    links: bindings::rb_node,
>
> Same question as with `rb_root`, do we need `Opaque`?

No.

> Otherwise the patch looks good.

Will you give a Reviewed-by conditional on the above changes?

Alice

Re: [PATCH v8 2/6] rust: rbtree: add red-black tree implementation backed by the C version

[Benno Lossin]

On 06.08.24 10:41, Alice Ryhl wrote:
> On Mon, Aug 5, 2024 at 9:02 PM Benno Lossin <benno.lossin@proton.me> wrote:
>> Otherwise the patch looks good.
> 
> Will you give a Reviewed-by conditional on the above changes?

Yes.

---
Cheers,
Benno

Re: [PATCH v8 5/6] rust: rbtree: add `RBTreeCursor`

[Benno Lossin]

On 06.08.24 10:24, Alice Ryhl wrote:
> On Mon, Aug 5, 2024 at 9:35 PM Benno Lossin <benno.lossin@proton.me> wrote:
>> On 27.07.24 22:30, Matt Gilbride wrote:
>>> +        // SAFETY: `best` is a non-null node so it is valid by the type invariants.
>>> +        let links = unsafe { addr_of_mut!((*best.as_ptr()).links) };
>>> +
>>> +        NonNull::new(links).map(|current| {
>>
>> Why would `links` be a null pointer? AFAIK it just came from `best`
>> which is non-null. (I don't know if we want to use `new_unchecked`
>> instead, but wanted to mention it)
> 
> It's never a null pointer in this branch. Do you prefer an extra
> unsafe block to call new_unchecked?

No need, doesn't seem like this part is hot and this is something that
field projections could solve.

>>> +            // INVARIANT:
>>> +            // - `current` is a valid node in the [`RBTree`] pointed to by `self`.
>>> +            // - Due to the type signature of this function, the returned [`RBTreeCursor`]
>>> +            //   borrows mutably from `self`.
>>> +            RBTreeCursor {
>>> +                current,
>>> +                tree: self,
>>> +            }
>>> +        })
>>> +    }
>>
>> [...]
>>
>>> +/// // Calling `remove_next` removes and returns the last element.
>>> +/// assert_eq!(cursor.remove_next().unwrap().to_key_value(), (30, 300));
>>> +///
>>> +/// # Ok::<(), Error>(())
>>> +/// ```
>>
>> I would put a newline here.
> 
> Ok.
> 
>>> +/// # Invariants
>>> +/// - `current` points to a node that is in the same [`RBTree`] as `tree`.
>>> +pub struct RBTreeCursor<'a, K, V> {
>>
>> I think we can name it just `Cursor`, since one can refer to it as
>> `rbtree::Cursor` and then it also follows the naming scheme for `Iter`
>> etc.
> 
> You are welcome to submit that as a follow-up change.
> 
>>> +    tree: &'a mut RBTree<K, V>,
>>> +    current: NonNull<bindings::rb_node>,
>>> +}
>>> +
>>> +// SAFETY: The [`RBTreeCursor`] gives out immutable references to K and mutable references to V,
>>> +// so it has the same thread safety requirements as mutable references.
>>> +unsafe impl<'a, K: Send, V: Send> Send for RBTreeCursor<'a, K, V> {}
>>
>> Again, do we want to use `K: Sync` here instead?
> 
> In this case, `K: Send` and `K: Sync` are both sufficient conditions,
> but `K: Send` will generally be less restrictive for the user.

What if `K = struct(RefCell<i32>, i32)` where only the second i32 is
used in `(Partial)Ord`? Then you can send `RBTreeCursor` to another
thread and call `borrow` there, even though `K: !Sync` (and the value
still lives on another thread).

---
Cheers,
Benno

Re: [PATCH v8 5/6] rust: rbtree: add `RBTreeCursor`

[Alice Ryhl]

On Tue, Aug 6, 2024 at 11:01 AM Benno Lossin <benno.lossin@proton.me> wrote:
>
> On 06.08.24 10:24, Alice Ryhl wrote:
> > On Mon, Aug 5, 2024 at 9:35 PM Benno Lossin <benno.lossin@proton.me> wrote:
> >> On 27.07.24 22:30, Matt Gilbride wrote:
> >>> +        // SAFETY: `best` is a non-null node so it is valid by the type invariants.
> >>> +        let links = unsafe { addr_of_mut!((*best.as_ptr()).links) };
> >>> +
> >>> +        NonNull::new(links).map(|current| {
> >>
> >> Why would `links` be a null pointer? AFAIK it just came from `best`
> >> which is non-null. (I don't know if we want to use `new_unchecked`
> >> instead, but wanted to mention it)
> >
> > It's never a null pointer in this branch. Do you prefer an extra
> > unsafe block to call new_unchecked?
>
> No need, doesn't seem like this part is hot and this is something that
> field projections could solve.
>
> >>> +            // INVARIANT:
> >>> +            // - `current` is a valid node in the [`RBTree`] pointed to by `self`.
> >>> +            // - Due to the type signature of this function, the returned [`RBTreeCursor`]
> >>> +            //   borrows mutably from `self`.
> >>> +            RBTreeCursor {
> >>> +                current,
> >>> +                tree: self,
> >>> +            }
> >>> +        })
> >>> +    }
> >>
> >> [...]
> >>
> >>> +/// // Calling `remove_next` removes and returns the last element.
> >>> +/// assert_eq!(cursor.remove_next().unwrap().to_key_value(), (30, 300));
> >>> +///
> >>> +/// # Ok::<(), Error>(())
> >>> +/// ```
> >>
> >> I would put a newline here.
> >
> > Ok.
> >
> >>> +/// # Invariants
> >>> +/// - `current` points to a node that is in the same [`RBTree`] as `tree`.
> >>> +pub struct RBTreeCursor<'a, K, V> {
> >>
> >> I think we can name it just `Cursor`, since one can refer to it as
> >> `rbtree::Cursor` and then it also follows the naming scheme for `Iter`
> >> etc.
> >
> > You are welcome to submit that as a follow-up change.
> >
> >>> +    tree: &'a mut RBTree<K, V>,
> >>> +    current: NonNull<bindings::rb_node>,
> >>> +}
> >>> +
> >>> +// SAFETY: The [`RBTreeCursor`] gives out immutable references to K and mutable references to V,
> >>> +// so it has the same thread safety requirements as mutable references.
> >>> +unsafe impl<'a, K: Send, V: Send> Send for RBTreeCursor<'a, K, V> {}
> >>
> >> Again, do we want to use `K: Sync` here instead?
> >
> > In this case, `K: Send` and `K: Sync` are both sufficient conditions,
> > but `K: Send` will generally be less restrictive for the user.
>
> What if `K = struct(RefCell<i32>, i32)` where only the second i32 is
> used in `(Partial)Ord`? Then you can send `RBTreeCursor` to another
> thread and call `borrow` there, even though `K: !Sync` (and the value
> still lives on another thread).

In that scenario, all immutable references to the key would be on the
same thread as the cursor. The cursor borrows the tree mutably, so
there can only be one cursor.

When using `K: Send`, it's basically like having `RBTreeCursor` store
mutable references to the key, but with an API that downgrades to
immutable reference when giving out access to the key.

Alice

Re: [PATCH v8 4/6] rust: rbtree: add mutable iterator

[Benno Lossin]

On 06.08.24 10:30, Alice Ryhl wrote:
> On Mon, Aug 5, 2024 at 9:22 PM Benno Lossin <benno.lossin@proton.me> wrote:
>>
>> On 27.07.24 22:30, Matt Gilbride wrote:
>>> From: Wedson Almeida Filho <wedsonaf@gmail.com>
>>>
>>> Add mutable Iterator implementation for `RBTree`,
>>> allowing iteration over (key, value) pairs in key order. Only values are
>>> mutable, as mutating keys implies modifying a node's position in the tree.
>>>
>>> Mutable iteration is used by the binder driver during shutdown to
>>> clean up the tree maintained by the "range allocator" [1].
>>>
>>> Link: https://lore.kernel.org/rust-for-linux/20231101-rust-binder-v1-6-08ba9197f637@google.com/ [1]
>>> Signed-off-by: Wedson Almeida Filho <wedsonaf@gmail.com>
>>> Signed-off-by: Matt Gilbride <mattgilbride@google.com>
>>> Reviewed-by: Alice Ryhl <aliceryhl@google.com>
>>> Tested-by: Alice Ryhl <aliceryhl@google.com>
>>> ---
>>>  rust/kernel/rbtree.rs | 98 ++++++++++++++++++++++++++++++++++++++++++++-------
>>>  1 file changed, 86 insertions(+), 12 deletions(-)
>>>
>>> diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs
>>> index d10074e4ac58..d7514ebadfa8 100644
>>> --- a/rust/kernel/rbtree.rs
>>> +++ b/rust/kernel/rbtree.rs
>>> @@ -197,8 +197,26 @@ pub fn iter(&self) -> Iter<'_, K, V> {
>>>          // INVARIANT: `bindings::rb_first` returns a valid pointer to a tree node given a valid pointer to a tree root.
>>
>> This INVARIANT is out of place, `Iter` doesn't have any INVARIANT any
>> more.
> 
> We can delete it.
> 
>>>          Iter {
>>>              _tree: PhantomData,
>>> -            // SAFETY: `self.root` is a valid pointer to the tree root.
>>> -            next: unsafe { bindings::rb_first(&self.root) },
>>> +            iter_raw: IterRaw {
>>
>> This `IterRaw` construction is missing an INVARIANT comment. I think you
>> can copy paste from below.
> 
> We can copy from below.
> 
>>> +                // SAFETY: by the invariants, all pointers are valid.
>>> +                next: unsafe { bindings::rb_first(&self.root) },
>>> +                _phantom: PhantomData,
>>> +            },
>>> +        }
>>> +    }
>>> +
>>> +    /// Returns a mutable iterator over the tree nodes, sorted by key.
>>> +    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
>>> +        IterMut {
>>> +            _tree: PhantomData,
>>> +            // INVARIANT:
>>> +            //   - `self.root` is a valid pointer to a tree root.
>>> +            //   - `bindings::rb_first` produces a valid pointer to a node given `root` is valid.
>>> +            iter_raw: IterRaw {
>>> +                // SAFETY: by the invariants, all pointers are valid.
>>> +                next: unsafe { bindings::rb_first(&self.root) },
>>
>> Does this really derive a mutable reference? Ie shouldn't this be:?
>>
>>     next: unsafe { bindings::rb_first(&mut self.root) },
> 
> Let's change this to:
> 
> next: unsafe { bindings::rb_first(ptr::from_mut(&mut self.root)) }
> 
> This way, the pointer will be derived from a mutable reference even if
> it becomes a `*const` through intermediate operations.

SGTM

> 
> 
>>> +                _phantom: PhantomData,
>>> +            },
>>>          }
>>>      }
>>>
>>> @@ -211,6 +229,11 @@ pub fn keys(&self) -> impl Iterator<Item = &'_ K> {
>>>      pub fn values(&self) -> impl Iterator<Item = &'_ V> {
>>>          self.iter().map(|(_, v)| v)
>>>      }
>>> +
>>> +    /// Returns a mutable iterator over the values of the nodes in the tree, sorted by key.
>>> +    pub fn values_mut(&mut self) -> impl Iterator<Item = &'_ mut V> {
>>> +        self.iter_mut().map(|(_, v)| v)
>>> +    }
>>>  }
>>>
>>>  impl<K, V> RBTree<K, V>
>>> @@ -414,13 +437,9 @@ fn into_iter(self) -> Self::IntoIter {
>>>  /// An iterator over the nodes of a [`RBTree`].
>>>  ///
>>>  /// Instances are created by calling [`RBTree::iter`].
>>> -///
>>> -/// # Invariants
>>> -/// - `self.next` is a valid pointer.
>>> -/// - `self.next` points to a node stored inside of a valid `RBTree`.
>>>  pub struct Iter<'a, K, V> {
>>>      _tree: PhantomData<&'a RBTree<K, V>>,
>>> -    next: *mut bindings::rb_node,
>>> +    iter_raw: IterRaw<K, V>,
>>>  }
>>>
>>>  // SAFETY: The [`Iter`] gives out immutable references to K and V, so it has the same
>>> @@ -434,21 +453,76 @@ unsafe impl<'a, K: Sync, V: Sync> Sync for Iter<'a, K, V> {}
>>>  impl<'a, K, V> Iterator for Iter<'a, K, V> {
>>>      type Item = (&'a K, &'a V);
>>>
>>> +    fn next(&mut self) -> Option<Self::Item> {
>>> +        self.iter_raw.next().map(|(k, v)|
>>> +            // SAFETY: Due to `self._tree`, `k` and `v` are valid for the lifetime of `'a`.
>>> +            unsafe { (&*k, &*v) })
>>
>> I don't really like the formatting here, can you move the SAFETY one
>> line upwards? It should format nicely then.
> 
> You suggested exactly the reverse formatting change on RBTreeCursor?

Do you mean on this version or in a previous one? If you mean in this
one, then I would argue that they are not "reverses" of each other. For
this instance I would prefer

        // SAFETY: Due to `self._tree`, `k` and `v` are valid for the lifetime of `'a`.
        self.iter_raw.next().map(|(k, v)| unsafe { (&*k, &*v) })

or

        self.iter_raw.next().map(|(k, v)| {
            // SAFETY: Due to `self._tree`, `k` and `v` are valid for the lifetime of `'a`.
            unsafe { (&*k, &*v) }
        })

I hope that this seems consistent, my motivation behind the suggestions 
are that I don't like the comment splitting the single line.

---
Cheers,
Benno

Re: [PATCH v8 5/6] rust: rbtree: add `RBTreeCursor`

[Benno Lossin]

On 06.08.24 11:04, Alice Ryhl wrote:
> On Tue, Aug 6, 2024 at 11:01 AM Benno Lossin <benno.lossin@proton.me> wrote:
>>
>> On 06.08.24 10:24, Alice Ryhl wrote:
>>> On Mon, Aug 5, 2024 at 9:35 PM Benno Lossin <benno.lossin@proton.me> wrote:
>>>> On 27.07.24 22:30, Matt Gilbride wrote:
>>>>> +    tree: &'a mut RBTree<K, V>,
>>>>> +    current: NonNull<bindings::rb_node>,
>>>>> +}
>>>>> +
>>>>> +// SAFETY: The [`RBTreeCursor`] gives out immutable references to K and mutable references to V,
>>>>> +// so it has the same thread safety requirements as mutable references.
>>>>> +unsafe impl<'a, K: Send, V: Send> Send for RBTreeCursor<'a, K, V> {}
>>>>
>>>> Again, do we want to use `K: Sync` here instead?
>>>
>>> In this case, `K: Send` and `K: Sync` are both sufficient conditions,
>>> but `K: Send` will generally be less restrictive for the user.
>>
>> What if `K = struct(RefCell<i32>, i32)` where only the second i32 is
>> used in `(Partial)Ord`? Then you can send `RBTreeCursor` to another
>> thread and call `borrow` there, even though `K: !Sync` (and the value
>> still lives on another thread).
> 
> In that scenario, all immutable references to the key would be on the
> same thread as the cursor. The cursor borrows the tree mutably, so
> there can only be one cursor.
> 
> When using `K: Send`, it's basically like having `RBTreeCursor` store
> mutable references to the key, but with an API that downgrades to
> immutable reference when giving out access to the key.

Ah that's true.

---
Cheers,
Benno