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         /* following fields used for testing augmented rbtree functionality */
  15         u32 val;
  16         u32 augmented;
  17 };
  18
  19 static struct rb_root root = RB_ROOT;
  20 static struct test_node nodes[NODES];
  21
  22 static struct rnd_state rnd;
  23
  24 static void insert(struct test_node *node, struct rb_root *root)
  25 {
  26         struct rb_node **new = &root->rb_node, *parent = NULL;
  27         u32 key = node->key;
  28
  29         while (*new) {
  30                 parent = *new;
  31                 if (key < rb_entry(parent, struct test_node, rb)->key)
  32                         new = &parent->rb_left;
  33                 else
  34                         new = &parent->rb_right;
  35         }
  36
  37         rb_link_node(&node->rb, parent, new);
  38         rb_insert_color(&node->rb, root);
  39 }
  40
  41 static inline void erase(struct test_node *node, struct rb_root *root)
  42 {
  43         rb_erase(&node->rb, root);
  44 }
  45
  46 static inline u32 augment_recompute(struct test_node *node)
  47 {
  48         u32 max = node->val, child_augmented;
  49         if (node->rb.rb_left) {
  50                 child_augmented = rb_entry(node->rb.rb_left, struct test_node,
  51                                            rb)->augmented;
  52                 if (max < child_augmented)
  53                         max = child_augmented;
  54         }
  55         if (node->rb.rb_right) {
  56                 child_augmented = rb_entry(node->rb.rb_right, struct test_node,
  57                                            rb)->augmented;
  58                 if (max < child_augmented)
  59                         max = child_augmented;
  60         }
  61         return max;
  62 }
  63
  64 static void augment_propagate(struct rb_node *rb, struct rb_node *stop)
  65 {
  66         while (rb != stop) {
  67                 struct test_node *node = rb_entry(rb, struct test_node, rb);
  68                 u32 augmented = augment_recompute(node);
  69                 if (node->augmented == augmented)
  70                         break;
  71                 node->augmented = augmented;
  72                 rb = rb_parent(&node->rb);
  73         }
  74 }
  75
  76 static void augment_copy(struct rb_node *rb_old, struct rb_node *rb_new)
  77 {
  78         struct test_node *old = rb_entry(rb_old, struct test_node, rb);
  79         struct test_node *new = rb_entry(rb_new, struct test_node, rb);
  80         new->augmented = old->augmented;
  81 }
  82
  83 static void augment_rotate(struct rb_node *rb_old, struct rb_node *rb_new)
  84 {
  85         struct test_node *old = rb_entry(rb_old, struct test_node, rb);
  86         struct test_node *new = rb_entry(rb_new, struct test_node, rb);
  87
  88         /* Rotation doesn't change subtree's augmented value */
  89         new->augmented = old->augmented;
  90         old->augmented = augment_recompute(old);
  91 }
  92
  93 static const struct rb_augment_callbacks augment_callbacks = {
  94         augment_propagate, augment_copy, augment_rotate
  95 };
  96
  97 static void insert_augmented(struct test_node *node, struct rb_root *root)
  98 {
  99         struct rb_node **new = &root->rb_node, *rb_parent = NULL;
 100         u32 key = node->key;
 101         u32 val = node->val;
 102         struct test_node *parent;
 103
 104         while (*new) {
 105                 rb_parent = *new;
 106                 parent = rb_entry(rb_parent, struct test_node, rb);
 107                 if (parent->augmented < val)
 108                         parent->augmented = val;
 109                 if (key < parent->key)
 110                         new = &parent->rb.rb_left;
 111                 else
 112                         new = &parent->rb.rb_right;
 113         }
 114
 115         node->augmented = val;
 116         rb_link_node(&node->rb, rb_parent, new);
 117         rb_insert_augmented(&node->rb, root, &augment_callbacks);
 118 }
 119
 120 static void erase_augmented(struct test_node *node, struct rb_root *root)
 121 {
 122         rb_erase_augmented(&node->rb, root, &augment_callbacks);
 123 }
 124
 125 static void init(void)
 126 {
 127         int i;
 128         for (i = 0; i < NODES; i++) {
 129                 nodes[i].key = prandom32(&rnd);
 130                 nodes[i].val = prandom32(&rnd);
 131         }
 132 }
 133
 134 static bool is_red(struct rb_node *rb)
 135 {
 136         return !(rb->__rb_parent_color & 1);
 137 }
 138
 139 static int black_path_count(struct rb_node *rb)
 140 {
 141         int count;
 142         for (count = 0; rb; rb = rb_parent(rb))
 143                 count += !is_red(rb);
 144         return count;
 145 }
 146
 147 static void check(int nr_nodes)
 148 {
 149         struct rb_node *rb;
 150         int count = 0;
 151         int blacks;
 152         u32 prev_key = 0;
 153
 154         for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
 155                 struct test_node *node = rb_entry(rb, struct test_node, rb);
 156                 WARN_ON_ONCE(node->key < prev_key);
 157                 WARN_ON_ONCE(is_red(rb) &&
 158                              (!rb_parent(rb) || is_red(rb_parent(rb))));
 159                 if (!count)
 160                         blacks = black_path_count(rb);
 161                 else
 162                         WARN_ON_ONCE((!rb->rb_left || !rb->rb_right) &&
 163                                      blacks != black_path_count(rb));
 164                 prev_key = node->key;
 165                 count++;
 166         }
 167         WARN_ON_ONCE(count != nr_nodes);
 168 }
 169
 170 static void check_augmented(int nr_nodes)
 171 {
 172         struct rb_node *rb;
 173
 174         check(nr_nodes);
 175         for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
 176                 struct test_node *node = rb_entry(rb, struct test_node, rb);
 177                 WARN_ON_ONCE(node->augmented != augment_recompute(node));
 178         }
 179 }
 180
 181 static int rbtree_test_init(void)
 182 {
 183         int i, j;
 184         cycles_t time1, time2, time;
 185
 186         printk(KERN_ALERT "rbtree testing");
 187
 188         prandom32_seed(&rnd, 3141592653589793238ULL);
 189         init();
 190
 191         time1 = get_cycles();
 192
 193         for (i = 0; i < PERF_LOOPS; i++) {
 194                 for (j = 0; j < NODES; j++)
 195                         insert(nodes + j, &root);
 196                 for (j = 0; j < NODES; j++)
 197                         erase(nodes + j, &root);
 198         }
 199
 200         time2 = get_cycles();
 201         time = time2 - time1;
 202
 203         time = div_u64(time, PERF_LOOPS);
 204         printk(" -> %llu cycles\n", (unsigned long long)time);
 205
 206         for (i = 0; i < CHECK_LOOPS; i++) {
 207                 init();
 208                 for (j = 0; j < NODES; j++) {
 209                         check(j);
 210                         insert(nodes + j, &root);
 211                 }
 212                 for (j = 0; j < NODES; j++) {
 213                         check(NODES - j);
 214                         erase(nodes + j, &root);
 215                 }
 216                 check(0);
 217         }
 218
 219         printk(KERN_ALERT "augmented rbtree testing");
 220
 221         init();
 222
 223         time1 = get_cycles();
 224
 225         for (i = 0; i < PERF_LOOPS; i++) {
 226                 for (j = 0; j < NODES; j++)
 227                         insert_augmented(nodes + j, &root);
 228                 for (j = 0; j < NODES; j++)
 229                         erase_augmented(nodes + j, &root);
 230         }
 231
 232         time2 = get_cycles();
 233         time = time2 - time1;
 234
 235         time = div_u64(time, PERF_LOOPS);
 236         printk(" -> %llu cycles\n", (unsigned long long)time);
 237
 238         for (i = 0; i < CHECK_LOOPS; i++) {
 239                 init();
 240                 for (j = 0; j < NODES; j++) {
 241                         check_augmented(j);
 242                         insert_augmented(nodes + j, &root);
 243                 }
 244                 for (j = 0; j < NODES; j++) {
 245                         check_augmented(NODES - j);
 246                         erase_augmented(nodes + j, &root);
 247                 }
 248                 check_augmented(0);
 249         }
 250
 251         return -EAGAIN; /* Fail will directly unload the module */
 252 }
 253
 254 static void rbtree_test_exit(void)
 255 {
 256         printk(KERN_ALERT "test exit\n");
 257 }
 258
 259 module_init(rbtree_test_init)
 260 module_exit(rbtree_test_exit)
 261
 262 MODULE_LICENSE("GPL");
 263 MODULE_AUTHOR("Michel Lespinasse");
 264 MODULE_DESCRIPTION("Red Black Tree test");