1/**************************************************************************/
2/* rb_map.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_MAP_H
32#define RB_MAP_H
33
34#include "core/error/error_macros.h"
35#include "core/os/memory.h"
36#include "core/templates/pair.h"
37
38// based on the very nice implementation of rb-trees by:
39// https://web.archive.org/web/20120507164830/https://web.mit.edu/~emin/www/source_code/red_black_tree/index.html
40
41template <class K, class V, class C = Comparator<K>, class A = DefaultAllocator>
42class RBMap {
43 enum Color {
44 RED,
45 BLACK
46 };
47 struct _Data;
48
49public:
50 class Element {
51 private:
52 friend class RBMap<K, V, C, A>;
53 int color = RED;
54 Element *right = nullptr;
55 Element *left = nullptr;
56 Element *parent = nullptr;
57 Element *_next = nullptr;
58 Element *_prev = nullptr;
59 KeyValue<K, V> _data;
60
61 public:
62 KeyValue<K, V> &key_value() { return _data; }
63 const KeyValue<K, V> &key_value() const { return _data; }
64
65 const Element *next() const {
66 return _next;
67 }
68 Element *next() {
69 return _next;
70 }
71 const Element *prev() const {
72 return _prev;
73 }
74 Element *prev() {
75 return _prev;
76 }
77 const K &key() const {
78 return _data.key;
79 }
80 V &value() {
81 return _data.value;
82 }
83 const V &value() const {
84 return _data.value;
85 }
86 V &get() {
87 return _data.value;
88 }
89 const V &get() const {
90 return _data.value;
91 }
92 Element(const KeyValue<K, V> &p_data) :
93 _data(p_data) {}
94 };
95
96 typedef KeyValue<K, V> ValueType;
97
98 struct Iterator {
99 _FORCE_INLINE_ KeyValue<K, V> &operator*() const {
100 return E->key_value();
101 }
102 _FORCE_INLINE_ KeyValue<K, V> *operator->() const { return &E->key_value(); }
103 _FORCE_INLINE_ Iterator &operator++() {
104 E = E->next();
105 return *this;
106 }
107 _FORCE_INLINE_ Iterator &operator--() {
108 E = E->prev();
109 return *this;
110 }
111
112 _FORCE_INLINE_ bool operator==(const Iterator &b) const { return E == b.E; }
113 _FORCE_INLINE_ bool operator!=(const Iterator &b) const { return E != b.E; }
114 explicit operator bool() const {
115 return E != nullptr;
116 }
117 Iterator(Element *p_E) { E = p_E; }
118 Iterator() {}
119 Iterator(const Iterator &p_it) { E = p_it.E; }
120
121 private:
122 Element *E = nullptr;
123 };
124
125 struct ConstIterator {
126 _FORCE_INLINE_ const KeyValue<K, V> &operator*() const {
127 return E->key_value();
128 }
129 _FORCE_INLINE_ const KeyValue<K, V> *operator->() const { return &E->key_value(); }
130 _FORCE_INLINE_ ConstIterator &operator++() {
131 E = E->next();
132 return *this;
133 }
134 _FORCE_INLINE_ ConstIterator &operator--() {
135 E = E->prev();
136 return *this;
137 }
138
139 _FORCE_INLINE_ bool operator==(const ConstIterator &b) const { return E == b.E; }
140 _FORCE_INLINE_ bool operator!=(const ConstIterator &b) const { return E != b.E; }
141 explicit operator bool() const {
142 return E != nullptr;
143 }
144 ConstIterator(const Element *p_E) { E = p_E; }
145 ConstIterator() {}
146 ConstIterator(const ConstIterator &p_it) { E = p_it.E; }
147
148 private:
149 const Element *E = nullptr;
150 };
151
152 _FORCE_INLINE_ Iterator begin() {
153 return Iterator(front());
154 }
155 _FORCE_INLINE_ Iterator end() {
156 return Iterator(nullptr);
157 }
158
159#if 0
160 //to use when replacing find()
161 _FORCE_INLINE_ Iterator find(const K &p_key) {
162 return Iterator(find(p_key));
163 }
164#endif
165 _FORCE_INLINE_ void remove(const Iterator &p_iter) {
166 return erase(p_iter.E);
167 }
168
169 _FORCE_INLINE_ ConstIterator begin() const {
170 return ConstIterator(front());
171 }
172 _FORCE_INLINE_ ConstIterator end() const {
173 return ConstIterator(nullptr);
174 }
175
176#if 0
177 //to use when replacing find()
178 _FORCE_INLINE_ ConstIterator find(const K &p_key) const {
179 return ConstIterator(find(p_key));
180 }
181#endif
182private:
183 struct _Data {
184 Element *_root = nullptr;
185 Element *_nil = nullptr;
186 int size_cache = 0;
187
188 _FORCE_INLINE_ _Data() {
189#ifdef GLOBALNIL_DISABLED
190 _nil = memnew_allocator(Element, A);
191 _nil->parent = _nil->left = _nil->right = _nil;
192 _nil->color = BLACK;
193#else
194 _nil = (Element *)&_GlobalNilClass::_nil;
195#endif
196 }
197
198 void _create_root() {
199 _root = memnew_allocator(Element(KeyValue<K, V>(K(), V())), A);
200 _root->parent = _root->left = _root->right = _nil;
201 _root->color = BLACK;
202 }
203
204 void _free_root() {
205 if (_root) {
206 memdelete_allocator<Element, A>(_root);
207 _root = nullptr;
208 }
209 }
210
211 ~_Data() {
212 _free_root();
213
214#ifdef GLOBALNIL_DISABLED
215 memdelete_allocator<Element, A>(_nil);
216#endif
217 }
218 };
219
220 _Data _data;
221
222 inline void _set_color(Element *p_node, int p_color) {
223 ERR_FAIL_COND(p_node == _data._nil && p_color == RED);
224 p_node->color = p_color;
225 }
226
227 inline void _rotate_left(Element *p_node) {
228 Element *r = p_node->right;
229 p_node->right = r->left;
230 if (r->left != _data._nil) {
231 r->left->parent = p_node;
232 }
233 r->parent = p_node->parent;
234 if (p_node == p_node->parent->left) {
235 p_node->parent->left = r;
236 } else {
237 p_node->parent->right = r;
238 }
239
240 r->left = p_node;
241 p_node->parent = r;
242 }
243
244 inline void _rotate_right(Element *p_node) {
245 Element *l = p_node->left;
246 p_node->left = l->right;
247 if (l->right != _data._nil) {
248 l->right->parent = p_node;
249 }
250 l->parent = p_node->parent;
251 if (p_node == p_node->parent->right) {
252 p_node->parent->right = l;
253 } else {
254 p_node->parent->left = l;
255 }
256
257 l->right = p_node;
258 p_node->parent = l;
259 }
260
261 inline Element *_successor(Element *p_node) const {
262 Element *node = p_node;
263
264 if (node->right != _data._nil) {
265 node = node->right;
266 while (node->left != _data._nil) { /* returns the minimum of the right subtree of node */
267 node = node->left;
268 }
269 return node;
270 } else {
271 while (node == node->parent->right) {
272 node = node->parent;
273 }
274
275 if (node->parent == _data._root) {
276 return nullptr; // No successor, as p_node = last node
277 }
278 return node->parent;
279 }
280 }
281
282 inline Element *_predecessor(Element *p_node) const {
283 Element *node = p_node;
284
285 if (node->left != _data._nil) {
286 node = node->left;
287 while (node->right != _data._nil) { /* returns the minimum of the left subtree of node */
288 node = node->right;
289 }
290 return node;
291 } else {
292 while (node == node->parent->left) {
293 node = node->parent;
294 }
295
296 if (node == _data._root) {
297 return nullptr; // No predecessor, as p_node = first node
298 }
299 return node->parent;
300 }
301 }
302
303 Element *_find(const K &p_key) const {
304 Element *node = _data._root->left;
305 C less;
306
307 while (node != _data._nil) {
308 if (less(p_key, node->_data.key)) {
309 node = node->left;
310 } else if (less(node->_data.key, p_key)) {
311 node = node->right;
312 } else {
313 return node; // found
314 }
315 }
316
317 return nullptr;
318 }
319
320 Element *_find_closest(const K &p_key) const {
321 Element *node = _data._root->left;
322 Element *prev = nullptr;
323 C less;
324
325 while (node != _data._nil) {
326 prev = node;
327
328 if (less(p_key, node->_data.key)) {
329 node = node->left;
330 } else if (less(node->_data.key, p_key)) {
331 node = node->right;
332 } else {
333 return node; // found
334 }
335 }
336
337 if (prev == nullptr) {
338 return nullptr; // tree empty
339 }
340
341 if (less(p_key, prev->_data.key)) {
342 prev = prev->_prev;
343 }
344
345 return prev;
346 }
347
348 void _insert_rb_fix(Element *p_new_node) {
349 Element *node = p_new_node;
350 Element *nparent = node->parent;
351 Element *ngrand_parent = nullptr;
352
353 while (nparent->color == RED) {
354 ngrand_parent = nparent->parent;
355
356 if (nparent == ngrand_parent->left) {
357 if (ngrand_parent->right->color == RED) {
358 _set_color(nparent, BLACK);
359 _set_color(ngrand_parent->right, BLACK);
360 _set_color(ngrand_parent, RED);
361 node = ngrand_parent;
362 nparent = node->parent;
363 } else {
364 if (node == nparent->right) {
365 _rotate_left(nparent);
366 node = nparent;
367 nparent = node->parent;
368 }
369 _set_color(nparent, BLACK);
370 _set_color(ngrand_parent, RED);
371 _rotate_right(ngrand_parent);
372 }
373 } else {
374 if (ngrand_parent->left->color == RED) {
375 _set_color(nparent, BLACK);
376 _set_color(ngrand_parent->left, BLACK);
377 _set_color(ngrand_parent, RED);
378 node = ngrand_parent;
379 nparent = node->parent;
380 } else {
381 if (node == nparent->left) {
382 _rotate_right(nparent);
383 node = nparent;
384 nparent = node->parent;
385 }
386 _set_color(nparent, BLACK);
387 _set_color(ngrand_parent, RED);
388 _rotate_left(ngrand_parent);
389 }
390 }
391 }
392
393 _set_color(_data._root->left, BLACK);
394 }
395
396 Element *_insert(const K &p_key, const V &p_value) {
397 Element *new_parent = _data._root;
398 Element *node = _data._root->left;
399 C less;
400
401 while (node != _data._nil) {
402 new_parent = node;
403
404 if (less(p_key, node->_data.key)) {
405 node = node->left;
406 } else if (less(node->_data.key, p_key)) {
407 node = node->right;
408 } else {
409 node->_data.value = p_value;
410 return node; // Return existing node with new value
411 }
412 }
413
414 typedef KeyValue<K, V> KV;
415 Element *new_node = memnew_allocator(Element(KV(p_key, p_value)), A);
416 new_node->parent = new_parent;
417 new_node->right = _data._nil;
418 new_node->left = _data._nil;
419
420 //new_node->data=_data;
421
422 if (new_parent == _data._root || less(p_key, new_parent->_data.key)) {
423 new_parent->left = new_node;
424 } else {
425 new_parent->right = new_node;
426 }
427
428 new_node->_next = _successor(new_node);
429 new_node->_prev = _predecessor(new_node);
430 if (new_node->_next) {
431 new_node->_next->_prev = new_node;
432 }
433 if (new_node->_prev) {
434 new_node->_prev->_next = new_node;
435 }
436
437 _data.size_cache++;
438 _insert_rb_fix(new_node);
439 return new_node;
440 }
441
442 void _erase_fix_rb(Element *p_node) {
443 Element *root = _data._root->left;
444 Element *node = _data._nil;
445 Element *sibling = p_node;
446 Element *parent = sibling->parent;
447
448 while (node != root) { // If red node found, will exit at a break
449 if (sibling->color == RED) {
450 _set_color(sibling, BLACK);
451 _set_color(parent, RED);
452 if (sibling == parent->right) {
453 sibling = sibling->left;
454 _rotate_left(parent);
455 } else {
456 sibling = sibling->right;
457 _rotate_right(parent);
458 }
459 }
460 if ((sibling->left->color == BLACK) && (sibling->right->color == BLACK)) {
461 _set_color(sibling, RED);
462 if (parent->color == RED) {
463 _set_color(parent, BLACK);
464 break;
465 } else { // loop: haven't found any red nodes yet
466 node = parent;
467 parent = node->parent;
468 sibling = (node == parent->left) ? parent->right : parent->left;
469 }
470 } else {
471 if (sibling == parent->right) {
472 if (sibling->right->color == BLACK) {
473 _set_color(sibling->left, BLACK);
474 _set_color(sibling, RED);
475 _rotate_right(sibling);
476 sibling = sibling->parent;
477 }
478 _set_color(sibling, parent->color);
479 _set_color(parent, BLACK);
480 _set_color(sibling->right, BLACK);
481 _rotate_left(parent);
482 break;
483 } else {
484 if (sibling->left->color == BLACK) {
485 _set_color(sibling->right, BLACK);
486 _set_color(sibling, RED);
487 _rotate_left(sibling);
488 sibling = sibling->parent;
489 }
490
491 _set_color(sibling, parent->color);
492 _set_color(parent, BLACK);
493 _set_color(sibling->left, BLACK);
494 _rotate_right(parent);
495 break;
496 }
497 }
498 }
499
500 ERR_FAIL_COND(_data._nil->color != BLACK);
501 }
502
503 void _erase(Element *p_node) {
504 Element *rp = ((p_node->left == _data._nil) || (p_node->right == _data._nil)) ? p_node : p_node->_next;
505 Element *node = (rp->left == _data._nil) ? rp->right : rp->left;
506
507 Element *sibling = nullptr;
508 if (rp == rp->parent->left) {
509 rp->parent->left = node;
510 sibling = rp->parent->right;
511 } else {
512 rp->parent->right = node;
513 sibling = rp->parent->left;
514 }
515
516 if (node->color == RED) {
517 node->parent = rp->parent;
518 _set_color(node, BLACK);
519 } else if (rp->color == BLACK && rp->parent != _data._root) {
520 _erase_fix_rb(sibling);
521 }
522
523 if (rp != p_node) {
524 ERR_FAIL_COND(rp == _data._nil);
525
526 rp->left = p_node->left;
527 rp->right = p_node->right;
528 rp->parent = p_node->parent;
529 rp->color = p_node->color;
530 if (p_node->left != _data._nil) {
531 p_node->left->parent = rp;
532 }
533 if (p_node->right != _data._nil) {
534 p_node->right->parent = rp;
535 }
536
537 if (p_node == p_node->parent->left) {
538 p_node->parent->left = rp;
539 } else {
540 p_node->parent->right = rp;
541 }
542 }
543
544 if (p_node->_next) {
545 p_node->_next->_prev = p_node->_prev;
546 }
547 if (p_node->_prev) {
548 p_node->_prev->_next = p_node->_next;
549 }
550
551 memdelete_allocator<Element, A>(p_node);
552 _data.size_cache--;
553 ERR_FAIL_COND(_data._nil->color == RED);
554 }
555
556 void _calculate_depth(Element *p_element, int &max_d, int d) const {
557 if (p_element == _data._nil) {
558 return;
559 }
560
561 _calculate_depth(p_element->left, max_d, d + 1);
562 _calculate_depth(p_element->right, max_d, d + 1);
563
564 if (d > max_d) {
565 max_d = d;
566 }
567 }
568
569 void _cleanup_tree(Element *p_element) {
570 if (p_element == _data._nil) {
571 return;
572 }
573
574 _cleanup_tree(p_element->left);
575 _cleanup_tree(p_element->right);
576 memdelete_allocator<Element, A>(p_element);
577 }
578
579 void _copy_from(const RBMap &p_map) {
580 clear();
581 // not the fastest way, but safeset to write.
582 for (Element *I = p_map.front(); I; I = I->next()) {
583 insert(I->key(), I->value());
584 }
585 }
586
587public:
588 const Element *find(const K &p_key) const {
589 if (!_data._root) {
590 return nullptr;
591 }
592
593 const Element *res = _find(p_key);
594 return res;
595 }
596
597 Element *find(const K &p_key) {
598 if (!_data._root) {
599 return nullptr;
600 }
601
602 Element *res = _find(p_key);
603 return res;
604 }
605
606 const Element *find_closest(const K &p_key) const {
607 if (!_data._root) {
608 return nullptr;
609 }
610
611 const Element *res = _find_closest(p_key);
612 return res;
613 }
614
615 Element *find_closest(const K &p_key) {
616 if (!_data._root) {
617 return nullptr;
618 }
619
620 Element *res = _find_closest(p_key);
621 return res;
622 }
623
624 bool has(const K &p_key) const {
625 return find(p_key) != nullptr;
626 }
627
628 Element *insert(const K &p_key, const V &p_value) {
629 if (!_data._root) {
630 _data._create_root();
631 }
632 return _insert(p_key, p_value);
633 }
634
635 void erase(Element *p_element) {
636 if (!_data._root || !p_element) {
637 return;
638 }
639
640 _erase(p_element);
641 if (_data.size_cache == 0 && _data._root) {
642 _data._free_root();
643 }
644 }
645
646 bool erase(const K &p_key) {
647 if (!_data._root) {
648 return false;
649 }
650
651 Element *e = find(p_key);
652 if (!e) {
653 return false;
654 }
655
656 _erase(e);
657 if (_data.size_cache == 0 && _data._root) {
658 _data._free_root();
659 }
660 return true;
661 }
662
663 const V &operator[](const K &p_key) const {
664 CRASH_COND(!_data._root);
665 const Element *e = find(p_key);
666 CRASH_COND(!e);
667 return e->_data.value;
668 }
669
670 V &operator[](const K &p_key) {
671 if (!_data._root) {
672 _data._create_root();
673 }
674
675 Element *e = find(p_key);
676 if (!e) {
677 e = insert(p_key, V());
678 }
679
680 return e->_data.value;
681 }
682
683 Element *front() const {
684 if (!_data._root) {
685 return nullptr;
686 }
687
688 Element *e = _data._root->left;
689 if (e == _data._nil) {
690 return nullptr;
691 }
692
693 while (e->left != _data._nil) {
694 e = e->left;
695 }
696
697 return e;
698 }
699
700 Element *back() const {
701 if (!_data._root) {
702 return nullptr;
703 }
704
705 Element *e = _data._root->left;
706 if (e == _data._nil) {
707 return nullptr;
708 }
709
710 while (e->right != _data._nil) {
711 e = e->right;
712 }
713
714 return e;
715 }
716
717 inline bool is_empty() const {
718 return _data.size_cache == 0;
719 }
720 inline int size() const {
721 return _data.size_cache;
722 }
723
724 int calculate_depth() const {
725 // used for debug mostly
726 if (!_data._root) {
727 return 0;
728 }
729
730 int max_d = 0;
731 _calculate_depth(_data._root->left, max_d, 0);
732 return max_d;
733 }
734
735 void clear() {
736 if (!_data._root) {
737 return;
738 }
739
740 _cleanup_tree(_data._root->left);
741 _data._root->left = _data._nil;
742 _data.size_cache = 0;
743 _data._free_root();
744 }
745
746 void operator=(const RBMap &p_map) {
747 _copy_from(p_map);
748 }
749
750 RBMap(const RBMap &p_map) {
751 _copy_from(p_map);
752 }
753
754 _FORCE_INLINE_ RBMap() {}
755
756 ~RBMap() {
757 clear();
758 }
759};
760
761#endif // RB_MAP_H
762