Linux kernel mirror (for testing)
git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git
kernel
os
linux
rust: rbtree: add iterator
- Add Iterator implementation for `RBTree`, allowing
iteration over (key, value) pairs in key order.
- Add individual `keys()` and `values()` functions to iterate over keys
or values alone.
- Update doctests to use iteration instead of explicitly getting items.
Iteration is needed by the binder driver to enumerate all values in a
tree for oneway spam detection [1].
Link: https://lore.kernel.org/rust-for-linux/20231101-rust-binder-v1-17-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: Benno Lossin <benno.lossin@proton.me>
Reviewed-by: Boqun Feng <boqun.feng@gmail.com>
Signed-off-by: Matt Gilbride <mattgilbride@google.com>
Link: https://lore.kernel.org/r/20240822-b4-rbtree-v12-2-014561758a57@google.com
Signed-off-by: Miguel Ojeda <ojeda@kernel.org>
authored by
Wedson Almeida Filho
and committed by
Miguel Ojeda
e601f1bb
a0d13aac
+112
-18
+112
-18
rust/kernel/rbtree.rs
+112
-18
rust/kernel/rbtree.rs
···
42
42
/// assert_eq!(tree.get(&30).unwrap(), &300);
43
43
/// }
44
44
///
45
+
/// // Iterate over the nodes we just inserted.
46
+
/// {
47
+
/// let mut iter = tree.iter();
48
+
/// assert_eq!(iter.next().unwrap(), (&10, &100));
49
+
/// assert_eq!(iter.next().unwrap(), (&20, &200));
50
+
/// assert_eq!(iter.next().unwrap(), (&30, &300));
51
+
/// assert!(iter.next().is_none());
52
+
/// }
53
+
///
54
+
/// // Print all elements.
55
+
/// for (key, value) in &tree {
56
+
/// pr_info!("{} = {}\n", key, value);
57
+
/// }
58
+
///
45
59
/// // Replace one of the elements.
46
60
/// tree.try_create_and_insert(10, 1000, flags::GFP_KERNEL)?;
47
61
///
48
62
/// // Check that the tree reflects the replacement.
49
63
/// {
50
-
/// assert_eq!(tree.get(&10).unwrap(), &1000);
51
-
/// assert_eq!(tree.get(&20).unwrap(), &200);
52
-
/// assert_eq!(tree.get(&30).unwrap(), &300);
64
+
/// let mut iter = tree.iter();
65
+
/// assert_eq!(iter.next().unwrap(), (&10, &1000));
66
+
/// assert_eq!(iter.next().unwrap(), (&20, &200));
67
+
/// assert_eq!(iter.next().unwrap(), (&30, &300));
68
+
/// assert!(iter.next().is_none());
53
69
/// }
54
70
///
55
71
/// // Change the value of one of the elements.
···
73
57
///
74
58
/// // Check that the tree reflects the update.
75
59
/// {
76
-
/// assert_eq!(tree.get(&10).unwrap(), &1000);
77
-
/// assert_eq!(tree.get(&20).unwrap(), &200);
78
-
/// assert_eq!(tree.get(&30).unwrap(), &3000);
60
+
/// let mut iter = tree.iter();
61
+
/// assert_eq!(iter.next().unwrap(), (&10, &1000));
62
+
/// assert_eq!(iter.next().unwrap(), (&20, &200));
63
+
/// assert_eq!(iter.next().unwrap(), (&30, &3000));
64
+
/// assert!(iter.next().is_none());
79
65
/// }
80
66
///
81
67
/// // Remove an element.
···
85
67
///
86
68
/// // Check that the tree reflects the removal.
87
69
/// {
88
-
/// assert_eq!(tree.get(&10), None);
89
-
/// assert_eq!(tree.get(&20).unwrap(), &200);
90
-
/// assert_eq!(tree.get(&30).unwrap(), &3000);
70
+
/// let mut iter = tree.iter();
71
+
/// assert_eq!(iter.next().unwrap(), (&20, &200));
72
+
/// assert_eq!(iter.next().unwrap(), (&30, &3000));
73
+
/// assert!(iter.next().is_none());
91
74
/// }
92
75
///
93
76
/// # Ok::<(), Error>(())
···
128
109
///
129
110
/// // Check the nodes we just inserted.
130
111
/// {
131
-
/// assert_eq!(tree.get(&10).unwrap(), &100);
132
-
/// assert_eq!(tree.get(&20).unwrap(), &200);
133
-
/// assert_eq!(tree.get(&30).unwrap(), &300);
112
+
/// let mut iter = tree.iter();
113
+
/// assert_eq!(iter.next().unwrap(), (&10, &100));
114
+
/// assert_eq!(iter.next().unwrap(), (&20, &200));
115
+
/// assert_eq!(iter.next().unwrap(), (&30, &300));
116
+
/// assert!(iter.next().is_none());
134
117
/// }
135
118
///
136
119
/// // Remove a node, getting back ownership of it.
···
140
119
///
141
120
/// // Check that the tree reflects the removal.
142
121
/// {
143
-
/// assert_eq!(tree.get(&10).unwrap(), &100);
144
-
/// assert_eq!(tree.get(&20).unwrap(), &200);
145
-
/// assert_eq!(tree.get(&30), None);
122
+
/// let mut iter = tree.iter();
123
+
/// assert_eq!(iter.next().unwrap(), (&10, &100));
124
+
/// assert_eq!(iter.next().unwrap(), (&20, &200));
125
+
/// assert!(iter.next().is_none());
146
126
/// }
147
127
///
148
128
/// // Create a preallocated reservation that we can re-use later.
···
155
133
///
156
134
/// // Check that the tree reflect the new insertion.
157
135
/// {
158
-
/// assert_eq!(tree.get(&10).unwrap(), &100);
159
-
/// assert_eq!(tree.get(&15).unwrap(), &150);
160
-
/// assert_eq!(tree.get(&20).unwrap(), &200);
136
+
/// let mut iter = tree.iter();
137
+
/// assert_eq!(iter.next().unwrap(), (&10, &100));
138
+
/// assert_eq!(iter.next().unwrap(), (&15, &150));
139
+
/// assert_eq!(iter.next().unwrap(), (&20, &200));
140
+
/// assert!(iter.next().is_none());
161
141
/// }
162
142
///
163
143
/// # Ok::<(), Error>(())
···
190
166
root: bindings::rb_root::default(),
191
167
_p: PhantomData,
192
168
}
169
+
}
170
+
171
+
/// Returns an iterator over the tree nodes, sorted by key.
172
+
pub fn iter(&self) -> Iter<'_, K, V> {
173
+
// INVARIANT: `bindings::rb_first` returns a valid pointer to a tree node given a valid pointer to a tree root.
174
+
Iter {
175
+
_tree: PhantomData,
176
+
// SAFETY: `self.root` is a valid pointer to the tree root.
177
+
next: unsafe { bindings::rb_first(&self.root) },
178
+
}
179
+
}
180
+
181
+
/// Returns an iterator over the keys of the nodes in the tree, in sorted order.
182
+
pub fn keys(&self) -> impl Iterator<Item = &'_ K> {
183
+
self.iter().map(|(k, _)| k)
184
+
}
185
+
186
+
/// Returns an iterator over the values of the nodes in the tree, sorted by key.
187
+
pub fn values(&self) -> impl Iterator<Item = &'_ V> {
188
+
self.iter().map(|(_, v)| v)
193
189
}
194
190
}
195
191
···
399
355
// SAFETY: `this` is valid per the loop invariant.
400
356
unsafe { drop(Box::from_raw(this.cast_mut())) };
401
357
}
358
+
}
359
+
}
360
+
361
+
impl<'a, K, V> IntoIterator for &'a RBTree<K, V> {
362
+
type Item = (&'a K, &'a V);
363
+
type IntoIter = Iter<'a, K, V>;
364
+
365
+
fn into_iter(self) -> Self::IntoIter {
366
+
self.iter()
367
+
}
368
+
}
369
+
370
+
/// An iterator over the nodes of a [`RBTree`].
371
+
///
372
+
/// Instances are created by calling [`RBTree::iter`].
373
+
///
374
+
/// # Invariants
375
+
/// - `self.next` is a valid pointer.
376
+
/// - `self.next` points to a node stored inside of a valid `RBTree`.
377
+
pub struct Iter<'a, K, V> {
378
+
_tree: PhantomData<&'a RBTree<K, V>>,
379
+
next: *mut bindings::rb_node,
380
+
}
381
+
382
+
// SAFETY: The [`Iter`] gives out immutable references to K and V, so it has the same
383
+
// thread safety requirements as immutable references.
384
+
unsafe impl<'a, K: Sync, V: Sync> Send for Iter<'a, K, V> {}
385
+
386
+
// SAFETY: The [`Iter`] gives out immutable references to K and V, so it has the same
387
+
// thread safety requirements as immutable references.
388
+
unsafe impl<'a, K: Sync, V: Sync> Sync for Iter<'a, K, V> {}
389
+
390
+
impl<'a, K, V> Iterator for Iter<'a, K, V> {
391
+
type Item = (&'a K, &'a V);
392
+
393
+
fn next(&mut self) -> Option<Self::Item> {
394
+
if self.next.is_null() {
395
+
return None;
396
+
}
397
+
398
+
// SAFETY: By the type invariant of `Iter`, `self.next` is a valid node in an `RBTree`,
399
+
// and by the type invariant of `RBTree`, all nodes point to the links field of `Node<K, V>` objects.
400
+
let cur = unsafe { container_of!(self.next, Node<K, V>, links) };
401
+
402
+
// SAFETY: `self.next` is a valid tree node by the type invariants.
403
+
self.next = unsafe { bindings::rb_next(self.next) };
404
+
405
+
// SAFETY: By the same reasoning above, it is safe to dereference the node. Additionally,
406
+
// it is ok to return a reference to members because the iterator must outlive it.
407
+
Some(unsafe { (&(*cur).key, &(*cur).value) })
402
408
}
403
409
}
404
410