1 // SPDX-License-Identifier: GPL-2.0
3 //! Red-black trees.
7 //! Reference: <https://docs.kernel.org/core-api/rbtree.html>
17 /// A red-black tree with owned nodes.
19 /// It is backed by the kernel C red-black trees.
44 /// // Iterate over the nodes we just inserted.
103 /// fn insert_test(tree: &SpinLock<RBTree<u32, u32>>) -> Result {
104 ///     // Pre-allocate node. This may fail (as it allocates memory).
148 /// // Create a preallocated reservation that we can re-use later.
169 /// Non-null parent/children pointers stored in instances of the `rb_node` C struct are always
186     pub fn new() -> Self {  in new()
196     pub fn is_empty(&self) -> bool {  in is_empty()
200     /// Returns an iterator over the tree nodes, sorted by key.
201     pub fn iter(&self) -> Iter<'_, K, V> {  in iter()
205             //   - `self.root` is a valid pointer to a tree root.  in iter()
206             //   - `bindings::rb_first` produces a valid pointer to a node given `root` is valid.  in iter()
215     /// Returns a mutable iterator over the tree nodes, sorted by key.
216     pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {  in iter_mut()
220             //   - `self.root` is a valid pointer to a tree root.  in iter_mut()
221             //   - `bindings::rb_first` produces a valid pointer to a node given `root` is valid.  in iter_mut()
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> {  in keys()
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> {  in values()
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> {  in values_mut()
245     /// Returns a cursor over the tree nodes, starting with the smallest key.
246     pub fn cursor_front(&mut self) -> Option<Cursor<'_, K, V>> {  in cursor_front()
249         let current = unsafe { bindings::rb_first(root) };  in cursor_front()  localVariable
250         NonNull::new(current).map(|current| {  in cursor_front()
252             // - `current` is a valid node in the [`RBTree`] pointed to by `self`.  in cursor_front()
254                 current,  in cursor_front()
260     /// Returns a cursor over the tree nodes, starting with the largest key.
261     pub fn cursor_back(&mut self) -> Option<Cursor<'_, K, V>> {  in cursor_back()
264         let current = unsafe { bindings::rb_last(root) };  in cursor_back()  localVariable
265         NonNull::new(current).map(|current| {  in cursor_back()
267             // - `current` is a valid node in the [`RBTree`] pointed to by `self`.  in cursor_back()
269                 current,  in cursor_back()
291     ) -> Result<Option<RBTreeNode<K, V>>> {  in try_create_and_insert()
301     pub fn insert(&mut self, node: RBTreeNode<K, V>) -> Option<RBTreeNode<K, V>> {  in insert()
311     fn raw_entry(&mut self, key: &K) -> RawEntry<'_, K, V> {  in raw_entry()
315         // - `node`: A pointer to an uninitialized node being inserted.  in raw_entry()
316         // - `parent`: A pointer to an existing node in the tree. One of its child pointers must be  in raw_entry()
319         // - `rb_link`: A pointer to either the left-child or right-child field of `parent`. This  in raw_entry()
340             // SAFETY: `node` is a non-null node so it is valid by the type invariants.  in raw_entry()
342                 // SAFETY: `curr` is a non-null node so it is valid by the type invariants.  in raw_entry()
344                 // SAFETY: `curr` is a non-null node so it is valid by the type invariants.  in raw_entry()
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> {  in entry()
373     pub fn find_mut(&mut self, key: &K) -> Option<OccupiedEntry<'_, K, V>> {  in find_mut()
381     pub fn get(&self, key: &K) -> Option<&V> {  in get()
384 …       // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`  in get()
387             // SAFETY: `this` is a non-null node so it is valid by the type invariants.  in get()
389                 // SAFETY: `node` is a non-null node so it is valid by the type invariants.  in get()
391                 // SAFETY: `node` is a non-null node so it is valid by the type invariants.  in get()
393                 // SAFETY: `node` is a non-null node so it is valid by the type invariants.  in get()
401     pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {  in get_mut()
408     pub fn remove_node(&mut self, key: &K) -> Option<RBTreeNode<K, V>> {  in remove_node()
415     pub fn remove(&mut self, key: &K) -> Option<V> {  in remove()
419     /// Returns a cursor over the tree nodes based on the given key.
424     pub fn cursor_lower_bound(&mut self, key: &K) -> Option<Cursor<'_, K, V>>  in cursor_lower_bound()
431 …       // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`  in cursor_lower_bound()
434             // SAFETY: `this` is a non-null node so it is valid by the type invariants.  in cursor_lower_bound()
436             // SAFETY: `node` is a non-null node so it is valid by the type invariants.  in cursor_lower_bound()
438             // SAFETY: `node` is a non-null node so it is valid by the type invariants.  in cursor_lower_bound()
452 …                        // SAFETY: `best` is a non-null node so it is valid by the type invariants.  in cursor_lower_bound()
467         // SAFETY: `best` is a non-null node so it is valid by the type invariants.  in cursor_lower_bound()
470         NonNull::new(links).map(|current| {  in cursor_lower_bound()
472             // - `current` is a valid node in the [`RBTree`] pointed to by `self`.  in cursor_lower_bound()
474                 current,  in cursor_lower_bound()
482     fn default() -> Self {  in default()
497             // Find out what the next node is before disposing of the current one.  in drop()
501             // INVARIANT: This is the destructor, so we break the type invariant during clean-up,  in drop()
510 /// A bidirectional cursor over the tree nodes, sorted by key.
515 /// The cursor allows us to iterate bidirectionally over key/value pairs in the tree.
530 /// let mut current = cursor.current();
531 /// assert_eq!(current, (&10, &100));
535 /// current = cursor.current();
536 /// assert_eq!(current, (&20, &200));
541 /// current = cursor.current();
542 /// assert_eq!(current, (&20, &200));
546 /// current = cursor.current();
547 /// assert_eq!(current, (&30, &300));
568 /// let current = cursor.current();
569 /// assert_eq!(current, (&30, &300));
602 /// let current = cursor.current();
603 /// assert_eq!(current, (&20, &200));
607 /// let current = cursor.current();
608 /// assert_eq!(current, (&30, &300));
633 /// // Get a mutable reference to the current value.
644 … allows node removal. The following examples demonstrate the behavior of removing the current node.
659 /// let mut current = cursor.current();
660 /// assert_eq!(current, (&10, &100));
663 /// // If a node exists after the current element, it is returned.
664 /// current = cursor.current();
665 /// assert_eq!(current, (&20, &200));
669 /// current = cursor.current();
670 /// assert_eq!(current, (&30, &300));
674 /// current = cursor.current();
675 /// assert_eq!(current, (&20, &200));
683 /// Nodes adjacent to the current node can also be removed.
698 /// let mut current = cursor.current();
699 /// assert_eq!(current, (&10, &100));
706 /// current = cursor.current();
707 /// assert_eq!(current, (&30, &300));
717 /// current = cursor.current();
718 /// assert_eq!(current, (&10, &100));
728 /// - `current` points to a node that is in the same [`RBTree`] as `tree`.
731     current: NonNull<bindings::rb_node>,  field
744     /// The current node
745     pub fn current(&self) -> (&K, &V) {  in current()  method
747         // - `self.current` is a valid node by the type invariants.  in current()
748         // - We have an immutable reference by the function signature.  in current()
749         unsafe { Self::to_key_value(self.current) }  in current()
752     /// The current node, with a mutable value
753     pub fn current_mut(&mut self) -> (&K, &mut V) {  in current_mut()
755         // - `self.current` is a valid node by the type invariants.  in current_mut()
756         // - We have an mutable reference by the function signature.  in current_mut()
757         unsafe { Self::to_key_value_mut(self.current) }  in current_mut()
760     /// Remove the current node from the tree.
765     pub fn remove_current(self) -> (Option<Self>, RBTreeNode<K, V>) {  in remove_current()
768         // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`  in remove_current()
770         let this = unsafe { container_of!(self.current.as_ptr(), Node<K, V>, links) };  in remove_current()
779         // - `current` is a valid node in the [`RBTree`] pointed to by `self.tree`.  in remove_current()
780         let cursor = next.or(prev).map(|current| Self {  in remove_current()
781             current,  in remove_current()
789     pub fn remove_prev(&mut self) -> Option<RBTreeNode<K, V>> {  in remove_prev()
794     pub fn remove_next(&mut self) -> Option<RBTreeNode<K, V>> {  in remove_next()
798     fn remove_neighbor(&mut self, direction: Direction) -> Option<RBTreeNode<K, V>> {  in remove_neighbor()
804 …       // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`  in remove_neighbor()
815     pub fn move_prev(self) -> Option<Self> {  in move_prev()
820     pub fn move_next(self) -> Option<Self> {  in move_next()
824     fn mv(self, direction: Direction) -> Option<Self> {  in mv()
826         // - `neighbor` is a valid node in the [`RBTree`] pointed to by `self.tree`.  in mv()
829             current: neighbor,  in mv()
834     pub fn peek_prev(&self) -> Option<(&K, &V)> {  in peek_prev()
839     pub fn peek_next(&self) -> Option<(&K, &V)> {  in peek_next()
843     fn peek(&self, direction: Direction) -> Option<(&K, &V)> {  in peek()
846             // - `neighbor` is a valid tree node.  in peek()
847             // - By the function signature, we have an immutable reference to `self`.  in peek()
853     pub fn peek_prev_mut(&mut self) -> Option<(&K, &mut V)> {  in peek_prev_mut()
858     pub fn peek_next_mut(&mut self) -> Option<(&K, &mut V)> {  in peek_next_mut()
862     fn peek_mut(&mut self, direction: Direction) -> Option<(&K, &mut V)> {  in peek_mut()
865             // - `neighbor` is a valid tree node.  in peek_mut()
866             // - By the function signature, we have a mutable reference to `self`.  in peek_mut()
871     fn get_neighbor_raw(&self, direction: Direction) -> Option<NonNull<bindings::rb_node>> {  in get_neighbor_raw()
872         // SAFETY: `self.current` is valid by the type invariants.  in get_neighbor_raw()
875                 Direction::Prev => bindings::rb_prev(self.current.as_ptr()),  in get_neighbor_raw()
876                 Direction::Next => bindings::rb_next(self.current.as_ptr()),  in get_neighbor_raw()
885     /// - `node` must be a valid pointer to a node in an [`RBTree`].
886     /// - The caller has immutable access to `node` for the duration of `'b`.
887     unsafe fn to_key_value<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, &'b V) {  in to_key_value()
896     /// - `node` must be a valid pointer to a node in an [`RBTree`].
897     /// - The caller has mutable access to `node` for the duration of `'b`.
898     unsafe fn to_key_value_mut<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, &'b mut V) {  in to_key_value_mut()
907     /// - `node` must be a valid pointer to a node in an [`RBTree`].
908     /// - The caller has immutable access to the key for the duration of `'b`.
909     unsafe fn to_key_value_raw<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, *mut V) {  in to_key_value_raw()
910         // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`  in to_key_value_raw()
913         // SAFETY: The passed `node` is the current node or a non-null neighbor,  in to_key_value_raw()
916         // SAFETY: The passed `node` is the current node or a non-null neighbor,  in to_key_value_raw()
935     fn into_iter(self) -> Self::IntoIter {  in into_iter()
940 /// An iterator over the nodes of a [`RBTree`].
959     fn next(&mut self) -> Option<Self::Item> {  in next()
969     fn into_iter(self) -> Self::IntoIter {  in into_iter()
974 /// A mutable iterator over the nodes of a [`RBTree`].
994     fn next(&mut self) -> Option<Self::Item> {  in next()
1001 /// A raw iterator over the nodes of a [`RBTree`].
1004 /// - `self.next` is a valid pointer.
1005 /// - `self.next` points to a node stored inside of a valid `RBTree`.
1008     _phantom: PhantomData<fn() -> (K, V)>,
1014     fn next(&mut self) -> Option<Self::Item> {  in next()
1031 /// A memory reservation for a red-black tree node.
1034 /// It contains the memory needed to hold a node that can be inserted into a red-black tree. One
1043     pub fn new(flags: Flags) -> Result<RBTreeNodeReservation<K, V>> {  in new()
1061     pub fn into_node(self, key: K, value: V) -> RBTreeNode<K, V> {  in into_node()
1074 /// A red-black tree node.
1085     pub fn new(key: K, value: V, flags: Flags) -> Result<RBTreeNode<K, V>> {  in new()
1090     pub fn to_key_value(self) -> (K, V) {  in to_key_value()
1108     /// It then becomes a reservation that can be re-initialised into a different node (i.e., with
1111     /// The existing key and value are dropped in-place as part of this operation, that is, memory
1113     pub fn into_reservation(self) -> RBTreeNodeReservation<K, V> {  in into_reservation()
1147 /// - `parent` may be null if the new node becomes the root.
1148 /// - `child_field_of_parent` is a valid pointer to the left-child or right-child of `parent`. If `…
1154     /// This points to the left-child or right-child field of `parent`, or `root` if `parent` is
1165     fn insert(self, node: RBTreeNode<K, V>) -> &'a mut V {  in insert()
1187     pub fn insert(self, value: V, reservation: RBTreeNodeReservation<K, V>) -> &'a mut V {  in insert()
1195 /// - `node_links` is a valid, non-null pointer to a tree node in `self.rbtree`
1204     pub fn get(&self) -> &V {  in get()
1206         // - `self.node_links` is a valid pointer to a node in the tree.  in get()
1207         // - We have shared access to the underlying tree, and can thus give out a shared reference.  in get()
1212     pub fn get_mut(&mut self) -> &mut V {  in get_mut()
1214         // - `self.node_links` is a valid pointer to a node in the tree.  in get_mut()
1215 …   // - We have exclusive access to the underlying tree, and can thus give out a mutable reference.  in get_mut()
1222     pub fn into_mut(self) -> &'a mut V {  in into_mut()
1224         // - `self.node_links` is a valid pointer to a node in the tree.  in into_mut()
1225 …// - This consumes the `&'a mut RBTree<K, V>`, therefore it can give out a mutable reference that …  in into_mut()
1230     pub fn remove_node(self) -> RBTreeNode<K, V> {  in remove_node()
1244     pub fn remove(self) -> V {  in remove()
1251     /// Swap the current node for the provided node.
1254     fn replace(self, node: RBTreeNode<K, V>) -> RBTreeNode<K, V> {  in replace()
1268         // - `self.node_ptr` produces a valid pointer to a node in the tree.  in replace()
1269         // - Now that we removed this entry from the tree, we can convert the node to a box.  in replace()