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 }
Linux® is a registered trademark of Linus Torvalds in the United States and other countries.
TOMOYO® is a registered trademark of NTT DATA CORPORATION.