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