[<prev] [next>] [<thread-prev] [thread-next>] [day] [month] [year] [list]
Message-ID: <309d3d27-62d0-42b0-b50a-40692a019b40@proton.me>
Date: Mon, 05 Aug 2024 19:35:05 +0000
From: Benno Lossin <benno.lossin@...ton.me>
To: Matt Gilbride <mattgilbride@...gle.com>, Miguel Ojeda <ojeda@...nel.org>, Alex Gaynor <alex.gaynor@...il.com>, Wedson Almeida Filho <wedsonaf@...il.com>, Boqun Feng <boqun.feng@...il.com>, Gary Guo <gary@...yguo.net>, Björn Roy Baron <bjorn3_gh@...tonmail.com>, Andreas Hindborg <a.hindborg@...sung.com>, Alice Ryhl <aliceryhl@...gle.com>, Greg Kroah-Hartman <gregkh@...uxfoundation.org>, Arve Hjønnevåg <arve@...roid.com>, Todd Kjos <tkjos@...roid.com>, Martijn Coenen <maco@...roid.com>, Joel Fernandes <joel@...lfernandes.org>, Carlos Llamas <cmllamas@...gle.com>, Suren Baghdasaryan <surenb@...gle.com>, Christian Brauner <brauner@...nel.org>
Cc: Rob Landley <rob@...dley.net>, Davidlohr Bueso <dave@...olabs.net>, Michel Lespinasse <michel@...pinasse.org>, rust-for-linux@...r.kernel.org, linux-kernel@...r.kernel.org
Subject: Re: [PATCH v8 5/6] rust: rbtree: add `RBTreeCursor`
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@...gle.com>
> Signed-off-by: Alice Ryhl <aliceryhl@...gle.com>
> Reviewed-by: Alice Ryhl <aliceryhl@...gle.com>
> Tested-by: Alice Ryhl <aliceryhl@...gle.com>
> Signed-off-by: Matt Gilbride <mattgilbride@...gle.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
>
Powered by blists - more mailing lists