1/*-------------------------------------------------------------------------
2 *
3 * rbtree.c
4 * implementation for PostgreSQL generic Red-Black binary tree package
5 * Adopted from http://algolist.manual.ru/ds/rbtree.php
6 *
7 * This code comes from Thomas Niemann's "Sorting and Searching Algorithms:
8 * a Cookbook".
9 *
10 * See http://www.cs.auckland.ac.nz/software/AlgAnim/niemann/s_man.htm for
11 * license terms: "Source code, when part of a software project, may be used
12 * freely without reference to the author."
13 *
14 * Red-black trees are a type of balanced binary tree wherein (1) any child of
15 * a red node is always black, and (2) every path from root to leaf traverses
16 * an equal number of black nodes. From these properties, it follows that the
17 * longest path from root to leaf is only about twice as long as the shortest,
18 * so lookups are guaranteed to run in O(lg n) time.
19 *
20 * Copyright (c) 2009-2019, PostgreSQL Global Development Group
21 *
22 * IDENTIFICATION
23 * src/backend/lib/rbtree.c
24 *
25 *-------------------------------------------------------------------------
26 */
27#include "postgres.h"
28
29#include "lib/rbtree.h"
30
31
32/*
33 * Colors of nodes (values of RBTNode.color)
34 */
35#define RBTBLACK (0)
36#define RBTRED (1)
37
38/*
39 * RBTree control structure
40 */
41struct RBTree
42{
43 RBTNode *root; /* root node, or RBTNIL if tree is empty */
44
45 /* Remaining fields are constant after rbt_create */
46
47 Size node_size; /* actual size of tree nodes */
48 /* The caller-supplied manipulation functions */
49 rbt_comparator comparator;
50 rbt_combiner combiner;
51 rbt_allocfunc allocfunc;
52 rbt_freefunc freefunc;
53 /* Passthrough arg passed to all manipulation functions */
54 void *arg;
55};
56
57/*
58 * all leafs are sentinels, use customized NIL name to prevent
59 * collision with system-wide constant NIL which is actually NULL
60 */
61#define RBTNIL (&sentinel)
62
63static RBTNode sentinel =
64{
65 RBTBLACK, RBTNIL, RBTNIL, NULL
66};
67
68
69/*
70 * rbt_create: create an empty RBTree
71 *
72 * Arguments are:
73 * node_size: actual size of tree nodes (> sizeof(RBTNode))
74 * The manipulation functions:
75 * comparator: compare two RBTNodes for less/equal/greater
76 * combiner: merge an existing tree entry with a new one
77 * allocfunc: allocate a new RBTNode
78 * freefunc: free an old RBTNode
79 * arg: passthrough pointer that will be passed to the manipulation functions
80 *
81 * Note that the combiner's righthand argument will be a "proposed" tree node,
82 * ie the input to rbt_insert, in which the RBTNode fields themselves aren't
83 * valid. Similarly, either input to the comparator may be a "proposed" node.
84 * This shouldn't matter since the functions aren't supposed to look at the
85 * RBTNode fields, only the extra fields of the struct the RBTNode is embedded
86 * in.
87 *
88 * The freefunc should just be pfree or equivalent; it should NOT attempt
89 * to free any subsidiary data, because the node passed to it may not contain
90 * valid data! freefunc can be NULL if caller doesn't require retail
91 * space reclamation.
92 *
93 * The RBTree node is palloc'd in the caller's memory context. Note that
94 * all contents of the tree are actually allocated by the caller, not here.
95 *
96 * Since tree contents are managed by the caller, there is currently not
97 * an explicit "destroy" operation; typically a tree would be freed by
98 * resetting or deleting the memory context it's stored in. You can pfree
99 * the RBTree node if you feel the urge.
100 */
101RBTree *
102rbt_create(Size node_size,
103 rbt_comparator comparator,
104 rbt_combiner combiner,
105 rbt_allocfunc allocfunc,
106 rbt_freefunc freefunc,
107 void *arg)
108{
109 RBTree *tree = (RBTree *) palloc(sizeof(RBTree));
110
111 Assert(node_size > sizeof(RBTNode));
112
113 tree->root = RBTNIL;
114 tree->node_size = node_size;
115 tree->comparator = comparator;
116 tree->combiner = combiner;
117 tree->allocfunc = allocfunc;
118 tree->freefunc = freefunc;
119
120 tree->arg = arg;
121
122 return tree;
123}
124
125/* Copy the additional data fields from one RBTNode to another */
126static inline void
127rbt_copy_data(RBTree *rbt, RBTNode *dest, const RBTNode *src)
128{
129 memcpy(dest + 1, src + 1, rbt->node_size - sizeof(RBTNode));
130}
131
132/**********************************************************************
133 * Search *
134 **********************************************************************/
135
136/*
137 * rbt_find: search for a value in an RBTree
138 *
139 * data represents the value to try to find. Its RBTNode fields need not
140 * be valid, it's the extra data in the larger struct that is of interest.
141 *
142 * Returns the matching tree entry, or NULL if no match is found.
143 */
144RBTNode *
145rbt_find(RBTree *rbt, const RBTNode *data)
146{
147 RBTNode *node = rbt->root;
148
149 while (node != RBTNIL)
150 {
151 int cmp = rbt->comparator(data, node, rbt->arg);
152
153 if (cmp == 0)
154 return node;
155 else if (cmp < 0)
156 node = node->left;
157 else
158 node = node->right;
159 }
160
161 return NULL;
162}
163
164/*
165 * rbt_leftmost: fetch the leftmost (smallest-valued) tree node.
166 * Returns NULL if tree is empty.
167 *
168 * Note: in the original implementation this included an unlink step, but
169 * that's a bit awkward. Just call rbt_delete on the result if that's what
170 * you want.
171 */
172RBTNode *
173rbt_leftmost(RBTree *rbt)
174{
175 RBTNode *node = rbt->root;
176 RBTNode *leftmost = rbt->root;
177
178 while (node != RBTNIL)
179 {
180 leftmost = node;
181 node = node->left;
182 }
183
184 if (leftmost != RBTNIL)
185 return leftmost;
186
187 return NULL;
188}
189
190/**********************************************************************
191 * Insertion *
192 **********************************************************************/
193
194/*
195 * Rotate node x to left.
196 *
197 * x's right child takes its place in the tree, and x becomes the left
198 * child of that node.
199 */
200static void
201rbt_rotate_left(RBTree *rbt, RBTNode *x)
202{
203 RBTNode *y = x->right;
204
205 /* establish x->right link */
206 x->right = y->left;
207 if (y->left != RBTNIL)
208 y->left->parent = x;
209
210 /* establish y->parent link */
211 if (y != RBTNIL)
212 y->parent = x->parent;
213 if (x->parent)
214 {
215 if (x == x->parent->left)
216 x->parent->left = y;
217 else
218 x->parent->right = y;
219 }
220 else
221 {
222 rbt->root = y;
223 }
224
225 /* link x and y */
226 y->left = x;
227 if (x != RBTNIL)
228 x->parent = y;
229}
230
231/*
232 * Rotate node x to right.
233 *
234 * x's left right child takes its place in the tree, and x becomes the right
235 * child of that node.
236 */
237static void
238rbt_rotate_right(RBTree *rbt, RBTNode *x)
239{
240 RBTNode *y = x->left;
241
242 /* establish x->left link */
243 x->left = y->right;
244 if (y->right != RBTNIL)
245 y->right->parent = x;
246
247 /* establish y->parent link */
248 if (y != RBTNIL)
249 y->parent = x->parent;
250 if (x->parent)
251 {
252 if (x == x->parent->right)
253 x->parent->right = y;
254 else
255 x->parent->left = y;
256 }
257 else
258 {
259 rbt->root = y;
260 }
261
262 /* link x and y */
263 y->right = x;
264 if (x != RBTNIL)
265 x->parent = y;
266}
267
268/*
269 * Maintain Red-Black tree balance after inserting node x.
270 *
271 * The newly inserted node is always initially marked red. That may lead to
272 * a situation where a red node has a red child, which is prohibited. We can
273 * always fix the problem by a series of color changes and/or "rotations",
274 * which move the problem progressively higher up in the tree. If one of the
275 * two red nodes is the root, we can always fix the problem by changing the
276 * root from red to black.
277 *
278 * (This does not work lower down in the tree because we must also maintain
279 * the invariant that every leaf has equal black-height.)
280 */
281static void
282rbt_insert_fixup(RBTree *rbt, RBTNode *x)
283{
284 /*
285 * x is always a red node. Initially, it is the newly inserted node. Each
286 * iteration of this loop moves it higher up in the tree.
287 */
288 while (x != rbt->root && x->parent->color == RBTRED)
289 {
290 /*
291 * x and x->parent are both red. Fix depends on whether x->parent is
292 * a left or right child. In either case, we define y to be the
293 * "uncle" of x, that is, the other child of x's grandparent.
294 *
295 * If the uncle is red, we flip the grandparent to red and its two
296 * children to black. Then we loop around again to check whether the
297 * grandparent still has a problem.
298 *
299 * If the uncle is black, we will perform one or two "rotations" to
300 * balance the tree. Either x or x->parent will take the
301 * grandparent's position in the tree and recolored black, and the
302 * original grandparent will be recolored red and become a child of
303 * that node. This always leaves us with a valid red-black tree, so
304 * the loop will terminate.
305 */
306 if (x->parent == x->parent->parent->left)
307 {
308 RBTNode *y = x->parent->parent->right;
309
310 if (y->color == RBTRED)
311 {
312 /* uncle is RBTRED */
313 x->parent->color = RBTBLACK;
314 y->color = RBTBLACK;
315 x->parent->parent->color = RBTRED;
316
317 x = x->parent->parent;
318 }
319 else
320 {
321 /* uncle is RBTBLACK */
322 if (x == x->parent->right)
323 {
324 /* make x a left child */
325 x = x->parent;
326 rbt_rotate_left(rbt, x);
327 }
328
329 /* recolor and rotate */
330 x->parent->color = RBTBLACK;
331 x->parent->parent->color = RBTRED;
332
333 rbt_rotate_right(rbt, x->parent->parent);
334 }
335 }
336 else
337 {
338 /* mirror image of above code */
339 RBTNode *y = x->parent->parent->left;
340
341 if (y->color == RBTRED)
342 {
343 /* uncle is RBTRED */
344 x->parent->color = RBTBLACK;
345 y->color = RBTBLACK;
346 x->parent->parent->color = RBTRED;
347
348 x = x->parent->parent;
349 }
350 else
351 {
352 /* uncle is RBTBLACK */
353 if (x == x->parent->left)
354 {
355 x = x->parent;
356 rbt_rotate_right(rbt, x);
357 }
358 x->parent->color = RBTBLACK;
359 x->parent->parent->color = RBTRED;
360
361 rbt_rotate_left(rbt, x->parent->parent);
362 }
363 }
364 }
365
366 /*
367 * The root may already have been black; if not, the black-height of every
368 * node in the tree increases by one.
369 */
370 rbt->root->color = RBTBLACK;
371}
372
373/*
374 * rbt_insert: insert a new value into the tree.
375 *
376 * data represents the value to insert. Its RBTNode fields need not
377 * be valid, it's the extra data in the larger struct that is of interest.
378 *
379 * If the value represented by "data" is not present in the tree, then
380 * we copy "data" into a new tree entry and return that node, setting *isNew
381 * to true.
382 *
383 * If the value represented by "data" is already present, then we call the
384 * combiner function to merge data into the existing node, and return the
385 * existing node, setting *isNew to false.
386 *
387 * "data" is unmodified in either case; it's typically just a local
388 * variable in the caller.
389 */
390RBTNode *
391rbt_insert(RBTree *rbt, const RBTNode *data, bool *isNew)
392{
393 RBTNode *current,
394 *parent,
395 *x;
396 int cmp;
397
398 /* find where node belongs */
399 current = rbt->root;
400 parent = NULL;
401 cmp = 0; /* just to prevent compiler warning */
402
403 while (current != RBTNIL)
404 {
405 cmp = rbt->comparator(data, current, rbt->arg);
406 if (cmp == 0)
407 {
408 /*
409 * Found node with given key. Apply combiner.
410 */
411 rbt->combiner(current, data, rbt->arg);
412 *isNew = false;
413 return current;
414 }
415 parent = current;
416 current = (cmp < 0) ? current->left : current->right;
417 }
418
419 /*
420 * Value is not present, so create a new node containing data.
421 */
422 *isNew = true;
423
424 x = rbt->allocfunc(rbt->arg);
425
426 x->color = RBTRED;
427
428 x->left = RBTNIL;
429 x->right = RBTNIL;
430 x->parent = parent;
431 rbt_copy_data(rbt, x, data);
432
433 /* insert node in tree */
434 if (parent)
435 {
436 if (cmp < 0)
437 parent->left = x;
438 else
439 parent->right = x;
440 }
441 else
442 {
443 rbt->root = x;
444 }
445
446 rbt_insert_fixup(rbt, x);
447
448 return x;
449}
450
451/**********************************************************************
452 * Deletion *
453 **********************************************************************/
454
455/*
456 * Maintain Red-Black tree balance after deleting a black node.
457 */
458static void
459rbt_delete_fixup(RBTree *rbt, RBTNode *x)
460{
461 /*
462 * x is always a black node. Initially, it is the former child of the
463 * deleted node. Each iteration of this loop moves it higher up in the
464 * tree.
465 */
466 while (x != rbt->root && x->color == RBTBLACK)
467 {
468 /*
469 * Left and right cases are symmetric. Any nodes that are children of
470 * x have a black-height one less than the remainder of the nodes in
471 * the tree. We rotate and recolor nodes to move the problem up the
472 * tree: at some stage we'll either fix the problem, or reach the root
473 * (where the black-height is allowed to decrease).
474 */
475 if (x == x->parent->left)
476 {
477 RBTNode *w = x->parent->right;
478
479 if (w->color == RBTRED)
480 {
481 w->color = RBTBLACK;
482 x->parent->color = RBTRED;
483
484 rbt_rotate_left(rbt, x->parent);
485 w = x->parent->right;
486 }
487
488 if (w->left->color == RBTBLACK && w->right->color == RBTBLACK)
489 {
490 w->color = RBTRED;
491
492 x = x->parent;
493 }
494 else
495 {
496 if (w->right->color == RBTBLACK)
497 {
498 w->left->color = RBTBLACK;
499 w->color = RBTRED;
500
501 rbt_rotate_right(rbt, w);
502 w = x->parent->right;
503 }
504 w->color = x->parent->color;
505 x->parent->color = RBTBLACK;
506 w->right->color = RBTBLACK;
507
508 rbt_rotate_left(rbt, x->parent);
509 x = rbt->root; /* Arrange for loop to terminate. */
510 }
511 }
512 else
513 {
514 RBTNode *w = x->parent->left;
515
516 if (w->color == RBTRED)
517 {
518 w->color = RBTBLACK;
519 x->parent->color = RBTRED;
520
521 rbt_rotate_right(rbt, x->parent);
522 w = x->parent->left;
523 }
524
525 if (w->right->color == RBTBLACK && w->left->color == RBTBLACK)
526 {
527 w->color = RBTRED;
528
529 x = x->parent;
530 }
531 else
532 {
533 if (w->left->color == RBTBLACK)
534 {
535 w->right->color = RBTBLACK;
536 w->color = RBTRED;
537
538 rbt_rotate_left(rbt, w);
539 w = x->parent->left;
540 }
541 w->color = x->parent->color;
542 x->parent->color = RBTBLACK;
543 w->left->color = RBTBLACK;
544
545 rbt_rotate_right(rbt, x->parent);
546 x = rbt->root; /* Arrange for loop to terminate. */
547 }
548 }
549 }
550 x->color = RBTBLACK;
551}
552
553/*
554 * Delete node z from tree.
555 */
556static void
557rbt_delete_node(RBTree *rbt, RBTNode *z)
558{
559 RBTNode *x,
560 *y;
561
562 /* This is just paranoia: we should only get called on a valid node */
563 if (!z || z == RBTNIL)
564 return;
565
566 /*
567 * y is the node that will actually be removed from the tree. This will
568 * be z if z has fewer than two children, or the tree successor of z
569 * otherwise.
570 */
571 if (z->left == RBTNIL || z->right == RBTNIL)
572 {
573 /* y has a RBTNIL node as a child */
574 y = z;
575 }
576 else
577 {
578 /* find tree successor */
579 y = z->right;
580 while (y->left != RBTNIL)
581 y = y->left;
582 }
583
584 /* x is y's only child */
585 if (y->left != RBTNIL)
586 x = y->left;
587 else
588 x = y->right;
589
590 /* Remove y from the tree. */
591 x->parent = y->parent;
592 if (y->parent)
593 {
594 if (y == y->parent->left)
595 y->parent->left = x;
596 else
597 y->parent->right = x;
598 }
599 else
600 {
601 rbt->root = x;
602 }
603
604 /*
605 * If we removed the tree successor of z rather than z itself, then move
606 * the data for the removed node to the one we were supposed to remove.
607 */
608 if (y != z)
609 rbt_copy_data(rbt, z, y);
610
611 /*
612 * Removing a black node might make some paths from root to leaf contain
613 * fewer black nodes than others, or it might make two red nodes adjacent.
614 */
615 if (y->color == RBTBLACK)
616 rbt_delete_fixup(rbt, x);
617
618 /* Now we can recycle the y node */
619 if (rbt->freefunc)
620 rbt->freefunc(y, rbt->arg);
621}
622
623/*
624 * rbt_delete: remove the given tree entry
625 *
626 * "node" must have previously been found via rbt_find or rbt_leftmost.
627 * It is caller's responsibility to free any subsidiary data attached
628 * to the node before calling rbt_delete. (Do *not* try to push that
629 * responsibility off to the freefunc, as some other physical node
630 * may be the one actually freed!)
631 */
632void
633rbt_delete(RBTree *rbt, RBTNode *node)
634{
635 rbt_delete_node(rbt, node);
636}
637
638/**********************************************************************
639 * Traverse *
640 **********************************************************************/
641
642static RBTNode *
643rbt_left_right_iterator(RBTreeIterator *iter)
644{
645 if (iter->last_visited == NULL)
646 {
647 iter->last_visited = iter->rbt->root;
648 while (iter->last_visited->left != RBTNIL)
649 iter->last_visited = iter->last_visited->left;
650
651 return iter->last_visited;
652 }
653
654 if (iter->last_visited->right != RBTNIL)
655 {
656 iter->last_visited = iter->last_visited->right;
657 while (iter->last_visited->left != RBTNIL)
658 iter->last_visited = iter->last_visited->left;
659
660 return iter->last_visited;
661 }
662
663 for (;;)
664 {
665 RBTNode *came_from = iter->last_visited;
666
667 iter->last_visited = iter->last_visited->parent;
668 if (iter->last_visited == NULL)
669 {
670 iter->is_over = true;
671 break;
672 }
673
674 if (iter->last_visited->left == came_from)
675 break; /* came from left sub-tree, return current
676 * node */
677
678 /* else - came from right sub-tree, continue to move up */
679 }
680
681 return iter->last_visited;
682}
683
684static RBTNode *
685rbt_right_left_iterator(RBTreeIterator *iter)
686{
687 if (iter->last_visited == NULL)
688 {
689 iter->last_visited = iter->rbt->root;
690 while (iter->last_visited->right != RBTNIL)
691 iter->last_visited = iter->last_visited->right;
692
693 return iter->last_visited;
694 }
695
696 if (iter->last_visited->left != RBTNIL)
697 {
698 iter->last_visited = iter->last_visited->left;
699 while (iter->last_visited->right != RBTNIL)
700 iter->last_visited = iter->last_visited->right;
701
702 return iter->last_visited;
703 }
704
705 for (;;)
706 {
707 RBTNode *came_from = iter->last_visited;
708
709 iter->last_visited = iter->last_visited->parent;
710 if (iter->last_visited == NULL)
711 {
712 iter->is_over = true;
713 break;
714 }
715
716 if (iter->last_visited->right == came_from)
717 break; /* came from right sub-tree, return current
718 * node */
719
720 /* else - came from left sub-tree, continue to move up */
721 }
722
723 return iter->last_visited;
724}
725
726/*
727 * rbt_begin_iterate: prepare to traverse the tree in any of several orders
728 *
729 * After calling rbt_begin_iterate, call rbt_iterate repeatedly until it
730 * returns NULL or the traversal stops being of interest.
731 *
732 * If the tree is changed during traversal, results of further calls to
733 * rbt_iterate are unspecified. Multiple concurrent iterators on the same
734 * tree are allowed.
735 *
736 * The iterator state is stored in the 'iter' struct. The caller should
737 * treat it as an opaque struct.
738 */
739void
740rbt_begin_iterate(RBTree *rbt, RBTOrderControl ctrl, RBTreeIterator *iter)
741{
742 /* Common initialization for all traversal orders */
743 iter->rbt = rbt;
744 iter->last_visited = NULL;
745 iter->is_over = (rbt->root == RBTNIL);
746
747 switch (ctrl)
748 {
749 case LeftRightWalk: /* visit left, then self, then right */
750 iter->iterate = rbt_left_right_iterator;
751 break;
752 case RightLeftWalk: /* visit right, then self, then left */
753 iter->iterate = rbt_right_left_iterator;
754 break;
755 default:
756 elog(ERROR, "unrecognized rbtree iteration order: %d", ctrl);
757 }
758}
759
760/*
761 * rbt_iterate: return the next node in traversal order, or NULL if no more
762 */
763RBTNode *
764rbt_iterate(RBTreeIterator *iter)
765{
766 if (iter->is_over)
767 return NULL;
768
769 return iter->iterate(iter);
770}
771