Messages in this thread | | | Date | Thu, 25 Apr 2024 21:58:09 +0000 | From | Benno Lossin <> | Subject | Re: [PATCH v3 3/5] rust: rbtree: add `RBTreeIteratorMut` |
| |
On 18.04.24 16:15, Matt Gilbride wrote: > 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 {
This is missing an INVARIANT comment.
> + _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> {
I think the names `Iter` and `IterMut` are more natural. That is what the collections in `std::collections` do. These are in the module `rbtree`, so you can refer to them as `rbtree::Iter`.
> + _tree: PhantomData<&'a RBTree<K, V>>,
This should have the type `PhantomData<&'a mut RBTree<K, V>>`.
> + next: *mut bindings::rb_node, > +}
You could create a common iterator type, since both `RBTreeIterator` and `RBTreeIteratorMut` are very similar.
How about a (private) `RawIter`:
struct RawIter<K, V> { next: *mut bindings::rb_node, _phantom: PhantomData<fn() -> (K, V)>, }
And implement `Iterator` with `Item = (*mut K, *mut V)` for `RawIter`.
Then you can change `Iter` to be:
pub struct Iter<'a, K, V> { raw_iter: RawIter<K, V>, _tree: PhantomData<&'a RBTree<K, V>>, }
-- Cheers, Benno
> + > +// 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 >
| |