projects/linux·blob master · refs

1433 lines · 53.1 KiB

1// SPDX-License-Identifier: GPL-2.0
2
3//! Red-black trees.
4//!
5//! C header: [`include/linux/rbtree.h`](srctree/include/linux/rbtree.h)
6//!
7//! Reference: <https://docs.kernel.org/core-api/rbtree.html>
8
9use crate::{alloc::Flags, bindings, container_of, error::Result, prelude::*};
10use core::{
11 cmp::{Ord, Ordering},
12 marker::PhantomData,
13 mem::MaybeUninit,
14 ptr::{addr_of_mut, from_mut, NonNull},
15};
16
17/// A red-black tree with owned nodes.
18///
19/// It is backed by the kernel C red-black trees.
20///
21/// # Examples
22///
23/// In the example below we do several operations on a tree. We note that insertions may fail if
24/// the system is out of memory.
25///
26/// ```
27/// use kernel::{alloc::flags, rbtree::{RBTree, RBTreeNode, RBTreeNodeReservation}};
28///
29/// // Create a new tree.
30/// let mut tree = RBTree::new();
31///
32/// // Insert three elements.
33/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
34/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
35/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
36///
37/// // Check the nodes we just inserted.
38/// {
39/// assert_eq!(tree.get(&10), Some(&100));
40/// assert_eq!(tree.get(&20), Some(&200));
41/// assert_eq!(tree.get(&30), Some(&300));
42/// }
43///
44/// // Iterate over the nodes we just inserted.
45/// {
46/// let mut iter = tree.iter();
47/// assert_eq!(iter.next(), Some((&10, &100)));
48/// assert_eq!(iter.next(), Some((&20, &200)));
49/// assert_eq!(iter.next(), Some((&30, &300)));
50/// assert!(iter.next().is_none());
51/// }
52///
53/// // Print all elements.
54/// for (key, value) in &tree {
55/// pr_info!("{} = {}\n", key, value);
56/// }
57///
58/// // Replace one of the elements.
59/// tree.try_create_and_insert(10, 1000, flags::GFP_KERNEL)?;
60///
61/// // Check that the tree reflects the replacement.
62/// {
63/// let mut iter = tree.iter();
64/// assert_eq!(iter.next(), Some((&10, &1000)));
65/// assert_eq!(iter.next(), Some((&20, &200)));
66/// assert_eq!(iter.next(), Some((&30, &300)));
67/// assert!(iter.next().is_none());
68/// }
69///
70/// // Change the value of one of the elements.
71/// *tree.get_mut(&30).unwrap() = 3000;
72///
73/// // Check that the tree reflects the update.
74/// {
75/// let mut iter = tree.iter();
76/// assert_eq!(iter.next(), Some((&10, &1000)));
77/// assert_eq!(iter.next(), Some((&20, &200)));
78/// assert_eq!(iter.next(), Some((&30, &3000)));
79/// assert!(iter.next().is_none());
80/// }
81///
82/// // Remove an element.
83/// tree.remove(&10);
84///
85/// // Check that the tree reflects the removal.
86/// {
87/// let mut iter = tree.iter();
88/// assert_eq!(iter.next(), Some((&20, &200)));
89/// assert_eq!(iter.next(), Some((&30, &3000)));
90/// assert!(iter.next().is_none());
91/// }
92///
93/// # Ok::<(), Error>(())
94/// ```
95///
96/// In the example below, we first allocate a node, acquire a spinlock, then insert the node into
97/// the tree. This is useful when the insertion context does not allow sleeping, for example, when
98/// holding a spinlock.
99///
100/// ```
101/// use kernel::{alloc::flags, rbtree::{RBTree, RBTreeNode}, sync::SpinLock};
102///
103/// fn insert_test(tree: &SpinLock<RBTree<u32, u32>>) -> Result {
104/// // Pre-allocate node. This may fail (as it allocates memory).
105/// let node = RBTreeNode::new(10, 100, flags::GFP_KERNEL)?;
106///
107/// // Insert node while holding the lock. It is guaranteed to succeed with no allocation
108/// // attempts.
109/// let mut guard = tree.lock();
110/// guard.insert(node);
111/// Ok(())
112/// }
113/// ```
114///
115/// In the example below, we reuse an existing node allocation from an element we removed.
116///
117/// ```
118/// use kernel::{alloc::flags, rbtree::{RBTree, RBTreeNodeReservation}};
119///
120/// // Create a new tree.
121/// let mut tree = RBTree::new();
122///
123/// // Insert three elements.
124/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
125/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
126/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
127///
128/// // Check the nodes we just inserted.
129/// {
130/// let mut iter = tree.iter();
131/// assert_eq!(iter.next(), Some((&10, &100)));
132/// assert_eq!(iter.next(), Some((&20, &200)));
133/// assert_eq!(iter.next(), Some((&30, &300)));
134/// assert!(iter.next().is_none());
135/// }
136///
137/// // Remove a node, getting back ownership of it.
138/// let existing = tree.remove(&30);
139///
140/// // Check that the tree reflects the removal.
141/// {
142/// let mut iter = tree.iter();
143/// assert_eq!(iter.next(), Some((&10, &100)));
144/// assert_eq!(iter.next(), Some((&20, &200)));
145/// assert!(iter.next().is_none());
146/// }
147///
148/// // Create a preallocated reservation that we can re-use later.
149/// let reservation = RBTreeNodeReservation::new(flags::GFP_KERNEL)?;
150///
151/// // Insert a new node into the tree, reusing the previous allocation. This is guaranteed to
152/// // succeed (no memory allocations).
153/// tree.insert(reservation.into_node(15, 150));
154///
155/// // Check that the tree reflect the new insertion.
156/// {
157/// let mut iter = tree.iter();
158/// assert_eq!(iter.next(), Some((&10, &100)));
159/// assert_eq!(iter.next(), Some((&15, &150)));
160/// assert_eq!(iter.next(), Some((&20, &200)));
161/// assert!(iter.next().is_none());
162/// }
163///
164/// # Ok::<(), Error>(())
165/// ```
166///
167/// # Invariants
168///
169/// Non-null parent/children pointers stored in instances of the `rb_node` C struct are always
170/// valid, and pointing to a field of our internal representation of a node.
171pub struct RBTree<K, V> {
172 root: bindings::rb_root,
173 _p: PhantomData<Node<K, V>>,
174}
175
176// SAFETY: An [`RBTree`] allows the same kinds of access to its values that a struct allows to its
177// fields, so we use the same Send condition as would be used for a struct with K and V fields.
178unsafe impl<K: Send, V: Send> Send for RBTree<K, V> {}
179
180// SAFETY: An [`RBTree`] allows the same kinds of access to its values that a struct allows to its
181// fields, so we use the same Sync condition as would be used for a struct with K and V fields.
182unsafe impl<K: Sync, V: Sync> Sync for RBTree<K, V> {}
183
184impl<K, V> RBTree<K, V> {
185 /// Creates a new and empty tree.
186 pub fn new() -> Self {
187 Self {
188 // INVARIANT: There are no nodes in the tree, so the invariant holds vacuously.
189 root: bindings::rb_root::default(),
190 _p: PhantomData,
191 }
192 }
193
194 /// Returns true if this tree is empty.
195 #[inline]
196 pub fn is_empty(&self) -> bool {
197 self.root.rb_node.is_null()
198 }
199
200 /// Returns an iterator over the tree nodes, sorted by key.
201 pub fn iter(&self) -> Iter<'_, K, V> {
202 Iter {
203 _tree: PhantomData,
204 // INVARIANT:
205 // - `self.root` is a valid pointer to a tree root.
206 // - `bindings::rb_first` produces a valid pointer to a node given `root` is valid.
207 iter_raw: IterRaw {
208 // SAFETY: by the invariants, all pointers are valid.
209 next: unsafe { bindings::rb_first(&self.root) },
210 _phantom: PhantomData,
211 },
212 }
213 }
214
215 /// Returns a mutable iterator over the tree nodes, sorted by key.
216 pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
217 IterMut {
218 _tree: PhantomData,
219 // INVARIANT:
220 // - `self.root` is a valid pointer to a tree root.
221 // - `bindings::rb_first` produces a valid pointer to a node given `root` is valid.
222 iter_raw: IterRaw {
223 // SAFETY: by the invariants, all pointers are valid.
224 next: unsafe { bindings::rb_first(from_mut(&mut self.root)) },
225 _phantom: PhantomData,
226 },
227 }
228 }
229
230 /// Returns an iterator over the keys of the nodes in the tree, in sorted order.
231 pub fn keys(&self) -> impl Iterator<Item = &'_ K> {
232 self.iter().map(|(k, _)| k)
233 }
234
235 /// Returns an iterator over the values of the nodes in the tree, sorted by key.
236 pub fn values(&self) -> impl Iterator<Item = &'_ V> {
237 self.iter().map(|(_, v)| v)
238 }
239
240 /// Returns a mutable iterator over the values of the nodes in the tree, sorted by key.
241 pub fn values_mut(&mut self) -> impl Iterator<Item = &'_ mut V> {
242 self.iter_mut().map(|(_, v)| v)
243 }
244
245 /// Returns a cursor over the tree nodes, starting with the smallest key.
246 pub fn cursor_front_mut(&mut self) -> Option<CursorMut<'_, K, V>> {
247 let root = addr_of_mut!(self.root);
248 // SAFETY: `self.root` is always a valid root node.
249 let current = unsafe { bindings::rb_first(root) };
250 NonNull::new(current).map(|current| {
251 // INVARIANT:
252 // - `current` is a valid node in the [`RBTree`] pointed to by `self`.
253 CursorMut {
254 current,
255 tree: self,
256 }
257 })
258 }
259
260 /// Returns an immutable cursor over the tree nodes, starting with the smallest key.
261 pub fn cursor_front(&self) -> Option<Cursor<'_, K, V>> {
262 let root = &raw const self.root;
263 // SAFETY: `self.root` is always a valid root node.
264 let current = unsafe { bindings::rb_first(root) };
265 NonNull::new(current).map(|current| {
266 // INVARIANT:
267 // - `current` is a valid node in the [`RBTree`] pointed to by `self`.
268 Cursor {
269 current,
270 _tree: PhantomData,
271 }
272 })
273 }
274
275 /// Returns a cursor over the tree nodes, starting with the largest key.
276 pub fn cursor_back_mut(&mut self) -> Option<CursorMut<'_, K, V>> {
277 let root = addr_of_mut!(self.root);
278 // SAFETY: `self.root` is always a valid root node.
279 let current = unsafe { bindings::rb_last(root) };
280 NonNull::new(current).map(|current| {
281 // INVARIANT:
282 // - `current` is a valid node in the [`RBTree`] pointed to by `self`.
283 CursorMut {
284 current,
285 tree: self,
286 }
287 })
288 }
289
290 /// Returns a cursor over the tree nodes, starting with the largest key.
291 pub fn cursor_back(&self) -> Option<Cursor<'_, K, V>> {
292 let root = &raw const self.root;
293 // SAFETY: `self.root` is always a valid root node.
294 let current = unsafe { bindings::rb_last(root) };
295 NonNull::new(current).map(|current| {
296 // INVARIANT:
297 // - `current` is a valid node in the [`RBTree`] pointed to by `self`.
298 Cursor {
299 current,
300 _tree: PhantomData,
301 }
302 })
303 }
304}
305
306impl<K, V> RBTree<K, V>
307where
308 K: Ord,
309{
310 /// Tries to insert a new value into the tree.
311 ///
312 /// It overwrites a node if one already exists with the same key and returns it (containing the
313 /// key/value pair). Returns [`None`] if a node with the same key didn't already exist.
314 ///
315 /// Returns an error if it cannot allocate memory for the new node.
316 pub fn try_create_and_insert(
317 &mut self,
318 key: K,
319 value: V,
320 flags: Flags,
321 ) -> Result<Option<RBTreeNode<K, V>>> {
322 Ok(self.insert(RBTreeNode::new(key, value, flags)?))
323 }
324
325 /// Inserts a new node into the tree.
326 ///
327 /// It overwrites a node if one already exists with the same key and returns it (containing the
328 /// key/value pair). Returns [`None`] if a node with the same key didn't already exist.
329 ///
330 /// This function always succeeds.
331 pub fn insert(&mut self, node: RBTreeNode<K, V>) -> Option<RBTreeNode<K, V>> {
332 match self.raw_entry(&node.node.key) {
333 RawEntry::Occupied(entry) => Some(entry.replace(node)),
334 RawEntry::Vacant(entry) => {
335 entry.insert(node);
336 None
337 }
338 }
339 }
340
341 fn raw_entry(&mut self, key: &K) -> RawEntry<'_, K, V> {
342 let raw_self: *mut RBTree<K, V> = self;
343 // The returned `RawEntry` is used to call either `rb_link_node` or `rb_replace_node`.
344 // The parameters of `bindings::rb_link_node` are as follows:
345 // - `node`: A pointer to an uninitialized node being inserted.
346 // - `parent`: A pointer to an existing node in the tree. One of its child pointers must be
347 // null, and `node` will become a child of `parent` by replacing that child pointer
348 // with a pointer to `node`.
349 // - `rb_link`: A pointer to either the left-child or right-child field of `parent`. This
350 // specifies which child of `parent` should hold `node` after this call. The
351 // value of `*rb_link` must be null before the call to `rb_link_node`. If the
352 // red/black tree is empty, then it’s also possible for `parent` to be null. In
353 // this case, `rb_link` is a pointer to the `root` field of the red/black tree.
354 //
355 // We will traverse the tree looking for a node that has a null pointer as its child,
356 // representing an empty subtree where we can insert our new node. We need to make sure
357 // that we preserve the ordering of the nodes in the tree. In each iteration of the loop
358 // we store `parent` and `child_field_of_parent`, and the new `node` will go somewhere
359 // in the subtree of `parent` that `child_field_of_parent` points at. Once
360 // we find an empty subtree, we can insert the new node using `rb_link_node`.
361 let mut parent = core::ptr::null_mut();
362 let mut child_field_of_parent: &mut *mut bindings::rb_node =
363 // SAFETY: `raw_self` is a valid pointer to the `RBTree` (created from `self` above).
364 unsafe { &mut (*raw_self).root.rb_node };
365 while !(*child_field_of_parent).is_null() {
366 let curr = *child_field_of_parent;
367 // SAFETY: All links fields we create are in a `Node<K, V>`.
368 let node = unsafe { container_of!(curr, Node<K, V>, links) };
369
370 // SAFETY: `node` is a non-null node so it is valid by the type invariants.
371 match key.cmp(unsafe { &(*node).key }) {
372 // SAFETY: `curr` is a non-null node so it is valid by the type invariants.
373 Ordering::Less => child_field_of_parent = unsafe { &mut (*curr).rb_left },
374 // SAFETY: `curr` is a non-null node so it is valid by the type invariants.
375 Ordering::Greater => child_field_of_parent = unsafe { &mut (*curr).rb_right },
376 Ordering::Equal => {
377 return RawEntry::Occupied(OccupiedEntry {
378 rbtree: self,
379 node_links: curr,
380 })
381 }
382 }
383 parent = curr;
384 }
385
386 RawEntry::Vacant(RawVacantEntry {
387 rbtree: raw_self,
388 parent,
389 child_field_of_parent,
390 _phantom: PhantomData,
391 })
392 }
393
394 /// Gets the given key's corresponding entry in the map for in-place manipulation.
395 pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
396 match self.raw_entry(&key) {
397 RawEntry::Occupied(entry) => Entry::Occupied(entry),
398 RawEntry::Vacant(entry) => Entry::Vacant(VacantEntry { raw: entry, key }),
399 }
400 }
401
402 /// Used for accessing the given node, if it exists.
403 pub fn find_mut(&mut self, key: &K) -> Option<OccupiedEntry<'_, K, V>> {
404 match self.raw_entry(key) {
405 RawEntry::Occupied(entry) => Some(entry),
406 RawEntry::Vacant(_entry) => None,
407 }
408 }
409
410 /// Returns a reference to the value corresponding to the key.
411 pub fn get(&self, key: &K) -> Option<&V> {
412 let mut node = self.root.rb_node;
413 while !node.is_null() {
414 // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`
415 // point to the links field of `Node<K, V>` objects.
416 let this = unsafe { container_of!(node, Node<K, V>, links) };
417
418 // SAFETY: `this` is a non-null node so it is valid by the type invariants.
419 let this_ref = unsafe { &*this };
420
421 // SAFETY: `node` is a non-null node so it is valid by the type invariants.
422 let node_ref = unsafe { &*node };
423
424 node = match key.cmp(&this_ref.key) {
425 Ordering::Less => node_ref.rb_left,
426 Ordering::Greater => node_ref.rb_right,
427 Ordering::Equal => return Some(&this_ref.value),
428 }
429 }
430 None
431 }
432
433 /// Returns a mutable reference to the value corresponding to the key.
434 pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
435 self.find_mut(key).map(|node| node.into_mut())
436 }
437
438 /// Removes the node with the given key from the tree.
439 ///
440 /// It returns the node that was removed if one exists, or [`None`] otherwise.
441 pub fn remove_node(&mut self, key: &K) -> Option<RBTreeNode<K, V>> {
442 self.find_mut(key).map(OccupiedEntry::remove_node)
443 }
444
445 /// Removes the node with the given key from the tree.
446 ///
447 /// It returns the value that was removed if one exists, or [`None`] otherwise.
448 pub fn remove(&mut self, key: &K) -> Option<V> {
449 self.find_mut(key).map(OccupiedEntry::remove)
450 }
451
452 /// Returns a cursor over the tree nodes based on the given key.
453 ///
454 /// If the given key exists, the cursor starts there.
455 /// Otherwise it starts with the first larger key in sort order.
456 /// If there is no larger key, it returns [`None`].
457 pub fn cursor_lower_bound_mut(&mut self, key: &K) -> Option<CursorMut<'_, K, V>>
458 where
459 K: Ord,
460 {
461 let best = self.find_best_match(key)?;
462
463 NonNull::new(best.as_ptr()).map(|current| {
464 // INVARIANT:
465 // - `current` is a valid node in the [`RBTree`] pointed to by `self`.
466 CursorMut {
467 current,
468 tree: self,
469 }
470 })
471 }
472
473 /// Returns a cursor over the tree nodes based on the given key.
474 ///
475 /// If the given key exists, the cursor starts there.
476 /// Otherwise it starts with the first larger key in sort order.
477 /// If there is no larger key, it returns [`None`].
478 pub fn cursor_lower_bound(&self, key: &K) -> Option<Cursor<'_, K, V>>
479 where
480 K: Ord,
481 {
482 let best = self.find_best_match(key)?;
483
484 NonNull::new(best.as_ptr()).map(|current| {
485 // INVARIANT:
486 // - `current` is a valid node in the [`RBTree`] pointed to by `self`.
487 Cursor {
488 current,
489 _tree: PhantomData,
490 }
491 })
492 }
493
494 fn find_best_match(&self, key: &K) -> Option<NonNull<bindings::rb_node>> {
495 let mut node = self.root.rb_node;
496 let mut best_key: Option<&K> = None;
497 let mut best_links: Option<NonNull<bindings::rb_node>> = None;
498 while !node.is_null() {
499 // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`
500 // point to the links field of `Node<K, V>` objects.
501 let this = unsafe { container_of!(node, Node<K, V>, links) };
502 // SAFETY: `this` is a non-null node so it is valid by the type invariants.
503 let this_key = unsafe { &(*this).key };
504
505 // SAFETY: `node` is a non-null node so it is valid by the type invariants.
506 let node_ref = unsafe { &*node };
507
508 match key.cmp(this_key) {
509 Ordering::Equal => {
510 // SAFETY: `this` is a non-null node so it is valid by the type invariants.
511 best_links = Some(unsafe { NonNull::new_unchecked(&mut (*this).links) });
512 break;
513 }
514 Ordering::Greater => {
515 node = node_ref.rb_right;
516 }
517 Ordering::Less => {
518 let is_better_match = match best_key {
519 None => true,
520 Some(best) => best > this_key,
521 };
522 if is_better_match {
523 best_key = Some(this_key);
524 // SAFETY: `this` is a non-null node so it is valid by the type invariants.
525 best_links = Some(unsafe { NonNull::new_unchecked(&mut (*this).links) });
526 }
527 node = node_ref.rb_left;
528 }
529 };
530 }
531 best_links
532 }
533}
534
535impl<K, V> Default for RBTree<K, V> {
536 fn default() -> Self {
537 Self::new()
538 }
539}
540
541impl<K, V> Drop for RBTree<K, V> {
542 fn drop(&mut self) {
543 // SAFETY: `root` is valid as it's embedded in `self` and we have a valid `self`.
544 let mut next = unsafe { bindings::rb_first_postorder(&self.root) };
545
546 // INVARIANT: The loop invariant is that all tree nodes from `next` in postorder are valid.
547 while !next.is_null() {
548 // SAFETY: All links fields we create are in a `Node<K, V>`.
549 let this = unsafe { container_of!(next, Node<K, V>, links) };
550
551 // Find out what the next node is before disposing of the current one.
552 // SAFETY: `next` and all nodes in postorder are still valid.
553 next = unsafe { bindings::rb_next_postorder(next) };
554
555 // INVARIANT: This is the destructor, so we break the type invariant during clean-up,
556 // but it is not observable. The loop invariant is still maintained.
557
558 // SAFETY: `this` is valid per the loop invariant.
559 unsafe { drop(KBox::from_raw(this)) };
560 }
561 }
562}
563
564/// A bidirectional mutable cursor over the tree nodes, sorted by key.
565///
566/// # Examples
567///
568/// In the following example, we obtain a cursor to the first element in the tree.
569/// The cursor allows us to iterate bidirectionally over key/value pairs in the tree.
570///
571/// ```
572/// use kernel::{alloc::flags, rbtree::RBTree};
573///
574/// // Create a new tree.
575/// let mut tree = RBTree::new();
576///
577/// // Insert three elements.
578/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
579/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
580/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
581///
582/// // Get a cursor to the first element.
583/// let mut cursor = tree.cursor_front_mut().unwrap();
584/// let mut current = cursor.current();
585/// assert_eq!(current, (&10, &100));
586///
587/// // Move the cursor, updating it to the 2nd element.
588/// cursor = cursor.move_next().unwrap();
589/// current = cursor.current();
590/// assert_eq!(current, (&20, &200));
591///
592/// // Peek at the next element without impacting the cursor.
593/// let next = cursor.peek_next().unwrap();
594/// assert_eq!(next, (&30, &300));
595/// current = cursor.current();
596/// assert_eq!(current, (&20, &200));
597///
598/// // Moving past the last element causes the cursor to return [`None`].
599/// cursor = cursor.move_next().unwrap();
600/// current = cursor.current();
601/// assert_eq!(current, (&30, &300));
602/// let cursor = cursor.move_next();
603/// assert!(cursor.is_none());
604///
605/// # Ok::<(), Error>(())
606/// ```
607///
608/// A cursor can also be obtained at the last element in the tree.
609///
610/// ```
611/// use kernel::{alloc::flags, rbtree::RBTree};
612///
613/// // Create a new tree.
614/// let mut tree = RBTree::new();
615///
616/// // Insert three elements.
617/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
618/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
619/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
620///
621/// let mut cursor = tree.cursor_back_mut().unwrap();
622/// let current = cursor.current();
623/// assert_eq!(current, (&30, &300));
624///
625/// # Ok::<(), Error>(())
626/// ```
627///
628/// Obtaining a cursor returns [`None`] if the tree is empty.
629///
630/// ```
631/// use kernel::rbtree::RBTree;
632///
633/// let mut tree: RBTree<u16, u16> = RBTree::new();
634/// assert!(tree.cursor_front_mut().is_none());
635///
636/// # Ok::<(), Error>(())
637/// ```
638///
639/// [`RBTree::cursor_lower_bound`] can be used to start at an arbitrary node in the tree.
640///
641/// ```
642/// use kernel::{alloc::flags, rbtree::RBTree};
643///
644/// // Create a new tree.
645/// let mut tree = RBTree::new();
646///
647/// // Insert five elements.
648/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
649/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
650/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
651/// tree.try_create_and_insert(40, 400, flags::GFP_KERNEL)?;
652/// tree.try_create_and_insert(50, 500, flags::GFP_KERNEL)?;
653///
654/// // If the provided key exists, a cursor to that key is returned.
655/// let cursor = tree.cursor_lower_bound(&20).unwrap();
656/// let current = cursor.current();
657/// assert_eq!(current, (&20, &200));
658///
659/// // If the provided key doesn't exist, a cursor to the first larger element in sort order is returned.
660/// let cursor = tree.cursor_lower_bound(&25).unwrap();
661/// let current = cursor.current();
662/// assert_eq!(current, (&30, &300));
663///
664/// // If there is no larger key, [`None`] is returned.
665/// let cursor = tree.cursor_lower_bound(&55);
666/// assert!(cursor.is_none());
667///
668/// # Ok::<(), Error>(())
669/// ```
670///
671/// The cursor allows mutation of values in the tree.
672///
673/// ```
674/// use kernel::{alloc::flags, rbtree::RBTree};
675///
676/// // Create a new tree.
677/// let mut tree = RBTree::new();
678///
679/// // Insert three elements.
680/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
681/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
682/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
683///
684/// // Retrieve a cursor.
685/// let mut cursor = tree.cursor_front_mut().unwrap();
686///
687/// // Get a mutable reference to the current value.
688/// let (k, v) = cursor.current_mut();
689/// *v = 1000;
690///
691/// // The updated value is reflected in the tree.
692/// let updated = tree.get(&10).unwrap();
693/// assert_eq!(updated, &1000);
694///
695/// # Ok::<(), Error>(())
696/// ```
697///
698/// It also allows node removal. The following examples demonstrate the behavior of removing the current node.
699///
700/// ```
701/// use kernel::{alloc::flags, rbtree::RBTree};
702///
703/// // Create a new tree.
704/// let mut tree = RBTree::new();
705///
706/// // Insert three elements.
707/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
708/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
709/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
710///
711/// // Remove the first element.
712/// let mut cursor = tree.cursor_front_mut().unwrap();
713/// let mut current = cursor.current();
714/// assert_eq!(current, (&10, &100));
715/// cursor = cursor.remove_current().0.unwrap();
716///
717/// // If a node exists after the current element, it is returned.
718/// current = cursor.current();
719/// assert_eq!(current, (&20, &200));
720///
721/// // Get a cursor to the last element, and remove it.
722/// cursor = tree.cursor_back_mut().unwrap();
723/// current = cursor.current();
724/// assert_eq!(current, (&30, &300));
725///
726/// // Since there is no next node, the previous node is returned.
727/// cursor = cursor.remove_current().0.unwrap();
728/// current = cursor.current();
729/// assert_eq!(current, (&20, &200));
730///
731/// // Removing the last element in the tree returns [`None`].
732/// assert!(cursor.remove_current().0.is_none());
733///
734/// # Ok::<(), Error>(())
735/// ```
736///
737/// Nodes adjacent to the current node can also be removed.
738///
739/// ```
740/// use kernel::{alloc::flags, rbtree::RBTree};
741///
742/// // Create a new tree.
743/// let mut tree = RBTree::new();
744///
745/// // Insert three elements.
746/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
747/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
748/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
749///
750/// // Get a cursor to the first element.
751/// let mut cursor = tree.cursor_front_mut().unwrap();
752/// let mut current = cursor.current();
753/// assert_eq!(current, (&10, &100));
754///
755/// // Calling `remove_prev` from the first element returns [`None`].
756/// assert!(cursor.remove_prev().is_none());
757///
758/// // Get a cursor to the last element.
759/// cursor = tree.cursor_back_mut().unwrap();
760/// current = cursor.current();
761/// assert_eq!(current, (&30, &300));
762///
763/// // Calling `remove_prev` removes and returns the middle element.
764/// assert_eq!(cursor.remove_prev().unwrap().to_key_value(), (20, 200));
765///
766/// // Calling `remove_next` from the last element returns [`None`].
767/// assert!(cursor.remove_next().is_none());
768///
769/// // Move to the first element
770/// cursor = cursor.move_prev().unwrap();
771/// current = cursor.current();
772/// assert_eq!(current, (&10, &100));
773///
774/// // Calling `remove_next` removes and returns the last element.
775/// assert_eq!(cursor.remove_next().unwrap().to_key_value(), (30, 300));
776///
777/// # Ok::<(), Error>(())
778///
779/// ```
780///
781/// # Invariants
782/// - `current` points to a node that is in the same [`RBTree`] as `tree`.
783pub struct CursorMut<'a, K, V> {
784 tree: &'a mut RBTree<K, V>,
785 current: NonNull<bindings::rb_node>,
786}
787
788/// A bidirectional immutable cursor over the tree nodes, sorted by key. This is a simpler
789/// variant of [`CursorMut`] that is basically providing read only access.
790///
791/// # Examples
792///
793/// In the following example, we obtain a cursor to the first element in the tree.
794/// The cursor allows us to iterate bidirectionally over key/value pairs in the tree.
795///
796/// ```
797/// use kernel::{alloc::flags, rbtree::RBTree};
798///
799/// // Create a new tree.
800/// let mut tree = RBTree::new();
801///
802/// // Insert three elements.
803/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
804/// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
805/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
806///
807/// // Get a cursor to the first element.
808/// let cursor = tree.cursor_front().unwrap();
809/// let current = cursor.current();
810/// assert_eq!(current, (&10, &100));
811///
812/// # Ok::<(), Error>(())
813/// ```
814pub struct Cursor<'a, K, V> {
815 _tree: PhantomData<&'a RBTree<K, V>>,
816 current: NonNull<bindings::rb_node>,
817}
818
819// SAFETY: The immutable cursor gives out shared access to `K` and `V` so if `K` and `V` can be
820// shared across threads, then it's safe to share the cursor.
821unsafe impl<'a, K: Sync, V: Sync> Send for Cursor<'a, K, V> {}
822
823// SAFETY: The immutable cursor gives out shared access to `K` and `V` so if `K` and `V` can be
824// shared across threads, then it's safe to share the cursor.
825unsafe impl<'a, K: Sync, V: Sync> Sync for Cursor<'a, K, V> {}
826
827impl<'a, K, V> Cursor<'a, K, V> {
828 /// The current node
829 pub fn current(&self) -> (&K, &V) {
830 // SAFETY:
831 // - `self.current` is a valid node by the type invariants.
832 // - We have an immutable reference by the function signature.
833 unsafe { Self::to_key_value(self.current) }
834 }
835
836 /// # Safety
837 ///
838 /// - `node` must be a valid pointer to a node in an [`RBTree`].
839 /// - The caller has immutable access to `node` for the duration of `'b`.
840 unsafe fn to_key_value<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, &'b V) {
841 // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`
842 // point to the links field of `Node<K, V>` objects.
843 let this = unsafe { container_of!(node.as_ptr(), Node<K, V>, links) };
844 // SAFETY: The passed `node` is the current node or a non-null neighbor,
845 // thus `this` is valid by the type invariants.
846 let k = unsafe { &(*this).key };
847 // SAFETY: The passed `node` is the current node or a non-null neighbor,
848 // thus `this` is valid by the type invariants.
849 let v = unsafe { &(*this).value };
850 (k, v)
851 }
852
853 /// Access the previous node without moving the cursor.
854 pub fn peek_prev(&self) -> Option<(&K, &V)> {
855 self.peek(Direction::Prev)
856 }
857
858 /// Access the next node without moving the cursor.
859 pub fn peek_next(&self) -> Option<(&K, &V)> {
860 self.peek(Direction::Next)
861 }
862
863 fn peek(&self, direction: Direction) -> Option<(&K, &V)> {
864 self.get_neighbor_raw(direction).map(|neighbor| {
865 // SAFETY:
866 // - `neighbor` is a valid tree node.
867 // - By the function signature, we have an immutable reference to `self`.
868 unsafe { Self::to_key_value(neighbor) }
869 })
870 }
871
872 fn get_neighbor_raw(&self, direction: Direction) -> Option<NonNull<bindings::rb_node>> {
873 // SAFETY: `self.current` is valid by the type invariants.
874 let neighbor = unsafe {
875 match direction {
876 Direction::Prev => bindings::rb_prev(self.current.as_ptr()),
877 Direction::Next => bindings::rb_next(self.current.as_ptr()),
878 }
879 };
880
881 NonNull::new(neighbor)
882 }
883}
884
885// SAFETY: The [`CursorMut`] has exclusive access to both `K` and `V`, so it is sufficient to
886// require them to be `Send`.
887// The cursor only gives out immutable references to the keys, but since it has exclusive access to
888// those same keys, `Send` is sufficient. `Sync` would be okay, but it is more restrictive to the
889// user.
890unsafe impl<'a, K: Send, V: Send> Send for CursorMut<'a, K, V> {}
891
892// SAFETY: The [`CursorMut`] gives out immutable references to `K` and mutable references to `V`,
893// so it has the same thread safety requirements as mutable references.
894unsafe impl<'a, K: Sync, V: Sync> Sync for CursorMut<'a, K, V> {}
895
896impl<'a, K, V> CursorMut<'a, K, V> {
897 /// The current node.
898 pub fn current(&self) -> (&K, &V) {
899 // SAFETY:
900 // - `self.current` is a valid node by the type invariants.
901 // - We have an immutable reference by the function signature.
902 unsafe { Self::to_key_value(self.current) }
903 }
904
905 /// The current node, with a mutable value
906 pub fn current_mut(&mut self) -> (&K, &mut V) {
907 // SAFETY:
908 // - `self.current` is a valid node by the type invariants.
909 // - We have an mutable reference by the function signature.
910 unsafe { Self::to_key_value_mut(self.current) }
911 }
912
913 /// Remove the current node from the tree.
914 ///
915 /// Returns a tuple where the first element is a cursor to the next node, if it exists,
916 /// else the previous node, else [`None`] (if the tree becomes empty). The second element
917 /// is the removed node.
918 pub fn remove_current(self) -> (Option<Self>, RBTreeNode<K, V>) {
919 let prev = self.get_neighbor_raw(Direction::Prev);
920 let next = self.get_neighbor_raw(Direction::Next);
921 // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`
922 // point to the links field of `Node<K, V>` objects.
923 let this = unsafe { container_of!(self.current.as_ptr(), Node<K, V>, links) };
924 // SAFETY: `this` is valid by the type invariants as described above.
925 let node = unsafe { KBox::from_raw(this) };
926 let node = RBTreeNode { node };
927 // SAFETY: The reference to the tree used to create the cursor outlives the cursor, so
928 // the tree cannot change. By the tree invariant, all nodes are valid.
929 unsafe { bindings::rb_erase(&mut (*this).links, addr_of_mut!(self.tree.root)) };
930
931 // INVARIANT:
932 // - `current` is a valid node in the [`RBTree`] pointed to by `self.tree`.
933 let cursor = next.or(prev).map(|current| Self {
934 current,
935 tree: self.tree,
936 });
937
938 (cursor, node)
939 }
940
941 /// Remove the previous node, returning it if it exists.
942 pub fn remove_prev(&mut self) -> Option<RBTreeNode<K, V>> {
943 self.remove_neighbor(Direction::Prev)
944 }
945
946 /// Remove the next node, returning it if it exists.
947 pub fn remove_next(&mut self) -> Option<RBTreeNode<K, V>> {
948 self.remove_neighbor(Direction::Next)
949 }
950
951 fn remove_neighbor(&mut self, direction: Direction) -> Option<RBTreeNode<K, V>> {
952 if let Some(neighbor) = self.get_neighbor_raw(direction) {
953 let neighbor = neighbor.as_ptr();
954 // SAFETY: The reference to the tree used to create the cursor outlives the cursor, so
955 // the tree cannot change. By the tree invariant, all nodes are valid.
956 unsafe { bindings::rb_erase(neighbor, addr_of_mut!(self.tree.root)) };
957 // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`
958 // point to the links field of `Node<K, V>` objects.
959 let this = unsafe { container_of!(neighbor, Node<K, V>, links) };
960 // SAFETY: `this` is valid by the type invariants as described above.
961 let node = unsafe { KBox::from_raw(this) };
962 return Some(RBTreeNode { node });
963 }
964 None
965 }
966
967 /// Move the cursor to the previous node, returning [`None`] if it doesn't exist.
968 pub fn move_prev(self) -> Option<Self> {
969 self.mv(Direction::Prev)
970 }
971
972 /// Move the cursor to the next node, returning [`None`] if it doesn't exist.
973 pub fn move_next(self) -> Option<Self> {
974 self.mv(Direction::Next)
975 }
976
977 fn mv(self, direction: Direction) -> Option<Self> {
978 // INVARIANT:
979 // - `neighbor` is a valid node in the [`RBTree`] pointed to by `self.tree`.
980 self.get_neighbor_raw(direction).map(|neighbor| Self {
981 tree: self.tree,
982 current: neighbor,
983 })
984 }
985
986 /// Access the previous node without moving the cursor.
987 pub fn peek_prev(&self) -> Option<(&K, &V)> {
988 self.peek(Direction::Prev)
989 }
990
991 /// Access the next node without moving the cursor.
992 pub fn peek_next(&self) -> Option<(&K, &V)> {
993 self.peek(Direction::Next)
994 }
995
996 fn peek(&self, direction: Direction) -> Option<(&K, &V)> {
997 self.get_neighbor_raw(direction).map(|neighbor| {
998 // SAFETY:
999 // - `neighbor` is a valid tree node.
1000 // - By the function signature, we have an immutable reference to `self`.
1001 unsafe { Self::to_key_value(neighbor) }
1002 })
1003 }
1004
1005 /// Access the previous node mutably without moving the cursor.
1006 pub fn peek_prev_mut(&mut self) -> Option<(&K, &mut V)> {
1007 self.peek_mut(Direction::Prev)
1008 }
1009
1010 /// Access the next node mutably without moving the cursor.
1011 pub fn peek_next_mut(&mut self) -> Option<(&K, &mut V)> {
1012 self.peek_mut(Direction::Next)
1013 }
1014
1015 fn peek_mut(&mut self, direction: Direction) -> Option<(&K, &mut V)> {
1016 self.get_neighbor_raw(direction).map(|neighbor| {
1017 // SAFETY:
1018 // - `neighbor` is a valid tree node.
1019 // - By the function signature, we have a mutable reference to `self`.
1020 unsafe { Self::to_key_value_mut(neighbor) }
1021 })
1022 }
1023
1024 fn get_neighbor_raw(&self, direction: Direction) -> Option<NonNull<bindings::rb_node>> {
1025 // SAFETY: `self.current` is valid by the type invariants.
1026 let neighbor = unsafe {
1027 match direction {
1028 Direction::Prev => bindings::rb_prev(self.current.as_ptr()),
1029 Direction::Next => bindings::rb_next(self.current.as_ptr()),
1030 }
1031 };
1032
1033 NonNull::new(neighbor)
1034 }
1035
1036 /// # Safety
1037 ///
1038 /// - `node` must be a valid pointer to a node in an [`RBTree`].
1039 /// - The caller has immutable access to `node` for the duration of `'b`.
1040 unsafe fn to_key_value<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, &'b V) {
1041 // SAFETY: the caller guarantees that `node` is a valid pointer in an `RBTree`.
1042 let (k, v) = unsafe { Self::to_key_value_raw(node) };
1043 // SAFETY: the caller guarantees immutable access to `node`.
1044 (k, unsafe { &*v })
1045 }
1046
1047 /// # Safety
1048 ///
1049 /// - `node` must be a valid pointer to a node in an [`RBTree`].
1050 /// - The caller has mutable access to `node` for the duration of `'b`.
1051 unsafe fn to_key_value_mut<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, &'b mut V) {
1052 // SAFETY: the caller guarantees that `node` is a valid pointer in an `RBTree`.
1053 let (k, v) = unsafe { Self::to_key_value_raw(node) };
1054 // SAFETY: the caller guarantees mutable access to `node`.
1055 (k, unsafe { &mut *v })
1056 }
1057
1058 /// # Safety
1059 ///
1060 /// - `node` must be a valid pointer to a node in an [`RBTree`].
1061 /// - The caller has immutable access to the key for the duration of `'b`.
1062 unsafe fn to_key_value_raw<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, *mut V) {
1063 // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`
1064 // point to the links field of `Node<K, V>` objects.
1065 let this = unsafe { container_of!(node.as_ptr(), Node<K, V>, links) };
1066 // SAFETY: The passed `node` is the current node or a non-null neighbor,
1067 // thus `this` is valid by the type invariants.
1068 let k = unsafe { &(*this).key };
1069 // SAFETY: The passed `node` is the current node or a non-null neighbor,
1070 // thus `this` is valid by the type invariants.
1071 let v = unsafe { addr_of_mut!((*this).value) };
1072 (k, v)
1073 }
1074}
1075
1076/// Direction for [`Cursor`] and [`CursorMut`] operations.
1077enum Direction {
1078 /// the node immediately before, in sort order
1079 Prev,
1080 /// the node immediately after, in sort order
1081 Next,
1082}
1083
1084impl<'a, K, V> IntoIterator for &'a RBTree<K, V> {
1085 type Item = (&'a K, &'a V);
1086 type IntoIter = Iter<'a, K, V>;
1087
1088 fn into_iter(self) -> Self::IntoIter {
1089 self.iter()
1090 }
1091}
1092
1093/// An iterator over the nodes of a [`RBTree`].
1094///
1095/// Instances are created by calling [`RBTree::iter`].
1096pub struct Iter<'a, K, V> {
1097 _tree: PhantomData<&'a RBTree<K, V>>,
1098 iter_raw: IterRaw<K, V>,
1099}
1100
1101// SAFETY: The [`Iter`] gives out immutable references to K and V, so it has the same
1102// thread safety requirements as immutable references.
1103unsafe impl<'a, K: Sync, V: Sync> Send for Iter<'a, K, V> {}
1104
1105// SAFETY: The [`Iter`] gives out immutable references to K and V, so it has the same
1106// thread safety requirements as immutable references.
1107unsafe impl<'a, K: Sync, V: Sync> Sync for Iter<'a, K, V> {}
1108
1109impl<'a, K, V> Iterator for Iter<'a, K, V> {
1110 type Item = (&'a K, &'a V);
1111
1112 fn next(&mut self) -> Option<Self::Item> {
1113 // SAFETY: Due to `self._tree`, `k` and `v` are valid for the lifetime of `'a`.
1114 self.iter_raw.next().map(|(k, v)| unsafe { (&*k, &*v) })
1115 }
1116}
1117
1118impl<'a, K, V> IntoIterator for &'a mut RBTree<K, V> {
1119 type Item = (&'a K, &'a mut V);
1120 type IntoIter = IterMut<'a, K, V>;
1121
1122 fn into_iter(self) -> Self::IntoIter {
1123 self.iter_mut()
1124 }
1125}
1126
1127/// A mutable iterator over the nodes of a [`RBTree`].
1128///
1129/// Instances are created by calling [`RBTree::iter_mut`].
1130pub struct IterMut<'a, K, V> {
1131 _tree: PhantomData<&'a mut RBTree<K, V>>,
1132 iter_raw: IterRaw<K, V>,
1133}
1134
1135// SAFETY: The [`IterMut`] has exclusive access to both `K` and `V`, so it is sufficient to require them to be `Send`.
1136// The iterator only gives out immutable references to the keys, but since the iterator has exclusive access to those same
1137// keys, `Send` is sufficient. `Sync` would be okay, but it is more restrictive to the user.
1138unsafe impl<'a, K: Send, V: Send> Send for IterMut<'a, K, V> {}
1139
1140// SAFETY: The [`IterMut`] gives out immutable references to K and mutable references to V, so it has the same
1141// thread safety requirements as mutable references.
1142unsafe impl<'a, K: Sync, V: Sync> Sync for IterMut<'a, K, V> {}
1143
1144impl<'a, K, V> Iterator for IterMut<'a, K, V> {
1145 type Item = (&'a K, &'a mut V);
1146
1147 fn next(&mut self) -> Option<Self::Item> {
1148 self.iter_raw.next().map(|(k, v)|
1149 // SAFETY: Due to `&mut self`, we have exclusive access to `k` and `v`, for the lifetime of `'a`.
1150 unsafe { (&*k, &mut *v) })
1151 }
1152}
1153
1154/// A raw iterator over the nodes of a [`RBTree`].
1155///
1156/// # Invariants
1157/// - `self.next` is a valid pointer.
1158/// - `self.next` points to a node stored inside of a valid `RBTree`.
1159struct IterRaw<K, V> {
1160 next: *mut bindings::rb_node,
1161 _phantom: PhantomData<fn() -> (K, V)>,
1162}
1163
1164impl<K, V> Iterator for IterRaw<K, V> {
1165 type Item = (*mut K, *mut V);
1166
1167 fn next(&mut self) -> Option<Self::Item> {
1168 if self.next.is_null() {
1169 return None;
1170 }
1171
1172 // SAFETY: By the type invariant of `IterRaw`, `self.next` is a valid node in an `RBTree`,
1173 // and by the type invariant of `RBTree`, all nodes point to the links field of `Node<K, V>` objects.
1174 let cur = unsafe { container_of!(self.next, Node<K, V>, links) };
1175
1176 // SAFETY: `self.next` is a valid tree node by the type invariants.
1177 self.next = unsafe { bindings::rb_next(self.next) };
1178
1179 // SAFETY: By the same reasoning above, it is safe to dereference the node.
1180 Some(unsafe { (addr_of_mut!((*cur).key), addr_of_mut!((*cur).value)) })
1181 }
1182}
1183
1184/// A memory reservation for a red-black tree node.
1185///
1186///
1187/// It contains the memory needed to hold a node that can be inserted into a red-black tree. One
1188/// can be obtained by directly allocating it ([`RBTreeNodeReservation::new`]).
1189pub struct RBTreeNodeReservation<K, V> {
1190 node: KBox<MaybeUninit<Node<K, V>>>,
1191}
1192
1193impl<K, V> RBTreeNodeReservation<K, V> {
1194 /// Allocates memory for a node to be eventually initialised and inserted into the tree via a
1195 /// call to [`RBTree::insert`].
1196 pub fn new(flags: Flags) -> Result<RBTreeNodeReservation<K, V>> {
1197 Ok(RBTreeNodeReservation {
1198 node: KBox::new_uninit(flags)?,
1199 })
1200 }
1201}
1202
1203// SAFETY: This doesn't actually contain K or V, and is just a memory allocation. Those can always
1204// be moved across threads.
1205unsafe impl<K, V> Send for RBTreeNodeReservation<K, V> {}
1206
1207// SAFETY: This doesn't actually contain K or V, and is just a memory allocation.
1208unsafe impl<K, V> Sync for RBTreeNodeReservation<K, V> {}
1209
1210impl<K, V> RBTreeNodeReservation<K, V> {
1211 /// Initialises a node reservation.
1212 ///
1213 /// It then becomes an [`RBTreeNode`] that can be inserted into a tree.
1214 pub fn into_node(self, key: K, value: V) -> RBTreeNode<K, V> {
1215 let node = KBox::write(
1216 self.node,
1217 Node {
1218 key,
1219 value,
1220 links: bindings::rb_node::default(),
1221 },
1222 );
1223 RBTreeNode { node }
1224 }
1225}
1226
1227/// A red-black tree node.
1228///
1229/// The node is fully initialised (with key and value) and can be inserted into a tree without any
1230/// extra allocations or failure paths.
1231pub struct RBTreeNode<K, V> {
1232 node: KBox<Node<K, V>>,
1233}
1234
1235impl<K, V> RBTreeNode<K, V> {
1236 /// Allocates and initialises a node that can be inserted into the tree via
1237 /// [`RBTree::insert`].
1238 pub fn new(key: K, value: V, flags: Flags) -> Result<RBTreeNode<K, V>> {
1239 Ok(RBTreeNodeReservation::new(flags)?.into_node(key, value))
1240 }
1241
1242 /// Get the key and value from inside the node.
1243 pub fn to_key_value(self) -> (K, V) {
1244 let node = KBox::into_inner(self.node);
1245
1246 (node.key, node.value)
1247 }
1248}
1249
1250// SAFETY: If K and V can be sent across threads, then it's also okay to send [`RBTreeNode`] across
1251// threads.
1252unsafe impl<K: Send, V: Send> Send for RBTreeNode<K, V> {}
1253
1254// SAFETY: If K and V can be accessed without synchronization, then it's also okay to access
1255// [`RBTreeNode`] without synchronization.
1256unsafe impl<K: Sync, V: Sync> Sync for RBTreeNode<K, V> {}
1257
1258impl<K, V> RBTreeNode<K, V> {
1259 /// Drop the key and value, but keep the allocation.
1260 ///
1261 /// It then becomes a reservation that can be re-initialised into a different node (i.e., with
1262 /// a different key and/or value).
1263 ///
1264 /// The existing key and value are dropped in-place as part of this operation, that is, memory
1265 /// may be freed (but only for the key/value; memory for the node itself is kept for reuse).
1266 pub fn into_reservation(self) -> RBTreeNodeReservation<K, V> {
1267 RBTreeNodeReservation {
1268 node: KBox::drop_contents(self.node),
1269 }
1270 }
1271}
1272
1273/// A view into a single entry in a map, which may either be vacant or occupied.
1274///
1275/// This enum is constructed from the [`RBTree::entry`].
1276///
1277/// [`entry`]: fn@RBTree::entry
1278pub enum Entry<'a, K, V> {
1279 /// This [`RBTree`] does not have a node with this key.
1280 Vacant(VacantEntry<'a, K, V>),
1281 /// This [`RBTree`] already has a node with this key.
1282 Occupied(OccupiedEntry<'a, K, V>),
1283}
1284
1285/// Like [`Entry`], except that it doesn't have ownership of the key.
1286enum RawEntry<'a, K, V> {
1287 Vacant(RawVacantEntry<'a, K, V>),
1288 Occupied(OccupiedEntry<'a, K, V>),
1289}
1290
1291/// A view into a vacant entry in a [`RBTree`]. It is part of the [`Entry`] enum.
1292pub struct VacantEntry<'a, K, V> {
1293 key: K,
1294 raw: RawVacantEntry<'a, K, V>,
1295}
1296
1297/// Like [`VacantEntry`], but doesn't hold on to the key.
1298///
1299/// # Invariants
1300/// - `parent` may be null if the new node becomes the root.
1301/// - `child_field_of_parent` is a valid pointer to the left-child or right-child of `parent`. If `parent` is
1302/// null, it is a pointer to the root of the [`RBTree`].
1303struct RawVacantEntry<'a, K, V> {
1304 rbtree: *mut RBTree<K, V>,
1305 /// The node that will become the parent of the new node if we insert one.
1306 parent: *mut bindings::rb_node,
1307 /// This points to the left-child or right-child field of `parent`, or `root` if `parent` is
1308 /// null.
1309 child_field_of_parent: *mut *mut bindings::rb_node,
1310 _phantom: PhantomData<&'a mut RBTree<K, V>>,
1311}
1312
1313impl<'a, K, V> RawVacantEntry<'a, K, V> {
1314 /// Inserts the given node into the [`RBTree`] at this entry.
1315 ///
1316 /// The `node` must have a key such that inserting it here does not break the ordering of this
1317 /// [`RBTree`].
1318 fn insert(self, node: RBTreeNode<K, V>) -> &'a mut V {
1319 let node = KBox::into_raw(node.node);
1320
1321 // SAFETY: `node` is valid at least until we call `KBox::from_raw`, which only happens when
1322 // the node is removed or replaced.
1323 let node_links = unsafe { addr_of_mut!((*node).links) };
1324
1325 // INVARIANT: We are linking in a new node, which is valid. It remains valid because we
1326 // "forgot" it with `KBox::into_raw`.
1327 // SAFETY: The type invariants of `RawVacantEntry` are exactly the safety requirements of `rb_link_node`.
1328 unsafe { bindings::rb_link_node(node_links, self.parent, self.child_field_of_parent) };
1329
1330 // SAFETY: All pointers are valid. `node` has just been inserted into the tree.
1331 unsafe { bindings::rb_insert_color(node_links, addr_of_mut!((*self.rbtree).root)) };
1332
1333 // SAFETY: The node is valid until we remove it from the tree.
1334 unsafe { &mut (*node).value }
1335 }
1336}
1337
1338impl<'a, K, V> VacantEntry<'a, K, V> {
1339 /// Inserts the given node into the [`RBTree`] at this entry.
1340 pub fn insert(self, value: V, reservation: RBTreeNodeReservation<K, V>) -> &'a mut V {
1341 self.raw.insert(reservation.into_node(self.key, value))
1342 }
1343}
1344
1345/// A view into an occupied entry in a [`RBTree`]. It is part of the [`Entry`] enum.
1346///
1347/// # Invariants
1348/// - `node_links` is a valid, non-null pointer to a tree node in `self.rbtree`
1349pub struct OccupiedEntry<'a, K, V> {
1350 rbtree: &'a mut RBTree<K, V>,
1351 /// The node that this entry corresponds to.
1352 node_links: *mut bindings::rb_node,
1353}
1354
1355impl<'a, K, V> OccupiedEntry<'a, K, V> {
1356 /// Gets a reference to the value in the entry.
1357 pub fn get(&self) -> &V {
1358 // SAFETY:
1359 // - `self.node_links` is a valid pointer to a node in the tree.
1360 // - We have shared access to the underlying tree, and can thus give out a shared reference.
1361 unsafe { &(*container_of!(self.node_links, Node<K, V>, links)).value }
1362 }
1363
1364 /// Gets a mutable reference to the value in the entry.
1365 pub fn get_mut(&mut self) -> &mut V {
1366 // SAFETY:
1367 // - `self.node_links` is a valid pointer to a node in the tree.
1368 // - We have exclusive access to the underlying tree, and can thus give out a mutable reference.
1369 unsafe { &mut (*(container_of!(self.node_links, Node<K, V>, links))).value }
1370 }
1371
1372 /// Converts the entry into a mutable reference to its value.
1373 ///
1374 /// If you need multiple references to the `OccupiedEntry`, see [`self#get_mut`].
1375 pub fn into_mut(self) -> &'a mut V {
1376 // SAFETY:
1377 // - `self.node_links` is a valid pointer to a node in the tree.
1378 // - This consumes the `&'a mut RBTree<K, V>`, therefore it can give out a mutable reference that lives for `'a`.
1379 unsafe { &mut (*(container_of!(self.node_links, Node<K, V>, links))).value }
1380 }
1381
1382 /// Remove this entry from the [`RBTree`].
1383 pub fn remove_node(self) -> RBTreeNode<K, V> {
1384 // SAFETY: The node is a node in the tree, so it is valid.
1385 unsafe { bindings::rb_erase(self.node_links, &mut self.rbtree.root) };
1386
1387 // INVARIANT: The node is being returned and the caller may free it, however, it was
1388 // removed from the tree. So the invariants still hold.
1389 RBTreeNode {
1390 // SAFETY: The node was a node in the tree, but we removed it, so we can convert it
1391 // back into a box.
1392 node: unsafe { KBox::from_raw(container_of!(self.node_links, Node<K, V>, links)) },
1393 }
1394 }
1395
1396 /// Takes the value of the entry out of the map, and returns it.
1397 pub fn remove(self) -> V {
1398 let rb_node = self.remove_node();
1399 let node = KBox::into_inner(rb_node.node);
1400
1401 node.value
1402 }
1403
1404 /// Swap the current node for the provided node.
1405 ///
1406 /// The key of both nodes must be equal.
1407 fn replace(self, node: RBTreeNode<K, V>) -> RBTreeNode<K, V> {
1408 let node = KBox::into_raw(node.node);
1409
1410 // SAFETY: `node` is valid at least until we call `KBox::from_raw`, which only happens when
1411 // the node is removed or replaced.
1412 let new_node_links = unsafe { addr_of_mut!((*node).links) };
1413
1414 // SAFETY: This updates the pointers so that `new_node_links` is in the tree where
1415 // `self.node_links` used to be.
1416 unsafe {
1417 bindings::rb_replace_node(self.node_links, new_node_links, &mut self.rbtree.root)
1418 };
1419
1420 // SAFETY:
1421 // - `self.node_ptr` produces a valid pointer to a node in the tree.
1422 // - Now that we removed this entry from the tree, we can convert the node to a box.
1423 let old_node = unsafe { KBox::from_raw(container_of!(self.node_links, Node<K, V>, links)) };
1424
1425 RBTreeNode { node: old_node }
1426 }
1427}
1428
1429struct Node<K, V> {
1430 links: bindings::rb_node,
1431 key: K,
1432 value: V,
1433}