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 <faith.ekstrand@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/24856>

diff --git a/src/util/rb_tree.h b/src/util/rb_tree.h
index 9466e73..ac44370 100644 (file)
--- 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 (file)
--- 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);
+}