Branch data Line data Source code
1 : : /*--------------------------------------------------------------------------
2 : : *
3 : : * test_rbtree.c
4 : : * Test correctness of red-black tree operations.
5 : : *
6 : : * Copyright (c) 2009-2026, PostgreSQL Global Development Group
7 : : *
8 : : * IDENTIFICATION
9 : : * src/test/modules/test_rbtree/test_rbtree.c
10 : : *
11 : : * -------------------------------------------------------------------------
12 : : */
13 : :
14 : : #include "postgres.h"
15 : :
16 : : #include "common/pg_prng.h"
17 : : #include "fmgr.h"
18 : : #include "lib/rbtree.h"
19 : : #include "utils/memutils.h"
20 : :
21 : 1 : PG_MODULE_MAGIC;
22 : :
23 : :
24 : : /*
25 : : * Our test trees store an integer key, and nothing else.
26 : : */
27 : : typedef struct IntRBTreeNode
28 : : {
29 : : RBTNode rbtnode;
30 : : int key;
31 : : } IntRBTreeNode;
32 : :
33 : :
34 : : /*
35 : : * Node comparator. We don't worry about overflow in the subtraction,
36 : : * since none of our test keys are negative.
37 : : */
38 : : static int
39 : 1340619 : irbt_cmp(const RBTNode *a, const RBTNode *b, void *arg)
40 : : {
41 : 1340619 : const IntRBTreeNode *ea = (const IntRBTreeNode *) a;
42 : 1340619 : const IntRBTreeNode *eb = (const IntRBTreeNode *) b;
43 : :
44 : 1340619 : return ea->key - eb->key;
45 : : }
46 : :
47 : : /*
48 : : * Node combiner. For testing purposes, just check that library doesn't
49 : : * try to combine unequal keys.
50 : : */
51 : : static void
52 : 6 : irbt_combine(RBTNode *existing, const RBTNode *newdata, void *arg)
53 : : {
54 : 6 : const IntRBTreeNode *eexist = (const IntRBTreeNode *) existing;
55 : 6 : const IntRBTreeNode *enew = (const IntRBTreeNode *) newdata;
56 : :
57 [ - + ]: 6 : if (eexist->key != enew->key)
58 [ # # ]: 0 : elog(ERROR, "red-black tree combines %d into %d",
59 : : enew->key, eexist->key);
60 : 6 : }
61 : :
62 : : /* Node allocator */
63 : : static RBTNode *
64 : 60000 : irbt_alloc(void *arg)
65 : : {
66 : 60000 : return (RBTNode *) palloc(sizeof(IntRBTreeNode));
67 : : }
68 : :
69 : : /* Node freer */
70 : : static void
71 : 15049 : irbt_free(RBTNode *node, void *arg)
72 : : {
73 : 15049 : pfree(node);
74 : 15049 : }
75 : :
76 : : /*
77 : : * Create a red-black tree using our support functions
78 : : */
79 : : static RBTree *
80 : 6 : create_int_rbtree(void)
81 : : {
82 : 6 : return rbt_create(sizeof(IntRBTreeNode),
83 : : irbt_cmp,
84 : : irbt_combine,
85 : : irbt_alloc,
86 : : irbt_free,
87 : : NULL);
88 : : }
89 : :
90 : : /*
91 : : * Generate a random permutation of the integers 0..size-1
92 : : */
93 : : static int *
94 : 6 : GetPermutation(int size)
95 : : {
96 : : int *permutation;
97 : : int i;
98 : :
99 : 6 : permutation = (int *) palloc(size * sizeof(int));
100 : :
101 : 6 : permutation[0] = 0;
102 : :
103 : : /*
104 : : * This is the "inside-out" variant of the Fisher-Yates shuffle algorithm.
105 : : * Notionally, we append each new value to the array and then swap it with
106 : : * a randomly-chosen array element (possibly including itself, else we
107 : : * fail to generate permutations with the last integer last). The swap
108 : : * step can be optimized by combining it with the insertion.
109 : : */
110 [ + + ]: 60000 : for (i = 1; i < size; i++)
111 : : {
112 : 59994 : int j = pg_prng_uint64_range(&pg_global_prng_state, 0, i);
113 : :
114 [ + + ]: 59994 : if (j < i) /* avoid fetching undefined data if j=i */
115 : 59945 : permutation[i] = permutation[j];
116 : 59994 : permutation[j] = i;
117 : : }
118 : :
119 : 6 : return permutation;
120 : : }
121 : :
122 : : /*
123 : : * Populate an empty RBTree with "size" integers having the values
124 : : * 0, step, 2*step, 3*step, ..., inserting them in random order
125 : : */
126 : : static void
127 : 6 : rbt_populate(RBTree *tree, int size, int step)
128 : : {
129 : 6 : int *permutation = GetPermutation(size);
130 : : IntRBTreeNode node;
131 : : bool isNew;
132 : : int i;
133 : :
134 : : /* Insert values. We don't expect any collisions. */
135 [ + + ]: 60006 : for (i = 0; i < size; i++)
136 : : {
137 : 60000 : node.key = step * permutation[i];
138 : 60000 : rbt_insert(tree, (RBTNode *) &node, &isNew);
139 [ - + ]: 60000 : if (!isNew)
140 [ # # ]: 0 : elog(ERROR, "unexpected !isNew result from rbt_insert");
141 : : }
142 : :
143 : : /*
144 : : * Re-insert the first value to make sure collisions work right. It's
145 : : * probably not useful to test that case over again for all the values.
146 : : */
147 [ + - ]: 6 : if (size > 0)
148 : : {
149 : 6 : node.key = step * permutation[0];
150 : 6 : rbt_insert(tree, (RBTNode *) &node, &isNew);
151 [ - + ]: 6 : if (isNew)
152 [ # # ]: 0 : elog(ERROR, "unexpected isNew result from rbt_insert");
153 : : }
154 : :
155 : 6 : pfree(permutation);
156 : 6 : }
157 : :
158 : : /*
159 : : * Check the correctness of left-right traversal.
160 : : * Left-right traversal is correct if all elements are
161 : : * visited in increasing order.
162 : : */
163 : : static void
164 : 1 : testleftright(int size)
165 : : {
166 : 1 : RBTree *tree = create_int_rbtree();
167 : : IntRBTreeNode *node;
168 : : RBTreeIterator iter;
169 : 1 : int lastKey = -1;
170 : 1 : int count = 0;
171 : :
172 : : /* check iteration over empty tree */
173 : 1 : rbt_begin_iterate(tree, LeftRightWalk, &iter);
174 [ - + ]: 1 : if (rbt_iterate(&iter) != NULL)
175 [ # # ]: 0 : elog(ERROR, "left-right walk over empty tree produced an element");
176 : :
177 : : /* fill tree with consecutive natural numbers */
178 : 1 : rbt_populate(tree, size, 1);
179 : :
180 : : /* iterate over the tree */
181 : 1 : rbt_begin_iterate(tree, LeftRightWalk, &iter);
182 : :
183 [ + + ]: 10001 : while ((node = (IntRBTreeNode *) rbt_iterate(&iter)) != NULL)
184 : : {
185 : : /* check that order is increasing */
186 [ - + ]: 10000 : if (node->key <= lastKey)
187 [ # # ]: 0 : elog(ERROR, "left-right walk gives elements not in sorted order");
188 : 10000 : lastKey = node->key;
189 : 10000 : count++;
190 : : }
191 : :
192 [ - + ]: 1 : if (lastKey != size - 1)
193 [ # # ]: 0 : elog(ERROR, "left-right walk did not reach end");
194 [ - + ]: 1 : if (count != size)
195 [ # # ]: 0 : elog(ERROR, "left-right walk missed some elements");
196 : 1 : }
197 : :
198 : : /*
199 : : * Check the correctness of right-left traversal.
200 : : * Right-left traversal is correct if all elements are
201 : : * visited in decreasing order.
202 : : */
203 : : static void
204 : 1 : testrightleft(int size)
205 : : {
206 : 1 : RBTree *tree = create_int_rbtree();
207 : : IntRBTreeNode *node;
208 : : RBTreeIterator iter;
209 : 1 : int lastKey = size;
210 : 1 : int count = 0;
211 : :
212 : : /* check iteration over empty tree */
213 : 1 : rbt_begin_iterate(tree, RightLeftWalk, &iter);
214 [ - + ]: 1 : if (rbt_iterate(&iter) != NULL)
215 [ # # ]: 0 : elog(ERROR, "right-left walk over empty tree produced an element");
216 : :
217 : : /* fill tree with consecutive natural numbers */
218 : 1 : rbt_populate(tree, size, 1);
219 : :
220 : : /* iterate over the tree */
221 : 1 : rbt_begin_iterate(tree, RightLeftWalk, &iter);
222 : :
223 [ + + ]: 10001 : while ((node = (IntRBTreeNode *) rbt_iterate(&iter)) != NULL)
224 : : {
225 : : /* check that order is decreasing */
226 [ - + ]: 10000 : if (node->key >= lastKey)
227 [ # # ]: 0 : elog(ERROR, "right-left walk gives elements not in sorted order");
228 : 10000 : lastKey = node->key;
229 : 10000 : count++;
230 : : }
231 : :
232 [ - + ]: 1 : if (lastKey != 0)
233 [ # # ]: 0 : elog(ERROR, "right-left walk did not reach end");
234 [ - + ]: 1 : if (count != size)
235 [ # # ]: 0 : elog(ERROR, "right-left walk missed some elements");
236 : 1 : }
237 : :
238 : : /*
239 : : * Check the correctness of the rbt_find operation by searching for
240 : : * both elements we inserted and elements we didn't.
241 : : */
242 : : static void
243 : 1 : testfind(int size)
244 : : {
245 : 1 : RBTree *tree = create_int_rbtree();
246 : : int i;
247 : :
248 : : /* Insert even integers from 0 to 2 * (size-1) */
249 : 1 : rbt_populate(tree, size, 2);
250 : :
251 : : /* Check that all inserted elements can be found */
252 [ + + ]: 10001 : for (i = 0; i < size; i++)
253 : : {
254 : : IntRBTreeNode node;
255 : : IntRBTreeNode *resultNode;
256 : :
257 : 10000 : node.key = 2 * i;
258 : 10000 : resultNode = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &node);
259 [ - + ]: 10000 : if (resultNode == NULL)
260 [ # # ]: 0 : elog(ERROR, "inserted element was not found");
261 [ - + ]: 10000 : if (node.key != resultNode->key)
262 [ # # ]: 0 : elog(ERROR, "find operation in rbtree gave wrong result");
263 : : }
264 : :
265 : : /*
266 : : * Check that not-inserted elements can not be found, being sure to try
267 : : * values before the first and after the last element.
268 : : */
269 [ + + ]: 10002 : for (i = -1; i <= 2 * size; i += 2)
270 : : {
271 : : IntRBTreeNode node;
272 : : IntRBTreeNode *resultNode;
273 : :
274 : 10001 : node.key = i;
275 : 10001 : resultNode = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &node);
276 [ - + ]: 10001 : if (resultNode != NULL)
277 [ # # ]: 0 : elog(ERROR, "not-inserted element was found");
278 : : }
279 : 1 : }
280 : :
281 : : /*
282 : : * Check the correctness of the rbt_find_less() and rbt_find_great() functions
283 : : * by searching for an equal key and iterating the lesser keys then the greater
284 : : * keys.
285 : : */
286 : : static void
287 : 1 : testfindltgt(int size)
288 : : {
289 : 1 : RBTree *tree = create_int_rbtree();
290 : :
291 : : /*
292 : : * Using the size as the random key to search wouldn't allow us to get at
293 : : * least one greater match, so we do size - 1
294 : : */
295 : 1 : int randomKey = pg_prng_uint64_range(&pg_global_prng_state, 0, size - 1);
296 : : bool keyDeleted;
297 : 1 : IntRBTreeNode searchNode = {.key = randomKey};
298 : : IntRBTreeNode *lteNode;
299 : : IntRBTreeNode *gteNode;
300 : : IntRBTreeNode *node;
301 : :
302 : : /* Insert natural numbers */
303 : 1 : rbt_populate(tree, size, 1);
304 : :
305 : : /*
306 : : * Since the search key is included in the naturals of the tree, we're
307 : : * sure to find an equal match
308 : : */
309 : 1 : lteNode = (IntRBTreeNode *) rbt_find_less(tree, (RBTNode *) &searchNode, true);
310 : 1 : gteNode = (IntRBTreeNode *) rbt_find_great(tree, (RBTNode *) &searchNode, true);
311 : :
312 [ + - - + ]: 1 : if (lteNode == NULL || lteNode->key != searchNode.key)
313 [ # # ]: 0 : elog(ERROR, "rbt_find_less() didn't find the equal key");
314 : :
315 [ + - - + ]: 1 : if (gteNode == NULL || gteNode->key != searchNode.key)
316 [ # # ]: 0 : elog(ERROR, "rbt_find_great() didn't find the equal key");
317 : :
318 [ - + ]: 1 : if (lteNode != gteNode)
319 [ # # ]: 0 : elog(ERROR, "rbt_find_less() and rbt_find_great() found different equal keys");
320 : :
321 : : /* Find the rest of the naturals lesser than the search key */
322 : 1 : keyDeleted = false;
323 [ + + ]: 8938 : for (; searchNode.key > 0; searchNode.key--)
324 : : {
325 : : /*
326 : : * Find the next key. If the current key is deleted, we can pass
327 : : * equal_match == true and still find the next one.
328 : : */
329 : 8937 : node = (IntRBTreeNode *) rbt_find_less(tree, (RBTNode *) &searchNode,
330 : : keyDeleted);
331 : :
332 : : /* ensure we find a lesser match */
333 [ + - - + ]: 8937 : if (!node || !(node->key < searchNode.key))
334 [ # # ]: 0 : elog(ERROR, "rbt_find_less() didn't find a lesser key");
335 : :
336 : : /* randomly delete the found key or leave it */
337 : 8937 : keyDeleted = (pg_prng_uint64_range(&pg_global_prng_state, 0, 1) == 1);
338 [ + + ]: 8937 : if (keyDeleted)
339 : 4498 : rbt_delete(tree, (RBTNode *) node);
340 : : }
341 : :
342 : : /* Find the rest of the naturals greater than the search key */
343 : 1 : keyDeleted = false;
344 [ + + ]: 1063 : for (searchNode.key = randomKey; searchNode.key < size - 1; searchNode.key++)
345 : : {
346 : : /*
347 : : * Find the next key. If the current key is deleted, we can pass
348 : : * equal_match == true and still find the next one.
349 : : */
350 : 1062 : node = (IntRBTreeNode *) rbt_find_great(tree, (RBTNode *) &searchNode,
351 : : keyDeleted);
352 : :
353 : : /* ensure we find a greater match */
354 [ + - - + ]: 1062 : if (!node || !(node->key > searchNode.key))
355 [ # # ]: 0 : elog(ERROR, "rbt_find_great() didn't find a greater key");
356 : :
357 : : /* randomly delete the found key or leave it */
358 : 1062 : keyDeleted = (pg_prng_uint64_range(&pg_global_prng_state, 0, 1) == 1);
359 [ + + ]: 1062 : if (keyDeleted)
360 : 551 : rbt_delete(tree, (RBTNode *) node);
361 : : }
362 : :
363 : : /* Check out of bounds searches find nothing */
364 : 1 : searchNode.key = -1;
365 : 1 : node = (IntRBTreeNode *) rbt_find_less(tree, (RBTNode *) &searchNode, true);
366 [ - + ]: 1 : if (node != NULL)
367 [ # # ]: 0 : elog(ERROR, "rbt_find_less() found non-inserted element");
368 : 1 : searchNode.key = 0;
369 : 1 : node = (IntRBTreeNode *) rbt_find_less(tree, (RBTNode *) &searchNode, false);
370 [ - + ]: 1 : if (node != NULL)
371 [ # # ]: 0 : elog(ERROR, "rbt_find_less() found non-inserted element");
372 : 1 : searchNode.key = size;
373 : 1 : node = (IntRBTreeNode *) rbt_find_great(tree, (RBTNode *) &searchNode, true);
374 [ - + ]: 1 : if (node != NULL)
375 [ # # ]: 0 : elog(ERROR, "rbt_find_great() found non-inserted element");
376 : 1 : searchNode.key = size - 1;
377 : 1 : node = (IntRBTreeNode *) rbt_find_great(tree, (RBTNode *) &searchNode, false);
378 [ - + ]: 1 : if (node != NULL)
379 [ # # ]: 0 : elog(ERROR, "rbt_find_great() found non-inserted element");
380 : 1 : }
381 : :
382 : : /*
383 : : * Check the correctness of the rbt_leftmost operation.
384 : : * This operation should always return the smallest element of the tree.
385 : : */
386 : : static void
387 : 1 : testleftmost(int size)
388 : : {
389 : 1 : RBTree *tree = create_int_rbtree();
390 : : IntRBTreeNode *result;
391 : :
392 : : /* Check that empty tree has no leftmost element */
393 [ - + ]: 1 : if (rbt_leftmost(tree) != NULL)
394 [ # # ]: 0 : elog(ERROR, "leftmost node of empty tree is not NULL");
395 : :
396 : : /* fill tree with consecutive natural numbers */
397 : 1 : rbt_populate(tree, size, 1);
398 : :
399 : : /* Check that leftmost element is the smallest one */
400 : 1 : result = (IntRBTreeNode *) rbt_leftmost(tree);
401 [ + - - + ]: 1 : if (result == NULL || result->key != 0)
402 [ # # ]: 0 : elog(ERROR, "rbt_leftmost gave wrong result");
403 : 1 : }
404 : :
405 : : /*
406 : : * Check the correctness of the rbt_delete operation.
407 : : */
408 : : static void
409 : 1 : testdelete(int size, int delsize)
410 : : {
411 : 1 : RBTree *tree = create_int_rbtree();
412 : : int *deleteIds;
413 : : bool *chosen;
414 : : int i;
415 : :
416 : : /* fill tree with consecutive natural numbers */
417 : 1 : rbt_populate(tree, size, 1);
418 : :
419 : : /* Choose unique ids to delete */
420 : 1 : deleteIds = (int *) palloc(delsize * sizeof(int));
421 : 1 : chosen = (bool *) palloc0(size * sizeof(bool));
422 : :
423 [ + + ]: 1001 : for (i = 0; i < delsize; i++)
424 : : {
425 : 1000 : int k = pg_prng_uint64_range(&pg_global_prng_state, 0, size - 1);
426 : :
427 [ + + ]: 1065 : while (chosen[k])
428 : 65 : k = (k + 1) % size;
429 : 1000 : deleteIds[i] = k;
430 : 1000 : chosen[k] = true;
431 : : }
432 : :
433 : : /* Delete elements */
434 [ + + ]: 1001 : for (i = 0; i < delsize; i++)
435 : : {
436 : : IntRBTreeNode find;
437 : : IntRBTreeNode *node;
438 : :
439 : 1000 : find.key = deleteIds[i];
440 : : /* Locate the node to be deleted */
441 : 1000 : node = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &find);
442 [ + - - + ]: 1000 : if (node == NULL || node->key != deleteIds[i])
443 [ # # ]: 0 : elog(ERROR, "expected element was not found during deleting");
444 : : /* Delete it */
445 : 1000 : rbt_delete(tree, (RBTNode *) node);
446 : : }
447 : :
448 : : /* Check that deleted elements are deleted */
449 [ + + ]: 10001 : for (i = 0; i < size; i++)
450 : : {
451 : : IntRBTreeNode node;
452 : : IntRBTreeNode *result;
453 : :
454 : 10000 : node.key = i;
455 : 10000 : result = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &node);
456 [ + + ]: 10000 : if (chosen[i])
457 : : {
458 : : /* Deleted element should be absent */
459 [ - + ]: 1000 : if (result != NULL)
460 [ # # ]: 0 : elog(ERROR, "deleted element still present in the rbtree");
461 : : }
462 : : else
463 : : {
464 : : /* Else it should be present */
465 [ + - - + ]: 9000 : if (result == NULL || result->key != i)
466 [ # # ]: 0 : elog(ERROR, "delete operation removed wrong rbtree value");
467 : : }
468 : : }
469 : :
470 : : /* Delete remaining elements, so as to exercise reducing tree to empty */
471 [ + + ]: 10001 : for (i = 0; i < size; i++)
472 : : {
473 : : IntRBTreeNode find;
474 : : IntRBTreeNode *node;
475 : :
476 [ + + ]: 10000 : if (chosen[i])
477 : 1000 : continue;
478 : 9000 : find.key = i;
479 : : /* Locate the node to be deleted */
480 : 9000 : node = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &find);
481 [ + - - + ]: 9000 : if (node == NULL || node->key != i)
482 [ # # ]: 0 : elog(ERROR, "expected element was not found during deleting");
483 : : /* Delete it */
484 : 9000 : rbt_delete(tree, (RBTNode *) node);
485 : : }
486 : :
487 : : /* Tree should now be empty */
488 [ - + ]: 1 : if (rbt_leftmost(tree) != NULL)
489 [ # # ]: 0 : elog(ERROR, "deleting all elements failed");
490 : :
491 : 1 : pfree(deleteIds);
492 : 1 : pfree(chosen);
493 : 1 : }
494 : :
495 : : /*
496 : : * SQL-callable entry point to perform all tests
497 : : *
498 : : * Argument is the number of entries to put in the trees
499 : : */
500 : 2 : PG_FUNCTION_INFO_V1(test_rb_tree);
501 : :
502 : : Datum
503 : 1 : test_rb_tree(PG_FUNCTION_ARGS)
504 : : {
505 : 1 : int size = PG_GETARG_INT32(0);
506 : :
507 [ + - - + ]: 1 : if (size <= 0 || size > MaxAllocSize / sizeof(int))
508 [ # # ]: 0 : elog(ERROR, "invalid size for test_rb_tree: %d", size);
509 : 1 : testleftright(size);
510 : 1 : testrightleft(size);
511 : 1 : testfind(size);
512 : 1 : testfindltgt(size);
513 : 1 : testleftmost(size);
514 [ + - ]: 1 : testdelete(size, Max(size / 10, 1));
515 : 1 : PG_RETURN_VOID();
516 : : }
|
