lib/interval_tree: add test case for interval_tree_iter_xxx() helpers

Verify interval_tree_iter_xxx() helpers could find intersection ranges
as expected.

[sfr@canb.auug.org.au: some of tools/ uses -Wno-unused-parameter]
Link: https://lkml.kernel.org/r/20250312113612.31ac808e@canb.auug.org.au
Link: https://lkml.kernel.org/r/20250310074938.26756-5-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 51863077d4ecf7633a649da6140d507ec14661cd..7821379e2c21746f2d17ddb964c48d8228732c68 100644 (file)
--- a/lib/interval_tree_test.c
+++ b/lib/interval_tree_test.c
@@ -5,6 +5,7 @@
 #include <linux/prandom.h>
 #include <linux/slab.h>
 #include <asm/timex.h>
+#include <linux/bitmap.h>
 
 #define __param(type, name, init, msg)         \
        static type name = init;                \
@@ -125,6 +126,73 @@ static int search_check(void)
        return 0;
 }
 
+static int intersection_range_check(void)
+{
+       int i, j, k;
+       unsigned long start, last;
+       struct interval_tree_node *node;
+       unsigned long *intxn1;
+       unsigned long *intxn2;
+
+       printk(KERN_ALERT "interval tree iteration\n");
+
+       intxn1 = bitmap_alloc(nnodes, GFP_KERNEL);
+       if (!intxn1) {
+               WARN_ON_ONCE("Failed to allocate intxn1\n");
+               return -ENOMEM;
+       }
+
+       intxn2 = bitmap_alloc(nnodes, GFP_KERNEL);
+       if (!intxn2) {
+               WARN_ON_ONCE("Failed to allocate intxn2\n");
+               bitmap_free(intxn1);
+               return -ENOMEM;
+       }
+
+       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);
+
+               /* 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;
+                       }
+
+                       /* Walk nodes to mark intersection nodes */
+                       bitmap_zero(intxn1, nnodes);
+                       for (j = 0; j < nnodes; j++) {
+                               node = nodes + j;
+
+                               if (start <= node->last && last >= node->start)
+                                       bitmap_set(intxn1, j, 1);
+                       }
+
+                       /* Iterate tree to clear intersection nodes */
+                       bitmap_zero(intxn2, nnodes);
+                       for (node = interval_tree_iter_first(&root, start, last); node;
+                            node = interval_tree_iter_next(node, start, last))
+                               bitmap_set(intxn2, node - nodes, 1);
+
+                       WARN_ON_ONCE(!bitmap_equal(intxn1, intxn2, nnodes));
+               }
+
+               for (j = 0; j < nnodes; j++)
+                       interval_tree_remove(nodes + j, &root);
+       }
+
+       bitmap_free(intxn1);
+       bitmap_free(intxn2);
+       return 0;
+}
+
 static int interval_tree_test_init(void)
 {
        nodes = kmalloc_array(nnodes, sizeof(struct interval_tree_node),
@@ -142,6 +210,7 @@ static int interval_tree_test_init(void)
 
        basic_check();
        search_check();
+       intersection_range_check();
 
        kfree(queries);
        kfree(nodes);

diff --git a/tools/include/linux/bitmap.h b/tools/include/linux/bitmap.h
index 2a7f260ef9dc5395aa0f5d885d325fa48f857289..d4d300040d019c8b2fd5ca553f1fc22eac600dd4 100644 (file)
--- a/tools/include/linux/bitmap.h
+++ b/tools/include/linux/bitmap.h
@@ -19,6 +19,7 @@ bool __bitmap_and(unsigned long *dst, const unsigned long *bitmap1,
                 const unsigned long *bitmap2, unsigned int bits);
 bool __bitmap_equal(const unsigned long *bitmap1,
                    const unsigned long *bitmap2, unsigned int bits);
+void __bitmap_set(unsigned long *map, unsigned int start, int len);
 void __bitmap_clear(unsigned long *map, unsigned int start, int len);
 bool __bitmap_intersects(const unsigned long *bitmap1,
                         const unsigned long *bitmap2, unsigned int bits);
@@ -79,6 +80,11 @@ static inline void bitmap_or(unsigned long *dst, const unsigned long *src1,
                __bitmap_or(dst, src1, src2, nbits);
 }
 
+static inline unsigned long *bitmap_alloc(unsigned int nbits, gfp_t flags __maybe_unused)
+{
+       return malloc(bitmap_size(nbits));
+}
+
 /**
  * bitmap_zalloc - Allocate bitmap
  * @nbits: Number of bits
@@ -150,6 +156,21 @@ static inline bool bitmap_intersects(const unsigned long *src1,
                return __bitmap_intersects(src1, src2, nbits);
 }
 
+static inline void bitmap_set(unsigned long *map, unsigned int start, unsigned int nbits)
+{
+       if (__builtin_constant_p(nbits) && nbits == 1)
+               __set_bit(start, map);
+       else if (small_const_nbits(start + nbits))
+               *map |= GENMASK(start + nbits - 1, start);
+       else if (__builtin_constant_p(start & BITMAP_MEM_MASK) &&
+                IS_ALIGNED(start, BITMAP_MEM_ALIGNMENT) &&
+                __builtin_constant_p(nbits & BITMAP_MEM_MASK) &&
+                IS_ALIGNED(nbits, BITMAP_MEM_ALIGNMENT))
+               memset((char *)map + start / 8, 0xff, nbits / 8);
+       else
+               __bitmap_set(map, start, nbits);
+}
+
 static inline void bitmap_clear(unsigned long *map, unsigned int start,
                               unsigned int nbits)
 {

diff --git a/tools/lib/bitmap.c b/tools/lib/bitmap.c
index 2178862bb114cbad1da670be63ad1fdb80836c0c..51255c69754d3aee414a4cb11a3a0ec99ea35e3d 100644 (file)
--- a/tools/lib/bitmap.c
+++ b/tools/lib/bitmap.c
@@ -101,6 +101,26 @@ bool __bitmap_intersects(const unsigned long *bitmap1,
        return false;
 }
 
+void __bitmap_set(unsigned long *map, unsigned int start, int len)
+{
+       unsigned long *p = map + BIT_WORD(start);
+       const unsigned int size = start + len;
+       int bits_to_set = BITS_PER_LONG - (start % BITS_PER_LONG);
+       unsigned long mask_to_set = BITMAP_FIRST_WORD_MASK(start);
+
+       while (len - bits_to_set >= 0) {
+               *p |= mask_to_set;
+               len -= bits_to_set;
+               bits_to_set = BITS_PER_LONG;
+               mask_to_set = ~0UL;
+               p++;
+       }
+       if (len) {
+               mask_to_set &= BITMAP_LAST_WORD_MASK(size);
+               *p |= mask_to_set;
+       }
+}
+
 void __bitmap_clear(unsigned long *map, unsigned int start, int len)
 {
        unsigned long *p = map + BIT_WORD(start);