1 /*
2 * Copyright (c) 2023 Institute of Parallel And Distributed Systems (IPADS), Shanghai Jiao Tong University (SJTU)
3 * Licensed under the Mulan PSL v2.
4 * You can use this software according to the terms and conditions of the Mulan PSL v2.
5 * You may obtain a copy of Mulan PSL v2 at:
6 * http://license.coscl.org.cn/MulanPSL2
7 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR
8 * IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR
9 * PURPOSE.
10 * See the Mulan PSL v2 for more details.
11 */
12 #include "common/rbtree.h"
13 #include "common/macro.h"
14 #include "common/types.h"
15 #include <common/rbtree_impl.h>
16
__rb_link_node(struct rb_node * node,struct rb_node * parent,struct rb_node ** child_link)17 static inline void __rb_link_node(struct rb_node *node, struct rb_node *parent,
18 struct rb_node **child_link)
19 {
20 node->__parent_color = 0;
21 rb_set_parent(node, parent);
22 node->left_child = node->right_child = NULL;
23 *child_link = node;
24 }
25
__rb_change_child(struct rb_root * root,struct rb_node * parent,struct rb_node * old,struct rb_node * new)26 static inline void __rb_change_child(struct rb_root *root,
27 struct rb_node *parent,
28 struct rb_node *old, struct rb_node *new)
29 {
30 if (parent) {
31 if (parent->left_child == old) {
32 parent->left_child = new;
33 } else {
34 parent->right_child = new;
35 }
36 } else {
37 root->root_node = new;
38 }
39 if (new) {
40 rb_set_parent(new, parent);
41 }
42 }
43
44 /*
45 * Reference: 邓俊辉. 数据结构: C++ 语言版. 清华大学出版社
46 * All double rotation variants (zip-zap, zap-zip, zip-zip, zap-zap...)
47 * can be unified to following connect34 operations. Because those
48 * operations are aimed at rebalancing three unbalanced nodes to the
49 * standard balanced state, like this:
50 * b
51 * / \
52 * a c
53 * / \ / \
54 * t1 t2 t3 t4
55 * There are three nodes and 4 corresponding subtrees, that's why this
56 * function is named as connect34.
57 * Differente kinds of double rotations only need to set arguments of
58 * this function according to their targeted tree patterns. See codes
59 * in __rb_insert_color and __rb_erase_color for concrete examples.
60 */
__connect34(struct rb_node * a,struct rb_node * b,struct rb_node * c,struct rb_node * t1,struct rb_node * t2,struct rb_node * t3,struct rb_node * t4)61 static inline void __connect34(struct rb_node *a, struct rb_node *b,
62 struct rb_node *c, struct rb_node *t1,
63 struct rb_node *t2, struct rb_node *t3,
64 struct rb_node *t4)
65 {
66 a->left_child = t1;
67 if (t1) {
68 rb_set_parent(t1, a);
69 }
70 a->right_child = t2;
71 if (t2) {
72 rb_set_parent(t2, a);
73 }
74 c->left_child = t3;
75 if (t3) {
76 rb_set_parent(t3, c);
77 }
78 c->right_child = t4;
79 if (t4) {
80 rb_set_parent(t4, c);
81 }
82 b->left_child = a;
83 b->right_child = c;
84 rb_set_parent(a, b);
85 rb_set_parent(c, b);
86 }
87
88 /*
89 * We use characters to represent a tree node. Lower case characters indicates
90 * a red node, and upper case ones indicate a black node. Characters enclosed in
91 * brackets represent any color. We assumed that every node has two subtrees,
92 * they can be actually subtree with identical blackheight, or NULL representing
93 * external nodes.
94 */
__rb_insert_color(struct rb_root * root,struct rb_node * node)95 static void __rb_insert_color(struct rb_root *root, struct rb_node *node)
96 {
97 struct rb_node *parent = rb_parent(node), *gparent = NULL, *uncle = NULL;
98
99 while (true) {
100 /*
101 * Loop invariant
102 * - @parent and @node is red
103 */
104 if (!parent) {
105 rb_set_color(node, RB_BLACK);
106 break;
107 }
108
109 if (rb_is_black(parent)) {
110 break;
111 }
112
113 // if double-red exists, then gparent must not be NULL.
114 gparent = rb_red_parent(parent);
115
116 if (parent == gparent->left_child) {
117 uncle = gparent->right_child;
118 if (!uncle || rb_is_black(uncle)) {
119 // uncle is black or NULL.
120 if (node == parent->left_child) {
121 /*
122 * Case-1-A
123 * Matched tree pattern:
124 * G P
125 * / \ / \
126 * p U -> x g
127 * / \
128 * x U
129 */
130
131 // __connect34 only reconnect node,
132 // parent and gparent, but we still have
133 // to connect new subtree root back to
134 // original tree, i.e, let parent
135 // inherit gparent's parent.
136 __rb_change_child(
137 root, rb_parent(gparent), gparent, parent);
138
139 __connect34(node,
140 parent,
141 gparent,
142 node->left_child,
143 node->right_child,
144 parent->right_child,
145 uncle);
146 rb_set_color(parent, RB_BLACK);
147 rb_set_color(gparent, RB_RED);
148
149 // gparent might be root node, but
150 // rotation here would alter root node
151 // update root node pointer stored in
152 // root
153 if (!rb_parent(parent)) {
154 root->root_node = parent;
155 }
156 } else {
157 /*
158 * Case-1-B
159 * Matched tree pattern:
160 * G X
161 * / \ / \
162 * p U -> p g
163 * \ \
164 * x U
165 */
166 __rb_change_child(root, rb_parent(gparent), gparent, node);
167
168 __connect34(parent,
169 node,
170 gparent,
171 parent->left_child,
172 node->left_child,
173 node->right_child,
174 uncle);
175 rb_set_color(node, RB_BLACK);
176 rb_set_color(gparent, RB_RED);
177
178 if (!rb_parent(node)) {
179 root->root_node = node;
180 }
181 }
182 break;
183 } else {
184 // uncle must be red.
185
186 /*
187 * Case-2
188 * Matched tree pattern:
189 * G g
190 * / \ / \
191 * p u -> P U
192 * / /
193 * x x
194 * g might be new source of double-red. So we
195 * set @node to g, and @parent to parent of g
196 * and continue iteration.
197 */
198
199 rb_set_color(gparent, RB_RED);
200 rb_set_color(parent, RB_BLACK);
201 rb_set_color(uncle, RB_BLACK);
202 node = gparent;
203 parent = rb_parent(gparent);
204 continue;
205 }
206 } else {
207 // mirror
208 uncle = gparent->left_child;
209 if (!uncle || rb_is_black(uncle)) {
210 if (node == parent->right_child) {
211 // mirror of Case-1-A
212 __rb_change_child(
213 root, rb_parent(gparent), gparent, parent);
214 __connect34(gparent,
215 parent,
216 node,
217 uncle,
218 parent->left_child,
219 node->left_child,
220 node->right_child);
221 rb_set_color(parent, RB_BLACK);
222 rb_set_color(gparent, RB_RED);
223 if (!rb_parent(parent)) {
224 root->root_node = parent;
225 }
226 break;
227 } else {
228 // mirror of Case-1-B
229 __rb_change_child(root, rb_parent(gparent), gparent, node);
230 __connect34(gparent,
231 node,
232 parent,
233 uncle,
234 node->left_child,
235 node->right_child,
236 parent->right_child);
237 rb_set_color(node, RB_BLACK);
238 rb_set_color(gparent, RB_RED);
239 if (!rb_parent(node)) {
240 root->root_node = node;
241 }
242 break;
243 }
244 } else {
245 // mirror of Case-2
246 rb_set_color(gparent, RB_RED);
247 rb_set_color(parent, RB_BLACK);
248 rb_set_color(uncle, RB_BLACK);
249 node = gparent;
250 parent = rb_parent(gparent);
251 continue;
252 }
253 }
254 }
255 }
256
rb_insert(struct rb_root * this,struct rb_node * data,less_func less)257 void rb_insert(struct rb_root *this, struct rb_node *data, less_func less)
258 {
259 struct rb_node **new_link = &this->root_node;
260 struct rb_node *parent = NULL, *cur = this->root_node;
261
262 while (cur) {
263 if (less(data, cur)) {
264 parent = cur;
265 new_link = &cur->left_child;
266 cur = cur->left_child;
267 } else {
268 parent = cur;
269 new_link = &cur->right_child;
270 cur = cur->right_child;
271 }
272 }
273
274 __rb_link_node(data, parent, new_link);
275 __rb_insert_color(this, data);
276 }
277
rb_replace_node(struct rb_root * this,struct rb_node * old,struct rb_node * new)278 void rb_replace_node(struct rb_root *this, struct rb_node *old,
279 struct rb_node *new)
280 {
281 *new = *old;
282
283 if (new->left_child) {
284 rb_set_parent(new->left_child, new);
285 }
286
287 if (new->right_child) {
288 rb_set_parent(new->right_child, new);
289 }
290
291 __rb_change_child(this, rb_parent(old), old, new);
292 }
293
__rb_erase(struct rb_root * this,struct rb_node * node,bool * is_left_child)294 static struct rb_node *__rb_erase(struct rb_root *this, struct rb_node *node,
295 bool *is_left_child)
296 {
297 struct rb_node *succ, *parent = rb_parent(node), *ret = NULL, *succ_child;
298
299 if (!node->left_child) {
300 // right_child might be NULL, if so and node is black, then
301 // double-black exists.
302 if (!node->right_child) {
303 // erase the only root_node requires no additional
304 // processing
305 if (rb_is_black(node) && parent) {
306 ret = parent;
307 *is_left_child = node == parent->left_child;
308 }
309 } else {
310 // If node is black and has only the right_child, the
311 // latter must be red and the former must be black, so
312 // we just set the right_child as black.
313 rb_set_color(node->right_child, RB_BLACK);
314 }
315 // Calling rb_replace_node here is not appropriate, because node
316 // might has no child.
317 __rb_change_child(this, parent, node, node->right_child);
318 } else if (!node->right_child) {
319 // node has at least one child(left child). If right_child is
320 // NULL, it's impossible for left_child to be black, otherwise
321 // the rbtree is illegal. So checking double-black or recoloring
322 // the left child are unnecessary here.
323 __rb_change_child(this, parent, node, node->left_child);
324 rb_set_color(node->left_child, RB_BLACK);
325 } else {
326 // find succ of node first.
327 succ = node->right_child;
328 if (!succ->left_child) {
329 succ_child = succ->right_child;
330 // check whether double-black exists. Logically, we
331 // should exchange node and succ first, then check
332 // against node and succ_child. But it's redundant,
333 // because node will be erased finally. So we check
334 // double-black using info of succ directly.
335 if (rb_is_black(succ)) {
336 if (succ_child && rb_is_red(succ_child)) {
337 rb_set_color(succ_child, RB_BLACK);
338 } else {
339 // due to succ is node->right_child,
340 // after adjustment, parent of
341 // succ_child is still succ.
342 ret = succ;
343 // succ_child will be located at
344 // original pos of succ, and succ must
345 // be right child of node here.
346 *is_left_child = false;
347 }
348 }
349 // succ of node is its right child. We just need to
350 // let succ inherit parent, color and **left child** of
351 // node.
352 __rb_change_child(this, parent, node, succ);
353 rb_set_color(succ, rb_color(node));
354 succ->left_child = node->left_child;
355 rb_set_parent(node->left_child, succ);
356 } else {
357 while (succ->left_child) {
358 succ = succ->left_child;
359 }
360 parent = rb_parent(succ);
361
362 succ_child = succ->right_child;
363
364 // check whether double-black exists
365 if (rb_is_black(succ)) {
366 if (succ_child && rb_is_red(succ_child)) {
367 rb_set_color(succ_child, RB_BLACK);
368 } else {
369 ret = parent;
370 *is_left_child = succ == parent->left_child;
371 }
372 }
373
374 // replace succ with its right_child
375 __rb_change_child(this, parent, succ, succ_child);
376
377 // replace node with succ, succ would inherit node's
378 // parent, color and **chilren**.
379 rb_replace_node(this, node, succ);
380 }
381 }
382
383 return ret;
384 }
385
__rb_erase_color(struct rb_root * root,struct rb_node * parent,bool is_left_child)386 static void __rb_erase_color(struct rb_root *root, struct rb_node *parent,
387 bool is_left_child)
388 {
389 struct rb_node *node = is_left_child ? parent->left_child :
390 parent->right_child;
391 struct rb_node *slibing, *nephew;
392
393 while (true) {
394 /*
395 * Loop invariants
396 * - @parent is not NULL(@node is not root node)
397 * - subtree of which @parent as root has 1 lower blackheight
398 */
399 if (!parent) {
400 break;
401 }
402 if (node == parent->left_child) {
403 // slibing must not be NULL.
404 slibing = parent->right_child;
405 if (rb_is_black(slibing)) {
406 if (slibing->right_child && rb_is_red(slibing->right_child)) {
407 nephew = slibing->right_child;
408 /*
409 * Case-1-A (x: @node, n: @nephew)
410 * Matched tree pattern:
411 * (p) (s)
412 * / \ / \
413 * X S -> P N
414 * / \ / \
415 * (s1) n X (s1)
416 */
417
418 __rb_change_child(root, rb_parent(parent), parent, slibing);
419
420 __connect34(parent,
421 slibing,
422 nephew,
423 node,
424 slibing->left_child,
425 nephew->left_child,
426 nephew->right_child);
427 rb_set_color(slibing, rb_color(parent));
428 rb_set_color(parent, RB_BLACK);
429 rb_set_color(nephew, RB_BLACK);
430
431 if (!rb_parent(slibing)) {
432 root->root_node = slibing;
433 }
434 break;
435 } else if (slibing->left_child
436 && rb_is_red(slibing->left_child)) {
437 nephew = slibing->left_child;
438 /*
439 * Case-1-B (x: @node, n: @nephew)
440 * Matched tree pattern:
441 * (p) (n)
442 * / \ / \
443 * X S -> P S
444 * / \ / \
445 * n (s1) X (s1)
446 */
447 __rb_change_child(root, rb_parent(parent), parent, nephew);
448
449 __connect34(parent,
450 nephew,
451 slibing,
452 node,
453 nephew->left_child,
454 nephew->right_child,
455 slibing->right_child);
456 rb_set_color(nephew, rb_color(parent));
457 rb_set_color(parent, RB_BLACK);
458
459 if (!rb_parent(nephew)) {
460 root->root_node = nephew;
461 }
462 break;
463 } else {
464 if (rb_is_black(parent)) {
465 /*
466 * Case-2 (x: @node, n: @nephew)
467 * Matched tree pattern:
468 * P P
469 * / \ / \
470 * X S -> X s
471 * / \ / \
472 * S1 S2 S1 S2
473 * Subtree P has 1 lower
474 * blackheight than slibings of
475 * P. (X has 1 lower blackheight
476 * than S before). So we let
477 * @node = P and continue
478 * iteration.
479 */
480 rb_set_color(slibing, RB_RED);
481 node = parent;
482 parent = rb_parent(node);
483 continue;
484 } else {
485 /*
486 * Case-3 (x: @node, n: @nephew)
487 * Matched tree pattern:
488 * p P
489 * / \ / \
490 * X S -> X s
491 * / \ / \
492 * S1 S2 S1 S2
493 */
494 rb_set_color(parent, RB_BLACK);
495 rb_set_color(slibing, RB_RED);
496 break;
497 }
498 }
499 } else {
500 /*
501 * Case-4 (x: @node, n: @nephew)
502 * Matched tree pattern:
503 * P S
504 * / \ / \
505 * X s -> p S2
506 * / \ / \
507 * S1 S2 X S1
508 * We do not fix rbtree here. However, after
509 * this rotation, X has a black slibing now, and
510 * will turn into Case-1 or Case-3 (depends on
511 * children of S1).
512 */
513 __rb_change_child(root, rb_parent(parent), parent, slibing);
514
515 parent->right_child = slibing->left_child;
516 if (slibing->left_child) {
517 rb_set_parent(slibing->left_child, parent);
518 }
519
520 slibing->left_child = parent;
521 rb_set_parent(parent, slibing);
522
523 rb_set_color(slibing, RB_BLACK);
524 rb_set_color(parent, RB_RED);
525
526 if (!rb_parent(slibing)) {
527 root->root_node = slibing;
528 }
529 continue;
530 }
531 } else {
532 slibing = parent->left_child;
533 if (rb_is_black(slibing)) {
534 if (slibing->left_child && rb_is_red(slibing->left_child)) {
535 // mirror of Case-1-A
536 nephew = slibing->left_child;
537
538 __rb_change_child(root, rb_parent(parent), parent, slibing);
539 __connect34(nephew,
540 slibing,
541 parent,
542 nephew->left_child,
543 nephew->right_child,
544 slibing->right_child,
545 node);
546 rb_set_color(slibing, rb_color(parent));
547 rb_set_color(nephew, RB_BLACK);
548 rb_set_color(parent, RB_BLACK);
549
550 if (!rb_parent(slibing)) {
551 root->root_node = slibing;
552 }
553 break;
554 } else if (slibing->right_child
555 && rb_is_red(slibing->right_child)) {
556 // mirror of Case-1-B
557 nephew = slibing->right_child;
558
559 __rb_change_child(root, rb_parent(parent), parent, nephew);
560
561 __connect34(slibing,
562 nephew,
563 parent,
564 slibing->left_child,
565 nephew->left_child,
566 nephew->right_child,
567 node);
568 rb_set_color(nephew, rb_color(parent));
569 rb_set_color(parent, RB_BLACK);
570
571 if (!rb_parent(nephew)) {
572 root->root_node = nephew;
573 }
574 break;
575 } else {
576 if (rb_is_black(parent)) {
577 // mirror of Case-2
578 rb_set_color(slibing, RB_RED);
579 node = parent;
580 parent = rb_parent(node);
581 continue;
582 } else {
583 // mirror of Case-3
584 rb_set_color(parent, RB_BLACK);
585 rb_set_color(slibing, RB_RED);
586 break;
587 }
588 }
589 } else {
590 // mirror of Case-4
591 __rb_change_child(root, rb_parent(parent), parent, slibing);
592 parent->left_child = slibing->right_child;
593 if (slibing->right_child) {
594 rb_set_parent(slibing->right_child, parent);
595 }
596
597 slibing->right_child = parent;
598 rb_set_parent(parent, slibing);
599
600 rb_set_color(slibing, RB_BLACK);
601 rb_set_color(parent, RB_RED);
602 continue;
603 }
604 }
605 }
606 }
607
rb_erase(struct rb_root * this,struct rb_node * node)608 void rb_erase(struct rb_root *this, struct rb_node *node)
609 {
610 bool is_left_child;
611 struct rb_node *rebalance = __rb_erase(this, node, &is_left_child);
612 if (rebalance) {
613 __rb_erase_color(this, rebalance, is_left_child);
614 }
615 }
616
rb_search(struct rb_root * this,const void * key,comp_key_func cmp)617 struct rb_node *rb_search(struct rb_root *this, const void *key,
618 comp_key_func cmp)
619 {
620 struct rb_node *cur = this->root_node;
621 int cmp_ret;
622 while (cur) {
623 cmp_ret = cmp(key, cur);
624 if (cmp_ret < 0) {
625 cur = cur->left_child;
626 } else if (cmp_ret > 0) {
627 cur = cur->right_child;
628 } else {
629 return cur;
630 }
631 }
632
633 return NULL;
634 }
635
rb_search_first(struct rb_root * this,const void * key,comp_key_func cmp)636 struct rb_node *rb_search_first(struct rb_root *this, const void *key,
637 comp_key_func cmp)
638 {
639 struct rb_node *cur = this->root_node, *ret = NULL;
640 int c;
641 while (cur) {
642 c = cmp(key, cur);
643 if (c <= 0) {
644 if (c == 0) {
645 ret = cur;
646 }
647 if (c != 0 && ret) {
648 break;
649 }
650 cur = cur->left_child;
651 } else {
652 cur = cur->right_child;
653 }
654 }
655
656 return ret;
657 }
658
rb_next(const struct rb_node * node)659 struct rb_node *rb_next(const struct rb_node *node)
660 {
661 struct rb_node *ret = NULL, *parent = NULL;
662
663 // If node has right_child, find its direct successor
664 if (node->right_child) {
665 ret = node->right_child;
666 while (ret->left_child) {
667 ret = ret->left_child;
668 }
669 return ret;
670 }
671
672 // If node has no right_child, backtrace to its parent.
673 // If node is left_child of its parent, then parent is our next node.
674 // If node is right_child of its parent, it indicates that parent is <=
675 // node, so we continue backtracing, until find a node which is
676 // left_child of its parent, or reach parent of root_node, i.e., NULL.
677 while ((parent = rb_parent(node)) && node == parent->right_child)
678 node = parent;
679
680 return parent;
681 }
682
rb_prev(const struct rb_node * node)683 struct rb_node *rb_prev(const struct rb_node *node)
684 {
685 // mirror of rb_next
686 struct rb_node *ret = NULL, *parent = NULL;
687
688 if (node->left_child) {
689 ret = node->left_child;
690 while (ret->right_child) {
691 ret = ret->right_child;
692 }
693 return ret;
694 }
695
696 while ((parent = rb_parent(node)) && node == parent->left_child)
697 node = parent;
698
699 return parent;
700 }
701
rb_first(const struct rb_root * this)702 struct rb_node *rb_first(const struct rb_root *this)
703 {
704 struct rb_node *cur = this->root_node;
705 if (!cur) {
706 return NULL;
707 }
708
709 while (cur->left_child) {
710 cur = cur->left_child;
711 }
712 return cur;
713 }
714
rb_last(const struct rb_root * this)715 struct rb_node *rb_last(const struct rb_root *this)
716 {
717 struct rb_node *cur = this->root_node;
718 if (!cur) {
719 return NULL;
720 }
721
722 while (cur->right_child) {
723 cur = cur->right_child;
724 }
725 return cur;
726 }
727
rb_next_match(const struct rb_node * node,const void * key,comp_key_func cmp)728 struct rb_node *rb_next_match(const struct rb_node *node, const void *key,
729 comp_key_func cmp)
730 {
731 struct rb_node *ret = rb_next(node);
732 if (ret && cmp(key, ret) != 0) {
733 ret = NULL;
734 }
735 return ret;
736 }