@@ -244,6 +244,36 @@ pub fn values(&self) -> impl Iterator<Item = &'_ V> {
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<RBTreeCursor<'_, 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) };
+ if current.is_null() {
+ return None;
+ }
+ Some(RBTreeCursor {
+ _tree: PhantomData,
+ root,
+ current,
+ })
+ }
+
+ /// Returns a cursor over the tree nodes, starting with the largest key.
+ pub fn cursor_back(&mut self) -> Option<RBTreeCursor<'_, 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) };
+ if current.is_null() {
+ return None;
+ }
+ Some(RBTreeCursor {
+ _tree: PhantomData,
+ root,
+ current,
+ })
+ }
}
impl<K, V> RBTree<K, V>
@@ -376,6 +406,59 @@ pub fn remove(&mut self, key: &K) -> Option<V> {
} = *node;
Some(value)
}
+
+ /// 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<RBTreeCursor<'_, 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: All links fields we create are in a `Node<K, V>`.
+ let this = unsafe { crate::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 };
+ if key == this_key {
+ return Some(RBTreeCursor {
+ _tree: PhantomData,
+ root: addr_of_mut!(self.root),
+ current: node,
+ });
+ } else {
+ node = if key > this_key {
+ right_child
+ } else {
+ 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);
+ }
+ left_child
+ }
+ };
+ }
+ best_match.map(|best| RBTreeCursor {
+ _tree: PhantomData,
+ root: addr_of_mut!(self.root),
+ // SAFETY: `best` is a non-null node so it is valid by the type invariants.
+ current: unsafe { addr_of_mut!((*best.as_ptr()).links) },
+ })
+ }
}
impl<K, V> Default for RBTree<K, V> {
@@ -406,6 +489,435 @@ fn drop(&mut self) {
}
}
+/// A bidirectional cursor over the tree nodes, sorted by key.
+///
+/// # Invariants
+///
+/// In instance of `RBTreeCursor` is only acquired from [`RBTree`].
+/// A reference to the tree used to create the cursor outlives the cursor, so
+/// the tree cannot change. By the tree invariant, all nodes are valid.
+///
+/// # Examples
+///
+/// In the following example, we obtain a cursor to the first element in the tree.
+/// The cursor allows us to iterate bidirectionally over key/value pairs in the tree.
+///
+/// ```
+/// use kernel::rbtree::RBTree;
+///
+/// // Create a new tree.
+/// let mut tree = RBTree::new();
+///
+/// // Insert three elements.
+/// tree.try_create_and_insert(10, 100)?;
+/// tree.try_create_and_insert(20, 200)?;
+/// tree.try_create_and_insert(30, 300)?;
+///
+/// // Get a cursor to the first element.
+/// let mut cursor = tree.cursor_front().unwrap();
+/// let mut current = cursor.current();
+/// assert_eq!(current, (&10, &100));
+///
+/// // Move the cursor, updating it to the 2nd element.
+/// cursor = cursor.move_next().unwrap();
+/// current = cursor.current();
+/// assert_eq!(current, (&20, &200));
+///
+/// // Peek at the next element without impacting the cursor.
+/// let next = cursor.peek_next().unwrap();
+/// assert_eq!(next, (&30, &300));
+/// current = cursor.current();
+/// assert_eq!(current, (&20, &200));
+///
+/// // Moving past the last element causes the cursor to return [`None`].
+/// cursor = cursor.move_next().unwrap();
+/// current = cursor.current();
+/// assert_eq!(current, (&30, &300));
+/// let cursor = cursor.move_next();
+/// assert!(cursor.is_none());
+///
+/// # Ok::<(), Error>(())
+/// ```
+///
+/// A cursor can also be obtained at the last element in the tree.
+///
+/// ```
+/// use kernel::rbtree::RBTree;
+///
+/// // Create a new tree.
+/// let mut tree = RBTree::new();
+///
+/// // Insert three elements.
+/// tree.try_create_and_insert(10, 100)?;
+/// tree.try_create_and_insert(20, 200)?;
+/// tree.try_create_and_insert(30, 300)?;
+///
+/// let mut cursor = tree.cursor_back().unwrap();
+/// let current = cursor.current();
+/// assert_eq!(current, (&30, &300));
+///
+/// # Ok::<(), Error>(())
+/// ```
+///
+/// Obtaining a cursor returns [`None`] if the tree is empty.
+///
+/// ```
+/// use kernel::rbtree::RBTree;
+///
+/// let mut tree: RBTree<u16, u16> = RBTree::new();
+/// assert!(tree.cursor_front().is_none());
+///
+/// # Ok::<(), Error>(())
+/// ```
+///
+/// [`RBTree::cursor_lower_bound`] can be used to start at an arbitrary node in the tree.
+///
+/// ```
+/// use kernel::rbtree::RBTree;
+///
+/// // Create a new tree.
+/// let mut tree = RBTree::new();
+///
+/// // Insert five elements.
+/// tree.try_create_and_insert(10, 100)?;
+/// tree.try_create_and_insert(20, 200)?;
+/// tree.try_create_and_insert(30, 300)?;
+/// tree.try_create_and_insert(40, 400)?;
+/// tree.try_create_and_insert(50, 500)?;
+///
+/// // If the provided key exists, a cursor to that key is returned.
+/// let cursor = tree.cursor_lower_bound(&20).unwrap();
+/// let current = cursor.current();
+/// assert_eq!(current, (&20, &200));
+///
+/// // If the provided key doesn't exist, a cursor to the first larger element in sort order is returned.
+/// let cursor = tree.cursor_lower_bound(&25).unwrap();
+/// let current = cursor.current();
+/// assert_eq!(current, (&30, &300));
+///
+/// // If there is no larger key, [`None`] is returned.
+/// let cursor = tree.cursor_lower_bound(&55);
+/// assert!(cursor.is_none());
+///
+/// # Ok::<(), Error>(())
+/// ```
+///
+/// The cursor allows mutation of values in the tree.
+///
+/// ```
+/// use kernel::rbtree::RBTree;
+///
+/// // Create a new tree.
+/// let mut tree = RBTree::new();
+///
+/// // Insert three elements.
+/// tree.try_create_and_insert(10, 100)?;
+/// tree.try_create_and_insert(20, 200)?;
+/// tree.try_create_and_insert(30, 300)?;
+///
+/// // Retrieve a cursor.
+/// let mut cursor = tree.cursor_front().unwrap();
+///
+/// // Get a mutable reference to the current value.
+/// let (k, v) = cursor.current_mut();
+/// *v = 1000;
+///
+/// // The updated value is reflected in the tree.
+/// let updated = tree.get(&10).unwrap();
+/// assert_eq!(updated, &1000);
+///
+/// # Ok::<(), Error>(())
+/// ```
+///
+/// It also allows node removal. The following examples demonstrate the behavior of removing the current node.
+///
+/// ```
+/// use kernel::rbtree::RBTree;
+///
+/// // Create a new tree.
+/// let mut tree = RBTree::new();
+///
+/// // Insert three elements.
+/// tree.try_create_and_insert(10, 100)?;
+/// tree.try_create_and_insert(20, 200)?;
+/// tree.try_create_and_insert(30, 300)?;
+///
+/// // Remove the first element.
+/// let mut cursor = tree.cursor_front().unwrap();
+/// let mut current = cursor.current();
+/// assert_eq!(current, (&10, &100));
+/// cursor = cursor.remove_current().unwrap();
+///
+/// // If a node exists after the current element, it is returned.
+/// current = cursor.current();
+/// assert_eq!(current, (&20, &200));
+///
+/// // Get a cursor to the last element, and remove it.
+/// cursor = tree.cursor_back().unwrap();
+/// current = cursor.current();
+/// assert_eq!(current, (&30, &300));
+///
+/// // Since there is no next node, the previous node is returned.
+/// cursor = cursor.remove_current().unwrap();
+/// current = cursor.current();
+/// assert_eq!(current, (&20, &200));
+///
+/// // Removing the last element in the tree returns [`None`].
+/// assert!(cursor.remove_current().is_none());
+///
+/// # Ok::<(), Error>(())
+/// ```
+///
+/// Nodes adjacent to the current node can also be removed.
+///
+/// ```
+/// use kernel::rbtree::RBTree;
+///
+/// // Create a new tree.
+/// let mut tree = RBTree::new();
+///
+/// // Insert three elements.
+/// tree.try_create_and_insert(10, 100)?;
+/// tree.try_create_and_insert(20, 200)?;
+/// tree.try_create_and_insert(30, 300)?;
+///
+/// // Get a cursor to the first element.
+/// let mut cursor = tree.cursor_front().unwrap();
+/// let mut current = cursor.current();
+/// assert_eq!(current, (&10, &100));
+///
+/// // Calling `remove_prev` from the first element returns [`None`].
+/// assert!(cursor.remove_prev().is_none());
+///
+/// // Get a cursor to the last element.
+/// cursor = tree.cursor_back().unwrap();
+/// current = cursor.current();
+/// assert_eq!(current, (&30, &300));
+///
+/// // Calling `remove_prev` removes and returns the middle element.
+/// assert_eq!(cursor.remove_prev().unwrap(), (20, 200));
+///
+/// // Calling `remove_next` from the last element returns [`None`].
+/// assert!(cursor.remove_next().is_none());
+///
+/// // Move to the first element
+/// cursor = cursor.move_prev().unwrap();
+/// current = cursor.current();
+/// assert_eq!(current, (&10, &100));
+///
+/// // Calling `remove_next` removes and returns the last element.
+/// assert_eq!(cursor.remove_next().unwrap(), (30, 300));
+///
+/// # Ok::<(), Error>(())
+/// ```
+pub struct RBTreeCursor<'a, K, V> {
+ _tree: PhantomData<&'a RBTree<K, V>>,
+ root: *mut bindings::rb_root,
+ current: *mut bindings::rb_node,
+}
+
+// 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<'a, K: Send, V: Send> Send for RBTreeCursor<'a, 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<'a, K: Sync, V: Sync> Sync for RBTreeCursor<'a, K, V> {}
+
+impl<'a, K, V> RBTreeCursor<'a, K, V> {
+ /// The current node
+ pub fn current(&self) -> (&K, &V) {
+ Self::to_key_value(self.current)
+ }
+
+ /// The current node, with a mutable value
+ pub fn current_mut(&mut self) -> (&K, &mut V) {
+ Self::to_key_value_mut(self.current)
+ }
+
+ /// Remove the current node from the tree.
+ ///
+ /// Returns a cursor to the next node, if it exists,
+ /// else the previous node. Returns [`None`] if the tree
+ /// becomes empty.
+ pub fn remove_current(mut self) -> Option<Self> {
+ let prev = self.get_neighbor_raw(Direction::Prev);
+ let next = self.get_neighbor_raw(Direction::Next);
+ // SAFETY: All links fields we create are in a `Node<K, V>`.
+ let this = unsafe { crate::container_of!(self.current, Node<K, V>, links) }.cast_mut();
+ // SAFETY: The reference to the tree used to create the cursor outlives the cursor, so
+ // the tree cannot change. By the tree invariant, all nodes are valid.
+ unsafe { bindings::rb_erase(&mut (*this).links, self.root) };
+
+ let current = match (prev, next) {
+ (_, Some(next)) => next,
+ (Some(prev), None) => prev,
+ (None, None) => {
+ return None;
+ }
+ };
+
+ Some(Self {
+ current,
+ _tree: self._tree,
+ root: self.root,
+ })
+ }
+
+ /// Remove the previous node, returning it if it exists.
+ pub fn remove_prev(&mut self) -> Option<(K, V)> {
+ self.remove_neighbor(Direction::Prev)
+ }
+
+ /// Remove the next node, returning it if it exists.
+ pub fn remove_next(&mut self) -> Option<(K, V)> {
+ self.remove_neighbor(Direction::Next)
+ }
+
+ fn remove_neighbor(&mut self, direction: Direction) -> Option<(K, V)> {
+ if let Some(neighbor) = self.get_neighbor_raw(direction) {
+ // SAFETY: All links fields we create are in a `Node<K, V>`.
+ let this = unsafe { crate::container_of!(neighbor, Node<K, V>, links) }.cast_mut();
+ // SAFETY: The reference to the tree used to create the cursor outlives the cursor, so
+ // the tree cannot change. By the tree invariant, all nodes are valid.
+ unsafe { bindings::rb_erase(&mut (*this).links, self.root) };
+ return Some(Self::to_key_value_owned(neighbor));
+ }
+ None
+ }
+
+ /// Move the cursor to the previous node, returning [`None`] if it doesn't exist.
+ pub fn move_prev(self) -> Option<Self> {
+ self.mv(Direction::Prev)
+ }
+
+ /// Move the cursor to the next node, returning [`None`] if it doesn't exist.
+ pub fn move_next(self) -> Option<Self> {
+ self.mv(Direction::Next)
+ }
+
+ fn mv(mut self, direction: Direction) -> Option<Self> {
+ self.get_neighbor_raw(direction).map(|neighbor| Self {
+ _tree: self._tree,
+ root: self.root,
+ current: neighbor,
+ })
+ }
+
+ /// Access the previous node without moving the cursor.
+ pub fn peek_prev(&self) -> Option<(&K, &V)> {
+ self.peek(Direction::Prev)
+ }
+
+ /// Access the previous node without moving the cursor.
+ pub fn peek_next(&self) -> Option<(&K, &V)> {
+ self.peek(Direction::Next)
+ }
+
+ fn peek(&self, direction: Direction) -> Option<(&K, &V)> {
+ // SAFETY: `self.current` is valid by the type invariants.
+ let neighbor = unsafe {
+ match direction {
+ Direction::Prev => bindings::rb_prev(self.current),
+ Direction::Next => bindings::rb_next(self.current),
+ }
+ };
+
+ if neighbor.is_null() {
+ return None;
+ }
+
+ Some(Self::to_key_value(neighbor))
+ }
+
+ /// Access the previous node mutably without moving the cursor.
+ pub fn peek_prev_mut(&mut self) -> Option<(&K, &mut V)> {
+ self.peek_mut(Direction::Prev)
+ }
+
+ /// Access the next node mutably without moving the cursor.
+ pub fn peek_next_mut(&mut self) -> Option<(&K, &mut V)> {
+ self.peek_mut(Direction::Next)
+ }
+
+ fn peek_mut(&mut self, direction: Direction) -> Option<(&K, &mut V)> {
+ // SAFETY: `self.current` is valid by the type invariants.
+ let neighbor = unsafe {
+ match direction {
+ Direction::Prev => bindings::rb_prev(self.current),
+ Direction::Next => bindings::rb_next(self.current),
+ }
+ };
+
+ if neighbor.is_null() {
+ return None;
+ }
+
+ Some(Self::to_key_value_mut(neighbor))
+ }
+
+ fn get_neighbor_raw(&mut self, direction: Direction) -> Option<*mut bindings::rb_node> {
+ // SAFETY: `self.current` is valid by the type invariants.
+ let neighbor = unsafe {
+ match direction {
+ Direction::Prev => bindings::rb_prev(self.current),
+ Direction::Next => bindings::rb_next(self.current),
+ }
+ };
+
+ if neighbor.is_null() {
+ return None;
+ }
+
+ Some(neighbor)
+ }
+
+ // This internal method should *only* be called with a valid pointer to a node.
+ fn to_key_value(node: *mut bindings::rb_node) -> (&'a K, &'a V) {
+ // SAFETY: All links fields we create are in a `Node<K, V>`.
+ let this = unsafe { crate::container_of!(node, Node<K, V>, links) };
+ // SAFETY: The passed `node` is the current node or a non-null neighbor,
+ // thus `this` is valid by the type invariants.
+ let k = unsafe { &(*this).key };
+ // SAFETY: The passed `node` is the current node or a non-null neighbor,
+ // thus `this` is valid by the type invariants.
+ let v = unsafe { &(*this).value };
+ (k, v)
+ }
+
+ // This internal method should *only* be called with a valid pointer to a node.
+ fn to_key_value_mut(node: *mut bindings::rb_node) -> (&'a K, &'a mut V) {
+ // SAFETY: All links fields we create are in a `Node<K, V>`.
+ let this = unsafe { crate::container_of!(node, Node<K, V>, links) }.cast_mut();
+ // SAFETY: The passed `node` is the current node or a non-null neighbor,
+ // thus `this` is valid by the type invariants.
+ let k = unsafe { &(*this).key };
+ // SAFETY: The passed `node` is the current node or a non-null neighbor,
+ // thus `this` is valid by the type invariants.
+ let v = unsafe { &mut (*this).value };
+ (k, v)
+ }
+
+ // This internal method should *only* be called with a valid pointer to a node *that is being removed*.
+ fn to_key_value_owned(node: *mut bindings::rb_node) -> (K, V) {
+ // SAFETY: All links fields we create are in a `Node<K, V>`.
+ let this = unsafe { crate::container_of!(node, Node<K, V>, links) }.cast_mut();
+ // SAFETY: The passed `node` is the current node or a non-null neighbor,
+ // thus `this` is valid by the type invariants.
+ let n = unsafe { Box::from_raw(this) };
+
+ (n.key, n.value)
+ }
+}
+
+/// Direction for [`RBTreeCursor`] operations.
+enum Direction {
+ /// the node immediately before, in sort order
+ Prev,
+ /// the node immediately after, in sort order
+ Next,
+}
+
impl<'a, K, V> IntoIterator for &'a RBTree<K, V> {
type Item = (&'a K, &'a V);
type IntoIter = RBTreeIterator<'a, K, V>;