1#ifdef TEST_RBIT
2
3#include <crypto/prng.h>
4#include <mem/alloc.h>
5#include <structures/rbit.h>
6#include <test.h>
7
8#define RBIT_N 256
9#define RBIT_OPS 4000
10#define RBIT_SEED 0xC0FFEEULL
11
12static size_t subtree_nodes(struct rbit_node *n) {
13 if (!n)
14 return 0;
15 return 1 + subtree_nodes(n: n->left) + subtree_nodes(n: n->right);
16}
17
18static int overlaps(struct interval a, struct interval b) {
19 return a.low <= b.high && b.low <= a.high;
20}
21
22TEST_DECLARE(rbit_order_search, .tier = TEST_TIER_UNIT) {
23 prng_seed(RBIT_SEED);
24 struct rbit tree;
25 rbit_init(rbit: &tree);
26
27 struct rbit_node *nodes =
28 kmalloc(sizeof(*nodes) * RBIT_N, ALLOC_FLAGS_ZERO);
29 TEST_ASSERT(nodes);
30
31 size_t low = 1;
32 for (size_t i = 0; i < RBIT_N; i++) {
33 low += 1 + (prng_next() % 64);
34 rbit_init_node(n: &nodes[i]);
35 nodes[i].interval.low = low;
36 nodes[i].interval.high = low + (prng_next() % 32);
37 rbit_insert(tree: &tree, new_node: &nodes[i]);
38 }
39
40 size_t prev = 0;
41 size_t count = 0;
42 struct rbit_node *it;
43 rbit_for_each(it, &tree) {
44 TEST_ASSERT(it->interval.low >= prev);
45 prev = it->interval.low;
46 count++;
47 }
48 TEST_ASSERT(count == RBIT_N);
49
50 for (size_t i = 0; i < RBIT_N; i++)
51 TEST_ASSERT(rbit_search(tree.root, nodes[i].interval) == &nodes[i]);
52
53 for (size_t i = 0; i < RBIT_N; i += 2)
54 rbit_delete(tree: &tree, z: &nodes[i]);
55 for (size_t i = 0; i < RBIT_N; i++) {
56 struct rbit_node *f = rbit_search(root: tree.root, iv: nodes[i].interval);
57 TEST_ASSERT((i % 2 == 0) ? (f == NULL) : (f == &nodes[i]));
58 }
59
60 for (size_t i = 1; i < RBIT_N; i += 2)
61 rbit_delete(tree: &tree, z: &nodes[i]);
62 TEST_ASSERT(rbit_empty(&tree));
63
64 return TEST_SUCCESS;
65}
66
67TEST_DECLARE(rbit_overlap, .tier = TEST_TIER_UNIT) {
68 prng_seed(RBIT_SEED + 1);
69 struct rbit tree;
70 rbit_init(rbit: &tree);
71
72 struct rbit_node *nodes =
73 kmalloc(sizeof(*nodes) * RBIT_N, ALLOC_FLAGS_ZERO);
74 TEST_ASSERT(nodes);
75
76 /* Random (possibly overlapping) intervals in a bounded space. */
77 for (size_t i = 0; i < RBIT_N; i++) {
78 size_t lo = prng_next() % 100000;
79 rbit_init_node(n: &nodes[i]);
80 nodes[i].interval.low = lo;
81 nodes[i].interval.high = lo + (prng_next() % 500);
82 rbit_insert(tree: &tree, new_node: &nodes[i]);
83 }
84
85 for (size_t q = 0; q < RBIT_OPS; q++) {
86 size_t lo = prng_next() % 100000;
87 struct interval iv = {.low = lo, .high = lo + (prng_next() % 500)};
88
89 bool brute = false;
90 for (size_t i = 0; i < RBIT_N; i++)
91 if (overlaps(a: nodes[i].interval, b: iv)) {
92 brute = true;
93 break;
94 }
95
96 struct rbit_node *res = rbit_overlap_search(root: tree.root, iv);
97 TEST_ASSERT((res != NULL) == brute);
98 if (res)
99 TEST_ASSERT(overlaps(res->interval, iv));
100 }
101
102 return TEST_SUCCESS;
103}
104
105struct count_node {
106 struct rbit_node node;
107 size_t subtree_count; /* maintained by the augment below */
108};
109
110static size_t cn_count(struct rbit_node *n) {
111 return n ? rbit_entry(n, struct count_node, node)->subtree_count : 0;
112}
113
114static bool count_augment(struct rbit_node *n) {
115 struct count_node *c = rbit_entry(n, struct count_node, node);
116 size_t old_max = n->max, old_cnt = c->subtree_count;
117
118 size_t mx = n->interval.high;
119 if (rbit_node_max(n: n->left) > mx)
120 mx = rbit_node_max(n: n->left);
121 if (rbit_node_max(n: n->right) > mx)
122 mx = rbit_node_max(n: n->right);
123
124 n->max = mx;
125 c->subtree_count = 1 + cn_count(n: n->left) + cn_count(n: n->right);
126 return n->max != old_max || c->subtree_count != old_cnt;
127}
128
129TEST_DECLARE(rbit_augment_hook, .tier = TEST_TIER_UNIT) {
130 prng_seed(RBIT_SEED + 2);
131 struct rbit tree;
132 rbit_init(rbit: &tree);
133 tree.augment = count_augment;
134
135 struct count_node *nodes =
136 kmalloc(sizeof(*nodes) * RBIT_N, ALLOC_FLAGS_ZERO);
137 TEST_ASSERT(nodes);
138 bool *live = kmalloc(sizeof(bool) * RBIT_N, ALLOC_FLAGS_ZERO);
139 TEST_ASSERT(live);
140
141 for (size_t i = 0; i < RBIT_N; i++) {
142 rbit_init_node(n: &nodes[i].node);
143 nodes[i].node.interval.low = i * 100 + 1;
144 nodes[i].node.interval.high = i * 100 + 50;
145 }
146
147 for (size_t op = 0; op < RBIT_OPS; op++) {
148 size_t i = prng_next() % RBIT_N;
149 if (live[i]) {
150 rbit_delete(tree: &tree, z: &nodes[i].node);
151 live[i] = false;
152 } else {
153 rbit_insert(tree: &tree, new_node: &nodes[i].node);
154 live[i] = true;
155 }
156
157 struct rbit_node *it;
158 rbit_for_each(it, &tree) {
159 struct count_node *c = rbit_entry(it, struct count_node, node);
160 TEST_ASSERT(c->subtree_count == subtree_nodes(it));
161 }
162 }
163
164 return TEST_SUCCESS;
165}
166
167#endif
168