Lines Matching +full:in +full:- +full:tree
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.
23 /// In the example below we do several operations on a tree. We note that insertions may fail if
29 /// // Create a new tree.
30 /// let mut tree = RBTree::new();
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)?;
39 /// assert_eq!(tree.get(&10).unwrap(), &100);
40 /// assert_eq!(tree.get(&20).unwrap(), &200);
41 /// assert_eq!(tree.get(&30).unwrap(), &300);
46 /// let mut iter = tree.iter();
54 /// for (key, value) in &tree {
59 /// tree.try_create_and_insert(10, 1000, flags::GFP_KERNEL)?;
61 /// // Check that the tree reflects the replacement.
63 /// let mut iter = tree.iter();
71 /// *tree.get_mut(&30).unwrap() = 3000;
73 /// // Check that the tree reflects the update.
75 /// let mut iter = tree.iter();
83 /// tree.remove(&10);
85 /// // Check that the tree reflects the removal.
87 /// let mut iter = tree.iter();
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
103 /// fn insert_test(tree: &SpinLock<RBTree<u32, u32>>) -> Result {
104 /// // Pre-allocate node. This may fail (as it allocates memory).
109 /// let mut guard = tree.lock();
115 /// In the example below, we reuse an existing node allocation from an element we removed.
120 /// // Create a new tree.
121 /// let mut tree = RBTree::new();
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)?;
130 /// let mut iter = tree.iter();
138 /// let existing = tree.remove(&30).unwrap();
140 /// // Check that the tree reflects the removal.
142 /// let mut iter = tree.iter();
148 /// // Create a preallocated reservation that we can re-use later.
151 /// // Insert a new node into the tree, reusing the previous allocation. This is guaranteed to
153 /// tree.insert(reservation.into_node(15, 150));
155 /// // Check that the tree reflect the new insertion.
157 /// let mut iter = tree.iter();
169 /// Non-null parent/children pointers stored in instances of the `rb_node` C struct are always
185 /// Creates a new and empty tree.
186 pub fn new() -> Self {
188 // INVARIANT: There are no nodes in the tree, so the invariant holds vacuously.
194 /// Returns an iterator over the tree nodes, sorted by key.
195 pub fn iter(&self) -> Iter<'_, K, V> {
199 // - `self.root` is a valid pointer to a tree root.
200 // - `bindings::rb_first` produces a valid pointer to a node given `root` is valid.
209 /// Returns a mutable iterator over the tree nodes, sorted by key.
210 pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
214 // - `self.root` is a valid pointer to a tree root.
215 // - `bindings::rb_first` produces a valid pointer to a node given `root` is valid.
224 /// Returns an iterator over the keys of the nodes in the tree, in sorted order.
225 pub fn keys(&self) -> impl Iterator<Item = &'_ K> {
229 /// Returns an iterator over the values of the nodes in the tree, sorted by key.
230 pub fn values(&self) -> impl Iterator<Item = &'_ V> {
234 /// Returns a mutable iterator over the values of the nodes in the tree, sorted by key.
235 pub fn values_mut(&mut self) -> impl Iterator<Item = &'_ mut V> {
239 /// Returns a cursor over the tree nodes, starting with the smallest key.
240 pub fn cursor_front(&mut self) -> Option<Cursor<'_, K, V>> {
246 // - `current` is a valid node in the [`RBTree`] pointed to by `self`.
249 tree: self,
254 /// Returns a cursor over the tree nodes, starting with the largest key.
255 pub fn cursor_back(&mut self) -> Option<Cursor<'_, K, V>> {
261 // - `current` is a valid node in the [`RBTree`] pointed to by `self`.
264 tree: self,
274 /// Tries to insert a new value into the tree.
285 ) -> Result<Option<RBTreeNode<K, V>>> {
289 /// Inserts a new node into the tree.
295 pub fn insert(&mut self, node: RBTreeNode<K, V>) -> Option<RBTreeNode<K, V>> {
305 fn raw_entry(&mut self, key: &K) -> RawEntry<'_, K, V> {
309 // - `node`: A pointer to an uninitialized node being inserted.
310 // - `parent`: A pointer to an existing node in the tree. One of its child pointers must be
313 // - `rb_link`: A pointer to either the left-child or right-child field of `parent`. This
316 // red/black tree is empty, then it’s also possible for `parent` to be null. In
317 // this case, `rb_link` is a pointer to the `root` field of the red/black tree.
319 // We will traverse the tree looking for a node that has a null pointer as its child,
321 // that we preserve the ordering of the nodes in the tree. In each iteration of the loop
323 // in the subtree of `parent` that `child_field_of_parent` points at. Once
331 // SAFETY: All links fields we create are in a `Node<K, V>`.
334 // SAFETY: `node` is a non-null node so it is valid by the type invariants.
336 // SAFETY: `curr` is a non-null node so it is valid by the type invariants.
338 // SAFETY: `curr` is a non-null node so it is valid by the type invariants.
358 /// Gets the given key's corresponding entry in the map for in-place manipulation.
359 pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
367 pub fn find_mut(&mut self, key: &K) -> Option<OccupiedEntry<'_, K, V>> {
375 pub fn get(&self, key: &K) -> Option<&V> {
378 // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`
381 // SAFETY: `this` is a non-null node so it is valid by the type invariants.
383 // SAFETY: `node` is a non-null node so it is valid by the type invariants.
385 // SAFETY: `node` is a non-null node so it is valid by the type invariants.
387 // SAFETY: `node` is a non-null node so it is valid by the type invariants.
395 pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
399 /// Removes the node with the given key from the tree.
402 pub fn remove_node(&mut self, key: &K) -> Option<RBTreeNode<K, V>> {
406 /// Removes the node with the given key from the tree.
409 pub fn remove(&mut self, key: &K) -> Option<V> {
413 /// Returns a cursor over the tree nodes based on the given key.
416 /// Otherwise it starts with the first larger key in sort order.
418 pub fn cursor_lower_bound(&mut self, key: &K) -> Option<Cursor<'_, K, V>>
425 // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`
428 // SAFETY: `this` is a non-null node so it is valid by the type invariants.
430 // SAFETY: `node` is a non-null node so it is valid by the type invariants.
432 // SAFETY: `node` is a non-null node so it is valid by the type invariants.
446 // SAFETY: `best` is a non-null node so it is valid by the type invariants.
461 // SAFETY: `best` is a non-null node so it is valid by the type invariants.
466 // - `current` is a valid node in the [`RBTree`] pointed to by `self`.
469 tree: self,
476 fn default() -> Self {
483 // SAFETY: `root` is valid as it's embedded in `self` and we have a valid `self`.
486 // INVARIANT: The loop invariant is that all tree nodes from `next` in postorder are valid.
488 // SAFETY: All links fields we create are in a `Node<K, V>`.
492 // SAFETY: `next` and all nodes in postorder are still valid.
495 // INVARIANT: This is the destructor, so we break the type invariant during clean-up,
504 /// A bidirectional cursor over the tree nodes, sorted by key.
508 /// In the following example, we obtain a cursor to the first element in the tree.
509 /// The cursor allows us to iterate bidirectionally over key/value pairs in the tree.
514 /// // Create a new tree.
515 /// let mut tree = RBTree::new();
518 /// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
519 /// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
520 /// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
523 /// let mut cursor = tree.cursor_front().unwrap();
548 /// A cursor can also be obtained at the last element in the tree.
553 /// // Create a new tree.
554 /// let mut tree = RBTree::new();
557 /// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
558 /// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
559 /// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
561 /// let mut cursor = tree.cursor_back().unwrap();
568 /// Obtaining a cursor returns [`None`] if the tree is empty.
573 /// let mut tree: RBTree<u16, u16> = RBTree::new();
574 /// assert!(tree.cursor_front().is_none());
579 /// [`RBTree::cursor_lower_bound`] can be used to start at an arbitrary node in the tree.
584 /// // Create a new tree.
585 /// let mut tree = RBTree::new();
588 /// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
589 /// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
590 /// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
591 /// tree.try_create_and_insert(40, 400, flags::GFP_KERNEL)?;
592 /// tree.try_create_and_insert(50, 500, flags::GFP_KERNEL)?;
595 /// let cursor = tree.cursor_lower_bound(&20).unwrap();
599 /// // If the provided key doesn't exist, a cursor to the first larger element in sort order is returned.
600 /// let cursor = tree.cursor_lower_bound(&25).unwrap();
605 /// let cursor = tree.cursor_lower_bound(&55);
611 /// The cursor allows mutation of values in the tree.
616 /// // Create a new tree.
617 /// let mut tree = RBTree::new();
620 /// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
621 /// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
622 /// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
625 /// let mut cursor = tree.cursor_front().unwrap();
631 /// // The updated value is reflected in the tree.
632 /// let updated = tree.get(&10).unwrap();
643 /// // Create a new tree.
644 /// let mut tree = RBTree::new();
647 /// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
648 /// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
649 /// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
652 /// let mut cursor = tree.cursor_front().unwrap();
662 /// cursor = tree.cursor_back().unwrap();
671 /// // Removing the last element in the tree returns [`None`].
682 /// // Create a new tree.
683 /// let mut tree = RBTree::new();
686 /// tree.try_create_and_insert(10, 100, flags::GFP_KERNEL)?;
687 /// tree.try_create_and_insert(20, 200, flags::GFP_KERNEL)?;
688 /// tree.try_create_and_insert(30, 300, flags::GFP_KERNEL)?;
691 /// let mut cursor = tree.cursor_front().unwrap();
699 /// cursor = tree.cursor_back().unwrap();
722 /// - `current` points to a node that is in the same [`RBTree`] as `tree`.
724 tree: &'a mut RBTree<K, V>,
739 pub fn current(&self) -> (&K, &V) {
741 // - `self.current` is a valid node by the type invariants.
742 // - We have an immutable reference by the function signature.
747 pub fn current_mut(&mut self) -> (&K, &mut V) {
749 // - `self.current` is a valid node by the type invariants.
750 // - We have an mutable reference by the function signature.
754 /// Remove the current node from the tree.
757 /// else the previous node, else [`None`] (if the tree becomes empty). The second element
759 pub fn remove_current(self) -> (Option<Self>, RBTreeNode<K, V>) {
762 // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`
768 // SAFETY: The reference to the tree used to create the cursor outlives the cursor, so
769 // the tree cannot change. By the tree invariant, all nodes are valid.
770 unsafe { bindings::rb_erase(&mut (*this).links, addr_of_mut!(self.tree.root)) };
782 // - `current` is a valid node in the [`RBTree`] pointed to by `self.tree`.
785 tree: self.tree,
792 pub fn remove_prev(&mut self) -> Option<RBTreeNode<K, V>> {
797 pub fn remove_next(&mut self) -> Option<RBTreeNode<K, V>> {
801 fn remove_neighbor(&mut self, direction: Direction) -> Option<RBTreeNode<K, V>> {
804 // SAFETY: The reference to the tree used to create the cursor outlives the cursor, so
805 // the tree cannot change. By the tree invariant, all nodes are valid.
806 unsafe { bindings::rb_erase(neighbor, addr_of_mut!(self.tree.root)) };
807 // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`
818 pub fn move_prev(self) -> Option<Self> {
823 pub fn move_next(self) -> Option<Self> {
827 fn mv(self, direction: Direction) -> Option<Self> {
829 // - `neighbor` is a valid node in the [`RBTree`] pointed to by `self.tree`.
831 tree: self.tree,
837 pub fn peek_prev(&self) -> Option<(&K, &V)> {
842 pub fn peek_next(&self) -> Option<(&K, &V)> {
846 fn peek(&self, direction: Direction) -> Option<(&K, &V)> {
849 // - `neighbor` is a valid tree node.
850 // - By the function signature, we have an immutable reference to `self`.
856 pub fn peek_prev_mut(&mut self) -> Option<(&K, &mut V)> {
861 pub fn peek_next_mut(&mut self) -> Option<(&K, &mut V)> {
865 fn peek_mut(&mut self, direction: Direction) -> Option<(&K, &mut V)> {
868 // - `neighbor` is a valid tree node.
869 // - By the function signature, we have a mutable reference to `self`.
874 fn get_neighbor_raw(&self, direction: Direction) -> Option<NonNull<bindings::rb_node>> {
888 /// - `node` must be a valid pointer to a node in an [`RBTree`].
889 /// - The caller has immutable access to `node` for the duration of 'b.
890 unsafe fn to_key_value<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, &'b V) {
891 // SAFETY: the caller guarantees that `node` is a valid pointer in an `RBTree`.
899 /// - `node` must be a valid pointer to a node in an [`RBTree`].
900 /// - The caller has mutable access to `node` for the duration of 'b.
901 unsafe fn to_key_value_mut<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, &'b mut V) {
902 // SAFETY: the caller guarantees that `node` is a valid pointer in an `RBTree`.
910 /// - `node` must be a valid pointer to a node in an [`RBTree`].
911 /// - The caller has immutable access to the key for the duration of 'b.
912 unsafe fn to_key_value_raw<'b>(node: NonNull<bindings::rb_node>) -> (&'b K, *mut V) {
913 // SAFETY: By the type invariant of `Self`, all non-null `rb_node` pointers stored in `self`
916 // SAFETY: The passed `node` is the current node or a non-null neighbor,
919 // SAFETY: The passed `node` is the current node or a non-null neighbor,
928 /// the node immediately before, in sort order
930 /// the node immediately after, in sort order
938 fn into_iter(self) -> Self::IntoIter {
962 fn next(&mut self) -> Option<Self::Item> {
972 fn into_iter(self) -> Self::IntoIter {
997 fn next(&mut self) -> Option<Self::Item> {
1007 /// - `self.next` is a valid pointer.
1008 /// - `self.next` points to a node stored inside of a valid `RBTree`.
1011 _phantom: PhantomData<fn() -> (K, V)>,
1017 fn next(&mut self) -> Option<Self::Item> {
1022 // SAFETY: By the type invariant of `IterRaw`, `self.next` is a valid node in an `RBTree`,
1026 // SAFETY: `self.next` is a valid tree node by the type invariants.
1034 /// A memory reservation for a red-black tree node.
1037 /// It contains the memory needed to hold a node that can be inserted into a red-black tree. One
1044 /// Allocates memory for a node to be eventually initialised and inserted into the tree via a
1046 pub fn new(flags: Flags) -> Result<RBTreeNodeReservation<K, V>> {
1063 /// It then becomes an [`RBTreeNode`] that can be inserted into a tree.
1064 pub fn into_node(self, key: K, value: V) -> RBTreeNode<K, V> {
1077 /// A red-black tree node.
1079 /// The node is fully initialised (with key and value) and can be inserted into a tree without any
1086 /// Allocates and initialises a node that can be inserted into the tree via
1088 pub fn new(key: K, value: V, flags: Flags) -> Result<RBTreeNode<K, V>> {
1093 pub fn to_key_value(self) -> (K, V) {
1111 /// It then becomes a reservation that can be re-initialised into a different node (i.e., with
1114 /// The existing key and value are dropped in-place as part of this operation, that is, memory
1116 pub fn into_reservation(self) -> RBTreeNodeReservation<K, V> {
1123 /// A view into a single entry in a map, which may either be vacant or occupied.
1141 /// A view into a vacant entry in a [`RBTree`]. It is part of the [`Entry`] enum.
1150 /// - `parent` may be null if the new node becomes the root.
1151 /// - `child_field_of_parent` is a valid pointer to the left-child or right-child of `parent`. If `parent` is
1157 /// This points to the left-child or right-child field of `parent`, or `root` if `parent` is
1168 fn insert(self, node: RBTreeNode<K, V>) -> &'a mut V {
1175 // INVARIANT: We are linking in a new node, which is valid. It remains valid because we
1180 // SAFETY: All pointers are valid. `node` has just been inserted into the tree.
1183 // SAFETY: The node is valid until we remove it from the tree.
1190 pub fn insert(self, value: V, reservation: RBTreeNodeReservation<K, V>) -> &'a mut V {
1195 /// A view into an occupied entry in a [`RBTree`]. It is part of the [`Entry`] enum.
1198 /// - `node_links` is a valid, non-null pointer to a tree node in `self.rbtree`
1206 /// Gets a reference to the value in the entry.
1207 pub fn get(&self) -> &V {
1209 // - `self.node_links` is a valid pointer to a node in the tree.
1210 // - We have shared access to the underlying tree, and can thus give out a shared reference.
1214 /// Gets a mutable reference to the value in the entry.
1215 pub fn get_mut(&mut self) -> &mut V {
1217 // - `self.node_links` is a valid pointer to a node in the tree.
1218 // - We have exclusive access to the underlying tree, and can thus give out a mutable reference.
1225 pub fn into_mut(self) -> &'a mut V {
1227 // - `self.node_links` is a valid pointer to a node in the tree.
1228 // - This consumes the `&'a mut RBTree<K, V>`, therefore it can give out a mutable reference that lives for `'a`.
1233 pub fn remove_node(self) -> RBTreeNode<K, V> {
1234 // SAFETY: The node is a node in the tree, so it is valid.
1238 // removed from the tree. So the invariants still hold.
1240 // SAFETY: The node was a node in the tree, but we removed it, so we can convert it
1249 pub fn remove(self) -> V {
1259 fn replace(self, node: RBTreeNode<K, V>) -> RBTreeNode<K, V> {
1266 // SAFETY: This updates the pointers so that `new_node_links` is in the tree where
1273 // - `self.node_ptr` produces a valid pointer to a node in the tree.
1274 // - Now that we removed this entry from the tree, we can convert the node to a box.