linux/lib/rbtree_test.c
<<
>>
Prefs
   1#include <linux/module.h>
   2#include <linux/rbtree_augmented.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
  10struct 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
  19static struct rb_root root = RB_ROOT;
  20static struct test_node nodes[NODES];
  21
  22static struct rnd_state rnd;
  23
  24static 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
  41static inline void erase(struct test_node *node, struct rb_root *root)
  42{
  43        rb_erase(&node->rb, root);
  44}
  45
  46static 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
  64RB_DECLARE_CALLBACKS(static, augment_callbacks, struct test_node, rb,
  65                     u32, augmented, augment_recompute)
  66
  67static void insert_augmented(struct test_node *node, struct rb_root *root)
  68{
  69        struct rb_node **new = &root->rb_node, *rb_parent = NULL;
  70        u32 key = node->key;
  71        u32 val = node->val;
  72        struct test_node *parent;
  73
  74        while (*new) {
  75                rb_parent = *new;
  76                parent = rb_entry(rb_parent, struct test_node, rb);
  77                if (parent->augmented < val)
  78                        parent->augmented = val;
  79                if (key < parent->key)
  80                        new = &parent->rb.rb_left;
  81                else
  82                        new = &parent->rb.rb_right;
  83        }
  84
  85        node->augmented = val;
  86        rb_link_node(&node->rb, rb_parent, new);
  87        rb_insert_augmented(&node->rb, root, &augment_callbacks);
  88}
  89
  90static void erase_augmented(struct test_node *node, struct rb_root *root)
  91{
  92        rb_erase_augmented(&node->rb, root, &augment_callbacks);
  93}
  94
  95static void init(void)
  96{
  97        int i;
  98        for (i = 0; i < NODES; i++) {
  99                nodes[i].key = prandom_u32_state(&rnd);
 100                nodes[i].val = prandom_u32_state(&rnd);
 101        }
 102}
 103
 104static bool is_red(struct rb_node *rb)
 105{
 106        return !(rb->__rb_parent_color & 1);
 107}
 108
 109static int black_path_count(struct rb_node *rb)
 110{
 111        int count;
 112        for (count = 0; rb; rb = rb_parent(rb))
 113                count += !is_red(rb);
 114        return count;
 115}
 116
 117static void check(int nr_nodes)
 118{
 119        struct rb_node *rb;
 120        int count = 0, blacks = 0;
 121        u32 prev_key = 0;
 122
 123        for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
 124                struct test_node *node = rb_entry(rb, struct test_node, rb);
 125                WARN_ON_ONCE(node->key < prev_key);
 126                WARN_ON_ONCE(is_red(rb) &&
 127                             (!rb_parent(rb) || is_red(rb_parent(rb))));
 128                if (!count)
 129                        blacks = black_path_count(rb);
 130                else
 131                        WARN_ON_ONCE((!rb->rb_left || !rb->rb_right) &&
 132                                     blacks != black_path_count(rb));
 133                prev_key = node->key;
 134                count++;
 135        }
 136
 137        WARN_ON_ONCE(count != nr_nodes);
 138        WARN_ON_ONCE(count < (1 << black_path_count(rb_last(&root))) - 1);
 139}
 140
 141static void check_augmented(int nr_nodes)
 142{
 143        struct rb_node *rb;
 144
 145        check(nr_nodes);
 146        for (rb = rb_first(&root); rb; rb = rb_next(rb)) {
 147                struct test_node *node = rb_entry(rb, struct test_node, rb);
 148                WARN_ON_ONCE(node->augmented != augment_recompute(node));
 149        }
 150}
 151
 152static int __init rbtree_test_init(void)
 153{
 154        int i, j;
 155        cycles_t time1, time2, time;
 156
 157        printk(KERN_ALERT "rbtree testing");
 158
 159        prandom_seed_state(&rnd, 3141592653589793238ULL);
 160        init();
 161
 162        time1 = get_cycles();
 163
 164        for (i = 0; i < PERF_LOOPS; i++) {
 165                for (j = 0; j < NODES; j++)
 166                        insert(nodes + j, &root);
 167                for (j = 0; j < NODES; j++)
 168                        erase(nodes + j, &root);
 169        }
 170
 171        time2 = get_cycles();
 172        time = time2 - time1;
 173
 174        time = div_u64(time, PERF_LOOPS);
 175        printk(" -> %llu cycles\n", (unsigned long long)time);
 176
 177        for (i = 0; i < CHECK_LOOPS; i++) {
 178                init();
 179                for (j = 0; j < NODES; j++) {
 180                        check(j);
 181                        insert(nodes + j, &root);
 182                }
 183                for (j = 0; j < NODES; j++) {
 184                        check(NODES - j);
 185                        erase(nodes + j, &root);
 186                }
 187                check(0);
 188        }
 189
 190        printk(KERN_ALERT "augmented rbtree testing");
 191
 192        init();
 193
 194        time1 = get_cycles();
 195
 196        for (i = 0; i < PERF_LOOPS; i++) {
 197                for (j = 0; j < NODES; j++)
 198                        insert_augmented(nodes + j, &root);
 199                for (j = 0; j < NODES; j++)
 200                        erase_augmented(nodes + j, &root);
 201        }
 202
 203        time2 = get_cycles();
 204        time = time2 - time1;
 205
 206        time = div_u64(time, PERF_LOOPS);
 207        printk(" -> %llu cycles\n", (unsigned long long)time);
 208
 209        for (i = 0; i < CHECK_LOOPS; i++) {
 210                init();
 211                for (j = 0; j < NODES; j++) {
 212                        check_augmented(j);
 213                        insert_augmented(nodes + j, &root);
 214                }
 215                for (j = 0; j < NODES; j++) {
 216                        check_augmented(NODES - j);
 217                        erase_augmented(nodes + j, &root);
 218                }
 219                check_augmented(0);
 220        }
 221
 222        return -EAGAIN; /* Fail will directly unload the module */
 223}
 224
 225static void __exit rbtree_test_exit(void)
 226{
 227        printk(KERN_ALERT "test exit\n");
 228}
 229
 230module_init(rbtree_test_init)
 231module_exit(rbtree_test_exit)
 232
 233MODULE_LICENSE("GPL");
 234MODULE_AUTHOR("Michel Lespinasse");
 235MODULE_DESCRIPTION("Red Black Tree test");
 236
lxr.linux.no kindly hosted by Redpill Linpro AS, provider of Linux consulting and operations services since 1995.