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
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(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
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    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
109static node_t *
110tree_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
137static unsigned
138tree_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
148static unsigned
149tree_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
159static void
160node_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
193static node_t *
194remove_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
204static node_t *
205remove_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
215TEST_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}
327TEST_END
328
329int
330main(void)
331{
332
333	return (test(
334	    test_rb_empty,
335	    test_rb_random));
336}
337