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