Linux kernel mirror (for testing)
git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git
kernel
os
linux
rust: rbtree: add mutable iterator
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>
Reviewed-by: Alice Ryhl <aliceryhl@google.com>
Tested-by: Alice Ryhl <aliceryhl@google.com>
Reviewed-by: Boqun Feng <boqun.feng@gmail.com>
Reviewed-by: Benno Lossin <benno.lossin@proton.me>
Signed-off-by: Matt Gilbride <mattgilbride@google.com>
Link: https://lore.kernel.org/r/20240822-b4-rbtree-v12-3-014561758a57@google.com
Signed-off-by: Miguel Ojeda <ojeda@kernel.org>
authored by
Wedson Almeida Filho
and committed by
Miguel Ojeda
cf5397d1
e601f1bb
+89
-14
+89
-14
rust/kernel/rbtree.rs
+89
-14
rust/kernel/rbtree.rs
···
12
12
cmp::{Ord, Ordering},
13
13
marker::PhantomData,
14
14
mem::MaybeUninit,
15
-
ptr::{addr_of_mut, NonNull},
15
+
ptr::{addr_of_mut, from_mut, NonNull},
16
16
};
17
17
18
18
/// A red-black tree with owned nodes.
···
194
194
195
195
/// Returns an iterator over the tree nodes, sorted by key.
196
196
pub fn iter(&self) -> Iter<'_, K, V> {
197
-
// INVARIANT: `bindings::rb_first` returns a valid pointer to a tree node given a valid pointer to a tree root.
198
197
Iter {
199
198
_tree: PhantomData,
200
-
// SAFETY: `self.root` is a valid pointer to the tree root.
201
-
next: unsafe { bindings::rb_first(&self.root) },
199
+
// INVARIANT:
200
+
// - `self.root` is a valid pointer to a tree root.
201
+
// - `bindings::rb_first` produces a valid pointer to a node given `root` is valid.
202
+
iter_raw: IterRaw {
203
+
// SAFETY: by the invariants, all pointers are valid.
204
+
next: unsafe { bindings::rb_first(&self.root) },
205
+
_phantom: PhantomData,
206
+
},
207
+
}
208
+
}
209
+
210
+
/// Returns a mutable iterator over the tree nodes, sorted by key.
211
+
pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
212
+
IterMut {
213
+
_tree: PhantomData,
214
+
// INVARIANT:
215
+
// - `self.root` is a valid pointer to a tree root.
216
+
// - `bindings::rb_first` produces a valid pointer to a node given `root` is valid.
217
+
iter_raw: IterRaw {
218
+
// SAFETY: by the invariants, all pointers are valid.
219
+
next: unsafe { bindings::rb_first(from_mut(&mut self.root)) },
220
+
_phantom: PhantomData,
221
+
},
202
222
}
203
223
}
204
224
···
230
210
/// Returns an iterator over the values of the nodes in the tree, sorted by key.
231
211
pub fn values(&self) -> impl Iterator<Item = &'_ V> {
232
212
self.iter().map(|(_, v)| v)
213
+
}
214
+
215
+
/// Returns a mutable iterator over the values of the nodes in the tree, sorted by key.
216
+
pub fn values_mut(&mut self) -> impl Iterator<Item = &'_ mut V> {
217
+
self.iter_mut().map(|(_, v)| v)
233
218
}
234
219
}
235
220
···
439
414
/// An iterator over the nodes of a [`RBTree`].
440
415
///
441
416
/// Instances are created by calling [`RBTree::iter`].
442
-
///
443
-
/// # Invariants
444
-
/// - `self.next` is a valid pointer.
445
-
/// - `self.next` points to a node stored inside of a valid `RBTree`.
446
417
pub struct Iter<'a, K, V> {
447
418
_tree: PhantomData<&'a RBTree<K, V>>,
448
-
next: *mut bindings::rb_node,
419
+
iter_raw: IterRaw<K, V>,
449
420
}
450
421
451
422
// SAFETY: The [`Iter`] gives out immutable references to K and V, so it has the same
···
456
435
type Item = (&'a K, &'a V);
457
436
458
437
fn next(&mut self) -> Option<Self::Item> {
438
+
// SAFETY: Due to `self._tree`, `k` and `v` are valid for the lifetime of `'a`.
439
+
self.iter_raw.next().map(|(k, v)| unsafe { (&*k, &*v) })
440
+
}
441
+
}
442
+
443
+
impl<'a, K, V> IntoIterator for &'a mut RBTree<K, V> {
444
+
type Item = (&'a K, &'a mut V);
445
+
type IntoIter = IterMut<'a, K, V>;
446
+
447
+
fn into_iter(self) -> Self::IntoIter {
448
+
self.iter_mut()
449
+
}
450
+
}
451
+
452
+
/// A mutable iterator over the nodes of a [`RBTree`].
453
+
///
454
+
/// Instances are created by calling [`RBTree::iter_mut`].
455
+
pub struct IterMut<'a, K, V> {
456
+
_tree: PhantomData<&'a mut RBTree<K, V>>,
457
+
iter_raw: IterRaw<K, V>,
458
+
}
459
+
460
+
// SAFETY: The [`IterMut`] has exclusive access to both `K` and `V`, so it is sufficient to require them to be `Send`.
461
+
// The iterator only gives out immutable references to the keys, but since the iterator has excusive access to those same
462
+
// keys, `Send` is sufficient. `Sync` would be okay, but it is more restrictive to the user.
463
+
unsafe impl<'a, K: Send, V: Send> Send for IterMut<'a, K, V> {}
464
+
465
+
// SAFETY: The [`IterMut`] gives out immutable references to K and mutable references to V, so it has the same
466
+
// thread safety requirements as mutable references.
467
+
unsafe impl<'a, K: Sync, V: Sync> Sync for IterMut<'a, K, V> {}
468
+
469
+
impl<'a, K, V> Iterator for IterMut<'a, K, V> {
470
+
type Item = (&'a K, &'a mut V);
471
+
472
+
fn next(&mut self) -> Option<Self::Item> {
473
+
self.iter_raw.next().map(|(k, v)|
474
+
// SAFETY: Due to `&mut self`, we have exclusive access to `k` and `v`, for the lifetime of `'a`.
475
+
unsafe { (&*k, &mut *v) })
476
+
}
477
+
}
478
+
479
+
/// A raw iterator over the nodes of a [`RBTree`].
480
+
///
481
+
/// # Invariants
482
+
/// - `self.next` is a valid pointer.
483
+
/// - `self.next` points to a node stored inside of a valid `RBTree`.
484
+
struct IterRaw<K, V> {
485
+
next: *mut bindings::rb_node,
486
+
_phantom: PhantomData<fn() -> (K, V)>,
487
+
}
488
+
489
+
impl<K, V> Iterator for IterRaw<K, V> {
490
+
type Item = (*mut K, *mut V);
491
+
492
+
fn next(&mut self) -> Option<Self::Item> {
459
493
if self.next.is_null() {
460
494
return None;
461
495
}
462
496
463
-
// SAFETY: By the type invariant of `Iter`, `self.next` is a valid node in an `RBTree`,
497
+
// SAFETY: By the type invariant of `IterRaw`, `self.next` is a valid node in an `RBTree`,
464
498
// and by the type invariant of `RBTree`, all nodes point to the links field of `Node<K, V>` objects.
465
-
let cur = unsafe { container_of!(self.next, Node<K, V>, links) };
499
+
let cur = unsafe { container_of!(self.next, Node<K, V>, links) }.cast_mut();
466
500
467
501
// SAFETY: `self.next` is a valid tree node by the type invariants.
468
502
self.next = unsafe { bindings::rb_next(self.next) };
469
503
470
-
// SAFETY: By the same reasoning above, it is safe to dereference the node. Additionally,
471
-
// it is ok to return a reference to members because the iterator must outlive it.
472
-
Some(unsafe { (&(*cur).key, &(*cur).value) })
504
+
// SAFETY: By the same reasoning above, it is safe to dereference the node.
505
+
Some(unsafe { (addr_of_mut!((*cur).key), addr_of_mut!((*cur).value)) })
473
506
}
474
507
}
475
508