Messages in this thread Patch in this message | | | Date | Thu, 18 Apr 2024 14:15:42 +0000 | Subject | [PATCH v3 3/5] rust: rbtree: add `RBTreeIteratorMut` | From | Matt Gilbride <> |
| |
From: Wedson Almeida Filho <wedsonaf@gmail.com>
Add mutable Iterator implementation (`RBTreeIteratorMut`) 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 | 64 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+)
diff --git a/rust/kernel/rbtree.rs b/rust/kernel/rbtree.rs index 2f836be7bdbe..50d440c9926d 100644 --- a/rust/kernel/rbtree.rs +++ b/rust/kernel/rbtree.rs @@ -222,6 +222,15 @@ pub fn iter(&self) -> RBTreeIterator<'_, K, V> { } } + /// Returns a mutable iterator over the tree nodes, sorted by key. + pub fn iter_mut(&mut self) -> RBTreeIteratorMut<'_, K, V> { + RBTreeIteratorMut { + _tree: PhantomData, + // SAFETY: `root` is valid as it's embedded in `self` and we have a valid `self`. + 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) @@ -231,6 +240,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> @@ -466,6 +480,56 @@ fn next(&mut self) -> Option<Self::Item> { } } +impl<'a, K, V> IntoIterator for &'a mut RBTree<K, V> { + type Item = (&'a K, &'a mut V); + type IntoIter = RBTreeIteratorMut<'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`]. +/// +/// # Invariants +/// - `self.next` is a valid pointer. +/// - `self.next` points to a node stored inside of a valid `RBTree`. +pub struct RBTreeIteratorMut<'a, K, V> { + _tree: PhantomData<&'a RBTree<K, V>>, + next: *mut bindings::rb_node, +} + +// SAFETY: The [`RBTreeIterator`] gives out mutable references to K and V, so it has the same +// thread safety requirements as mutable references. +unsafe impl<'a, K: Send, V: Send> Send for RBTreeIteratorMut<'a, K, V> {} + +// SAFETY: The [`RBTreeIterator`] gives out mutable references to K and V, so it has the same +// thread safety requirements as mutable references. +unsafe impl<'a, K: Sync, V: Sync> Sync for RBTreeIteratorMut<'a, K, V> {} + +impl<'a, K, V> Iterator for RBTreeIteratorMut<'a, K, V> { + type Item = (&'a K, &'a mut V); + + fn next(&mut self) -> Option<Self::Item> { + if self.next.is_null() { + return None; + } + + // 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 cur = 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, &mut (*cur).value) }) + } +} + /// 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 -- 2.44.0.769.g3c40516874-goog
| |