// SPDX-License-Identifier: GPL-2.0 //! Red-black trees.//!//! C header: [`include/linux/rbtree.h`](srctree/include/linux/rbtree.h)//!//! Reference: <https://docs.kernel.org/core-api/rbtree.html> use crate::{alloc::Flags, bindings, container_of, error::Result, prelude::*};use core::{ cmp::{Ord, Ordering}, marker::PhantomData, mem::MaybeUninit, ptr::{addr_of_mut, from_mut, NonNull},}; /// A red-black tree with owned nodes.////// It is backed by the kernel C red-black trees.////// # Examples////// In the example below we do several operations on a tree. We note that insertions may fail if/// the system is out of memory.////// ```/// use kernel::{alloc::flags, rbtree::{RBTree, RBTreeNode, RBTreeNodeReservation}};////// // Create a new tree./// let mut tree = RBTree::new();////// // Insert three elements./// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;////// // Check the nodes we just inserted./// {/// assert_eq!(tree.get(&10).unwrap(), &100);/// assert_eq!(tree.get(&20).unwrap(), &200);/// assert_eq!(tree.get(&30).unwrap(), &300);/// }////// // Iterate over the nodes we just inserted./// {/// let mut iter = tree.iter();/// assert_eq!(iter.next().unwrap(), (&10, &100));/// assert_eq!(iter.next().unwrap(), (&20, &200));/// assert_eq!(iter.next().unwrap(), (&30, &300));/// assert!(iter.next().is_none());/// }////// // Print all elements./// for (key, value) in &tree {/// pr_info!("{} = {}\n", key, value);/// }////// // Replace one of the elements./// tree.try_create_and_insert(10, 1000, flags::GFP_KERNEL)?;////// // Check that the tree reflects the replacement./// {/// let mut iter = tree.iter();/// assert_eq!(iter.next().unwrap(), (&10, &1000));/// assert_eq!(iter.next().unwrap(), (&20, &200));/// assert_eq!(iter.next().unwrap(), (&30, &300));/// assert!(iter.next().is_none());/// }////// // Change the value of one of the elements./// *tree.get_mut(&30).unwrap() = 3000;////// // Check that the tree reflects the update./// {/// let mut iter = tree.iter();/// assert_eq!(iter.next().unwrap(), (&10, &1000));/// assert_eq!(iter.next().unwrap(), (&20, &200));/// assert_eq!(iter.next().unwrap(), (&30, &3000));/// assert!(iter.next().is_none());/// }////// // Remove an element./// tree.remove(&10);////// // Check that the tree reflects the removal./// {/// let mut iter = tree.iter();/// assert_eq!(iter.next().unwrap(), (&20, &200));/// assert_eq!(iter.next().unwrap(), (&30, &3000));/// assert!(iter.next().is_none());/// }////// # Ok::<(), Error>(())/// ```////// In the example below, we first allocate a node, acquire a spinlock, then insert the node into/// the tree. This is useful when the insertion context does not allow sleeping, for example, when/// holding a spinlock.////// ```/// use kernel::{alloc::flags, rbtree::{RBTree, RBTreeNode}, sync::SpinLock};////// fn insert_test(tree: &SpinLock<RBTree<u32, u32>>) -> Result {/// // Pre-allocate node. This may fail (as it allocates memory)./// let node = RBTreeNode::new(10, 100, flags::GFP_KERNEL)?;////// // Insert node while holding the lock. It is guaranteed to succeed with no allocation/// // attempts./// let mut guard = tree.lock();/// guard.insert(node);/// Ok(())/// }/// ```////// In the example below, we reuse an existing node allocation from an element we removed.////// ```/// use kernel::{alloc::flags, rbtree::{RBTree, RBTreeNodeReservation}};////// // Create a new tree./// let mut tree = RBTree::new();////// // Insert three elements./// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;/// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;/// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;////// // Check the nodes we just inserted./// {/// let mut iter = tree.iter();/// assert_eq!(iter.next().unwrap(), (&10, &100));/// assert_eq!(iter.next().unwrap(), (&20, &200));/// assert_eq!(iter.next().unwrap(), (&30, &300));/// assert!(iter.next().is_none());/// }////// // Remove a node, getting back ownership of it./// let existing = tree.remove(&30).unwrap();////// // Check that the tree reflects the removal./// {/// let mut iter = tree.iter();/// assert_eq!(iter.next().unwrap(), (&10, &100));/// assert_eq!(iter.next().unwrap(), (&20, &200));/// assert!(iter.next().is_none());/// }////// // Create a preallocated reservation that we can re-use later./// let reservation = RBTreeNodeReservation::new(flags::GFP_KERNEL)?;////// // Insert a new node into the tree, reusing the previous allocation. This is guaranteed to/// // succeed (no memory allocations)./// tree.insert(reservation.into_node(15, 150));////// // Check that the tree reflect the new insertion./// {/// let mut iter = tree.iter();/// assert_eq!(iter.next().unwrap(), (&10, &100));/// assert_eq!(iter.next().unwrap(), (&15, &150));/// assert_eq!(iter.next().unwrap(), (&20, &200));/// assert!(iter.next().is_none());/// }////// # Ok::<(), Error>(())/// ```////// # Invariants////// Non-null parent/children pointers stored in instances of the `rb_node` C struct are always/// valid, and pointing to a field of our internal representation of a node.pub struct RBTree<K, V> { root: bindings::rb_root, _p: PhantomData<Node<K, V>>,} // SAFETY: An [`RBTree`] allows the same kinds of access to its values that a struct allows to its// fields, so we use the same Send condition as would be used for a struct with K and V fields.unsafe impl<K: Send, V: Send> Send for RBTree<K, V> {} // SAFETY: An [`RBTree`] allows the same kinds of access to its values that a struct allows to its// fields, so we use the same Sync condition as would be used for a struct with K and V fields.unsafe impl<K: Sync, V: Sync> Sync for RBTree<K, V> {} impl<K, V> RBTree<K, V> { /// Creates a new and empty tree. pub fn new() -> Self { Self { // INVARIANT: There are no nodes in the tree, so the invariant holds vacuously. root: bindings::rb_root::default(), _p: PhantomData, } } /// Returns an iterator over the tree nodes, sorted by key. pub fn iter(&self) -> Iter<'_, K, V> { Iter { _tree: PhantomData, // INVARIANT: // - `self.root` is a valid pointer to a tree root. // - `bindings::rb_first` produces a valid pointer to a node given `root` is valid. iter_raw: IterRaw { // SAFETY: by the invariants, all pointers are valid. next: unsafe { bindings::rb_first(&self.root) }, _phantom: PhantomData, }, } } /// Returns a mutable iterator over the tree nodes, sorted by key. pub fn iter_mut(&mut self) -> IterMut<'_, K, V> { IterMut { _tree: PhantomData, // INVARIANT: // - `self.root` is a valid pointer to a tree root. // - `bindings::rb_first` produces a valid pointer to a node given `root` is valid. iter_raw: IterRaw { // SAFETY: by the invariants, all pointers are valid. next: unsafe { bindings::rb_first(from_mut(&mut self.root)) }, _phantom: PhantomData, }, } } /// Returns an iterator over the keys of the nodes in the tree, in sorted order. pub fn keys(&self) -> impl Iterator<Item = &'_ K> { self.iter().map(|(k, _)| k) } /// Returns an iterator over the values of the nodes in the tree, sorted by key. pub fn values(&self) -> impl Iterator<Item = &'_ V> { self.iter().map(|(_, v)| v) } /// Returns a mutable iterator over the values of the nodes in the tree, sorted by key. pub fn values_mut(&mut self) -> impl Iterator<Item = &'_ mut V> { self.iter_mut().map(|(_, v)| v) } /// Returns a cursor over the tree nodes, starting with the smallest key. pub fn cursor_front(&mut self) -> Option<Cursor<'_, K, V>> { let root = addr_of_mut!(self.root); // SAFETY: `self.root` is always a valid root node let current = unsafe { bindings::rb_first(root) }; NonNull::new(current).map(|current| { // INVARIANT: // - `current` is a valid node in the [`RBTree`] pointed to by `self`. Cursor { current, tree: self, } }) } /// Returns a cursor over the tree nodes, starting with the largest key. pub fn cursor_back(&mut self) -> Option<Cursor<'_, K, V>> { let root = addr_of_mut!(self.root); // SAFETY: `self.root` is always a valid root node let current = unsafe { bindings::rb_last(root) }; NonNull::new(current).map(|current| { // INVARIANT: // - `current` is a valid node in the [`RBTree`] pointed to by `self`. Cursor { current, tree: self, } }) }} impl<K, V> RBTree<K, V>where K: Ord,{ /// Tries to insert a new value into the tree. /// /// It overwrites a node if one already exists with the same key and returns it (containing the /// key/value pair). Returns [`None`] if a node with the same key didn't already exist. /// /// Returns an error if it cannot allocate memory for the new node. pub fn try_create_and_insert( &mut self, key: K, value: V, flags: Flags, ) -> Result<Option<RBTreeNode<K, V>>> { Ok(self.insert(RBTreeNode::new(key, value, flags)?)) } /// Inserts a new node into the tree. /// /// It overwrites a node if one already exists with the same key and returns it (containing the /// key/value pair). Returns [`None`] if a node with the same key didn't already exist. /// /// This function always succeeds. pub fn insert(&mut self, node: RBTreeNode<K, V>) -> Option<RBTreeNode<K, V>> { match self.raw_entry(&node.node.key) { RawEntry::Occupied(entry) => Some(entry.replace(node)), RawEntry::Vacant(entry) => { entry.insert(node); None } } } fn raw_entry(&mut self, key: &K) -> RawEntry<'_, K, V> { let raw_self: *mut RBTree<K, V> = self; // The returned `RawEntry` is used to call either `rb_link_node` or `rb_replace_node`. // The parameters of `bindings::rb_link_node` are as follows: // - `node`: A pointer to an uninitialized node being inserted. // - `parent`: A pointer to an existing node in the tree. One of its child pointers must be // null, and `node` will become a child of `parent` by replacing that child pointer // with a pointer to `node`. // - `rb_link`: A pointer to either the left-child or right-child field of `parent`. This // specifies which child of `parent` should hold `node` after this call. The // value of `*rb_link` must be null before the call to `rb_link_node`. If the // red/black tree is empty, then it’s also possible for `parent` to be null. In // this case, `rb_link` is a pointer to the `root` field of the red/black tree. // // We will traverse the tree looking for a node that has a null pointer as its child, // representing an empty subtree where we can insert our new node. We need to make sure // that we preserve the ordering of the nodes in the tree. In each iteration of the loop // we store `parent` and `child_field_of_parent`, and the new `node` will go somewhere // in the subtree of `parent` that `child_field_of_parent` points at. Once // we find an empty subtree, we can insert the new node using `rb_link_node`. let mut parent = core::ptr::null_mut(); let mut child_field_of_parent: &mut *mut bindings::rb_node = // SAFETY: `raw_self` is a valid pointer to the `RBTree` (created from `self` above). unsafe { &mut (*raw_self).root.rb_node }; while !(*child_field_of_parent).is_null() { let curr = *child_field_of_parent; // SAFETY: All links fields we create are in a `Node<K, V>`. let node = unsafe { container_of!(curr, Node<K, V>, links) }; // SAFETY: `node` is a non-null node so it is valid by the type invariants. match key.cmp(unsafe { &(*node).key }) { // SAFETY: `curr` is a non-null node so it is valid by the type invariants. Ordering::Less => child_field_of_parent = unsafe { &mut (*curr).rb_left }, // SAFETY: `curr` is a non-null node so it is valid by the type invariants. Ordering::Greater => child_field_of_parent = unsafe { &mut (*curr).rb_right }, Ordering::Equal => { return RawEntry::Occupied(OccupiedEntry { rbtree: self, node_links: curr, }) } } parent = curr; } RawEntry::Vacant(RawVacantEntry { rbtree: raw_self, parent, child_field_of_parent, _phantom: PhantomData, }) } /// Gets the given key's corresponding entry in the map for in-place manipulation. pub fn entry(&mut self, key: K) -> Entry<'_, K, V> { match self.raw_entry(&key) { RawEntry::Occupied(entry) => Entry::Occupied(entry), RawEntry::Vacant(entry) => Entry::Vacant(VacantEntry { raw: entry, key }), } } /// Used for accessing the given node, if it exists. pub fn find_mut(&mut self, key: &K) -> Option<OccupiedEntry<'_, K, V>> { match self.raw_entry(key) { RawEntry::Occupied(entry) => Some(entry), RawEntry::Vacant(_entry) => None, } } /// Returns a reference to the value corresponding to the key. pub fn get(&self, key: &K) -> Option<&V> { let mut node = self.root.rb_node; while !node.is_null() { // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self` // point to the links field of `Node<K, V>` objects. let this = unsafe { container_of!(node, Node<K, V>, links) }; // SAFETY: `this` is a non-null node so it is valid by the type invariants. node = match key.cmp(unsafe { &(*this).key }) { // SAFETY: `node` is a non-null node so it is valid by the type invariants. Ordering::Less => unsafe { (*node).rb_left }, // SAFETY: `node` is a non-null node so it is valid by the type invariants. Ordering::Greater => unsafe { (*node).rb_right }, // SAFETY: `node` is a non-null node so it is valid by the type invariants. Ordering::Equal => return Some(unsafe { &(*this).value }), } } None } /// Returns a mutable reference to the value corresponding to the key. pub fn get_mut(&mut self, key: &K) -> Option<&mut V> { self.find_mut(key).map(|node| node.into_mut()) } /// Removes the node with the given key from the tree. /// /// It returns the node that was removed if one exists, or [`None`] otherwise. pub fn remove_node(&mut self, key: &K) -> Option<RBTreeNode<K, V>> { self.find_mut(key).map(OccupiedEntry::remove_node) } /// Removes the node with the given key from the tree. /// /// It returns the value that was removed if one exists, or [`None`] otherwise. pub fn remove(&mut self, key: &K) -> Option<V> { self.find_mut(key).map(OccupiedEntry::remove) } /// Returns a cursor over the tree nodes based on the given key. /// /// If the given key exists, the cursor starts there. /// Otherwise it starts with the first larger key in sort order. /// If there is no larger key, it returns [`None`]. pub fn cursor_lower_bound(&mut self, key: &K) -> Option<Cursor<'_, K, V>> where K: Ord, { let mut node = self.root.rb_node; let mut best_match: Option<NonNull<Node<K, V>>> = None; while !node.is_null() { // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self` // point to the links field of `Node<K, V>` objects. let this = unsafe { container_of!(node, Node<K, V>, links) }.cast_mut(); // SAFETY: `this` is a non-null node so it is valid by the type invariants. let this_key = unsafe { &(*this).key }; // SAFETY: `node` is a non-null node so it is valid by the type invariants. let left_child = unsafe { (*node).rb_left }; // SAFETY: `node` is a non-null node so it is valid by the type invariants. let right_child = unsafe { (*node).rb_right }; match key.cmp(this_key) { Ordering::Equal => { best_match = NonNull::new(this); break; } Ordering::Greater => { node = right_child; } Ordering::Less => { let is_better_match = match best_match { None => true, Some(best) => { // SAFETY: `best` is a non-null node so it is valid by the type invariants. let best_key = unsafe { &(*best.as_ptr()).key }; best_key > this_key } }; if is_better_match { best_match = NonNull::new(this); } node = left_child; } }; } let best = best_match?; // SAFETY: `best` is a non-null node so it is valid by the type invariants. let links = unsafe { addr_of_mut!((*best.as_ptr()).links) }; NonNull::new(links).map(|current| { // INVARIANT: // - `current` is a valid node in the [`RBTree`] pointed to by `self`. Cursor { current, tree: self, } }) }} impl<K, V> Default for RBTree<K, V> { fn default() -> Self { Self::new() }} impl<K, V> Drop for RBTree<K, V> { fn drop(&mut self) { // SAFETY: `root` is valid as it's embedded in `self` and we have a valid `self`. let mut next = unsafe { bindings::rb_first_postorder(&self.root) }; // INVARIANT: The loop invariant is that all tree nodes from `next` in postorder are valid. while !next.is_null() { // SAFETY: All links fields we create are in a `Node<K, V>`. let this = unsafe { container_of!(next, Node<K, V>, links) }; // Find out what the next node is before disposing of the current one. // SAFETY: `next` and all nodes in postorder are still valid. next = unsafe { bindings::rb_next_postorder(next) }; // INVARIANT: This is the destructor, so we break the type invariant during clean-up, // but it is not observable. The loop invariant is still maintained. // SAFETY: `this` is valid per the loop invariant. unsafe { drop(KBox::from_raw(this.cast_mut())) }; }Lines 1–500 of 1286. Load more
// SPDX-License-Identifier: GPL-2.0 //! Red-black trees. //! //! C header: [`include/linux/rbtree.h`](srctree/include/linux/rbtree.h) //! //! Reference: <https://docs.kernel.org/core-api/rbtree.html> use crate::{alloc::Flags, bindings, container_of, error::Result, prelude::*}; use core::{ cmp::{Ord, Ordering}, marker::PhantomData, mem::MaybeUninit, ptr::{addr_of_mut, from_mut, NonNull}, }; /// A red-black tree with owned nodes. /// /// It is backed by the kernel C red-black trees. /// /// # Examples /// /// In the example below we do several operations on a tree. We note that insertions may fail if /// the system is out of memory. /// /// ``` /// use kernel::{alloc::flags, rbtree::{RBTree, RBTreeNode, RBTreeNodeReservation}}; /// /// // Create a new tree. /// let mut tree = RBTree::new(); /// /// // Insert three elements. /// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; /// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; /// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; /// /// // Check the nodes we just inserted. /// { /// assert_eq!(tree.get(&10).unwrap(), &100); /// assert_eq!(tree.get(&20).unwrap(), &200); /// assert_eq!(tree.get(&30).unwrap(), &300); /// } /// /// // Iterate over the nodes we just inserted. /// { /// let mut iter = tree.iter(); /// assert_eq!(iter.next().unwrap(), (&10, &100)); /// assert_eq!(iter.next().unwrap(), (&20, &200)); /// assert_eq!(iter.next().unwrap(), (&30, &300)); /// assert!(iter.next().is_none()); /// } /// /// // Print all elements. /// for (key, value) in &tree { /// pr_info!("{} = {}\n", key, value); /// } /// /// // Replace one of the elements. /// tree.try_create_and_insert(10, 1000, flags::GFP_KERNEL)?; /// /// // Check that the tree reflects the replacement. /// { /// let mut iter = tree.iter(); /// assert_eq!(iter.next().unwrap(), (&10, &1000)); /// assert_eq!(iter.next().unwrap(), (&20, &200)); /// assert_eq!(iter.next().unwrap(), (&30, &300)); /// assert!(iter.next().is_none()); /// } /// /// // Change the value of one of the elements. /// *tree.get_mut(&30).unwrap() = 3000; /// /// // Check that the tree reflects the update. /// { /// let mut iter = tree.iter(); /// assert_eq!(iter.next().unwrap(), (&10, &1000)); /// assert_eq!(iter.next().unwrap(), (&20, &200)); /// assert_eq!(iter.next().unwrap(), (&30, &3000)); /// assert!(iter.next().is_none()); /// } /// /// // Remove an element. /// tree.remove(&10); /// /// // Check that the tree reflects the removal. /// { /// let mut iter = tree.iter(); /// assert_eq!(iter.next().unwrap(), (&20, &200)); /// assert_eq!(iter.next().unwrap(), (&30, &3000)); /// assert!(iter.next().is_none()); /// } /// /// # Ok::<(), Error>(()) /// ``` /// /// In the example below, we first allocate a node, acquire a spinlock, then insert the node into /// the tree. This is useful when the insertion context does not allow sleeping, for example, when /// holding a spinlock. /// /// ``` /// use kernel::{alloc::flags, rbtree::{RBTree, RBTreeNode}, sync::SpinLock}; /// /// fn insert_test(tree: &SpinLock<RBTree<u32, u32>>) -> Result { /// // Pre-allocate node. This may fail (as it allocates memory). /// let node = RBTreeNode::new(10, 100, flags::GFP_KERNEL)?; /// /// // Insert node while holding the lock. It is guaranteed to succeed with no allocation /// // attempts. /// let mut guard = tree.lock(); /// guard.insert(node); /// Ok(()) /// } /// ``` /// /// In the example below, we reuse an existing node allocation from an element we removed. /// /// ``` /// use kernel::{alloc::flags, rbtree::{RBTree, RBTreeNodeReservation}}; /// /// // Create a new tree. /// let mut tree = RBTree::new(); /// /// // Insert three elements. /// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?; /// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?; /// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?; /// /// // Check the nodes we just inserted. /// { /// let mut iter = tree.iter(); /// assert_eq!(iter.next().unwrap(), (&10, &100)); /// assert_eq!(iter.next().unwrap(), (&20, &200)); /// assert_eq!(iter.next().unwrap(), (&30, &300)); /// assert!(iter.next().is_none()); /// } /// /// // Remove a node, getting back ownership of it. /// let existing = tree.remove(&30).unwrap(); /// /// // Check that the tree reflects the removal. /// { /// let mut iter = tree.iter(); /// assert_eq!(iter.next().unwrap(), (&10, &100)); /// assert_eq!(iter.next().unwrap(), (&20, &200)); /// assert!(iter.next().is_none()); /// } /// /// // Create a preallocated reservation that we can re-use later. /// let reservation = RBTreeNodeReservation::new(flags::GFP_KERNEL)?; /// /// // Insert a new node into the tree, reusing the previous allocation. This is guaranteed to /// // succeed (no memory allocations). /// tree.insert(reservation.into_node(15, 150)); /// /// // Check that the tree reflect the new insertion. /// { /// let mut iter = tree.iter(); /// assert_eq!(iter.next().unwrap(), (&10, &100)); /// assert_eq!(iter.next().unwrap(), (&15, &150)); /// assert_eq!(iter.next().unwrap(), (&20, &200)); /// assert!(iter.next().is_none()); /// } /// /// # Ok::<(), Error>(()) /// ``` /// /// # Invariants /// /// Non-null parent/children pointers stored in instances of the `rb_node` C struct are always /// valid, and pointing to a field of our internal representation of a node. pub struct RBTree<K, V> { root: bindings::rb_root, _p: PhantomData<Node<K, V>>, } // SAFETY: An [`RBTree`] allows the same kinds of access to its values that a struct allows to its // fields, so we use the same Send condition as would be used for a struct with K and V fields. unsafe impl<K: Send, V: Send> Send for RBTree<K, V> {} // SAFETY: An [`RBTree`] allows the same kinds of access to its values that a struct allows to its // fields, so we use the same Sync condition as would be used for a struct with K and V fields. unsafe impl<K: Sync, V: Sync> Sync for RBTree<K, V> {} impl<K, V> RBTree<K, V> { /// Creates a new and empty tree. pub fn new() -> Self { Self { // INVARIANT: There are no nodes in the tree, so the invariant holds vacuously. root: bindings::rb_root::default(), _p: PhantomData, } } /// Returns an iterator over the tree nodes, sorted by key. pub fn iter(&self) -> Iter<'_, K, V> { Iter { _tree: PhantomData, // INVARIANT: // - `self.root` is a valid pointer to a tree root. // - `bindings::rb_first` produces a valid pointer to a node given `root` is valid. iter_raw: IterRaw { // SAFETY: by the invariants, all pointers are valid. next: unsafe { bindings::rb_first(&self.root) }, _phantom: PhantomData, }, } } /// Returns a mutable iterator over the tree nodes, sorted by key. pub fn iter_mut(&mut self) -> IterMut<'_, K, V> { IterMut { _tree: PhantomData, // INVARIANT: // - `self.root` is a valid pointer to a tree root. // - `bindings::rb_first` produces a valid pointer to a node given `root` is valid. iter_raw: IterRaw { // SAFETY: by the invariants, all pointers are valid. next: unsafe { bindings::rb_first(from_mut(&mut self.root)) }, _phantom: PhantomData, }, } } /// Returns an iterator over the keys of the nodes in the tree, in sorted order. pub fn keys(&self) -> impl Iterator<Item = &'_ K> { self.iter().map(|(k, _)| k) } /// Returns an iterator over the values of the nodes in the tree, sorted by key. pub fn values(&self) -> impl Iterator<Item = &'_ V> { self.iter().map(|(_, v)| v) } /// Returns a mutable iterator over the values of the nodes in the tree, sorted by key. pub fn values_mut(&mut self) -> impl Iterator<Item = &'_ mut V> { self.iter_mut().map(|(_, v)| v) } /// Returns a cursor over the tree nodes, starting with the smallest key. pub fn cursor_front(&mut self) -> Option<Cursor<'_, K, V>> { let root = addr_of_mut!(self.root); // SAFETY: `self.root` is always a valid root node let current = unsafe { bindings::rb_first(root) }; NonNull::new(current).map(|current| { // INVARIANT: // - `current` is a valid node in the [`RBTree`] pointed to by `self`. Cursor { current, tree: self, } }) } /// Returns a cursor over the tree nodes, starting with the largest key. pub fn cursor_back(&mut self) -> Option<Cursor<'_, K, V>> { let root = addr_of_mut!(self.root); // SAFETY: `self.root` is always a valid root node let current = unsafe { bindings::rb_last(root) }; NonNull::new(current).map(|current| { // INVARIANT: // - `current` is a valid node in the [`RBTree`] pointed to by `self`. Cursor { current, tree: self, } }) } } impl<K, V> RBTree<K, V> where K: Ord, { /// Tries to insert a new value into the tree. /// /// It overwrites a node if one already exists with the same key and returns it (containing the /// key/value pair). Returns [`None`] if a node with the same key didn't already exist. /// /// Returns an error if it cannot allocate memory for the new node. pub fn try_create_and_insert( &mut self, key: K, value: V, flags: Flags, ) -> Result<Option<RBTreeNode<K, V>>> { Ok(self.insert(RBTreeNode::new(key, value, flags)?)) } /// Inserts a new node into the tree. /// /// It overwrites a node if one already exists with the same key and returns it (containing the /// key/value pair). Returns [`None`] if a node with the same key didn't already exist. /// /// This function always succeeds. pub fn insert(&mut self, node: RBTreeNode<K, V>) -> Option<RBTreeNode<K, V>> { match self.raw_entry(&node.node.key) { RawEntry::Occupied(entry) => Some(entry.replace(node)), RawEntry::Vacant(entry) => { entry.insert(node); None } } } fn raw_entry(&mut self, key: &K) -> RawEntry<'_, K, V> { let raw_self: *mut RBTree<K, V> = self; // The returned `RawEntry` is used to call either `rb_link_node` or `rb_replace_node`. // The parameters of `bindings::rb_link_node` are as follows: // - `node`: A pointer to an uninitialized node being inserted. // - `parent`: A pointer to an existing node in the tree. One of its child pointers must be // null, and `node` will become a child of `parent` by replacing that child pointer // with a pointer to `node`. // - `rb_link`: A pointer to either the left-child or right-child field of `parent`. This // specifies which child of `parent` should hold `node` after this call. The // value of `*rb_link` must be null before the call to `rb_link_node`. If the // red/black tree is empty, then it’s also possible for `parent` to be null. In // this case, `rb_link` is a pointer to the `root` field of the red/black tree. // // We will traverse the tree looking for a node that has a null pointer as its child, // representing an empty subtree where we can insert our new node. We need to make sure // that we preserve the ordering of the nodes in the tree. In each iteration of the loop // we store `parent` and `child_field_of_parent`, and the new `node` will go somewhere // in the subtree of `parent` that `child_field_of_parent` points at. Once // we find an empty subtree, we can insert the new node using `rb_link_node`. let mut parent = core::ptr::null_mut(); let mut child_field_of_parent: &mut *mut bindings::rb_node = // SAFETY: `raw_self` is a valid pointer to the `RBTree` (created from `self` above). unsafe { &mut (*raw_self).root.rb_node }; while !(*child_field_of_parent).is_null() { let curr = *child_field_of_parent; // SAFETY: All links fields we create are in a `Node<K, V>`. let node = unsafe { container_of!(curr, Node<K, V>, links) }; // SAFETY: `node` is a non-null node so it is valid by the type invariants. match key.cmp(unsafe { &(*node).key }) { // SAFETY: `curr` is a non-null node so it is valid by the type invariants. Ordering::Less => child_field_of_parent = unsafe { &mut (*curr).rb_left }, // SAFETY: `curr` is a non-null node so it is valid by the type invariants. Ordering::Greater => child_field_of_parent = unsafe { &mut (*curr).rb_right }, Ordering::Equal => { return RawEntry::Occupied(OccupiedEntry { rbtree: self, node_links: curr, }) } } parent = curr; } RawEntry::Vacant(RawVacantEntry { rbtree: raw_self, parent, child_field_of_parent, _phantom: PhantomData, }) } /// Gets the given key's corresponding entry in the map for in-place manipulation. pub fn entry(&mut self, key: K) -> Entry<'_, K, V> { match self.raw_entry(&key) { RawEntry::Occupied(entry) => Entry::Occupied(entry), RawEntry::Vacant(entry) => Entry::Vacant(VacantEntry { raw: entry, key }), } } /// Used for accessing the given node, if it exists. pub fn find_mut(&mut self, key: &K) -> Option<OccupiedEntry<'_, K, V>> { match self.raw_entry(key) { RawEntry::Occupied(entry) => Some(entry), RawEntry::Vacant(_entry) => None, } } /// Returns a reference to the value corresponding to the key. pub fn get(&self, key: &K) -> Option<&V> { let mut node = self.root.rb_node; while !node.is_null() { // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self` // point to the links field of `Node<K, V>` objects. let this = unsafe { container_of!(node, Node<K, V>, links) }; // SAFETY: `this` is a non-null node so it is valid by the type invariants. node = match key.cmp(unsafe { &(*this).key }) { // SAFETY: `node` is a non-null node so it is valid by the type invariants. Ordering::Less => unsafe { (*node).rb_left }, // SAFETY: `node` is a non-null node so it is valid by the type invariants. Ordering::Greater => unsafe { (*node).rb_right }, // SAFETY: `node` is a non-null node so it is valid by the type invariants. Ordering::Equal => return Some(unsafe { &(*this).value }), } } None } /// Returns a mutable reference to the value corresponding to the key. pub fn get_mut(&mut self, key: &K) -> Option<&mut V> { self.find_mut(key).map(|node| node.into_mut()) } /// Removes the node with the given key from the tree. /// /// It returns the node that was removed if one exists, or [`None`] otherwise. pub fn remove_node(&mut self, key: &K) -> Option<RBTreeNode<K, V>> { self.find_mut(key).map(OccupiedEntry::remove_node) } /// Removes the node with the given key from the tree. /// /// It returns the value that was removed if one exists, or [`None`] otherwise. pub fn remove(&mut self, key: &K) -> Option<V> { self.find_mut(key).map(OccupiedEntry::remove) } /// Returns a cursor over the tree nodes based on the given key. /// /// If the given key exists, the cursor starts there. /// Otherwise it starts with the first larger key in sort order. /// If there is no larger key, it returns [`None`]. pub fn cursor_lower_bound(&mut self, key: &K) -> Option<Cursor<'_, K, V>> where K: Ord, { let mut node = self.root.rb_node; let mut best_match: Option<NonNull<Node<K, V>>> = None; while !node.is_null() { // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self` // point to the links field of `Node<K, V>` objects. let this = unsafe { container_of!(node, Node<K, V>, links) }.cast_mut(); // SAFETY: `this` is a non-null node so it is valid by the type invariants. let this_key = unsafe { &(*this).key }; // SAFETY: `node` is a non-null node so it is valid by the type invariants. let left_child = unsafe { (*node).rb_left }; // SAFETY: `node` is a non-null node so it is valid by the type invariants. let right_child = unsafe { (*node).rb_right }; match key.cmp(this_key) { Ordering::Equal => { best_match = NonNull::new(this); break; } Ordering::Greater => { node = right_child; } Ordering::Less => { let is_better_match = match best_match { None => true, Some(best) => { // SAFETY: `best` is a non-null node so it is valid by the type invariants. let best_key = unsafe { &(*best.as_ptr()).key }; best_key > this_key } }; if is_better_match { best_match = NonNull::new(this); } node = left_child; } }; } let best = best_match?; // SAFETY: `best` is a non-null node so it is valid by the type invariants. let links = unsafe { addr_of_mut!((*best.as_ptr()).links) }; NonNull::new(links).map(|current| { // INVARIANT: // - `current` is a valid node in the [`RBTree`] pointed to by `self`. Cursor { current, tree: self, } }) } } impl<K, V> Default for RBTree<K, V> { fn default() -> Self { Self::new() } } impl<K, V> Drop for RBTree<K, V> { fn drop(&mut self) { // SAFETY: `root` is valid as it's embedded in `self` and we have a valid `self`. let mut next = unsafe { bindings::rb_first_postorder(&self.root) }; // INVARIANT: The loop invariant is that all tree nodes from `next` in postorder are valid. while !next.is_null() { // SAFETY: All links fields we create are in a `Node<K, V>`. let this = unsafe { container_of!(next, Node<K, V>, links) }; // Find out what the next node is before disposing of the current one. // SAFETY: `next` and all nodes in postorder are still valid. next = unsafe { bindings::rb_next_postorder(next) }; // INVARIANT: This is the destructor, so we break the type invariant during clean-up, // but it is not observable. The loop invariant is still maintained. // SAFETY: `this` is valid per the loop invariant. unsafe { drop(KBox::from_raw(this.cast_mut())) }; }