[PATCH bpf-next 3/3] selftests/bpf: libarena: Add rb_next()

[Changwoo Min <changwoo@igalia.com>]

Add rb_next(), the in-order successor through parent pointers.

Add a selftest that walks a tree with repeated keys and checks the walk
against the order in which the head leaves the tree.

Signed-off-by: Changwoo Min <changwoo@igalia.com>
---
 .../bpf/libarena/include/libarena/rbtree.h    | 23 ++++++++
 .../bpf/libarena/selftests/test_rbtree.bpf.c  | 54 +++++++++++++++++++
 2 files changed, 77 insertions(+)

diff --git a/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h b/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h
index a8dd8c2e7596..956b4a24b52a 100644
--- a/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h
+++ b/tools/testing/selftests/bpf/libarena/include/libarena/rbtree.h
@@ -103,6 +103,29 @@ void rb_node_free(struct rbnode __arena *rbnode);
 
 int rb_integrity_check(struct rbtree __arena *rbtree);
 
+/*
+ * Return the in-order successor of @node, or NULL if @node is the last one.
+ * Walking from rb_first() with it visits the tree in key order, equal keys in
+ * insertion order. Call it under the writer's lock.
+ */
+static __always_inline struct rbnode __arena *rb_next(struct rbnode __arena *node)
+{
+	struct rbnode __arena *parent;
+
+	if (node->right) {
+		node = node->right;
+		while (node->left && can_loop)
+			node = node->left;
+		return node;
+	}
+
+	/* Climb until we leave a left subtree; its parent comes next. */
+	while ((parent = node->parent) && node == parent->right && can_loop)
+		node = parent;
+
+	return parent;
+}
+
 /*
  * Lockless reads of the head and the node count. They need not be called
  * under the writer's lock, so they are advisory: the tree may change before
diff --git a/tools/testing/selftests/bpf/libarena/selftests/test_rbtree.bpf.c b/tools/testing/selftests/bpf/libarena/selftests/test_rbtree.bpf.c
index 57c4395c046a..784cc9edea17 100644
--- a/tools/testing/selftests/bpf/libarena/selftests/test_rbtree.bpf.c
+++ b/tools/testing/selftests/bpf/libarena/selftests/test_rbtree.bpf.c
@@ -425,6 +425,60 @@ __weak int test_rbtree_first_update(void)
 	return rb_destroy(rbtree);
 }
 
+#define NEXT_NR_NODES	64
+
+static struct rbnode __arena *next_nodes[NEXT_NR_NODES];
+
+SEC("syscall")
+__weak int test_rbtree_next(void)
+{
+	struct rbtree __arena *rbtree;
+	struct rbnode __arena *node;
+	u64 seed = 0xcafe;
+	int i;
+
+	rbtree = rb_create(RB_NOALLOC, RB_DUPLICATE);
+	if (!rbtree)
+		return 1;
+
+	/* Pseudo-random keys with repeats, so the walk crosses equal keys. */
+	for (i = 0; i < NEXT_NR_NODES && can_loop; i++) {
+		/* LCG using Knuth's MMIX constants for determinism. */
+		seed = seed * 6364136223846793005ULL + 1442695040888963407ULL;
+		node = rb_node_alloc((seed >> 33) % 16, i);
+		if (!node)
+			return 2;
+
+		if (rb_insert_node(rbtree, node))
+			return 3;
+	}
+
+	node = rb_first(rbtree);
+	for (i = 0; i < NEXT_NR_NODES && can_loop; i++) {
+		if (!node)
+			return 4;
+
+		next_nodes[i] = node;
+		node = rb_next(node);
+	}
+
+	if (node)
+		return 5;
+
+	for (i = 0; i < NEXT_NR_NODES && can_loop; i++) {
+		node = rb_first(rbtree);
+		if (node != next_nodes[i])
+			return 6;
+
+		if (rb_remove_node(rbtree, node))
+			return 7;
+
+		rb_node_free(node);
+	}
+
+	return rb_destroy(rbtree);
+}
+
 struct rb_embed {
 	u64 pad;
 	struct rbtree tree;
-- 
2.56.0