checkasm: Use a self-balancing tree

Tested functions are internally kept in a binary search tree for efficient
lookups. The downside of the current implementation is that the tree quickly
becomes unbalanced which causes an unneccessary amount of comparisons between
nodes. Improve this by changing the tree into a self-balancing left-leaning
red-black tree with a worst case lookup/insertion time complexity of O(log n).

Significantly reduces the recursion depth and makes the tests run around 10%
faster overall. The relative performance improvement compared to the existing
non-balanced tree will also most likely increase as more tests are added.

Signed-off-by: Anton Khirnov <anton@khirnov.net>

1 files changed, 47 insertions, 12 deletions

diff --git a/tests/checkasm/checkasm.c b/tests/checkasm/checkasm.c
index 3d47a43d00..013e197b43 100644
--- a/tests/checkasm/checkasm.c
+++ b/tests/checkasm/checkasm.c
@@ -125,6 +125,7 @@ typedef struct CheckasmFuncVersion {
 typedef struct CheckasmFunc {
     struct CheckasmFunc *child[2];
     CheckasmFuncVersion versions;
+    uint8_t color; /* 0 = red, 1 = black */
     char name[1];
 } CheckasmFunc;
 
@@ -287,24 +288,57 @@ static int cmp_func_names(const char *a, const char *b)
     return (digit_diff = av_isdigit(*a) - av_isdigit(*b)) ? digit_diff : ascii_diff;
 }
 
+/* Perform a tree rotation in the specified direction and return the new root */
+static CheckasmFunc *rotate_tree(CheckasmFunc *f, int dir)
+{
+    CheckasmFunc *r = f->child[dir^1];
+    f->child[dir^1] = r->child[dir];
+    r->child[dir] = f;
+    r->color = f->color;
+    f->color = 0;
+    return r;
+}
+
+#define is_red(f) ((f) && !(f)->color)
+
+/* Balance a left-leaning red-black tree at the specified node */
+static void balance_tree(CheckasmFunc **root)
+{
+    CheckasmFunc *f = *root;
+
+    if (is_red(f->child[0]) && is_red(f->child[1])) {
+        f->color ^= 1;
+        f->child[0]->color = f->child[1]->color = 1;
+    }
+
+    if (!is_red(f->child[0]) && is_red(f->child[1]))
+        *root = rotate_tree(f, 0); /* Rotate left */
+    else if (is_red(f->child[0]) && is_red(f->child[0]->child[0]))
+        *root = rotate_tree(f, 1); /* Rotate right */
+}
+
 /* Get a node with the specified name, creating it if it doesn't exist */
-static CheckasmFunc *get_func(const char *name, int length)
+static CheckasmFunc *get_func(CheckasmFunc **root, const char *name)
 {
-    CheckasmFunc *f, **f_ptr = &state.funcs;
+    CheckasmFunc *f = *root;
 
-    /* Search the tree for a matching node */
-    while ((f = *f_ptr)) {
+    if (f) {
+        /* Search the tree for a matching node */
         int cmp = cmp_func_names(name, f->name);
-        if (!cmp)
-            return f;
+        if (cmp) {
+            f = get_func(&f->child[cmp > 0], name);
 
-        f_ptr = &f->child[(cmp > 0)];
+            /* Rebalance the tree on the way up if a new node was inserted */
+            if (!f->versions.func)
+                balance_tree(root);
+        }
+    } else {
+        /* Allocate and insert a new node into the tree */
+        int name_length = strlen(name);
+        f = *root = checkasm_malloc(sizeof(CheckasmFunc) + name_length);
+        memcpy(f->name, name, name_length + 1);
     }
 
-    /* Allocate and insert a new node into the tree */
-    f = *f_ptr = checkasm_malloc(sizeof(CheckasmFunc) + length);
-    memcpy(f->name, name, length+1);
-
     return f;
 }
 
@@ -405,7 +439,8 @@ void *checkasm_check_func(void *func, const char *name, ...)
     if (!func || name_length <= 0 || name_length >= sizeof(name_buf))
         return NULL;
 
-    state.current_func = get_func(name_buf, name_length);
+    state.current_func = get_func(&state.funcs, name_buf);
+    state.funcs->color = 1;
     v = &state.current_func->versions;
 
     if (v->func) {