kernel/
rbtree.rs

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