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