1 // SPDX-License-Identifier: GPL-2.0
2
3 //! Red-black trees.
4 //!
5 //! C header: [`include/linux/rbtree.h`](srctr
6 //!
7 //! Reference: <https://docs.kernel.org/core-a
8
9 use crate::{alloc::Flags, bindings, container_
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 tre
21 ///
22 /// # Examples
23 ///
24 /// In the example below we do several operati
25 /// the system is out of memory.
26 ///
27 /// ```
28 /// use kernel::{alloc::flags, rbtree::{RBTree
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:
35 /// tree.try_create_and_insert(10, 100, flags:
36 /// tree.try_create_and_insert(30, 300, flags:
37 ///
38 /// // Check the nodes we just inserted.
39 /// {
40 /// assert_eq!(tree.get(&10).unwrap(), &10
41 /// assert_eq!(tree.get(&20).unwrap(), &20
42 /// assert_eq!(tree.get(&30).unwrap(), &30
43 /// }
44 ///
45 /// // Iterate over the nodes we just inserted
46 /// {
47 /// let mut iter = tree.iter();
48 /// assert_eq!(iter.next().unwrap(), (&10,
49 /// assert_eq!(iter.next().unwrap(), (&20,
50 /// assert_eq!(iter.next().unwrap(), (&30,
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
61 ///
62 /// // Check that the tree reflects the replac
63 /// {
64 /// let mut iter = tree.iter();
65 /// assert_eq!(iter.next().unwrap(), (&10,
66 /// assert_eq!(iter.next().unwrap(), (&20,
67 /// assert_eq!(iter.next().unwrap(), (&30,
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,
78 /// assert_eq!(iter.next().unwrap(), (&20,
79 /// assert_eq!(iter.next().unwrap(), (&30,
80 /// assert!(iter.next().is_none());
81 /// }
82 ///
83 /// // Remove an element.
84 /// tree.remove(&10);
85 ///
86 /// // Check that the tree reflects the remova
87 /// {
88 /// let mut iter = tree.iter();
89 /// assert_eq!(iter.next().unwrap(), (&20,
90 /// assert_eq!(iter.next().unwrap(), (&30,
91 /// assert!(iter.next().is_none());
92 /// }
93 ///
94 /// # Ok::<(), Error>(())
95 /// ```
96 ///
97 /// In the example below, we first allocate a
98 /// the tree. This is useful when the insertio
99 /// holding a spinlock.
100 ///
101 /// ```
102 /// use kernel::{alloc::flags, rbtree::{RBTree
103 ///
104 /// fn insert_test(tree: &SpinLock<RBTree<u32,
105 /// // Pre-allocate node. This may fail (a
106 /// let node = RBTreeNode::new(10, 100, fl
107 ///
108 /// // Insert node while holding the lock.
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
117 ///
118 /// ```
119 /// use kernel::{alloc::flags, rbtree::{RBTree
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:
126 /// tree.try_create_and_insert(10, 100, flags:
127 /// tree.try_create_and_insert(30, 300, flags:
128 ///
129 /// // Check the nodes we just inserted.
130 /// {
131 /// let mut iter = tree.iter();
132 /// assert_eq!(iter.next().unwrap(), (&10,
133 /// assert_eq!(iter.next().unwrap(), (&20,
134 /// assert_eq!(iter.next().unwrap(), (&30,
135 /// assert!(iter.next().is_none());
136 /// }
137 ///
138 /// // Remove a node, getting back ownership o
139 /// let existing = tree.remove(&30).unwrap();
140 ///
141 /// // Check that the tree reflects the remova
142 /// {
143 /// let mut iter = tree.iter();
144 /// assert_eq!(iter.next().unwrap(), (&10,
145 /// assert_eq!(iter.next().unwrap(), (&20,
146 /// assert!(iter.next().is_none());
147 /// }
148 ///
149 /// // Create a preallocated reservation that
150 /// let reservation = RBTreeNodeReservation::n
151 ///
152 /// // Insert a new node into the tree, reusin
153 /// // succeed (no memory allocations).
154 /// tree.insert(reservation.into_node(15, 150)
155 ///
156 /// // Check that the tree reflect the new ins
157 /// {
158 /// let mut iter = tree.iter();
159 /// assert_eq!(iter.next().unwrap(), (&10,
160 /// assert_eq!(iter.next().unwrap(), (&15,
161 /// assert_eq!(iter.next().unwrap(), (&20,
162 /// assert!(iter.next().is_none());
163 /// }
164 ///
165 /// # Ok::<(), Error>(())
166 /// ```
167 ///
168 /// # Invariants
169 ///
170 /// Non-null parent/children pointers stored i
171 /// valid, and pointing to a field of our inte
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
178 // fields, so we use the same Send condition a
179 unsafe impl<K: Send, V: Send> Send for RBTree<
180
181 // SAFETY: An [`RBTree`] allows the same kinds
182 // fields, so we use the same Sync condition a
183 unsafe impl<K: Sync, V: Sync> Sync for RBTree<
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 i
190 root: bindings::rb_root::default()
191 _p: PhantomData,
192 }
193 }
194
195 /// Returns an iterator over the tree node
196 pub fn iter(&self) -> Iter<'_, K, V> {
197 Iter {
198 _tree: PhantomData,
199 // INVARIANT:
200 // - `self.root` is a valid poin
201 // - `bindings::rb_first` produc
202 iter_raw: IterRaw {
203 // SAFETY: by the invariants,
204 next: unsafe { bindings::rb_fi
205 _phantom: PhantomData,
206 },
207 }
208 }
209
210 /// Returns a mutable iterator over the tr
211 pub fn iter_mut(&mut self) -> IterMut<'_,
212 IterMut {
213 _tree: PhantomData,
214 // INVARIANT:
215 // - `self.root` is a valid poin
216 // - `bindings::rb_first` produc
217 iter_raw: IterRaw {
218 // SAFETY: by the invariants,
219 next: unsafe { bindings::rb_fi
220 _phantom: PhantomData,
221 },
222 }
223 }
224
225 /// Returns an iterator over the keys of t
226 pub fn keys(&self) -> impl Iterator<Item =
227 self.iter().map(|(k, _)| k)
228 }
229
230 /// Returns an iterator over the values of
231 pub fn values(&self) -> impl Iterator<Item
232 self.iter().map(|(_, v)| v)
233 }
234
235 /// Returns a mutable iterator over the va
236 pub fn values_mut(&mut self) -> impl Itera
237 self.iter_mut().map(|(_, v)| v)
238 }
239
240 /// Returns a cursor over the tree nodes,
241 pub fn cursor_front(&mut self) -> Option<C
242 let root = addr_of_mut!(self.root);
243 // SAFETY: `self.root` is always a val
244 let current = unsafe { bindings::rb_fi
245 NonNull::new(current).map(|current| {
246 // INVARIANT:
247 // - `current` is a valid node in
248 Cursor {
249 current,
250 tree: self,
251 }
252 })
253 }
254
255 /// Returns a cursor over the tree nodes,
256 pub fn cursor_back(&mut self) -> Option<Cu
257 let root = addr_of_mut!(self.root);
258 // SAFETY: `self.root` is always a val
259 let current = unsafe { bindings::rb_la
260 NonNull::new(current).map(|current| {
261 // INVARIANT:
262 // - `current` is a valid node in
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 t
276 ///
277 /// It overwrites a node if one already ex
278 /// key/value pair). Returns [`None`] if a
279 ///
280 /// Returns an error if it cannot allocate
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, va
288 }
289
290 /// Inserts a new node into the tree.
291 ///
292 /// It overwrites a node if one already ex
293 /// key/value pair). Returns [`None`] if a
294 ///
295 /// This function always succeeds.
296 pub fn insert(&mut self, node: RBTreeNode<
297 match self.raw_entry(&node.node.key) {
298 RawEntry::Occupied(entry) => Some(
299 RawEntry::Vacant(entry) => {
300 entry.insert(node);
301 None
302 }
303 }
304 }
305
306 fn raw_entry(&mut self, key: &K) -> RawEnt
307 let raw_self: *mut RBTree<K, V> = self
308 // The returned `RawEntry` is used to
309 // The parameters of `bindings::rb_lin
310 // - `node`: A pointer to an uninitial
311 // - `parent`: A pointer to an existin
312 // null, and `node` will beco
313 // with a pointer to `node`.
314 // - `rb_link`: A pointer to either th
315 // specifies which child of `
316 // value of `*rb_link` must b
317 // red/black tree is empty, t
318 // this case, `rb_link` is a
319 //
320 // We will traverse the tree looking f
321 // representing an empty subtree where
322 // that we preserve the ordering of th
323 // we store `parent` and `child_field_
324 // in the subtree of `parent` that `ch
325 // we find an empty subtree, we can in
326 let mut parent = core::ptr::null_mut()
327 let mut child_field_of_parent: &mut *m
328 // SAFETY: `raw_self` is a valid p
329 unsafe { &mut (*raw_self).root.rb_
330 while !(*child_field_of_parent).is_nul
331 let curr = *child_field_of_parent;
332 // SAFETY: All links fields we cre
333 let node = unsafe { container_of!(
334
335 // SAFETY: `node` is a non-null no
336 match key.cmp(unsafe { &(*node).ke
337 // SAFETY: `curr` is a non-nul
338 Ordering::Less => child_field_
339 // SAFETY: `curr` is a non-nul
340 Ordering::Greater => child_fie
341 Ordering::Equal => {
342 return RawEntry::Occupied(
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 ent
360 pub fn entry(&mut self, key: K) -> Entry<'
361 match self.raw_entry(&key) {
362 RawEntry::Occupied(entry) => Entry
363 RawEntry::Vacant(entry) => Entry::
364 }
365 }
366
367 /// Used for accessing the given node, if
368 pub fn find_mut(&mut self, key: &K) -> Opt
369 match self.raw_entry(key) {
370 RawEntry::Occupied(entry) => Some(
371 RawEntry::Vacant(_entry) => None,
372 }
373 }
374
375 /// Returns a reference to the value corre
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 o
380 // point to the links field of `No
381 let this = unsafe { container_of!(
382 // SAFETY: `this` is a non-null no
383 node = match key.cmp(unsafe { &(*t
384 // SAFETY: `node` is a non-nul
385 Ordering::Less => unsafe { (*n
386 // SAFETY: `node` is a non-nul
387 Ordering::Greater => unsafe {
388 // SAFETY: `node` is a non-nul
389 Ordering::Equal => return Some
390 }
391 }
392 None
393 }
394
395 /// Returns a mutable reference to the val
396 pub fn get_mut(&mut self, key: &K) -> Opti
397 self.find_mut(key).map(|node| node.int
398 }
399
400 /// Removes the node with the given key fr
401 ///
402 /// It returns the node that was removed i
403 pub fn remove_node(&mut self, key: &K) ->
404 self.find_mut(key).map(OccupiedEntry::
405 }
406
407 /// Removes the node with the given key fr
408 ///
409 /// It returns the value that was removed
410 pub fn remove(&mut self, key: &K) -> Optio
411 self.find_mut(key).map(OccupiedEntry::
412 }
413
414 /// Returns a cursor over the tree nodes b
415 ///
416 /// If the given key exists, the cursor st
417 /// Otherwise it starts with the first lar
418 /// If there is no larger key, it returns
419 pub fn cursor_lower_bound(&mut self, key:
420 where
421 K: Ord,
422 {
423 let mut node = self.root.rb_node;
424 let mut best_match: Option<NonNull<Nod
425 while !node.is_null() {
426 // SAFETY: By the type invariant o
427 // point to the links field of `No
428 let this = unsafe { container_of!(
429 // SAFETY: `this` is a non-null no
430 let this_key = unsafe { &(*this).k
431 // SAFETY: `node` is a non-null no
432 let left_child = unsafe { (*node).
433 // SAFETY: `node` is a non-null no
434 let right_child = unsafe { (*node)
435 match key.cmp(this_key) {
436 Ordering::Equal => {
437 best_match = NonNull::new(
438 break;
439 }
440 Ordering::Greater => {
441 node = right_child;
442 }
443 Ordering::Less => {
444 let is_better_match = matc
445 None => true,
446 Some(best) => {
447 // SAFETY: `best`
448 let best_key = uns
449 best_key > this_ke
450 }
451 };
452 if is_better_match {
453 best_match = NonNull::
454 }
455 node = left_child;
456 }
457 };
458 }
459
460 let best = best_match?;
461
462 // SAFETY: `best` is a non-null node s
463 let links = unsafe { addr_of_mut!((*be
464
465 NonNull::new(links).map(|current| {
466 // INVARIANT:
467 // - `current` is a valid node in
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 emb
485 let mut next = unsafe { bindings::rb_f
486
487 // INVARIANT: The loop invariant is th
488 while !next.is_null() {
489 // SAFETY: All links fields we cre
490 let this = unsafe { container_of!(
491
492 // Find out what the next node is
493 // SAFETY: `next` and all nodes in
494 next = unsafe { bindings::rb_next_
495
496 // INVARIANT: This is the destruct
497 // but it is not observable. The l
498
499 // SAFETY: `this` is valid per the
500 unsafe { drop(Box::from_raw(this.c
501 }
502 }
503 }
504
505 /// A bidirectional cursor over the tree nodes
506 ///
507 /// # Examples
508 ///
509 /// In the following example, we obtain a curs
510 /// The cursor allows us to iterate bidirectio
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:
520 /// tree.try_create_and_insert(20, 200, flags:
521 /// tree.try_create_and_insert(30, 300, flags:
522 ///
523 /// // Get a cursor to the first element.
524 /// let mut cursor = tree.cursor_front().unwra
525 /// let mut current = cursor.current();
526 /// assert_eq!(current, (&10, &100));
527 ///
528 /// // Move the cursor, updating it to the 2nd
529 /// cursor = cursor.move_next().unwrap();
530 /// current = cursor.current();
531 /// assert_eq!(current, (&20, &200));
532 ///
533 /// // Peek at the next element without impact
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
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
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:
559 /// tree.try_create_and_insert(20, 200, flags:
560 /// tree.try_create_and_insert(30, 300, flags:
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
570 ///
571 /// ```
572 /// use kernel::rbtree::RBTree;
573 ///
574 /// let mut tree: RBTree<u16, u16> = RBTree::n
575 /// assert!(tree.cursor_front().is_none());
576 ///
577 /// # Ok::<(), Error>(())
578 /// ```
579 ///
580 /// [`RBTree::cursor_lower_bound`] can be used
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:
590 /// tree.try_create_and_insert(20, 200, flags:
591 /// tree.try_create_and_insert(30, 300, flags:
592 /// tree.try_create_and_insert(40, 400, flags:
593 /// tree.try_create_and_insert(50, 500, flags:
594 ///
595 /// // If the provided key exists, a cursor to
596 /// let cursor = tree.cursor_lower_bound(&20).
597 /// let current = cursor.current();
598 /// assert_eq!(current, (&20, &200));
599 ///
600 /// // If the provided key doesn't exist, a cu
601 /// let cursor = tree.cursor_lower_bound(&25).
602 /// let current = cursor.current();
603 /// assert_eq!(current, (&30, &300));
604 ///
605 /// // If there is no larger key, [`None`] is
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 th
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:
622 /// tree.try_create_and_insert(20, 200, flags:
623 /// tree.try_create_and_insert(30, 300, flags:
624 ///
625 /// // Retrieve a cursor.
626 /// let mut cursor = tree.cursor_front().unwra
627 ///
628 /// // Get a mutable reference to the current
629 /// let (k, v) = cursor.current_mut();
630 /// *v = 1000;
631 ///
632 /// // The updated value is reflected in the t
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
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:
649 /// tree.try_create_and_insert(20, 200, flags:
650 /// tree.try_create_and_insert(30, 300, flags:
651 ///
652 /// // Remove the first element.
653 /// let mut cursor = tree.cursor_front().unwra
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 elem
659 /// current = cursor.current();
660 /// assert_eq!(current, (&20, &200));
661 ///
662 /// // Get a cursor to the last element, and r
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 previo
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 r
673 /// assert!(cursor.remove_current().0.is_none(
674 ///
675 /// # Ok::<(), Error>(())
676 /// ```
677 ///
678 /// Nodes adjacent to the current node can als
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:
688 /// tree.try_create_and_insert(20, 200, flags:
689 /// tree.try_create_and_insert(30, 300, flags:
690 ///
691 /// // Get a cursor to the first element.
692 /// let mut cursor = tree.cursor_front().unwra
693 /// let mut current = cursor.current();
694 /// assert_eq!(current, (&10, &100));
695 ///
696 /// // Calling `remove_prev` from the first el
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 retur
705 /// assert_eq!(cursor.remove_prev().unwrap().t
706 ///
707 /// // Calling `remove_next` from the last ele
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 retur
716 /// assert_eq!(cursor.remove_next().unwrap().t
717 ///
718 /// # Ok::<(), Error>(())
719 ///
720 /// ```
721 ///
722 /// # Invariants
723 /// - `current` points to a node that is in th
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
730 // The cursor only gives out immutable referen
731 // keys, `Send` is sufficient. `Sync` would be
732 unsafe impl<'a, K: Send, V: Send> Send for Cur
733
734 // SAFETY: The [`Cursor`] gives out immutable
735 // so it has the same thread safety requiremen
736 unsafe impl<'a, K: Sync, V: Sync> Sync for Cur
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
743 // - We have an immutable reference by
744 unsafe { Self::to_key_value(self.curre
745 }
746
747 /// The current node, with a mutable value
748 pub fn current_mut(&mut self) -> (&K, &mut
749 // SAFETY:
750 // - `self.current` is a valid node by
751 // - We have an mutable reference by t
752 unsafe { Self::to_key_value_mut(self.c
753 }
754
755 /// Remove the current node from the tree.
756 ///
757 /// Returns a tuple where the first elemen
758 /// else the previous node, else [`None`]
759 /// is the removed node.
760 pub fn remove_current(self) -> (Option<Sel
761 let prev = self.get_neighbor_raw(Direc
762 let next = self.get_neighbor_raw(Direc
763 // SAFETY: By the type invariant of `S
764 // point to the links field of `Node<K
765 let this = unsafe { container_of!(self
766 // SAFETY: `this` is valid by the type
767 let node = unsafe { Box::from_raw(this
768 let node = RBTreeNode { node };
769 // SAFETY: The reference to the tree u
770 // the tree cannot change. By the tree
771 unsafe { bindings::rb_erase(&mut (*thi
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
784 Some(Self {
785 current,
786 tree: self.tree,
787 }),
788 node,
789 )
790 }
791
792 /// Remove the previous node, returning it
793 pub fn remove_prev(&mut self) -> Option<RB
794 self.remove_neighbor(Direction::Prev)
795 }
796
797 /// Remove the next node, returning it if
798 pub fn remove_next(&mut self) -> Option<RB
799 self.remove_neighbor(Direction::Next)
800 }
801
802 fn remove_neighbor(&mut self, direction: D
803 if let Some(neighbor) = self.get_neigh
804 let neighbor = neighbor.as_ptr();
805 // SAFETY: The reference to the tr
806 // the tree cannot change. By the
807 unsafe { bindings::rb_erase(neighb
808 // SAFETY: By the type invariant o
809 // point to the links field of `No
810 let this = unsafe { container_of!(
811 // SAFETY: `this` is valid by the
812 let node = unsafe { Box::from_raw(
813 return Some(RBTreeNode { node });
814 }
815 None
816 }
817
818 /// Move the cursor to the previous node,
819 pub fn move_prev(self) -> Option<Self> {
820 self.mv(Direction::Prev)
821 }
822
823 /// Move the cursor to the next node, retu
824 pub fn move_next(self) -> Option<Self> {
825 self.mv(Direction::Next)
826 }
827
828 fn mv(self, direction: Direction) -> Optio
829 // INVARIANT:
830 // - `neighbor` is a valid node in the
831 self.get_neighbor_raw(direction).map(|
832 tree: self.tree,
833 current: neighbor,
834 })
835 }
836
837 /// Access the previous node without movin
838 pub fn peek_prev(&self) -> Option<(&K, &V)
839 self.peek(Direction::Prev)
840 }
841
842 /// Access the previous node without movin
843 pub fn peek_next(&self) -> Option<(&K, &V)
844 self.peek(Direction::Next)
845 }
846
847 fn peek(&self, direction: Direction) -> Op
848 self.get_neighbor_raw(direction).map(|
849 // SAFETY:
850 // - `neighbor` is a valid tree no
851 // - By the function signature, we
852 unsafe { Self::to_key_value(neighb
853 })
854 }
855
856 /// Access the previous node mutably witho
857 pub fn peek_prev_mut(&mut self) -> Option<
858 self.peek_mut(Direction::Prev)
859 }
860
861 /// Access the next node mutably without m
862 pub fn peek_next_mut(&mut self) -> Option<
863 self.peek_mut(Direction::Next)
864 }
865
866 fn peek_mut(&mut self, direction: Directio
867 self.get_neighbor_raw(direction).map(|
868 // SAFETY:
869 // - `neighbor` is a valid tree no
870 // - By the function signature, we
871 unsafe { Self::to_key_value_mut(ne
872 })
873 }
874
875 fn get_neighbor_raw(&self, direction: Dire
876 // SAFETY: `self.current` is valid by
877 let neighbor = unsafe {
878 match direction {
879 Direction::Prev => bindings::r
880 Direction::Next => bindings::r
881 }
882 };
883
884 NonNull::new(neighbor)
885 }
886
887 /// SAFETY:
888 /// - `node` must be a valid pointer to a
889 /// - The caller has immutable access to `
890 unsafe fn to_key_value<'b>(node: NonNull<b
891 // SAFETY: the caller guarantees that
892 let (k, v) = unsafe { Self::to_key_val
893 // SAFETY: the caller guarantees immut
894 (k, unsafe { &*v })
895 }
896
897 /// SAFETY:
898 /// - `node` must be a valid pointer to a
899 /// - The caller has mutable access to `no
900 unsafe fn to_key_value_mut<'b>(node: NonNu
901 // SAFETY: the caller guarantees that
902 let (k, v) = unsafe { Self::to_key_val
903 // SAFETY: the caller guarantees mutab
904 (k, unsafe { &mut *v })
905 }
906
907 /// SAFETY:
908 /// - `node` must be a valid pointer to a
909 /// - The caller has immutable access to t
910 unsafe fn to_key_value_raw<'b>(node: NonNu
911 // SAFETY: By the type invariant of `S
912 // point to the links field of `Node<K
913 let this = unsafe { container_of!(node
914 // SAFETY: The passed `node` is the cu
915 // thus `this` is valid by the type in
916 let k = unsafe { &(*this).key };
917 // SAFETY: The passed `node` is the cu
918 // thus `this` is valid by the type in
919 let v = unsafe { addr_of_mut!((*this).
920 (k, v)
921 }
922 }
923
924 /// Direction for [`Cursor`] operations.
925 enum Direction {
926 /// the node immediately before, in sort o
927 Prev,
928 /// the node immediately after, in sort or
929 Next,
930 }
931
932 impl<'a, K, V> IntoIterator for &'a RBTree<K,
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:
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 re
950 // thread safety requirements as immutable ref
951 unsafe impl<'a, K: Sync, V: Sync> Send for Ite
952
953 // SAFETY: The [`Iter`] gives out immutable re
954 // thread safety requirements as immutable ref
955 unsafe impl<'a, K: Sync, V: Sync> Sync for Ite
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` an
962 self.iter_raw.next().map(|(k, v)| unsa
963 }
964 }
965
966 impl<'a, K, V> IntoIterator for &'a mut RBTree
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 [`R
976 ///
977 /// Instances are created by calling [`RBTree:
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 acces
984 // The iterator only gives out immutable refer
985 // keys, `Send` is sufficient. `Sync` would be
986 unsafe impl<'a, K: Send, V: Send> Send for Ite
987
988 // SAFETY: The [`IterMut`] gives out immutable
989 // thread safety requirements as mutable refer
990 unsafe impl<'a, K: Sync, V: Sync> Sync for Ite
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
998 unsafe { (&*k, &mut *v) })
999 }
1000 }
1001
1002 /// A raw iterator over the nodes of a [`RBTr
1003 ///
1004 /// # Invariants
1005 /// - `self.next` is a valid pointer.
1006 /// - `self.next` points to a node stored ins
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 `
1021 // and by the type invariant of `RBTr
1022 let cur = unsafe { container_of!(self
1023
1024 // SAFETY: `self.next` is a valid tre
1025 self.next = unsafe { bindings::rb_nex
1026
1027 // SAFETY: By the same reasoning abov
1028 Some(unsafe { (addr_of_mut!((*cur).ke
1029 }
1030 }
1031
1032 /// A memory reservation for a red-black tree
1033 ///
1034 ///
1035 /// It contains the memory needed to hold a n
1036 /// can be obtained by directly allocating it
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 eve
1043 /// call to [`RBTree::insert`].
1044 pub fn new(flags: Flags) -> Result<RBTree
1045 Ok(RBTreeNodeReservation {
1046 node: <Box<_> as BoxExt<_>>::new_
1047 })
1048 }
1049 }
1050
1051 // SAFETY: This doesn't actually contain K or
1052 // be moved across threads.
1053 unsafe impl<K, V> Send for RBTreeNodeReservat
1054
1055 // SAFETY: This doesn't actually contain K or
1056 unsafe impl<K, V> Sync for RBTreeNodeReservat
1057
1058 impl<K, V> RBTreeNodeReservation<K, V> {
1059 /// Initialises a node reservation.
1060 ///
1061 /// It then becomes an [`RBTreeNode`] tha
1062 pub fn into_node(mut self, key: K, value:
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_
1070 RBTreeNode { node }
1071 }
1072 }
1073
1074 /// A red-black tree node.
1075 ///
1076 /// The node is fully initialised (with key a
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
1084 /// [`RBTree::insert`].
1085 pub fn new(key: K, value: V, flags: Flags
1086 Ok(RBTreeNodeReservation::new(flags)?
1087 }
1088
1089 /// Get the key and value from inside the
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 thre
1096 // threads.
1097 unsafe impl<K: Send, V: Send> Send for RBTree
1098
1099 // SAFETY: If K and V can be accessed without
1100 // [`RBTreeNode`] without synchronization.
1101 unsafe impl<K: Sync, V: Sync> Sync for RBTree
1102
1103 impl<K, V> RBTreeNode<K, V> {
1104 /// Drop the key and value, but keep the
1105 ///
1106 /// It then becomes a reservation that ca
1107 /// a different key and/or value).
1108 ///
1109 /// The existing key and value are droppe
1110 /// may be freed (but only for the key/va
1111 pub fn into_reservation(self) -> RBTreeNo
1112 RBTreeNodeReservation {
1113 node: Box::drop_contents(self.nod
1114 }
1115 }
1116 }
1117
1118 /// A view into a single entry in a map, whic
1119 ///
1120 /// This enum is constructed from the [`RBTre
1121 ///
1122 /// [`entry`]: fn@RBTree::entry
1123 pub enum Entry<'a, K, V> {
1124 /// This [`RBTree`] does not have a node
1125 Vacant(VacantEntry<'a, K, V>),
1126 /// This [`RBTree`] already has a node wi
1127 Occupied(OccupiedEntry<'a, K, V>),
1128 }
1129
1130 /// Like [`Entry`], except that it doesn't ha
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`
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
1143 ///
1144 /// # Invariants
1145 /// - `parent` may be null if the new node be
1146 /// - `child_field_of_parent` is a valid poin
1147 /// null, it is a pointer to the root of
1148 struct RawVacantEntry<'a, K, V> {
1149 rbtree: *mut RBTree<K, V>,
1150 /// The node that will become the parent
1151 parent: *mut bindings::rb_node,
1152 /// This points to the left-child or righ
1153 /// null.
1154 child_field_of_parent: *mut *mut bindings
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 [`RBT
1160 ///
1161 /// The `node` must have a key such that
1162 /// [`RBTree`].
1163 fn insert(self, node: RBTreeNode<K, V>) -
1164 let node = Box::into_raw(node.node);
1165
1166 // SAFETY: `node` is valid at least u
1167 // the node is removed or replaced.
1168 let node_links = unsafe { addr_of_mut
1169
1170 // INVARIANT: We are linking in a new
1171 // "forgot" it with `Box::into_raw`.
1172 // SAFETY: The type invariants of `Ra
1173 unsafe { bindings::rb_link_node(node_
1174
1175 // SAFETY: All pointers are valid. `n
1176 unsafe { bindings::rb_insert_color(no
1177
1178 // SAFETY: The node is valid until we
1179 unsafe { &mut (*node).value }
1180 }
1181 }
1182
1183 impl<'a, K, V> VacantEntry<'a, K, V> {
1184 /// Inserts the given node into the [`RBT
1185 pub fn insert(self, value: V, reservation
1186 self.raw.insert(reservation.into_node
1187 }
1188 }
1189
1190 /// A view into an occupied entry in a [`RBTr
1191 ///
1192 /// # Invariants
1193 /// - `node_links` is a valid, non-null point
1194 pub struct OccupiedEntry<'a, K, V> {
1195 rbtree: &'a mut RBTree<K, V>,
1196 /// The node that this entry corresponds
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
1202 pub fn get(&self) -> &V {
1203 // SAFETY:
1204 // - `self.node_links` is a valid poi
1205 // - We have shared access to the und
1206 unsafe { &(*container_of!(self.node_l
1207 }
1208
1209 /// Gets a mutable reference to the value
1210 pub fn get_mut(&mut self) -> &mut V {
1211 // SAFETY:
1212 // - `self.node_links` is a valid poi
1213 // - We have exclusive access to the
1214 unsafe { &mut (*(container_of!(self.n
1215 }
1216
1217 /// Converts the entry into a mutable ref
1218 ///
1219 /// If you need multiple references to th
1220 pub fn into_mut(self) -> &'a mut V {
1221 // SAFETY:
1222 // - `self.node_links` is a valid poi
1223 // - This consumes the `&'a mut RBTre
1224 unsafe { &mut (*(container_of!(self.n
1225 }
1226
1227 /// Remove this entry from the [`RBTree`]
1228 pub fn remove_node(self) -> RBTreeNode<K,
1229 // SAFETY: The node is a node in the
1230 unsafe { bindings::rb_erase(self.node
1231
1232 // INVARIANT: The node is being retur
1233 // removed from the tree. So the inva
1234 RBTreeNode {
1235 // SAFETY: The node was a node in
1236 // back into a box.
1237 node: unsafe {
1238 Box::from_raw(container_of!(s
1239 },
1240 }
1241 }
1242
1243 /// Takes the value of the entry out of t
1244 pub fn remove(self) -> V {
1245 self.remove_node().node.value
1246 }
1247
1248 /// Swap the current node for the provide
1249 ///
1250 /// The key of both nodes must be equal.
1251 fn replace(self, node: RBTreeNode<K, V>)
1252 let node = Box::into_raw(node.node);
1253
1254 // SAFETY: `node` is valid at least u
1255 // the node is removed or replaced.
1256 let new_node_links = unsafe { addr_of
1257
1258 // SAFETY: This updates the pointers
1259 // `self.node_links` used to be.
1260 unsafe {
1261 bindings::rb_replace_node(self.no
1262 };
1263
1264 // SAFETY:
1265 // - `self.node_ptr` produces a valid
1266 // - Now that we removed this entry f
1267 let old_node =
1268 unsafe { Box::from_raw(container_
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.