Messages in this thread | | | Date | Thu, 25 Apr 2024 22:20:27 +0000 | From | Benno Lossin <> | Subject | Re: [PATCH v3 4/5] rust: rbtree: add `RBTreeCursor` |
| |
On 18.04.24 16:15, Matt Gilbride wrote: > @@ -400,6 +442,72 @@ 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 };
Since you have this pattern multiple times, I think you could have a single function that walks the tree and takes care of most of the `unsafe` stuff. A good starting point might be this:
unsafe fn walk<F, R>(node: *mut bindings::rb_node, dir: F) -> R where /* this, key */ F: FnMut(*mut bindings::rb_node, &K) -> Either<Direction, R>;
> + if key == this_key { > + // INVARIANT: > + // - `self.root` and `node` are valid pointers. > + // - `self.root` is the root of an [`RBTree`]. > + // - `node` is a valid node in an [`RBTree`]. > + // - Due to the type signature of this function, the returned [`RBTreeCursor`] > + // borrows from `self`. > + return Some(RBTreeCursor { > + _tree: PhantomData, > + root: addr_of_mut!(self.root), > + current: node, > + }); > + } 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 > + } > + }; > + } > + // INVARIANT: > + // - `self.root` and `best` are valid pointers. > + // - `self.root` is the root of an [`RBTree`]. > + // - `best` is a valid node in an [`RBTree`]. > + // - Due to the type signature of this function, the returned [`RBTreeCursor`] > + // borrows from `self`. > + best_match.map(|best| RBTreeCursor { > + _tree: PhantomData, > + root: addr_of_mut!(self.root), > + // SAFETY: `best` is a non-null node so it is valid by the type invariants. > + current: unsafe { addr_of_mut!((*best.as_ptr()).links) }, > + }) > + } > } > > impl<K, V> Default for RBTree<K, V> {
[...]
> +/// # Invariants > +/// - `root` and `current` are valid pointers. > +/// - `root` points to the `root` node of an [`RBTree`]. > +/// - `current` points to a node that is in the same [`RBTree`] that `root` is pointing to. > +/// - A cursor must borrow the [`RBTree`] containing `root` and `current` mutably. > +pub struct RBTreeCursor<'a, K, V> { > + _tree: PhantomData<&'a RBTree<K, V>>,
Why is this not `&'a mut RBTree<K, V>`?
> + root: *mut bindings::rb_root, > + current: *mut 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 cursor to the next node, if it exists, > + /// else the previous node. Returns [`None`] if the tree > + /// becomes empty. > + pub fn remove_current(self) -> Option<Self> { > + 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, Node<K, V>, links) }.cast_mut(); > + // 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, self.root) }; > + > + let current = match (prev, next) { > + (_, Some(next)) => next, > + (Some(prev), None) => prev, > + (None, None) => { > + return None; > + } > + }; > + > + // INVARIANT: > + // - `self.root` and `current` are valid pointers. > + // - `self.root` is the root of an [`RBTree`]. > + // - `current` is a valid node in an [`RBTree`]. > + // - 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, > + root: self.root, > + }) > + } > + > + /// Remove the previous node, returning it if it exists. > + pub fn remove_prev(&mut self) -> Option<(K, V)> {
Why do these functions not return `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<(K, V)> { > + self.remove_neighbor(Direction::Next) > + } > + > + fn remove_neighbor(&mut self, direction: Direction) -> Option<(K, V)> { > + if let Some(neighbor) = self.get_neighbor_raw(direction) { > + // 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, self.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 n = unsafe { Box::from_raw(this) }; > + return Some((n.key, n.value)); > + } > + 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> {
Does it hurt to name this `move`?
> + // INVARIANT: > + // - `self.root` and `neighbor` are valid pointers. > + // - `self.root` is the root of an [`RBTree`]. > + // - `neighbor` is a valid node in an [`RBTree`]. > + // - 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, > + root: self.root, > + current: 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 'a. > + unsafe fn to_key_value(node: *mut bindings::rb_node) -> (&'a K, &'a 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, Node<K, V>, links) }; > + // 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 { &(*this).value }; > + (k, 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 'a. > + unsafe fn to_key_value_mut(node: *mut bindings::rb_node) -> (&'a K, &'a 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, 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 { &mut (*this).value }; > + (k, v) > + }
You can create a single function that does the `container_of!` stuff and that returns `(*mut K, *mut V)` and implement these two in terms of that one.
-- Cheers, Benno
> +} > + > +/// 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 = RBTreeIterator<'a, K, V>; > > -- > 2.44.0.769.g3c40516874-goog >
| |