map.h 15 KB

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