1/**************************************************************************/
2/* rb_set.h */
3/**************************************************************************/
4/* This file is part of: */
5/* GODOT ENGINE */
6/* https://godotengine.org */
7/**************************************************************************/
8/* Copyright (c) 2014-present Godot Engine contributors (see AUTHORS.md). */
9/* Copyright (c) 2007-2014 Juan Linietsky, Ariel Manzur. */
10/* */
11/* Permission is hereby granted, free of charge, to any person obtaining */
12/* a copy of this software and associated documentation files (the */
13/* "Software"), to deal in the Software without restriction, including */
14/* without limitation the rights to use, copy, modify, merge, publish, */
15/* distribute, sublicense, and/or sell copies of the Software, and to */
16/* permit persons to whom the Software is furnished to do so, subject to */
17/* the following conditions: */
18/* */
19/* The above copyright notice and this permission notice shall be */
20/* included in all copies or substantial portions of the Software. */
21/* */
22/* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, */
23/* EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF */
24/* MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. */
25/* IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY */
26/* CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, */
27/* TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE */
28/* SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. */
29/**************************************************************************/
30
31#ifndef RB_SET_H
32#define RB_SET_H
33
34#include "core/os/memory.h"
35#include "core/typedefs.h"
36
37// based on the very nice implementation of rb-trees by:
38// https://web.archive.org/web/20120507164830/https://web.mit.edu/~emin/www/source_code/red_black_tree/index.html
39
40template <class T, class C = Comparator<T>, class A = DefaultAllocator>
41class RBSet {
42 enum Color {
43 RED,
44 BLACK
45 };
46 struct _Data;
47
48public:
49 class Element {
50 private:
51 friend class RBSet<T, C, A>;
52 int color = RED;
53 Element *right = nullptr;
54 Element *left = nullptr;
55 Element *parent = nullptr;
56 Element *_next = nullptr;
57 Element *_prev = nullptr;
58 T value;
59 //_Data *data;
60
61 public:
62 const Element *next() const {
63 return _next;
64 }
65 Element *next() {
66 return _next;
67 }
68 const Element *prev() const {
69 return _prev;
70 }
71 Element *prev() {
72 return _prev;
73 }
74 T &get() {
75 return value;
76 }
77 const T &get() const {
78 return value;
79 };
80 Element() {}
81 };
82
83 typedef T ValueType;
84
85 struct Iterator {
86 _FORCE_INLINE_ T &operator*() const {
87 return E->get();
88 }
89 _FORCE_INLINE_ T *operator->() const { return &E->get(); }
90 _FORCE_INLINE_ Iterator &operator++() {
91 E = E->next();
92 return *this;
93 }
94 _FORCE_INLINE_ Iterator &operator--() {
95 E = E->prev();
96 return *this;
97 }
98
99 _FORCE_INLINE_ bool operator==(const Iterator &b) const { return E == b.E; }
100 _FORCE_INLINE_ bool operator!=(const Iterator &b) const { return E != b.E; }
101
102 explicit operator bool() const { return E != nullptr; }
103 Iterator(Element *p_E) { E = p_E; }
104 Iterator() {}
105 Iterator(const Iterator &p_it) { E = p_it.E; }
106
107 private:
108 Element *E = nullptr;
109 };
110
111 struct ConstIterator {
112 _FORCE_INLINE_ const T &operator*() const {
113 return E->get();
114 }
115 _FORCE_INLINE_ const T *operator->() const { return &E->get(); }
116 _FORCE_INLINE_ ConstIterator &operator++() {
117 E = E->next();
118 return *this;
119 }
120 _FORCE_INLINE_ ConstIterator &operator--() {
121 E = E->prev();
122 return *this;
123 }
124
125 _FORCE_INLINE_ bool operator==(const ConstIterator &b) const { return E == b.E; }
126 _FORCE_INLINE_ bool operator!=(const ConstIterator &b) const { return E != b.E; }
127
128 _FORCE_INLINE_ ConstIterator(const Element *p_E) { E = p_E; }
129 _FORCE_INLINE_ ConstIterator() {}
130 _FORCE_INLINE_ ConstIterator(const ConstIterator &p_it) { E = p_it.E; }
131
132 explicit operator bool() const { return E != nullptr; }
133
134 private:
135 const Element *E = nullptr;
136 };
137
138 _FORCE_INLINE_ Iterator begin() {
139 return Iterator(front());
140 }
141 _FORCE_INLINE_ Iterator end() {
142 return Iterator(nullptr);
143 }
144
145#if 0
146 //to use when replacing find()
147 _FORCE_INLINE_ Iterator find(const K &p_key) {
148 return Iterator(find(p_key));
149 }
150#endif
151
152 _FORCE_INLINE_ ConstIterator begin() const {
153 return ConstIterator(front());
154 }
155 _FORCE_INLINE_ ConstIterator end() const {
156 return ConstIterator(nullptr);
157 }
158
159#if 0
160 //to use when replacing find()
161 _FORCE_INLINE_ ConstIterator find(const K &p_key) const {
162 return ConstIterator(find(p_key));
163 }
164#endif
165private:
166 struct _Data {
167 Element *_root = nullptr;
168 Element *_nil = nullptr;
169 int size_cache = 0;
170
171 _FORCE_INLINE_ _Data() {
172#ifdef GLOBALNIL_DISABLED
173 _nil = memnew_allocator(Element, A);
174 _nil->parent = _nil->left = _nil->right = _nil;
175 _nil->color = BLACK;
176#else
177 _nil = (Element *)&_GlobalNilClass::_nil;
178#endif
179 }
180
181 void _create_root() {
182 _root = memnew_allocator(Element, A);
183 _root->parent = _root->left = _root->right = _nil;
184 _root->color = BLACK;
185 }
186
187 void _free_root() {
188 if (_root) {
189 memdelete_allocator<Element, A>(_root);
190 _root = nullptr;
191 }
192 }
193
194 ~_Data() {
195 _free_root();
196
197#ifdef GLOBALNIL_DISABLED
198 memdelete_allocator<Element, A>(_nil);
199#endif
200 }
201 };
202
203 _Data _data;
204
205 inline void _set_color(Element *p_node, int p_color) {
206 ERR_FAIL_COND(p_node == _data._nil && p_color == RED);
207 p_node->color = p_color;
208 }
209
210 inline void _rotate_left(Element *p_node) {
211 Element *r = p_node->right;
212 p_node->right = r->left;
213 if (r->left != _data._nil) {
214 r->left->parent = p_node;
215 }
216 r->parent = p_node->parent;
217 if (p_node == p_node->parent->left) {
218 p_node->parent->left = r;
219 } else {
220 p_node->parent->right = r;
221 }
222
223 r->left = p_node;
224 p_node->parent = r;
225 }
226
227 inline void _rotate_right(Element *p_node) {
228 Element *l = p_node->left;
229 p_node->left = l->right;
230 if (l->right != _data._nil) {
231 l->right->parent = p_node;
232 }
233 l->parent = p_node->parent;
234 if (p_node == p_node->parent->right) {
235 p_node->parent->right = l;
236 } else {
237 p_node->parent->left = l;
238 }
239
240 l->right = p_node;
241 p_node->parent = l;
242 }
243
244 inline Element *_successor(Element *p_node) const {
245 Element *node = p_node;
246
247 if (node->right != _data._nil) {
248 node = node->right;
249 while (node->left != _data._nil) { /* returns the minimum of the right subtree of node */
250 node = node->left;
251 }
252 return node;
253 } else {
254 while (node == node->parent->right) {
255 node = node->parent;
256 }
257
258 if (node->parent == _data._root) {
259 return nullptr; // No successor, as p_node = last node
260 }
261 return node->parent;
262 }
263 }
264
265 inline Element *_predecessor(Element *p_node) const {
266 Element *node = p_node;
267
268 if (node->left != _data._nil) {
269 node = node->left;
270 while (node->right != _data._nil) { /* returns the minimum of the left subtree of node */
271 node = node->right;
272 }
273 return node;
274 } else {
275 while (node == node->parent->left) {
276 node = node->parent;
277 }
278
279 if (node == _data._root) {
280 return nullptr; // No predecessor, as p_node = first node.
281 }
282 return node->parent;
283 }
284 }
285
286 Element *_find(const T &p_value) const {
287 Element *node = _data._root->left;
288 C less;
289
290 while (node != _data._nil) {
291 if (less(p_value, node->value)) {
292 node = node->left;
293 } else if (less(node->value, p_value)) {
294 node = node->right;
295 } else {
296 return node; // found
297 }
298 }
299
300 return nullptr;
301 }
302
303 Element *_lower_bound(const T &p_value) const {
304 Element *node = _data._root->left;
305 Element *prev = nullptr;
306 C less;
307
308 while (node != _data._nil) {
309 prev = node;
310
311 if (less(p_value, node->value)) {
312 node = node->left;
313 } else if (less(node->value, p_value)) {
314 node = node->right;
315 } else {
316 return node; // found
317 }
318 }
319
320 if (prev == nullptr) {
321 return nullptr; // tree empty
322 }
323
324 if (less(prev->value, p_value)) {
325 prev = prev->_next;
326 }
327
328 return prev;
329 }
330
331 void _insert_rb_fix(Element *p_new_node) {
332 Element *node = p_new_node;
333 Element *nparent = node->parent;
334 Element *ngrand_parent = nullptr;
335
336 while (nparent->color == RED) {
337 ngrand_parent = nparent->parent;
338
339 if (nparent == ngrand_parent->left) {
340 if (ngrand_parent->right->color == RED) {
341 _set_color(nparent, BLACK);
342 _set_color(ngrand_parent->right, BLACK);
343 _set_color(ngrand_parent, RED);
344 node = ngrand_parent;
345 nparent = node->parent;
346 } else {
347 if (node == nparent->right) {
348 _rotate_left(nparent);
349 node = nparent;
350 nparent = node->parent;
351 }
352 _set_color(nparent, BLACK);
353 _set_color(ngrand_parent, RED);
354 _rotate_right(ngrand_parent);
355 }
356 } else {
357 if (ngrand_parent->left->color == RED) {
358 _set_color(nparent, BLACK);
359 _set_color(ngrand_parent->left, BLACK);
360 _set_color(ngrand_parent, RED);
361 node = ngrand_parent;
362 nparent = node->parent;
363 } else {
364 if (node == nparent->left) {
365 _rotate_right(nparent);
366 node = nparent;
367 nparent = node->parent;
368 }
369 _set_color(nparent, BLACK);
370 _set_color(ngrand_parent, RED);
371 _rotate_left(ngrand_parent);
372 }
373 }
374 }
375
376 _set_color(_data._root->left, BLACK);
377 }
378
379 Element *_insert(const T &p_value) {
380 Element *new_parent = _data._root;
381 Element *node = _data._root->left;
382 C less;
383
384 while (node != _data._nil) {
385 new_parent = node;
386
387 if (less(p_value, node->value)) {
388 node = node->left;
389 } else if (less(node->value, p_value)) {
390 node = node->right;
391 } else {
392 return node; // Return existing node
393 }
394 }
395
396 Element *new_node = memnew_allocator(Element, A);
397 new_node->parent = new_parent;
398 new_node->right = _data._nil;
399 new_node->left = _data._nil;
400 new_node->value = p_value;
401 //new_node->data=_data;
402
403 if (new_parent == _data._root || less(p_value, new_parent->value)) {
404 new_parent->left = new_node;
405 } else {
406 new_parent->right = new_node;
407 }
408
409 new_node->_next = _successor(new_node);
410 new_node->_prev = _predecessor(new_node);
411 if (new_node->_next) {
412 new_node->_next->_prev = new_node;
413 }
414 if (new_node->_prev) {
415 new_node->_prev->_next = new_node;
416 }
417
418 _data.size_cache++;
419 _insert_rb_fix(new_node);
420 return new_node;
421 }
422
423 void _erase_fix_rb(Element *p_node) {
424 Element *root = _data._root->left;
425 Element *node = _data._nil;
426 Element *sibling = p_node;
427 Element *parent = sibling->parent;
428
429 while (node != root) { // If red node found, will exit at a break
430 if (sibling->color == RED) {
431 _set_color(sibling, BLACK);
432 _set_color(parent, RED);
433 if (sibling == parent->right) {
434 sibling = sibling->left;
435 _rotate_left(parent);
436 } else {
437 sibling = sibling->right;
438 _rotate_right(parent);
439 }
440 }
441 if ((sibling->left->color == BLACK) && (sibling->right->color == BLACK)) {
442 _set_color(sibling, RED);
443 if (parent->color == RED) {
444 _set_color(parent, BLACK);
445 break;
446 } else { // loop: haven't found any red nodes yet
447 node = parent;
448 parent = node->parent;
449 sibling = (node == parent->left) ? parent->right : parent->left;
450 }
451 } else {
452 if (sibling == parent->right) {
453 if (sibling->right->color == BLACK) {
454 _set_color(sibling->left, BLACK);
455 _set_color(sibling, RED);
456 _rotate_right(sibling);
457 sibling = sibling->parent;
458 }
459 _set_color(sibling, parent->color);
460 _set_color(parent, BLACK);
461 _set_color(sibling->right, BLACK);
462 _rotate_left(parent);
463 break;
464 } else {
465 if (sibling->left->color == BLACK) {
466 _set_color(sibling->right, BLACK);
467 _set_color(sibling, RED);
468 _rotate_left(sibling);
469 sibling = sibling->parent;
470 }
471
472 _set_color(sibling, parent->color);
473 _set_color(parent, BLACK);
474 _set_color(sibling->left, BLACK);
475 _rotate_right(parent);
476 break;
477 }
478 }
479 }
480
481 ERR_FAIL_COND(_data._nil->color != BLACK);
482 }
483
484 void _erase(Element *p_node) {
485 Element *rp = ((p_node->left == _data._nil) || (p_node->right == _data._nil)) ? p_node : p_node->_next;
486 Element *node = (rp->left == _data._nil) ? rp->right : rp->left;
487
488 Element *sibling = nullptr;
489 if (rp == rp->parent->left) {
490 rp->parent->left = node;
491 sibling = rp->parent->right;
492 } else {
493 rp->parent->right = node;
494 sibling = rp->parent->left;
495 }
496
497 if (node->color == RED) {
498 node->parent = rp->parent;
499 _set_color(node, BLACK);
500 } else if (rp->color == BLACK && rp->parent != _data._root) {
501 _erase_fix_rb(sibling);
502 }
503
504 if (rp != p_node) {
505 ERR_FAIL_COND(rp == _data._nil);
506
507 rp->left = p_node->left;
508 rp->right = p_node->right;
509 rp->parent = p_node->parent;
510 rp->color = p_node->color;
511 if (p_node->left != _data._nil) {
512 p_node->left->parent = rp;
513 }
514 if (p_node->right != _data._nil) {
515 p_node->right->parent = rp;
516 }
517
518 if (p_node == p_node->parent->left) {
519 p_node->parent->left = rp;
520 } else {
521 p_node->parent->right = rp;
522 }
523 }
524
525 if (p_node->_next) {
526 p_node->_next->_prev = p_node->_prev;
527 }
528 if (p_node->_prev) {
529 p_node->_prev->_next = p_node->_next;
530 }
531
532 memdelete_allocator<Element, A>(p_node);
533 _data.size_cache--;
534 ERR_FAIL_COND(_data._nil->color == RED);
535 }
536
537 void _calculate_depth(Element *p_element, int &max_d, int d) const {
538 if (p_element == _data._nil) {
539 return;
540 }
541
542 _calculate_depth(p_element->left, max_d, d + 1);
543 _calculate_depth(p_element->right, max_d, d + 1);
544
545 if (d > max_d) {
546 max_d = d;
547 }
548 }
549
550 void _cleanup_tree(Element *p_element) {
551 if (p_element == _data._nil) {
552 return;
553 }
554
555 _cleanup_tree(p_element->left);
556 _cleanup_tree(p_element->right);
557 memdelete_allocator<Element, A>(p_element);
558 }
559
560 void _copy_from(const RBSet &p_set) {
561 clear();
562 // not the fastest way, but safeset to write.
563 for (Element *I = p_set.front(); I; I = I->next()) {
564 insert(I->get());
565 }
566 }
567
568public:
569 const Element *find(const T &p_value) const {
570 if (!_data._root) {
571 return nullptr;
572 }
573
574 const Element *res = _find(p_value);
575 return res;
576 }
577
578 Element *find(const T &p_value) {
579 if (!_data._root) {
580 return nullptr;
581 }
582
583 Element *res = _find(p_value);
584 return res;
585 }
586
587 Element *lower_bound(const T &p_value) const {
588 if (!_data._root) {
589 return nullptr;
590 }
591 return _lower_bound(p_value);
592 }
593
594 bool has(const T &p_value) const {
595 return find(p_value) != nullptr;
596 }
597
598 Element *insert(const T &p_value) {
599 if (!_data._root) {
600 _data._create_root();
601 }
602 return _insert(p_value);
603 }
604
605 void erase(Element *p_element) {
606 if (!_data._root || !p_element) {
607 return;
608 }
609
610 _erase(p_element);
611 if (_data.size_cache == 0 && _data._root) {
612 _data._free_root();
613 }
614 }
615
616 bool erase(const T &p_value) {
617 if (!_data._root) {
618 return false;
619 }
620
621 Element *e = find(p_value);
622 if (!e) {
623 return false;
624 }
625
626 _erase(e);
627 if (_data.size_cache == 0 && _data._root) {
628 _data._free_root();
629 }
630 return true;
631 }
632
633 Element *front() const {
634 if (!_data._root) {
635 return nullptr;
636 }
637
638 Element *e = _data._root->left;
639 if (e == _data._nil) {
640 return nullptr;
641 }
642
643 while (e->left != _data._nil) {
644 e = e->left;
645 }
646
647 return e;
648 }
649
650 Element *back() const {
651 if (!_data._root) {
652 return nullptr;
653 }
654
655 Element *e = _data._root->left;
656 if (e == _data._nil) {
657 return nullptr;
658 }
659
660 while (e->right != _data._nil) {
661 e = e->right;
662 }
663
664 return e;
665 }
666
667 inline bool is_empty() const {
668 return _data.size_cache == 0;
669 }
670 inline int size() const {
671 return _data.size_cache;
672 }
673
674 int calculate_depth() const {
675 // used for debug mostly
676 if (!_data._root) {
677 return 0;
678 }
679
680 int max_d = 0;
681 _calculate_depth(_data._root->left, max_d, 0);
682 return max_d;
683 }
684
685 void clear() {
686 if (!_data._root) {
687 return;
688 }
689
690 _cleanup_tree(_data._root->left);
691 _data._root->left = _data._nil;
692 _data.size_cache = 0;
693 _data._free_root();
694 }
695
696 void operator=(const RBSet &p_set) {
697 _copy_from(p_set);
698 }
699
700 RBSet(const RBSet &p_set) {
701 _copy_from(p_set);
702 }
703
704 _FORCE_INLINE_ RBSet() {}
705
706 ~RBSet() {
707 clear();
708 }
709};
710
711#endif // RB_SET_H
712