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 1337779 : irbt_cmp(const RBTNode *a, const RBTNode *b, void *arg)
40 : {
41 1337779 : const IntRBTreeNode *ea = (const IntRBTreeNode *) a;
42 1337779 : const IntRBTreeNode *eb = (const IntRBTreeNode *) b;
43 :
44 1337779 : 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 15002 : irbt_free(RBTNode *node, void *arg)
72 : {
73 15002 : pfree(node);
74 15002 : }
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 59939 : 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 6004 : 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 6003 : node = (IntRBTreeNode *) rbt_find_less(tree, (RBTNode *) &searchNode,
330 : keyDeleted);
331 :
332 : /* ensure we find a lesser match */
333 6003 : 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 6003 : keyDeleted = (pg_prng_uint64_range(&pg_global_prng_state, 0, 1) == 1);
338 6003 : if (keyDeleted)
339 3026 : rbt_delete(tree, (RBTNode *) node);
340 : }
341 :
342 : /* Find the rest of the naturals greater than the search key */
343 1 : keyDeleted = false;
344 3997 : 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 3996 : node = (IntRBTreeNode *) rbt_find_great(tree, (RBTNode *) &searchNode,
351 : keyDeleted);
352 :
353 : /* ensure we find a greater match */
354 3996 : 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 3996 : keyDeleted = (pg_prng_uint64_range(&pg_global_prng_state, 0, 1) == 1);
359 3996 : if (keyDeleted)
360 1976 : 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 1053 : while (chosen[k])
428 53 : 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 : }
|
