git: 2c545cf3b063 - main - rb_tree: test rank balance

The branch main has been updated by dougm:

URL: https://cgit.FreeBSD.org/src/commit/?id=2c545cf3b06310e248dd4427f31e73f0bc1ad504

commit 2c545cf3b06310e248dd4427f31e73f0bc1ad504
Author:     Doug Moore <dougm@FreeBSD.org>
AuthorDate: 2022-09-08 02:40:05 +0000
Commit:     Doug Moore <dougm@FreeBSD.org>
CommitDate: 2022-09-08 02:40:05 +0000

    rb_tree: test rank balance
    
    With _RB_DIAGNOSTIC defined, provide an RB_RANK method to compute the
    rank of a node in an rb-tree, if the subtree rooted at that node is
    rank-balanced, and -1 otherwise.
    
    In rb_test, rewrite a bit to avoid malloc/free and nondeterministic
    running times because of randomness. Allocate all the nodes on the
    stack, and shuffle a set of keys to get randomness for the testing.
    
    Add a rank-balance check for the completed tree.
    
    Reviewed by:    markj
    MFC after:      3 weeks
    Differential Revision:  https://reviews.freebsd.org/D36484
---
 sys/sys/tree.h          | 37 ++++++++++++++++++++++++++++
 tests/sys/sys/rb_test.c | 65 ++++++++++++++++++++++++++-----------------------
 2 files changed, 71 insertions(+), 31 deletions(-)

diff --git a/sys/sys/tree.h b/sys/sys/tree.h
index 971562a5c121..74f1f23d3576 100644
--- a/sys/sys/tree.h
+++ b/sys/sys/tree.h
@@ -410,12 +410,17 @@ struct {								\
 	RB_SET_PARENT(elm, tmp, field);					\
 } while (/*CONSTCOND*/ 0)
 
+#if defined(_KERNEL) && defined(DIAGNOSTIC) && !defined(_RB_DIAGNOSTIC)
+#define _RB_DIAGNOSTIC 1
+#endif
+
 /* Generates prototypes and inline functions */
 #define	RB_PROTOTYPE(name, type, field, cmp)				\
 	RB_PROTOTYPE_INTERNAL(name, type, field, cmp,)
 #define	RB_PROTOTYPE_STATIC(name, type, field, cmp)			\
 	RB_PROTOTYPE_INTERNAL(name, type, field, cmp, __unused static)
 #define RB_PROTOTYPE_INTERNAL(name, type, field, cmp, attr)		\
+	RB_PROTOTYPE_RANK(name, type, attr)				\
 	RB_PROTOTYPE_INSERT_COLOR(name, type, attr);			\
 	RB_PROTOTYPE_REMOVE_COLOR(name, type, attr);			\
 	RB_PROTOTYPE_INSERT(name, type, attr);				\
