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 0 : 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 0 : irbt_cmp(const RBTNode *a, const RBTNode *b, void *arg)
40 : {
41 0 : const IntRBTreeNode *ea = (const IntRBTreeNode *) a;
42 0 : const IntRBTreeNode *eb = (const IntRBTreeNode *) b;
43 :
44 0 : return ea->key - eb->key;
45 0 : }
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 0 : irbt_combine(RBTNode *existing, const RBTNode *newdata, void *arg)
53 : {
54 0 : const IntRBTreeNode *eexist = (const IntRBTreeNode *) existing;
55 0 : const IntRBTreeNode *enew = (const IntRBTreeNode *) newdata;
56 :
57 0 : if (eexist->key != enew->key)
58 0 : elog(ERROR, "red-black tree combines %d into %d",
59 : enew->key, eexist->key);
60 0 : }
61 :
62 : /* Node allocator */
63 : static RBTNode *
64 0 : irbt_alloc(void *arg)
65 : {
66 0 : return (RBTNode *) palloc(sizeof(IntRBTreeNode));
67 : }
68 :
69 : /* Node freer */
70 : static void
71 0 : irbt_free(RBTNode *node, void *arg)
72 : {
73 0 : pfree(node);
74 0 : }
75 :
76 : /*
77 : * Create a red-black tree using our support functions
78 : */
79 : static RBTree *
80 0 : create_int_rbtree(void)
81 : {
82 0 : 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 0 : GetPermutation(int size)
95 : {
96 0 : int *permutation;
97 0 : int i;
98 :
99 0 : permutation = (int *) palloc(size * sizeof(int));
100 :
101 0 : 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 0 : for (i = 1; i < size; i++)
111 : {
112 0 : int j = pg_prng_uint64_range(&pg_global_prng_state, 0, i);
113 :
114 0 : if (j < i) /* avoid fetching undefined data if j=i */
115 0 : permutation[i] = permutation[j];
116 0 : permutation[j] = i;
117 0 : }
118 :
119 0 : return permutation;
120 0 : }
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 0 : rbt_populate(RBTree *tree, int size, int step)
128 : {
129 0 : int *permutation = GetPermutation(size);
130 0 : IntRBTreeNode node;
131 0 : bool isNew;
132 0 : int i;
133 :
134 : /* Insert values. We don't expect any collisions. */
135 0 : for (i = 0; i < size; i++)
136 : {
137 0 : node.key = step * permutation[i];
138 0 : rbt_insert(tree, (RBTNode *) &node, &isNew);
139 0 : if (!isNew)
140 0 : elog(ERROR, "unexpected !isNew result from rbt_insert");
141 0 : }
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 0 : if (size > 0)
148 : {
149 0 : node.key = step * permutation[0];
150 0 : rbt_insert(tree, (RBTNode *) &node, &isNew);
151 0 : if (isNew)
152 0 : elog(ERROR, "unexpected isNew result from rbt_insert");
153 0 : }
154 :
155 0 : pfree(permutation);
156 0 : }
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 0 : testleftright(int size)
165 : {
166 0 : RBTree *tree = create_int_rbtree();
167 0 : IntRBTreeNode *node;
168 0 : RBTreeIterator iter;
169 0 : int lastKey = -1;
170 0 : int count = 0;
171 :
172 : /* check iteration over empty tree */
173 0 : rbt_begin_iterate(tree, LeftRightWalk, &iter);
174 0 : 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 0 : rbt_populate(tree, size, 1);
179 :
180 : /* iterate over the tree */
181 0 : rbt_begin_iterate(tree, LeftRightWalk, &iter);
182 :
183 0 : while ((node = (IntRBTreeNode *) rbt_iterate(&iter)) != NULL)
184 : {
185 : /* check that order is increasing */
186 0 : if (node->key <= lastKey)
187 0 : elog(ERROR, "left-right walk gives elements not in sorted order");
188 0 : lastKey = node->key;
189 0 : count++;
190 : }
191 :
192 0 : if (lastKey != size - 1)
193 0 : elog(ERROR, "left-right walk did not reach end");
194 0 : if (count != size)
195 0 : elog(ERROR, "left-right walk missed some elements");
196 0 : }
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 0 : testrightleft(int size)
205 : {
206 0 : RBTree *tree = create_int_rbtree();
207 0 : IntRBTreeNode *node;
208 0 : RBTreeIterator iter;
209 0 : int lastKey = size;
210 0 : int count = 0;
211 :
212 : /* check iteration over empty tree */
213 0 : rbt_begin_iterate(tree, RightLeftWalk, &iter);
214 0 : 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 0 : rbt_populate(tree, size, 1);
219 :
220 : /* iterate over the tree */
221 0 : rbt_begin_iterate(tree, RightLeftWalk, &iter);
222 :
223 0 : while ((node = (IntRBTreeNode *) rbt_iterate(&iter)) != NULL)
224 : {
225 : /* check that order is decreasing */
226 0 : if (node->key >= lastKey)
227 0 : elog(ERROR, "right-left walk gives elements not in sorted order");
228 0 : lastKey = node->key;
229 0 : count++;
230 : }
231 :
232 0 : if (lastKey != 0)
233 0 : elog(ERROR, "right-left walk did not reach end");
234 0 : if (count != size)
235 0 : elog(ERROR, "right-left walk missed some elements");
236 0 : }
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 0 : testfind(int size)
244 : {
245 0 : RBTree *tree = create_int_rbtree();
246 0 : int i;
247 :
248 : /* Insert even integers from 0 to 2 * (size-1) */
249 0 : rbt_populate(tree, size, 2);
250 :
251 : /* Check that all inserted elements can be found */
252 0 : for (i = 0; i < size; i++)
253 : {
254 0 : IntRBTreeNode node;
255 0 : IntRBTreeNode *resultNode;
256 :
257 0 : node.key = 2 * i;
258 0 : resultNode = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &node);
259 0 : if (resultNode == NULL)
260 0 : elog(ERROR, "inserted element was not found");
261 0 : if (node.key != resultNode->key)
262 0 : elog(ERROR, "find operation in rbtree gave wrong result");
263 0 : }
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 0 : for (i = -1; i <= 2 * size; i += 2)
270 : {
271 0 : IntRBTreeNode node;
272 0 : IntRBTreeNode *resultNode;
273 :
274 0 : node.key = i;
275 0 : resultNode = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &node);
276 0 : if (resultNode != NULL)
277 0 : elog(ERROR, "not-inserted element was found");
278 0 : }
279 0 : }
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 0 : testfindltgt(int size)
288 : {
289 0 : 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 0 : int randomKey = pg_prng_uint64_range(&pg_global_prng_state, 0, size - 1);
296 0 : bool keyDeleted;
297 0 : IntRBTreeNode searchNode = {.key = randomKey};
298 0 : IntRBTreeNode *lteNode;
299 0 : IntRBTreeNode *gteNode;
300 0 : IntRBTreeNode *node;
301 :
302 : /* Insert natural numbers */
303 0 : 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 0 : lteNode = (IntRBTreeNode *) rbt_find_less(tree, (RBTNode *) &searchNode, true);
310 0 : gteNode = (IntRBTreeNode *) rbt_find_great(tree, (RBTNode *) &searchNode, true);
311 :
312 0 : if (lteNode == NULL || lteNode->key != searchNode.key)
313 0 : elog(ERROR, "rbt_find_less() didn't find the equal key");
314 :
315 0 : if (gteNode == NULL || gteNode->key != searchNode.key)
316 0 : elog(ERROR, "rbt_find_great() didn't find the equal key");
317 :
318 0 : 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 0 : keyDeleted = false;
323 0 : 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 0 : node = (IntRBTreeNode *) rbt_find_less(tree, (RBTNode *) &searchNode,
330 0 : keyDeleted);
331 :
332 : /* ensure we find a lesser match */
333 0 : 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 0 : keyDeleted = (pg_prng_uint64_range(&pg_global_prng_state, 0, 1) == 1);
338 0 : if (keyDeleted)
339 0 : rbt_delete(tree, (RBTNode *) node);
340 0 : }
341 :
342 : /* Find the rest of the naturals greater than the search key */
343 0 : keyDeleted = false;
344 0 : 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 0 : node = (IntRBTreeNode *) rbt_find_great(tree, (RBTNode *) &searchNode,
351 0 : keyDeleted);
352 :
353 : /* ensure we find a greater match */
354 0 : 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 0 : keyDeleted = (pg_prng_uint64_range(&pg_global_prng_state, 0, 1) == 1);
359 0 : if (keyDeleted)
360 0 : rbt_delete(tree, (RBTNode *) node);
361 0 : }
362 :
363 : /* Check out of bounds searches find nothing */
364 0 : searchNode.key = -1;
365 0 : node = (IntRBTreeNode *) rbt_find_less(tree, (RBTNode *) &searchNode, true);
366 0 : if (node != NULL)
367 0 : elog(ERROR, "rbt_find_less() found non-inserted element");
368 0 : searchNode.key = 0;
369 0 : node = (IntRBTreeNode *) rbt_find_less(tree, (RBTNode *) &searchNode, false);
370 0 : if (node != NULL)
371 0 : elog(ERROR, "rbt_find_less() found non-inserted element");
372 0 : searchNode.key = size;
373 0 : node = (IntRBTreeNode *) rbt_find_great(tree, (RBTNode *) &searchNode, true);
374 0 : if (node != NULL)
375 0 : elog(ERROR, "rbt_find_great() found non-inserted element");
376 0 : searchNode.key = size - 1;
377 0 : node = (IntRBTreeNode *) rbt_find_great(tree, (RBTNode *) &searchNode, false);
378 0 : if (node != NULL)
379 0 : elog(ERROR, "rbt_find_great() found non-inserted element");
380 0 : }
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 0 : testleftmost(int size)
388 : {
389 0 : RBTree *tree = create_int_rbtree();
390 0 : IntRBTreeNode *result;
391 :
392 : /* Check that empty tree has no leftmost element */
393 0 : 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 0 : rbt_populate(tree, size, 1);
398 :
399 : /* Check that leftmost element is the smallest one */
400 0 : result = (IntRBTreeNode *) rbt_leftmost(tree);
401 0 : if (result == NULL || result->key != 0)
402 0 : elog(ERROR, "rbt_leftmost gave wrong result");
403 0 : }
404 :
405 : /*
406 : * Check the correctness of the rbt_delete operation.
407 : */
408 : static void
409 0 : testdelete(int size, int delsize)
410 : {
411 0 : RBTree *tree = create_int_rbtree();
412 0 : int *deleteIds;
413 0 : bool *chosen;
414 0 : int i;
415 :
416 : /* fill tree with consecutive natural numbers */
417 0 : rbt_populate(tree, size, 1);
418 :
419 : /* Choose unique ids to delete */
420 0 : deleteIds = (int *) palloc(delsize * sizeof(int));
421 0 : chosen = (bool *) palloc0(size * sizeof(bool));
422 :
423 0 : for (i = 0; i < delsize; i++)
424 : {
425 0 : int k = pg_prng_uint64_range(&pg_global_prng_state, 0, size - 1);
426 :
427 0 : while (chosen[k])
428 0 : k = (k + 1) % size;
429 0 : deleteIds[i] = k;
430 0 : chosen[k] = true;
431 0 : }
432 :
433 : /* Delete elements */
434 0 : for (i = 0; i < delsize; i++)
435 : {
436 0 : IntRBTreeNode find;
437 0 : IntRBTreeNode *node;
438 :
439 0 : find.key = deleteIds[i];
440 : /* Locate the node to be deleted */
441 0 : node = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &find);
442 0 : if (node == NULL || node->key != deleteIds[i])
443 0 : elog(ERROR, "expected element was not found during deleting");
444 : /* Delete it */
445 0 : rbt_delete(tree, (RBTNode *) node);
446 0 : }
447 :
448 : /* Check that deleted elements are deleted */
449 0 : for (i = 0; i < size; i++)
450 : {
451 0 : IntRBTreeNode node;
452 0 : IntRBTreeNode *result;
453 :
454 0 : node.key = i;
455 0 : result = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &node);
456 0 : if (chosen[i])
457 : {
458 : /* Deleted element should be absent */
459 0 : if (result != NULL)
460 0 : elog(ERROR, "deleted element still present in the rbtree");
461 0 : }
462 : else
463 : {
464 : /* Else it should be present */
465 0 : if (result == NULL || result->key != i)
466 0 : elog(ERROR, "delete operation removed wrong rbtree value");
467 : }
468 0 : }
469 :
470 : /* Delete remaining elements, so as to exercise reducing tree to empty */
471 0 : for (i = 0; i < size; i++)
472 : {
473 0 : IntRBTreeNode find;
474 0 : IntRBTreeNode *node;
475 :
476 0 : if (chosen[i])
477 0 : continue;
478 0 : find.key = i;
479 : /* Locate the node to be deleted */
480 0 : node = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &find);
481 0 : if (node == NULL || node->key != i)
482 0 : elog(ERROR, "expected element was not found during deleting");
483 : /* Delete it */
484 0 : rbt_delete(tree, (RBTNode *) node);
485 0 : }
486 :
487 : /* Tree should now be empty */
488 0 : if (rbt_leftmost(tree) != NULL)
489 0 : elog(ERROR, "deleting all elements failed");
490 :
491 0 : pfree(deleteIds);
492 0 : pfree(chosen);
493 0 : }
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 0 : PG_FUNCTION_INFO_V1(test_rb_tree);
501 :
502 : Datum
503 0 : test_rb_tree(PG_FUNCTION_ARGS)
504 : {
505 0 : int size = PG_GETARG_INT32(0);
506 :
507 0 : if (size <= 0 || size > MaxAllocSize / sizeof(int))
508 0 : elog(ERROR, "invalid size for test_rb_tree: %d", size);
509 0 : testleftright(size);
510 0 : testrightleft(size);
511 0 : testfind(size);
512 0 : testfindltgt(size);
513 0 : testleftmost(size);
514 0 : testdelete(size, Max(size / 10, 1));
515 0 : PG_RETURN_VOID();
516 0 : }
|
