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