lib/interval_tree: add test case for span iteration

Verify interval_tree_span_iter_xxx() helpers works as expected.

Link: https://lkml.kernel.org/r/20250310074938.26756-6-richard.weiyang@gmail.com
Signed-off-by: Wei Yang <richard.weiyang@gmail.com>
Cc: Matthew Wilcox <willy@infradead.org>
Cc: Michel Lespinasse <michel@lespinasse.org>
Cc: Jason Gunthorpe <jgg@nvidia.com>
Signed-off-by: Andrew Morton <akpm@linux-foundation.org>

diff --git a/lib/interval_tree_test.c b/lib/interval_tree_test.c
index 7821379e2c21746f2d17ddb964c48d8228732c68..5fd62656f42e33e9831839875d5dd66054a82c9d 100644 (file)
--- a/lib/interval_tree_test.c
+++ b/lib/interval_tree_test.c
@@ -6,6 +6,7 @@
 #include <linux/slab.h>
 #include <asm/timex.h>
 #include <linux/bitmap.h>
+#include <linux/maple_tree.h>
 
 #define __param(type, name, init, msg)         \
        static type name = init;                \
@@ -193,6 +194,121 @@ static int intersection_range_check(void)
        return 0;
 }
 
+#ifdef CONFIG_INTERVAL_TREE_SPAN_ITER
+/*
+ * Helper function to get span of current position from maple tree point of
+ * view.
+ */
+static void mas_cur_span(struct ma_state *mas, struct interval_tree_span_iter *state)
+{
+       unsigned long cur_start;
+       unsigned long cur_last;
+       int is_hole;
+
+       if (mas->status == ma_overflow)
+               return;
+
+       /* walk to current position */
+       state->is_hole = mas_walk(mas) ? 0 : 1;
+
+       cur_start = mas->index < state->first_index ?
+                       state->first_index : mas->index;
+
+       /* whether we have followers */
+       do {
+
+               cur_last = mas->last > state->last_index ?
+                               state->last_index : mas->last;
+
+               is_hole = mas_next_range(mas, state->last_index) ? 0 : 1;
+
+       } while (mas->status != ma_overflow && is_hole == state->is_hole);
+
+       if (state->is_hole) {
+               state->start_hole = cur_start;
+               state->last_hole = cur_last;
+       } else {
+               state->start_used = cur_start;
+               state->last_used = cur_last;
+       }
+
+       /* advance position for next round */
+       if (mas->status != ma_overflow)
+               mas_set(mas, cur_last + 1);
+}
+
+static int span_iteration_check(void)
+{
+       int i, j, k;
+       unsigned long start, last;
+       struct interval_tree_span_iter span, mas_span;
+
+       DEFINE_MTREE(tree);
+
+       MA_STATE(mas, &tree, 0, 0);
+
+       printk(KERN_ALERT "interval tree span iteration\n");
+
+       for (i = 0; i < search_loops; i++) {
+               /* Initialize interval tree for each round */
+               init();
+               for (j = 0; j < nnodes; j++)
+                       interval_tree_insert(nodes + j, &root);
+
+               /* Put all the range into maple tree */
+               mt_init_flags(&tree, MT_FLAGS_ALLOC_RANGE);
+               mt_set_in_rcu(&tree);
+
+               for (j = 0; j < nnodes; j++)
+                       WARN_ON_ONCE(mtree_store_range(&tree, nodes[j].start,
+                                       nodes[j].last, nodes + j, GFP_KERNEL));
+
+               /* Let's try nsearches different ranges */
+               for (k = 0; k < nsearches; k++) {
+                       /* Try whole range once */
+                       if (!k) {
+                               start = 0UL;
+                               last = ULONG_MAX;
+                       } else {
+                               last = (prandom_u32_state(&rnd) >> 4) % max_endpoint;
+                               start = (prandom_u32_state(&rnd) >> 4) % last;
+                       }
+
+                       mas_span.first_index = start;
+                       mas_span.last_index = last;
+                       mas_span.is_hole = -1;
+                       mas_set(&mas, start);
+
+                       interval_tree_for_each_span(&span, &root, start, last) {
+                               mas_cur_span(&mas, &mas_span);
+
+                               WARN_ON_ONCE(span.is_hole != mas_span.is_hole);
+
+                               if (span.is_hole) {
+                                       WARN_ON_ONCE(span.start_hole != mas_span.start_hole);
+                                       WARN_ON_ONCE(span.last_hole != mas_span.last_hole);
+                               } else {
+                                       WARN_ON_ONCE(span.start_used != mas_span.start_used);
+                                       WARN_ON_ONCE(span.last_used != mas_span.last_used);
+                               }
+                       }
+
+               }
+
+               WARN_ON_ONCE(mas.status != ma_overflow);
+
+               /* Cleanup maple tree for each round */
+               mtree_destroy(&tree);
+               /* Cleanup interval tree for each round */
+               for (j = 0; j < nnodes; j++)
+                       interval_tree_remove(nodes + j, &root);
+       }
+       return 0;
+}
+#else
+static inline int span_iteration_check(void) {return 0; }
+#endif
+
 static int interval_tree_test_init(void)
 {
        nodes = kmalloc_array(nnodes, sizeof(struct interval_tree_node),
@@ -211,6 +327,7 @@ static int interval_tree_test_init(void)
        basic_check();
        search_check();
        intersection_range_check();
+       span_iteration_check();
 
        kfree(queries);
        kfree(nodes);

diff --git a/tools/testing/rbtree/Makefile b/tools/testing/rbtree/Makefile
index bac6931b499d0d16cf3d6d51d734aa09bfee784f..d7bbae2af4c7803304e806294ce9c72ebd3b8ecc 100644 (file)
--- a/tools/testing/rbtree/Makefile
+++ b/tools/testing/rbtree/Makefile
@@ -3,7 +3,7 @@
 .PHONY: clean
 
 TARGETS = rbtree_test interval_tree_test
-OFILES = $(LIBS) rbtree-shim.o interval_tree-shim.o
+OFILES = $(SHARED_OFILES) rbtree-shim.o interval_tree-shim.o maple-shim.o
 DEPS = ../../../include/linux/rbtree.h \
        ../../../include/linux/rbtree_types.h \
        ../../../include/linux/rbtree_augmented.h \
@@ -25,7 +25,9 @@ $(TARGETS):   $(OFILES)
 rbtree-shim.o: $(DEPS)
 rbtree_test.o:  ../../../lib/rbtree_test.c
 interval_tree-shim.o: $(DEPS)
+interval_tree-shim.o: CFLAGS += -DCONFIG_INTERVAL_TREE_SPAN_ITER
 interval_tree_test.o:  ../../../lib/interval_tree_test.c
+interval_tree_test.o: CFLAGS += -DCONFIG_INTERVAL_TREE_SPAN_ITER
 
 clean:
-       $(RM) $(TARGETS) *.o generated/*
+       $(RM) $(TARGETS) *.o radix-tree.c idr.c generated/*

diff --git a/tools/testing/rbtree/interval_tree_test.c b/tools/testing/rbtree/interval_tree_test.c
index 63775b831c1c4faa328dfe8d1c03926c6cc37de9..49bc5b534330903bbbcb1735bb86e25350510add 100644 (file)
--- a/tools/testing/rbtree/interval_tree_test.c
+++ b/tools/testing/rbtree/interval_tree_test.c
@@ -6,6 +6,7 @@
 #include <linux/math64.h>
 #include <linux/kern_levels.h>
 #include "shared.h"
+#include "maple-shared.h"
 
 #include "../../../lib/interval_tree_test.c"
 
@@ -51,6 +52,7 @@ int main(int argc, char **argv)
                        usage();
        }
 
+       maple_tree_init();
        interval_tree_tests();
        return 0;
 }