lib/rbtree_test.c

   1 #include <linux/module.h>
   2 #include <linux/rbtree.h>
   3 #include <linux/random.h>
   4 #include <asm/timex.h>
   5
   6 #define NODES       100
   7 #define PERF_LOOPS  100000
   8 #define CHECK_LOOPS 100
   9
  10 struct test_node {
  11         struct rb_node rb;
  12         u32 key;
  13 };
  14
  15 static struct rb_root root = RB_ROOT;
  16 static struct test_node nodes[NODES];
  17
  18 static struct rnd_state rnd;
  19
  20 static void insert(struct test_node *node, struct rb_root *root)
  21 {
  22         struct rb_node **new = &root->rb_node, *parent = NULL;
  23
  24         while (*new) {
  25                 parent = *new;
  26                 if (node->key < rb_entry(parent, struct test_node, rb)->key)
  27                         new = &parent->rb_left;
  28                 else
  29                         new = &parent->rb_right;
  30         }
  31
  32         rb_link_node(&node->rb, parent, new);
  33         rb_insert_color(&node->rb, root);
  34 }
  35
  36 static inline void erase(struct test_node *node, struct rb_root *root)
  37 {
  38         rb_erase(&node->rb, root);
  39 }
  40
  41 static void init(void)
  42 {
  43         int i;
  44         for (i = 0; i < NODES; i++)
  45                 nodes[i].key = prandom32(&rnd);
  46 }
  47
  48 static bool is_red(struct rb_node *rb)
  49 {
  50         return !(rb->__rb_parent_color & 1);
  51 }
  52
  53 static int black_path_count(struct rb_node *rb)
  54 {
  55         int count;
  56         for (count = 0; rb; rb = rb_parent(rb))
  57                 count += !is_red(rb);
  58         return count;
  59 }
  60
  61 static void check(int nr_nodes)
  62 {
  63         struct rb_node *rb;
  64         int count = 0;
  65         int blacks;
  66         u32 prev_key = 0;
  67
  68         for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
  69                 struct test_node *node = rb_entry(rb, struct test_node, rb);
  70                 WARN_ON_ONCE(node->key < prev_key);
  71                 WARN_ON_ONCE(is_red(rb) &&
  72                              (!rb_parent(rb) || is_red(rb_parent(rb))));
  73                 if (!count)
  74                         blacks = black_path_count(rb);
  75                 else
  76                         WARN_ON_ONCE((!rb->rb_left || !rb->rb_right) &&
  77                                      blacks != black_path_count(rb));
  78                 prev_key = node->key;
  79                 count++;
  80         }
  81         WARN_ON_ONCE(count != nr_nodes);
  82 }
  83
  84 static int rbtree_test_init(void)
  85 {
  86         int i, j;
  87         cycles_t time1, time2, time;
  88
  89         printk(KERN_ALERT "rbtree testing");
  90
  91         prandom32_seed(&rnd, 3141592653589793238);
  92         init();
  93
  94         time1 = get_cycles();
  95
  96         for (i = 0; i < PERF_LOOPS; i++) {
  97                 for (j = 0; j < NODES; j++)
  98                         insert(nodes + j, &root);
  99                 for (j = 0; j < NODES; j++)
 100                         erase(nodes + j, &root);
 101         }
 102
 103         time2 = get_cycles();
 104         time = time2 - time1;
 105
 106         time = div_u64(time, PERF_LOOPS);
 107         printk(" -> %llu cycles\n", (unsigned long long)time);
 108
 109         for (i = 0; i < CHECK_LOOPS; i++) {
 110                 init();
 111                 for (j = 0; j < NODES; j++) {
 112                         check(j);
 113                         insert(nodes + j, &root);
 114                 }
 115                 for (j = 0; j < NODES; j++) {
 116                         check(NODES - j);
 117                         erase(nodes + j, &root);
 118                 }
 119                 check(0);
 120         }
 121
 122         return -EAGAIN; /* Fail will directly unload the module */
 123 }
 124
 125 static void rbtree_test_exit(void)
 126 {
 127         printk(KERN_ALERT "test exit\n");
 128 }
 129
 130 module_init(rbtree_test_init)
 131 module_exit(rbtree_test_exit)
 132
 133 MODULE_LICENSE("GPL");
 134 MODULE_AUTHOR("Michel Lespinasse");
 135 MODULE_DESCRIPTION("Red Black Tree test");