@@ -426,6 +431,12 @@ struct {								\
 	RB_PROTOTYPE_PREV(name, type, attr);				\
 	RB_PROTOTYPE_MINMAX(name, type, attr);				\
 	RB_PROTOTYPE_REINSERT(name, type, attr);
+#ifdef _RB_DIAGNOSTIC
+#define RB_PROTOTYPE_RANK(name, type, attr)				\
+	attr int name##_RB_RANK(struct type *);
+#else
+#define RB_PROTOTYPE_RANK(name, type, attr)
+#endif
 #define RB_PROTOTYPE_INSERT_COLOR(name, type, attr)			\
 	attr void name##_RB_INSERT_COLOR(struct name *, struct type *)
 #define RB_PROTOTYPE_REMOVE_COLOR(name, type, attr)			\
@@ -456,6 +467,7 @@ struct {								\
 #define	RB_GENERATE_STATIC(name, type, field, cmp)			\
 	RB_GENERATE_INTERNAL(name, type, field, cmp, __unused static)
 #define RB_GENERATE_INTERNAL(name, type, field, cmp, attr)		\
+	RB_GENERATE_RANK(name, type, field, attr)			\
 	RB_GENERATE_INSERT_COLOR(name, type, field, attr)		\
 	RB_GENERATE_REMOVE_COLOR(name, type, field, attr)		\
 	RB_GENERATE_INSERT(name, type, field, cmp, attr)		\
@@ -467,6 +479,31 @@ struct {								\
 	RB_GENERATE_MINMAX(name, type, field, attr)			\
 	RB_GENERATE_REINSERT(name, type, field, cmp, attr)
 
+#ifdef _RB_DIAGNOSTIC
+#define RB_GENERATE_RANK(name, type, field, attr)			\
+attr int								\
+name##_RB_RANK(struct type *elm)					\
+{									\
+	struct type *left, *right;					\
+	int left_rank, right_rank;					\
+	__uintptr_t bits;						\
+									\
+	if (elm == NULL)						\
+		return (0);						\
+	bits = RB_BITS(elm, field);					\
+	left = RB_LEFT(elm, field);					\
+	left_rank = ((bits & RB_RED_L) ? 2 : 1) + name##_RB_RANK(left);	\
+	right = RB_RIGHT(elm, field);					\
+	right_rank = ((bits & RB_RED_R) ? 2 : 1) + name##_RB_RANK(right); \
+	if (left_rank != right_rank ||					\
+	    (left_rank == 2 && left == NULL && right == NULL))		\
+		return (-1);						\
+	return (left_rank);						\
+}
+#else
+#define RB_GENERATE_RANK(name, type, field, attr)
+#endif
+
 #define RB_GENERATE_INSERT_COLOR(name, type, field, attr)		\
 attr void								\
 name##_RB_INSERT_COLOR(struct name *head, struct type *elm)		\
diff --git a/tests/sys/sys/rb_test.c b/tests/sys/sys/rb_test.c
index c558ad6098cf..149fbe052bde 100644
--- a/tests/sys/sys/rb_test.c
+++ b/tests/sys/sys/rb_test.c
@@ -29,6 +29,7 @@
  */
 #include <sys/types.h>
 
+#define _RB_DIAGNOSTIC 1
 #include <sys/tree.h>
 #include <stdlib.h>
 
@@ -54,52 +55,54 @@ RB_PROTOTYPE(tree, node, node, compare);
 RB_GENERATE(tree, node, node, compare);
 
 #define ITER 150
-#define MIN 5
-#define MAX 5000
 
 ATF_TC_WITHOUT_HEAD(rb_test);
 ATF_TC_BODY(rb_test, tc)
 {
-	struct node *tmp, *ins;
-	int i, max, min;
+	struct node *tmp, *ins, store[ITER];
+	int i, j, k, max, min;
 
-	max = min = 42; /* pacify gcc */
+	min = ITER;
+	max = -1;
 
 	RB_INIT(&root);
 
+	/* Initialize keys */
+	for (i = 0; i < ITER; i++)
+		store[i].key = i;
+
+	/* Randomly shuffle keys */
+	for (i = 0; i < ITER; i++) {
+		j = i + arc4random_uniform(ITER - i);
+		k = store[j].key;
+		store[j].key = store[i].key;
+		store[i].key = k;
+	}
+
 	for (i = 0; i < ITER; i++) {
-		tmp = malloc(sizeof(struct node));
-		ATF_REQUIRE_MSG(tmp != NULL, "malloc failed");
-		do {
-			tmp->key = arc4random_uniform(MAX-MIN);
-			tmp->key += MIN;
-		} while (RB_FIND(tree, &root, tmp) != NULL);
-		if (i == 0)
-			max = min = tmp->key;
-		else {
-			if (tmp->key > max)
-				max = tmp->key;
-			if (tmp->key < min)
-				min = tmp->key;
+		for (j = 0; j < i; ++j) {
+			tmp = &store[j];
+			ATF_REQUIRE_EQ(tmp, RB_FIND(tree, &root, tmp));
 		}
+		tmp = &store[i];
+		if (tmp->key > max)
+			max = tmp->key;
+		if (tmp->key < min)
+			min = tmp->key;
 		ATF_REQUIRE_EQ(NULL, RB_INSERT(tree, &root, tmp));
+		ins = RB_MIN(tree, &root);
+		ATF_REQUIRE_MSG(ins != NULL, "RB_MIN error");
+		ATF_CHECK_EQ(min, ins->key);
+		ins = RB_MAX(tree, &root);
+		ATF_REQUIRE_MSG(ins != NULL, "RB_MAX error");
+		ATF_CHECK_EQ(max, ins->key);
 	}
-
-	ins = RB_MIN(tree, &root);
-	ATF_REQUIRE_MSG(ins != NULL, "RB_MIN error");
-	ATF_CHECK_EQ(min, ins->key);
-	tmp = ins;
-	ins = RB_MAX(tree, &root);
-	ATF_REQUIRE_MSG(ins != NULL, "RB_MAX error");
-	ATF_CHECK_EQ(max, ins->key);
-
-	ATF_CHECK_EQ(tmp, RB_REMOVE(tree, &root, tmp));
-
-	for (i = 0; i < ITER - 1; i++) {
+	tmp = RB_ROOT(&root);
+	ATF_REQUIRE_MSG(tree_RB_RANK(tmp) >= 0, "RB rank balance error");
+	for (i = 0; i < ITER; i++) {
 		tmp = RB_ROOT(&root);
 		ATF_REQUIRE_MSG(tmp != NULL, "RB_ROOT error");
 		ATF_CHECK_EQ(tmp, RB_REMOVE(tree, &root, tmp));
-		free(tmp);
 	}
 }