1 // SPDX-License-Identifier: CDDL-1.0
2 /*
3 * This file and its contents are supplied under the terms of the
4 * Common Development and Distribution License ("CDDL"), version 1.0.
5 * You may only use this file in accordance with the terms of version
6 * 1.0 of the CDDL.
7 *
8 * A full copy of the text of the CDDL should have accompanied this
9 * source. A copy of the CDDL is also available via the Internet at
10 * https://opensource.org/license/CDDL-1.0.
11 */
12
13 /*
14 * Copyright (c) 2026, Christos Longros.
15 */
16
17 #include <stddef.h>
18 #include <stdlib.h>
19 #include <sys/types.h>
20 #include <sys/sysmacros.h>
21 #include <sys/avl.h>
22 #include <sys/btree.h>
23
24 #include "unit.h"
25
26 #define DRAIN_COUNT (64 * 1024)
27
28 /* ========== */
29
30 /*
31 * The B-Tree stores arbitrary sortable data; these tests use plain uint64_t
32 * values. Elements are kept in sorted order, so a comparison function must
33 * return -1, 0, or +1 for less-than, equal, and greater-than.
34 */
35 static int
u64_compare(const void * a,const void * b)36 u64_compare(const void *a, const void *b)
37 {
38 const uint64_t x = *(const uint64_t *)a;
39 const uint64_t y = *(const uint64_t *)b;
40
41 return (TREE_CMP(x, y));
42 }
43
44 /*
45 * Create a tree of uint64_t values.
46 */
47 static void
btree_create_u64(zfs_btree_t * bt)48 btree_create_u64(zfs_btree_t *bt)
49 {
50 zfs_btree_create(bt, u64_compare, NULL, sizeof (uint64_t));
51 }
52
53 /*
54 * Cross-check the B-Tree against an AVL tree holding the same
55 * values, used as reference.
56 */
57 typedef struct int_node {
58 avl_node_t node;
59 uint64_t data;
60 } int_node_t;
61
62 static int
avl_u64_compare(const void * v1,const void * v2)63 avl_u64_compare(const void *v1, const void *v2)
64 {
65 const int_node_t *n1 = v1;
66 const int_node_t *n2 = v2;
67
68 return (TREE_CMP(n1->data, n2->data));
69 }
70
71 /* ========== */
72
73 /* A new tree is empty: no nodes, and nothing to walk. */
74 static MunitResult
test_btree_empty(const MunitParameter params[],void * data)75 test_btree_empty(const MunitParameter params[], void *data)
76 {
77 (void) params, (void) data;
78
79 zfs_btree_t bt;
80 zfs_btree_index_t idx;
81
82 btree_create_u64(&bt);
83 unit_zero(zfs_btree_numnodes(&bt));
84 unit_true(zfs_btree_first(&bt, &idx) == NULL);
85 zfs_btree_clear(&bt);
86 zfs_btree_destroy(&bt);
87
88 return (MUNIT_OK);
89 }
90
91 /* Added values can be found; the count tracks them; absent values are not. */
92 static MunitResult
test_btree_add_find(const MunitParameter params[],void * data)93 test_btree_add_find(const MunitParameter params[], void *data)
94 {
95 (void) params, (void) data;
96
97 zfs_btree_t bt;
98 zfs_btree_index_t idx;
99 btree_create_u64(&bt);
100
101 uint64_t vals[] = { 50, 10, 30 };
102 for (size_t i = 0; i < ARRAY_SIZE(vals); i++)
103 zfs_btree_add(&bt, &vals[i]);
104
105 /* The tree now holds exactly the values we added. */
106 unit_eq(zfs_btree_numnodes(&bt), ARRAY_SIZE(vals));
107
108 /* Each value is found; find() returns a pointer to the stored copy. */
109 for (size_t i = 0; i < ARRAY_SIZE(vals); i++) {
110 uint64_t *found = zfs_btree_find(&bt, &vals[i], &idx);
111 unit_true(found != NULL);
112 unit_eq(*found, vals[i]);
113 }
114
115 /* A value that was never added is not in the tree. */
116 uint64_t absent = 99;
117 unit_true(zfs_btree_find(&bt, &absent, &idx) == NULL);
118
119 zfs_btree_clear(&bt);
120 zfs_btree_destroy(&bt);
121 return (MUNIT_OK);
122 }
123
124 /* Removing a value drops it from the tree and lowers the count. */
125 static MunitResult
test_btree_remove(const MunitParameter params[],void * data)126 test_btree_remove(const MunitParameter params[], void *data)
127 {
128 (void) params, (void) data;
129
130 zfs_btree_t bt;
131 zfs_btree_index_t idx;
132 btree_create_u64(&bt);
133
134 uint64_t vals[] = { 10, 20, 30 };
135 for (size_t i = 0; i < ARRAY_SIZE(vals); i++)
136 zfs_btree_add(&bt, &vals[i]);
137
138 uint64_t gone = 20;
139 zfs_btree_remove(&bt, &gone);
140
141 unit_eq(zfs_btree_numnodes(&bt), 2);
142 unit_true(zfs_btree_find(&bt, &gone, &idx) == NULL);
143
144 /* The values we kept are still present. */
145 uint64_t keep1 = 10, keep2 = 30;
146 unit_true(zfs_btree_find(&bt, &keep1, &idx) != NULL);
147 unit_true(zfs_btree_find(&bt, &keep2, &idx) != NULL);
148
149 zfs_btree_clear(&bt);
150 zfs_btree_destroy(&bt);
151 return (MUNIT_OK);
152 }
153
154 /* Values inserted out of order are walked back in ascending order. */
155 static MunitResult
test_btree_walk(const MunitParameter params[],void * data)156 test_btree_walk(const MunitParameter params[], void *data)
157 {
158 (void) params, (void) data;
159
160 zfs_btree_t bt;
161 zfs_btree_index_t idx;
162 btree_create_u64(&bt);
163
164 uint64_t vals[] = { 50, 10, 40, 20, 30 };
165 for (size_t i = 0; i < ARRAY_SIZE(vals); i++)
166 zfs_btree_add(&bt, &vals[i]);
167
168 /* first()/next() yield the elements smallest-to-largest. */
169 uint64_t prev = 0;
170 uint64_t count = 0;
171 for (uint64_t *p = zfs_btree_first(&bt, &idx); p != NULL;
172 p = zfs_btree_next(&bt, &idx, &idx)) {
173 unit_gt(*p, prev); /* strictly increasing */
174 prev = *p;
175 count++;
176 }
177 unit_eq(count, ARRAY_SIZE(vals));
178
179 zfs_btree_clear(&bt);
180 zfs_btree_destroy(&bt);
181 return (MUNIT_OK);
182 }
183
184 /* Verify that zfs_btree_find() works correctly when passed a NULL index. */
185 static MunitResult
test_btree_find_without_index(const MunitParameter params[],void * data)186 test_btree_find_without_index(const MunitParameter params[], void *data)
187 {
188 (void) params, (void) data;
189
190 zfs_btree_t bt;
191 btree_create_u64(&bt);
192
193 uint64_t i = 12345;
194 zfs_btree_add(&bt, &i);
195
196 uint64_t *p = zfs_btree_find(&bt, &i, NULL);
197 unit_true(p != NULL);
198 unit_eq(*p, i);
199
200 uint64_t absent = i + 1;
201 unit_true(zfs_btree_find(&bt, &absent, NULL) == NULL);
202
203 zfs_btree_clear(&bt);
204 zfs_btree_destroy(&bt);
205 return (MUNIT_OK);
206 }
207
208 /*
209 * Fill a B-Tree and an AVL tree with the same random values, then drain
210 * both from alternating ends, checking they stay identical at every step.
211 */
212 static MunitResult
test_btree_drain(const MunitParameter params[],void * data)213 test_btree_drain(const MunitParameter params[], void *data)
214 {
215 (void) params, (void) data;
216
217 zfs_btree_t bt;
218 avl_tree_t avl;
219 zfs_btree_index_t bt_idx = {0};
220 avl_index_t avl_idx = {0};
221 int_node_t *node;
222
223 btree_create_u64(&bt);
224 avl_create(&avl, avl_u64_compare, sizeof (int_node_t),
225 offsetof(int_node_t, node));
226
227 /* Fill both trees with the same data. */
228 for (int i = 0; i < DRAIN_COUNT; i++) {
229 uint64_t randval = unit_rand_uint64();
230 if (zfs_btree_find(&bt, &randval, &bt_idx) != NULL)
231 continue;
232 zfs_btree_add_idx(&bt, &randval, &bt_idx);
233
234 node = malloc(sizeof (int_node_t));
235 unit_true(node != NULL);
236 node->data = randval;
237
238 /* New to the btree, so the avl must not have it either. */
239 unit_true(avl_find(&avl, node, &avl_idx) == NULL);
240 avl_insert(&avl, node, avl_idx);
241 }
242
243 /* Remove from alternating ends, comparing the trees as we go. */
244 while (avl_numnodes(&avl) != 0) {
245 uint64_t *bt_data;
246
247 unit_eq(zfs_btree_numnodes(&bt), avl_numnodes(&avl));
248 if (avl_numnodes(&avl) % 2 == 0) {
249 node = avl_first(&avl);
250 bt_data = zfs_btree_first(&bt, &bt_idx);
251 } else {
252 node = avl_last(&avl);
253 bt_data = zfs_btree_last(&bt, &bt_idx);
254 }
255 unit_eq(*bt_data, node->data);
256 zfs_btree_remove_idx(&bt, &bt_idx);
257 avl_remove(&avl, node);
258 free(node);
259
260 if (avl_numnodes(&avl) == 0)
261 break;
262
263 /* Both ends still agree after the removal. */
264 uint64_t *bt_lo = zfs_btree_first(&bt, NULL);
265 uint64_t *bt_hi = zfs_btree_last(&bt, NULL);
266 int_node_t *avl_lo = avl_first(&avl);
267 int_node_t *avl_hi = avl_last(&avl);
268 unit_eq(*bt_lo, avl_lo->data);
269 unit_eq(*bt_hi, avl_hi->data);
270 }
271 unit_zero(zfs_btree_numnodes(&bt));
272
273 avl_destroy(&avl);
274 zfs_btree_clear(&bt);
275 zfs_btree_destroy(&bt);
276 return (MUNIT_OK);
277 }
278
279 /* ========== */
280
281 static const MunitTest btree_tests[] = {
282 UNIT_TEST("empty", test_btree_empty),
283 UNIT_TEST("add_find", test_btree_add_find),
284 UNIT_TEST("remove", test_btree_remove),
285 UNIT_TEST("walk", test_btree_walk),
286 UNIT_TEST("find_without_index", test_btree_find_without_index),
287 UNIT_TEST("drain", test_btree_drain),
288 { 0 },
289 };
290
291 static const MunitSuite btree_test_suite = {
292 "btree.",
293 btree_tests,
294 NULL,
295 1,
296 MUNIT_SUITE_OPTION_NONE,
297 };
298
299 int
main(int argc,char ** argv)300 main(int argc, char **argv)
301 {
302 zfs_btree_init();
303 int ret = munit_suite_main(&btree_test_suite, NULL, argc, argv);
304 zfs_btree_fini();
305 return (ret);
306 }
307