From 2ff38260b9a794a84b059cb355f5ffb255739f8e Mon Sep 17 00:00:00 2001 From: Ian Romanick Date: Wed, 23 Aug 2023 12:08:21 -0700 Subject: [PATCH] util/rb-tree: Return the actual first node from rb_tree_search Previously rb_tree_search would return the first node encountered, but that may not be the first node that would be encoutnered by, say, rb_tree_foreach. Two test cases are added. v2: Add more curly braces. Suggested by Faith. Reviewed-by: Faith Ekstrand Part-of: --- src/util/rb_tree.h | 53 ++++++++++++++++++++++---------------- src/util/tests/rb_tree_test.cpp | 56 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 88 insertions(+), 21 deletions(-) diff --git a/src/util/rb_tree.h b/src/util/rb_tree.h index 9466e73..ac44370 100644 --- a/src/util/rb_tree.h +++ b/src/util/rb_tree.h @@ -81,6 +81,18 @@ rb_tree_is_empty(const struct rb_tree *T) return T->root == NULL; } +/** Get the first (left-most) node in the tree or NULL */ +struct rb_node *rb_tree_first(struct rb_tree *T); + +/** Get the last (right-most) node in the tree or NULL */ +struct rb_node *rb_tree_last(struct rb_tree *T); + +/** Get the next node (to the right) in the tree or NULL */ +struct rb_node *rb_node_next(struct rb_node *node); + +/** Get the next previous (to the left) in the tree or NULL */ +struct rb_node *rb_node_prev(struct rb_node *node); + /** Retrieve the data structure containing a node * * \param type The type of the containing data structure @@ -170,12 +182,23 @@ rb_tree_search(struct rb_tree *T, const void *key, struct rb_node *x = T->root; while (x != NULL) { int c = cmp(x, key); - if (c < 0) + if (c < 0) { x = x->left; - else if (c > 0) + } else if (c > 0) { x = x->right; - else - return x; + } else { + /* x is the first *encountered* node matching the key. There may + * be other nodes in the left subtree that also match the key. + */ + while (true) { + struct rb_node *prev = rb_node_prev(x); + + if (prev == NULL || cmp(prev, key) != 0) + return x; + + x = prev; + } + } } return x; @@ -184,11 +207,11 @@ rb_tree_search(struct rb_tree *T, const void *key, /** Sloppily search the tree for a node * * This function searches the tree for a given node. If a node with a - * matching key exists, that first matching node found will be returned. - * If no node with an exactly matching key exists, the node returned will - * be either the right-most node comparing less than \p key or the - * right-most node comparing greater than \p key. If the tree is empty, - * NULL is returned. + * matching key exists, that first encountered matching node found (there may + * be other matching nodes in the left subtree) will be returned. If no node + * with an exactly matching key exists, the node returned will be either the + * right-most node comparing less than \p key or the right-most node comparing + * greater than \p key. If the tree is empty, NULL is returned. * * \param T The red-black tree to search * @@ -219,18 +242,6 @@ rb_tree_search_sloppy(struct rb_tree *T, const void *key, return y; } -/** Get the first (left-most) node in the tree or NULL */ -struct rb_node *rb_tree_first(struct rb_tree *T); - -/** Get the last (right-most) node in the tree or NULL */ -struct rb_node *rb_tree_last(struct rb_tree *T); - -/** Get the next node (to the right) in the tree or NULL */ -struct rb_node *rb_node_next(struct rb_node *node); - -/** Get the next previous (to the left) in the tree or NULL */ -struct rb_node *rb_node_prev(struct rb_node *node); - #define rb_node_next_or_null(n) ((n) == NULL ? NULL : rb_node_next(n)) #define rb_node_prev_or_null(n) ((n) == NULL ? NULL : rb_node_prev(n)) diff --git a/src/util/tests/rb_tree_test.cpp b/src/util/tests/rb_tree_test.cpp index 7bda246..2676bd5 100644 --- a/src/util/tests/rb_tree_test.cpp +++ b/src/util/tests/rb_tree_test.cpp @@ -227,3 +227,59 @@ TEST(RBTreeTest, InsertAndSearch) validate_search(&tree, i + 1, ARRAY_SIZE(test_numbers) - 1); } } + +TEST(RBTreeTest, FindFirst) +{ + struct rb_test_node nodes[100]; + struct rb_tree tree; + + rb_tree_init(&tree); + + const int x = 13; + + for (unsigned i = 0; i < ARRAY_SIZE(nodes); i++) { + nodes[i].key = x; + rb_tree_insert(&tree, &nodes[i].node, rb_test_node_cmp); + } + + struct rb_node *n = rb_tree_search(&tree, &x, rb_test_node_cmp_void); + + ASSERT_NE(nullptr, n); + EXPECT_EQ(nullptr, rb_node_prev(n)); + EXPECT_EQ(n, rb_tree_first(&tree)); +} + +TEST(RBTreeTest, FindFirstOfMiddle) +{ + struct rb_test_node nodes[3 * 33]; + struct rb_tree tree; + + rb_tree_init(&tree); + + unsigned i; + for (i = 0; i < ARRAY_SIZE(nodes) / 3; i++) { + nodes[i].key = i * 13; + rb_tree_insert(&tree, &nodes[i].node, rb_test_node_cmp); + } + + const int x = i * 13; + for (/* empty */; i < 2 * ARRAY_SIZE(nodes) / 3; i++) { + nodes[i].key = x; + rb_tree_insert(&tree, &nodes[i].node, rb_test_node_cmp); + } + + for (/* empty */; i < ARRAY_SIZE(nodes); i++) { + nodes[i].key = i * 13; + rb_tree_insert(&tree, &nodes[i].node, rb_test_node_cmp); + } + + struct rb_node *n = rb_tree_search(&tree, &x, rb_test_node_cmp_void); + + ASSERT_NE(nullptr, n); + + struct rb_node *prev = rb_node_prev(n); + + ASSERT_NE(nullptr, prev); + + EXPECT_NE(rb_test_node_cmp(prev, n), 0); +} -- 2.7.4