1 #include "test/jemalloc_test.h"
2 
3 #define	rbtn_black_height(a_type, a_field, a_rbt, r_height) do {	\
4     a_type *rbp_bh_t;							\
5     for (rbp_bh_t = (a_rbt)->rbt_root, (r_height) = 0;			\
6       rbp_bh_t != &(a_rbt)->rbt_nil;					\
7       rbp_bh_t = rbtn_left_get(a_type, a_field, rbp_bh_t)) {		\
8 	if (!rbtn_red_get(a_type, a_field, rbp_bh_t)) {			\
9 	    (r_height)++;						\
10 	}								\
11     }									\
12 } while (0)
13 
14 typedef struct node_s node_t;
15 
16 struct node_s {
17 #define	NODE_MAGIC 0x9823af7e
18 	uint32_t magic;
19 	rb_node(node_t) link;
20 	uint64_t key;
21 };
22 
23 static int
node_cmp(node_t * a,node_t * b)24 node_cmp(node_t *a, node_t *b) {
25 	int ret;
26 
27 	assert_u32_eq(a->magic, NODE_MAGIC, "Bad magic");
28 	assert_u32_eq(b->magic, NODE_MAGIC, "Bad magic");
29 
30 	ret = (a->key > b->key) - (a->key < b->key);
31 	if (ret == 0) {
32 		/*
33 		 * Duplicates are not allowed in the tree, so force an
34 		 * arbitrary ordering for non-identical items with equal keys.
35 		 */
36 		ret = (((uintptr_t)a) > ((uintptr_t)b))
37 		    - (((uintptr_t)a) < ((uintptr_t)b));
38 	}
39 	return (ret);
40 }
41 
42 typedef rb_tree(node_t) tree_t;
43 rb_gen(static, tree_, tree_t, node_t, link, node_cmp);
44 
TEST_BEGIN(test_rb_empty)45 TEST_BEGIN(test_rb_empty)
46 {
47 	tree_t tree;
48 	node_t key;
49 
50 	tree_new(&tree);
51 
52 	assert_true(tree_empty(&tree), "Tree should be empty");
53 	assert_ptr_null(tree_first(&tree), "Unexpected node");
54 	assert_ptr_null(tree_last(&tree), "Unexpected node");
55 
56 	key.key = 0;
57 	key.magic = NODE_MAGIC;
58 	assert_ptr_null(tree_search(&tree, &key), "Unexpected node");
59 
60 	key.key = 0;
61 	key.magic = NODE_MAGIC;
62 	assert_ptr_null(tree_nsearch(&tree, &key), "Unexpected node");
63 
64 	key.key = 0;
65 	key.magic = NODE_MAGIC;
66 	assert_ptr_null(tree_psearch(&tree, &key), "Unexpected node");
67 }
68 TEST_END
69 
70 static unsigned
tree_recurse(node_t * node,unsigned black_height,unsigned black_depth,node_t * nil)71 tree_recurse(node_t *node, unsigned black_height, unsigned black_depth,
72     node_t *nil)
73 {
74 	unsigned ret = 0;
75 	node_t *left_node = rbtn_left_get(node_t, link, node);
76 	node_t *right_node = rbtn_right_get(node_t, link, node);
77 
78 	if (!rbtn_red_get(node_t, link, node))
79 		black_depth++;
80 
81 	/* Red nodes must be interleaved with black nodes. */
82 	if (rbtn_red_get(node_t, link, node)) {
83 		assert_false(rbtn_red_get(node_t, link, left_node),
84 		    "Node should be black");
85 		assert_false(rbtn_red_get(node_t, link, right_node),
86 		    "Node should be black");
87 	}
88 
89 	if (node == nil)
90 		return (ret);
91 	/* Self. */
92 	assert_u32_eq(node->magic, NODE_MAGIC, "Bad magic");
93 
94 	/* Left subtree. */
95 	if (left_node != nil)
96 		ret += tree_recurse(left_node, black_height, black_depth, nil);
97 	else
98 		ret += (black_depth != black_height);
99 
100 	/* Right subtree. */
101 	if (right_node != nil)
102 		ret += tree_recurse(right_node, black_height, black_depth, nil);
103 	else
104 		ret += (black_depth != black_height);
105 
106 	return (ret);
107 }
108 
109 static node_t *
tree_iterate_cb(tree_t * tree,node_t * node,void * data)110 tree_iterate_cb(tree_t *tree, node_t *node, void *data)
111 {
112 	unsigned *i = (unsigned *)data;
113 	node_t *search_node;
114 
115 	assert_u32_eq(node->magic, NODE_MAGIC, "Bad magic");
116 
117 	/* Test rb_search(). */
118 	search_node = tree_search(tree, node);
119 	assert_ptr_eq(search_node, node,
120 	    "tree_search() returned unexpected node");
121 
122 	/* Test rb_nsearch(). */
123 	search_node = tree_nsearch(tree, node);
124 	assert_ptr_eq(search_node, node,
125 	    "tree_nsearch() returned unexpected node");
126 
127 	/* Test rb_psearch(). */
128 	search_node = tree_psearch(tree, node);
129 	assert_ptr_eq(search_node, node,
130 	    "tree_psearch() returned unexpected node");
131 
132 	(*i)++;
133 
134 	return (NULL);
135 }
136 
137 static unsigned
tree_iterate(tree_t * tree)138 tree_iterate(tree_t *tree)
139 {
140 	unsigned i;
141 
142 	i = 0;
143 	tree_iter(tree, NULL, tree_iterate_cb, (void *)&i);
144 
145 	return (i);
146 }
147 
148 static unsigned
tree_iterate_reverse(tree_t * tree)149 tree_iterate_reverse(tree_t *tree)
150 {
151 	unsigned i;
152 
153 	i = 0;
154 	tree_reverse_iter(tree, NULL, tree_iterate_cb, (void *)&i);
155 
156 	return (i);
157 }
158 
159 static void
node_remove(tree_t * tree,node_t * node,unsigned nnodes)160 node_remove(tree_t *tree, node_t *node, unsigned nnodes)
161 {
162 	node_t *search_node;
163 	unsigned black_height, imbalances;
164 
165 	tree_remove(tree, node);
166 
167 	/* Test rb_nsearch(). */
168 	search_node = tree_nsearch(tree, node);
169 	if (search_node != NULL) {
170 		assert_u64_ge(search_node->key, node->key,
171 		    "Key ordering error");
172 	}
173 
174 	/* Test rb_psearch(). */
175 	search_node = tree_psearch(tree, node);
176 	if (search_node != NULL) {
177 		assert_u64_le(search_node->key, node->key,
178 		    "Key ordering error");
179 	}
180 
181 	node->magic = 0;
182 
183 	rbtn_black_height(node_t, link, tree, black_height);
184 	imbalances = tree_recurse(tree->rbt_root, black_height, 0,
185 	    &(tree->rbt_nil));
186 	assert_u_eq(imbalances, 0, "Tree is unbalanced");
187 	assert_u_eq(tree_iterate(tree), nnodes-1,
188 	    "Unexpected node iteration count");
189 	assert_u_eq(tree_iterate_reverse(tree), nnodes-1,
190 	    "Unexpected node iteration count");
191 }
192 
193 static node_t *
remove_iterate_cb(tree_t * tree,node_t * node,void * data)194 remove_iterate_cb(tree_t *tree, node_t *node, void *data)
195 {
196 	unsigned *nnodes = (unsigned *)data;
197 	node_t *ret = tree_next(tree, node);
198 
199 	node_remove(tree, node, *nnodes);
200 
201 	return (ret);
202 }
203 
204 static node_t *
remove_reverse_iterate_cb(tree_t * tree,node_t * node,void * data)205 remove_reverse_iterate_cb(tree_t *tree, node_t *node, void *data)
206 {
207 	unsigned *nnodes = (unsigned *)data;
208 	node_t *ret = tree_prev(tree, node);
209 
210 	node_remove(tree, node, *nnodes);
211 
212 	return (ret);
213 }
214 
TEST_BEGIN(test_rb_random)215 TEST_BEGIN(test_rb_random)
216 {
217 #define	NNODES 25
218 #define	NBAGS 250
219 #define	SEED 42
220 	sfmt_t *sfmt;
221 	uint64_t bag[NNODES];
222 	tree_t tree;
223 	node_t nodes[NNODES];
224 	unsigned i, j, k, black_height, imbalances;
225 
226 	sfmt = init_gen_rand(SEED);
227 	for (i = 0; i < NBAGS; i++) {
228 		switch (i) {
229 		case 0:
230 			/* Insert in order. */
231 			for (j = 0; j < NNODES; j++)
232 				bag[j] = j;
233 			break;
234 		case 1:
235 			/* Insert in reverse order. */
236 			for (j = 0; j < NNODES; j++)
237 				bag[j] = NNODES - j - 1;
238 			break;
239 		default:
240 			for (j = 0; j < NNODES; j++)
241 				bag[j] = gen_rand64_range(sfmt, NNODES);
242 		}
243 
244 		for (j = 1; j <= NNODES; j++) {
245 			/* Initialize tree and nodes. */
246 			tree_new(&tree);
247 			tree.rbt_nil.magic = 0;
248 			for (k = 0; k < j; k++) {
249 				nodes[k].magic = NODE_MAGIC;
250 				nodes[k].key = bag[k];
251 			}
252 
253 			/* Insert nodes. */
254 			for (k = 0; k < j; k++) {
255 				tree_insert(&tree, &nodes[k]);
256 
257 				rbtn_black_height(node_t, link, &tree,
258 				    black_height);
259 				imbalances = tree_recurse(tree.rbt_root,
260 				    black_height, 0, &(tree.rbt_nil));
261 				assert_u_eq(imbalances, 0,
262 				    "Tree is unbalanced");
263 
264 				assert_u_eq(tree_iterate(&tree), k+1,
265 				    "Unexpected node iteration count");
266 				assert_u_eq(tree_iterate_reverse(&tree), k+1,
267 				    "Unexpected node iteration count");
268 
269 				assert_false(tree_empty(&tree),
270 				    "Tree should not be empty");
271 				assert_ptr_not_null(tree_first(&tree),
272 				    "Tree should not be empty");
273 				assert_ptr_not_null(tree_last(&tree),
274 				    "Tree should not be empty");
275 
276 				tree_next(&tree, &nodes[k]);
277 				tree_prev(&tree, &nodes[k]);
278 			}
279 
280 			/* Remove nodes. */
281 			switch (i % 4) {
282 			case 0:
283 				for (k = 0; k < j; k++)
284 					node_remove(&tree, &nodes[k], j - k);
285 				break;
286 			case 1:
287 				for (k = j; k > 0; k--)
288 					node_remove(&tree, &nodes[k-1], k);
289 				break;
290 			case 2: {
291 				node_t *start;
292 				unsigned nnodes = j;
293 
294 				start = NULL;
295 				do {
296 					start = tree_iter(&tree, start,
297 					    remove_iterate_cb, (void *)&nnodes);
298 					nnodes--;
299 				} while (start != NULL);
300 				assert_u_eq(nnodes, 0,
301 				    "Removal terminated early");
302 				break;
303 			} case 3: {
304 				node_t *start;
305 				unsigned nnodes = j;
306 
307 				start = NULL;
308 				do {
309 					start = tree_reverse_iter(&tree, start,
310 					    remove_reverse_iterate_cb,
311 					    (void *)&nnodes);
312 					nnodes--;
313 				} while (start != NULL);
314 				assert_u_eq(nnodes, 0,
315 				    "Removal terminated early");
316 				break;
317 			} default:
318 				not_reached();
319 			}
320 		}
321 	}
322 	fini_gen_rand(sfmt);
323 #undef NNODES
324 #undef NBAGS
325 #undef SEED
326 }
327 TEST_END
328 
329 int
main(void)330 main(void)
331 {
332 
333 	return (test(
334 	    test_rb_empty,
335 	    test_rb_random));
336 }
337