Line data Source code
1 : /*--------------------------------------------------------------------------
2 : *
3 : * test_rbtree.c
4 : * Test correctness of red-black tree operations.
5 : *
6 : * Copyright (c) 2009-2021, 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 "fmgr.h"
17 : #include "lib/rbtree.h"
18 : #include "utils/memutils.h"
19 :
20 2 : PG_MODULE_MAGIC;
21 :
22 :
23 : /*
24 : * Our test trees store an integer key, and nothing else.
25 : */
26 : typedef struct IntRBTreeNode
27 : {
28 : RBTNode rbtnode;
29 : int key;
30 : } IntRBTreeNode;
31 :
32 :
33 : /*
34 : * Node comparator. We don't worry about overflow in the subtraction,
35 : * since none of our test keys are negative.
36 : */
37 : static int
38 2172338 : irbt_cmp(const RBTNode *a, const RBTNode *b, void *arg)
39 : {
40 2172338 : const IntRBTreeNode *ea = (const IntRBTreeNode *) a;
41 2172338 : const IntRBTreeNode *eb = (const IntRBTreeNode *) b;
42 :
43 2172338 : return ea->key - eb->key;
44 : }
45 :
46 : /*
47 : * Node combiner. For testing purposes, just check that library doesn't
48 : * try to combine unequal keys.
49 : */
50 : static void
51 10 : irbt_combine(RBTNode *existing, const RBTNode *newdata, void *arg)
52 : {
53 10 : const IntRBTreeNode *eexist = (const IntRBTreeNode *) existing;
54 10 : const IntRBTreeNode *enew = (const IntRBTreeNode *) newdata;
55 :
56 10 : if (eexist->key != enew->key)
57 0 : elog(ERROR, "red-black tree combines %d into %d",
58 : enew->key, eexist->key);
59 10 : }
60 :
61 : /* Node allocator */
62 : static RBTNode *
63 100000 : irbt_alloc(void *arg)
64 : {
65 100000 : return (RBTNode *) palloc(sizeof(IntRBTreeNode));
66 : }
67 :
68 : /* Node freer */
69 : static void
70 20000 : irbt_free(RBTNode *node, void *arg)
71 : {
72 20000 : pfree(node);
73 20000 : }
74 :
75 : /*
76 : * Create a red-black tree using our support functions
77 : */
78 : static RBTree *
79 10 : create_int_rbtree(void)
80 : {
81 10 : return rbt_create(sizeof(IntRBTreeNode),
82 : irbt_cmp,
83 : irbt_combine,
84 : irbt_alloc,
85 : irbt_free,
86 : NULL);
87 : }
88 :
89 : /*
90 : * Generate a random permutation of the integers 0..size-1
91 : */
92 : static int *
93 10 : GetPermutation(int size)
94 : {
95 : int *permutation;
96 : int i;
97 :
98 10 : permutation = (int *) palloc(size * sizeof(int));
99 :
100 10 : permutation[0] = 0;
101 :
102 : /*
103 : * This is the "inside-out" variant of the Fisher-Yates shuffle algorithm.
104 : * Notionally, we append each new value to the array and then swap it with
105 : * a randomly-chosen array element (possibly including itself, else we
106 : * fail to generate permutations with the last integer last). The swap
107 : * step can be optimized by combining it with the insertion.
108 : */
109 100000 : for (i = 1; i < size; i++)
110 : {
111 99990 : int j = random() % (i + 1);
112 :
113 99990 : if (j < i) /* avoid fetching undefined data if j=i */
114 99916 : permutation[i] = permutation[j];
115 99990 : permutation[j] = i;
116 : }
117 :
118 10 : return permutation;
119 : }
120 :
121 : /*
122 : * Populate an empty RBTree with "size" integers having the values
123 : * 0, step, 2*step, 3*step, ..., inserting them in random order
124 : */
125 : static void
126 10 : rbt_populate(RBTree *tree, int size, int step)
127 : {
128 10 : int *permutation = GetPermutation(size);
129 : IntRBTreeNode node;
130 : bool isNew;
131 : int i;
132 :
133 : /* Insert values. We don't expect any collisions. */
134 100010 : for (i = 0; i < size; i++)
135 : {
136 100000 : node.key = step * permutation[i];
137 100000 : rbt_insert(tree, (RBTNode *) &node, &isNew);
138 100000 : if (!isNew)
139 0 : elog(ERROR, "unexpected !isNew result from rbt_insert");
140 : }
141 :
142 : /*
143 : * Re-insert the first value to make sure collisions work right. It's
144 : * probably not useful to test that case over again for all the values.
145 : */
146 10 : if (size > 0)
147 : {
148 10 : node.key = step * permutation[0];
149 10 : rbt_insert(tree, (RBTNode *) &node, &isNew);
150 10 : if (isNew)
151 0 : elog(ERROR, "unexpected isNew result from rbt_insert");
152 : }
153 :
154 10 : pfree(permutation);
155 10 : }
156 :
157 : /*
158 : * Check the correctness of left-right traversal.
159 : * Left-right traversal is correct if all elements are
160 : * visited in increasing order.
161 : */
162 : static void
163 2 : testleftright(int size)
164 : {
165 2 : RBTree *tree = create_int_rbtree();
166 : IntRBTreeNode *node;
167 : RBTreeIterator iter;
168 2 : int lastKey = -1;
169 2 : int count = 0;
170 :
171 : /* check iteration over empty tree */
172 2 : rbt_begin_iterate(tree, LeftRightWalk, &iter);
173 2 : if (rbt_iterate(&iter) != NULL)
174 0 : elog(ERROR, "left-right walk over empty tree produced an element");
175 :
176 : /* fill tree with consecutive natural numbers */
177 2 : rbt_populate(tree, size, 1);
178 :
179 : /* iterate over the tree */
180 2 : rbt_begin_iterate(tree, LeftRightWalk, &iter);
181 :
182 20002 : while ((node = (IntRBTreeNode *) rbt_iterate(&iter)) != NULL)
183 : {
184 : /* check that order is increasing */
185 20000 : if (node->key <= lastKey)
186 0 : elog(ERROR, "left-right walk gives elements not in sorted order");
187 20000 : lastKey = node->key;
188 20000 : count++;
189 : }
190 :
191 2 : if (lastKey != size - 1)
192 0 : elog(ERROR, "left-right walk did not reach end");
193 2 : if (count != size)
194 0 : elog(ERROR, "left-right walk missed some elements");
195 2 : }
196 :
197 : /*
198 : * Check the correctness of right-left traversal.
199 : * Right-left traversal is correct if all elements are
200 : * visited in decreasing order.
201 : */
202 : static void
203 2 : testrightleft(int size)
204 : {
205 2 : RBTree *tree = create_int_rbtree();
206 : IntRBTreeNode *node;
207 : RBTreeIterator iter;
208 2 : int lastKey = size;
209 2 : int count = 0;
210 :
211 : /* check iteration over empty tree */
212 2 : rbt_begin_iterate(tree, RightLeftWalk, &iter);
213 2 : if (rbt_iterate(&iter) != NULL)
214 0 : elog(ERROR, "right-left walk over empty tree produced an element");
215 :
216 : /* fill tree with consecutive natural numbers */
217 2 : rbt_populate(tree, size, 1);
218 :
219 : /* iterate over the tree */
220 2 : rbt_begin_iterate(tree, RightLeftWalk, &iter);
221 :
222 20002 : while ((node = (IntRBTreeNode *) rbt_iterate(&iter)) != NULL)
223 : {
224 : /* check that order is decreasing */
225 20000 : if (node->key >= lastKey)
226 0 : elog(ERROR, "right-left walk gives elements not in sorted order");
227 20000 : lastKey = node->key;
228 20000 : count++;
229 : }
230 :
231 2 : if (lastKey != 0)
232 0 : elog(ERROR, "right-left walk did not reach end");
233 2 : if (count != size)
234 0 : elog(ERROR, "right-left walk missed some elements");
235 2 : }
236 :
237 : /*
238 : * Check the correctness of the rbt_find operation by searching for
239 : * both elements we inserted and elements we didn't.
240 : */
241 : static void
242 2 : testfind(int size)
243 : {
244 2 : RBTree *tree = create_int_rbtree();
245 : int i;
246 :
247 : /* Insert even integers from 0 to 2 * (size-1) */
248 2 : rbt_populate(tree, size, 2);
249 :
250 : /* Check that all inserted elements can be found */
251 20002 : for (i = 0; i < size; i++)
252 : {
253 : IntRBTreeNode node;
254 : IntRBTreeNode *resultNode;
255 :
256 20000 : node.key = 2 * i;
257 20000 : resultNode = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &node);
258 20000 : if (resultNode == NULL)
259 0 : elog(ERROR, "inserted element was not found");
260 20000 : if (node.key != resultNode->key)
261 0 : elog(ERROR, "find operation in rbtree gave wrong result");
262 : }
263 :
264 : /*
265 : * Check that not-inserted elements can not be found, being sure to try
266 : * values before the first and after the last element.
267 : */
268 20004 : for (i = -1; i <= 2 * size; i += 2)
269 : {
270 : IntRBTreeNode node;
271 : IntRBTreeNode *resultNode;
272 :
273 20002 : node.key = i;
274 20002 : resultNode = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &node);
275 20002 : if (resultNode != NULL)
276 0 : elog(ERROR, "not-inserted element was found");
277 : }
278 2 : }
279 :
280 : /*
281 : * Check the correctness of the rbt_leftmost operation.
282 : * This operation should always return the smallest element of the tree.
283 : */
284 : static void
285 2 : testleftmost(int size)
286 : {
287 2 : RBTree *tree = create_int_rbtree();
288 : IntRBTreeNode *result;
289 :
290 : /* Check that empty tree has no leftmost element */
291 2 : if (rbt_leftmost(tree) != NULL)
292 0 : elog(ERROR, "leftmost node of empty tree is not NULL");
293 :
294 : /* fill tree with consecutive natural numbers */
295 2 : rbt_populate(tree, size, 1);
296 :
297 : /* Check that leftmost element is the smallest one */
298 2 : result = (IntRBTreeNode *) rbt_leftmost(tree);
299 2 : if (result == NULL || result->key != 0)
300 0 : elog(ERROR, "rbt_leftmost gave wrong result");
301 2 : }
302 :
303 : /*
304 : * Check the correctness of the rbt_delete operation.
305 : */
306 : static void
307 2 : testdelete(int size, int delsize)
308 : {
309 2 : RBTree *tree = create_int_rbtree();
310 : int *deleteIds;
311 : bool *chosen;
312 : int i;
313 :
314 : /* fill tree with consecutive natural numbers */
315 2 : rbt_populate(tree, size, 1);
316 :
317 : /* Choose unique ids to delete */
318 2 : deleteIds = (int *) palloc(delsize * sizeof(int));
319 2 : chosen = (bool *) palloc0(size * sizeof(bool));
320 :
321 2002 : for (i = 0; i < delsize; i++)
322 : {
323 2000 : int k = random() % size;
324 :
325 2092 : while (chosen[k])
326 92 : k = (k + 1) % size;
327 2000 : deleteIds[i] = k;
328 2000 : chosen[k] = true;
329 : }
330 :
331 : /* Delete elements */
332 2002 : for (i = 0; i < delsize; i++)
333 : {
334 : IntRBTreeNode find;
335 : IntRBTreeNode *node;
336 :
337 2000 : find.key = deleteIds[i];
338 : /* Locate the node to be deleted */
339 2000 : node = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &find);
340 2000 : if (node == NULL || node->key != deleteIds[i])
341 0 : elog(ERROR, "expected element was not found during deleting");
342 : /* Delete it */
343 2000 : rbt_delete(tree, (RBTNode *) node);
344 : }
345 :
346 : /* Check that deleted elements are deleted */
347 20002 : for (i = 0; i < size; i++)
348 : {
349 : IntRBTreeNode node;
350 : IntRBTreeNode *result;
351 :
352 20000 : node.key = i;
353 20000 : result = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &node);
354 20000 : if (chosen[i])
355 : {
356 : /* Deleted element should be absent */
357 2000 : if (result != NULL)
358 0 : elog(ERROR, "deleted element still present in the rbtree");
359 : }
360 : else
361 : {
362 : /* Else it should be present */
363 18000 : if (result == NULL || result->key != i)
364 0 : elog(ERROR, "delete operation removed wrong rbtree value");
365 : }
366 : }
367 :
368 : /* Delete remaining elements, so as to exercise reducing tree to empty */
369 20002 : for (i = 0; i < size; i++)
370 : {
371 : IntRBTreeNode find;
372 : IntRBTreeNode *node;
373 :
374 20000 : if (chosen[i])
375 2000 : continue;
376 18000 : find.key = i;
377 : /* Locate the node to be deleted */
378 18000 : node = (IntRBTreeNode *) rbt_find(tree, (RBTNode *) &find);
379 18000 : if (node == NULL || node->key != i)
380 0 : elog(ERROR, "expected element was not found during deleting");
381 : /* Delete it */
382 18000 : rbt_delete(tree, (RBTNode *) node);
383 : }
384 :
385 : /* Tree should now be empty */
386 2 : if (rbt_leftmost(tree) != NULL)
387 0 : elog(ERROR, "deleting all elements failed");
388 :
389 2 : pfree(deleteIds);
390 2 : pfree(chosen);
391 2 : }
392 :
393 : /*
394 : * SQL-callable entry point to perform all tests
395 : *
396 : * Argument is the number of entries to put in the trees
397 : */
398 4 : PG_FUNCTION_INFO_V1(test_rb_tree);
399 :
400 : Datum
401 2 : test_rb_tree(PG_FUNCTION_ARGS)
402 : {
403 2 : int size = PG_GETARG_INT32(0);
404 :
405 2 : if (size <= 0 || size > MaxAllocSize / sizeof(int))
406 0 : elog(ERROR, "invalid size for test_rb_tree: %d", size);
407 2 : testleftright(size);
408 2 : testrightleft(size);
409 2 : testfind(size);
410 2 : testleftmost(size);
411 2 : testdelete(size, Max(size / 10, 1));
412 2 : PG_RETURN_VOID();
413 : }
|