rb_set.h 16 KB

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