Refactor core/map.h to be similar to core/set.h

1 files changed, 158 insertions, 161 deletions

diff --git a/core/map.h b/core/map.h
index a37d898a9c..f01062ebed 100644
--- a/core/map.h
+++ b/core/map.h
@@ -31,6 +31,7 @@
 #define MAP_H
 
 #include "set.h"
+
 /**
 	@author Juan Linietsky <reduzio@gmail.com>
 */
@@ -52,7 +53,6 @@ public:
 
 	private:
 		friend class Map<K, V, C, A>;
-		//Color color;
 		int color;
 		Element *right;
 		Element *left;
@@ -61,7 +61,6 @@ public:
 		Element *_prev;
 		K _key;
 		V _value;
-
 		//_Data *data;
 
 	public:
@@ -147,7 +146,6 @@ private:
 #ifdef GLOBALNIL_DISABLED
 			memdelete_allocator<Element, A>(_nil);
 #endif
-			//memdelete_allocator<Element,A>(_root);
 		}
 	};
 
@@ -158,6 +156,7 @@ private:
 		ERR_FAIL_COND(p_node == _data._nil && p_color == RED);
 		p_node->color = p_color;
 	}
+
 	inline void _rotate_left(Element *p_node) {
 
 		Element *r = p_node->right;
@@ -206,8 +205,9 @@ private:
 			while (node == node->parent->right) {
 				node = node->parent;
 			}
+
 			if (node->parent == _data._root)
-				return NULL;
+				return NULL; // No successor, as p_node = last node
 			return node->parent;
 		}
 	}
@@ -225,10 +225,11 @@ private:
 		} else {
 
 			while (node == node->parent->left) {
-				if (node->parent == _data._root)
-					return NULL;
 				node = node->parent;
 			}
+
+			if (node == _data._root)
+				return NULL; // No predecessor, as p_node = first node
 			return node->parent;
 		}
 	}
@@ -239,16 +240,15 @@ private:
 		C less;
 
 		while (node != _data._nil) {
-
 			if (less(p_key, node->_key))
 				node = node->left;
 			else if (less(node->_key, p_key))
 				node = node->right;
 			else
-				break; // found
+				return node; // found
 		}
 
-		return (node != _data._nil) ? node : NULL;
+		return NULL;
 	}
 
 	Element *_find_closest(const K &p_key) const {
@@ -265,24 +265,68 @@ private:
 			else if (less(node->_key, p_key))
 				node = node->right;
 			else
-				break; // found
+				return node; // found
 		}
 
-		if (node == _data._nil) {
-			if (prev == NULL)
-				return NULL;
-			if (less(p_key, prev->_key)) {
+		if (prev == NULL)
+			return NULL; // tree empty
 
-				prev = prev->_prev;
-			}
+		if (less(p_key, prev->_key))
+			prev = prev->_prev;
 
-			return prev;
+		return prev;
+	}
 
-		} else
-			return node;
+	void _insert_rb_fix(Element *p_new_node) {
+
+		Element *node = p_new_node;
+		Element *nparent = node->parent;
+		Element *ngrand_parent;
+
+		while (nparent->color == RED) {
+			ngrand_parent = nparent->parent;
+
+			if (nparent == ngrand_parent->left) {
+				if (ngrand_parent->right->color == RED) {
+					_set_color(nparent, BLACK);
+					_set_color(ngrand_parent->right, BLACK);
+					_set_color(ngrand_parent, RED);
+					node = ngrand_parent;
+					nparent = node->parent;
+				} else {
+					if (node == nparent->right) {
+						_rotate_left(nparent);
+						node = nparent;
+						nparent = node->parent;
+					}
+					_set_color(nparent, BLACK);
+					_set_color(ngrand_parent, RED);
+					_rotate_right(ngrand_parent);
+				}
+			} else {
+				if (ngrand_parent->left->color == RED) {
+					_set_color(nparent, BLACK);
+					_set_color(ngrand_parent->left, BLACK);
+					_set_color(ngrand_parent, RED);
+					node = ngrand_parent;
+					nparent = node->parent;
+				} else {
+					if (node == nparent->left) {
+						_rotate_right(nparent);
+						node = nparent;
+						nparent = node->parent;
+					}
+					_set_color(nparent, BLACK);
+					_set_color(ngrand_parent, RED);
+					_rotate_left(ngrand_parent);
+				}
+			}
+		}
+
+		_set_color(_data._root->left, BLACK);
 	}
 
-	Element *_insert(const K &p_key, bool &r_exists) {
+	Element *_insert(const K &p_key, const V &p_value) {
 
 		Element *new_parent = _data._root;
 		Element *node = _data._root->left;
@@ -297,27 +341,25 @@ private:
 			else if (less(node->_key, p_key))
 				node = node->right;
 			else {
-				r_exists = true;
-				return node;
+				node->_value = p_value;
+				return node; // Return existing node with new value
 			}
 		}
 
 		Element *new_node = memnew_allocator(Element, A);
-
 		new_node->parent = new_parent;
 		new_node->right = _data._nil;
 		new_node->left = _data._nil;
 		new_node->_key = p_key;
+		new_node->_value = p_value;
 		//new_node->data=_data;
-		if (new_parent == _data._root || less(p_key, new_parent->_key)) {
 
+		if (new_parent == _data._root || less(p_key, new_parent->_key)) {
 			new_parent->left = new_node;
 		} else {
 			new_parent->right = new_node;
 		}
 
-		r_exists = false;
-
 		new_node->_next = _successor(new_node);
 		new_node->_prev = _predecessor(new_node);
 		if (new_node->_next)
@@ -325,168 +367,113 @@ private:
 		if (new_node->_prev)
 			new_node->_prev->_next = new_node;
 
-		return new_node;
-	}
-
-	Element *_insert_rb(const K &p_key, const V &p_value) {
-
-		bool exists = false;
-		Element *new_node = _insert(p_key, exists);
-
-		if (new_node) {
-			new_node->_value = p_value;
-		}
-		if (exists)
-			return new_node;
-
-		Element *node = new_node;
 		_data.size_cache++;
-
-		while (node->parent->color == RED) {
-
-			if (node->parent == node->parent->parent->left) {
-
-				Element *aux = node->parent->parent->right;
-
-				if (aux->color == RED) {
-					_set_color(node->parent, BLACK);
-					_set_color(aux, BLACK);
-					_set_color(node->parent->parent, RED);
-					node = node->parent->parent;
-				} else {
-					if (node == node->parent->right) {
-						node = node->parent;
-						_rotate_left(node);
-					}
-					_set_color(node->parent, BLACK);
-					_set_color(node->parent->parent, RED);
-					_rotate_right(node->parent->parent);
-				}
-			} else {
-				Element *aux = node->parent->parent->left;
-
-				if (aux->color == RED) {
-					_set_color(node->parent, BLACK);
-					_set_color(aux, BLACK);
-					_set_color(node->parent->parent, RED);
-					node = node->parent->parent;
-				} else {
-					if (node == node->parent->left) {
-						node = node->parent;
-						_rotate_right(node);
-					}
-					_set_color(node->parent, BLACK);
-					_set_color(node->parent->parent, RED);
-					_rotate_left(node->parent->parent);
-				}
-			}
-		}
-		_set_color(_data._root->left, BLACK);
+		_insert_rb_fix(new_node);
 		return new_node;
 	}
 
-	void _erase_fix(Element *p_node) {
+	void _erase_fix_rb(Element *p_node) {
 
 		Element *root = _data._root->left;
-		Element *node = p_node;
-
-		while ((node->color == BLACK) && (root != node)) {
-			if (node == node->parent->left) {
-				Element *aux = node->parent->right;
-				if (aux->color == RED) {
-					_set_color(aux, BLACK);
-					_set_color(node->parent, RED);
-					_rotate_left(node->parent);
-					aux = node->parent->right;
-				}
-				if ((aux->right->color == BLACK) && (aux->left->color == BLACK)) {
-					_set_color(aux, RED);
-					node = node->parent;
+		Element *node = _data._nil;
+		Element *sibling = p_node;
+		Element *parent = sibling->parent;
+
+		while (node != root) { // If red node found, will exit at a break
+			if (sibling->color == RED) {
+				_set_color(sibling, BLACK);
+				_set_color(parent, RED);
+				if (sibling == parent->right) {
+					sibling = sibling->left;
+					_rotate_left(parent);
 				} else {
-					if (aux->right->color == BLACK) {
-						_set_color(aux->left, BLACK);
-						_set_color(aux, RED);
-						_rotate_right(aux);
-						aux = node->parent->right;
-					}
-					_set_color(aux, node->parent->color);
-					_set_color(node->parent, BLACK);
-					_set_color(aux->right, BLACK);
-					_rotate_left(node->parent);
-					node = root; /* this is to exit while loop */
+					sibling = sibling->right;
+					_rotate_right(parent);
 				}
-			} else { /* the code below is has left and right switched from above */
-				Element *aux = node->parent->left;
-				if (aux->color == RED) {
-					_set_color(aux, BLACK);
-					_set_color(node->parent, RED);
-					_rotate_right(node->parent);
-					aux = node->parent->left;
+			}
+			if ((sibling->left->color == BLACK) && (sibling->right->color == BLACK)) {
+				_set_color(sibling, RED);
+				if (parent->color == RED) {
+					_set_color(parent, BLACK);
+					break;
+				} else { // loop: haven't found any red nodes yet
+					node = parent;
+					parent = node->parent;
+					sibling = (node == parent->left) ? parent->right : parent->left;
 				}
-				if ((aux->right->color == BLACK) && (aux->left->color == BLACK)) {
-					_set_color(aux, RED);
-					node = node->parent;
+			} else {
+				if (sibling == parent->right) {
+					if (sibling->right->color == BLACK) {
+						_set_color(sibling->left, BLACK);
+						_set_color(sibling, RED);
+						_rotate_right(sibling);
+						sibling = sibling->parent;
+					}
+					_set_color(sibling, parent->color);
+					_set_color(parent, BLACK);
+					_set_color(sibling->right, BLACK);
+					_rotate_left(parent);
+					break;
 				} else {
-					if (aux->left->color == BLACK) {
-						_set_color(aux->right, BLACK);
-						_set_color(aux, RED);
-						_rotate_left(aux);
-						aux = node->parent->left;
+					if (sibling->left->color == BLACK) {
+						_set_color(sibling->right, BLACK);
+						_set_color(sibling, RED);
+						_rotate_left(sibling);
+						sibling = sibling->parent;
 					}
-					_set_color(aux, node->parent->color);
-					_set_color(node->parent, BLACK);
-					_set_color(aux->left, BLACK);
-					_rotate_right(node->parent);
-					node = root;
+
+					_set_color(sibling, parent->color);
+					_set_color(parent, BLACK);
+					_set_color(sibling->left, BLACK);
+					_rotate_right(parent);
+					break;
 				}
 			}
 		}
 
-		_set_color(node, BLACK);
-
 		ERR_FAIL_COND(_data._nil->color != BLACK);
 	}
 
 	void _erase(Element *p_node) {
 
-		Element *rp = ((p_node->left == _data._nil) || (p_node->right == _data._nil)) ? p_node : _successor(p_node);
-		if (!rp)
-			rp = _data._nil;
+		Element *rp = ((p_node->left == _data._nil) || (p_node->right == _data._nil)) ? p_node : p_node->_next;
 		Element *node = (rp->left == _data._nil) ? rp->right : rp->left;
 		node->parent = rp->parent;
 
-		if (_data._root == node->parent) {
-			_data._root->left = node;
+		Element *sibling;
+		if (rp == rp->parent->left) {
+			rp->parent->left = node;
+			sibling = rp->parent->right;
 		} else {
-			if (rp == rp->parent->left) {
-				rp->parent->left = node;
-			} else {
-				rp->parent->right = node;
-			}
+			rp->parent->right = node;
+			sibling = rp->parent->left;
+		}
+
+		if (node->color == RED) {
+			node->parent = rp->parent;
+			_set_color(node, BLACK);
+		} else if (rp->color == BLACK && rp->parent != _data._root) {
+			_erase_fix_rb(sibling);
 		}
 
 		if (rp != p_node) {
 
 			ERR_FAIL_COND(rp == _data._nil);
 
-			if (rp->color == BLACK)
-				_erase_fix(node);
-
 			rp->left = p_node->left;
 			rp->right = p_node->right;
 			rp->parent = p_node->parent;
 			rp->color = p_node->color;
-			p_node->left->parent = rp;
-			p_node->right->parent = rp;
+			if (p_node->left != _data._nil)
+				p_node->left->parent = rp;
+			if (p_node->right != _data._nil)
+				p_node->right->parent = rp;
 
 			if (p_node == p_node->parent->left) {
 				p_node->parent->left = rp;
 			} else {
 				p_node->parent->right = rp;
 			}
-		} else {
-			if (p_node->color == BLACK)
-				_erase_fix(node);
 		}
 
 		if (p_node->_next)
@@ -501,11 +488,12 @@ private:
 
 	void _calculate_depth(Element *p_element, int &max_d, int d) const {
 
-		if (p_element == _data._nil) {
+		if (p_element == _data._nil)
 			return;
-		}
+
 		_calculate_depth(p_element->left, max_d, d + 1);
 		_calculate_depth(p_element->right, max_d, d + 1);
+
 		if (d > max_d)
 			max_d = d;
 	}
@@ -544,6 +532,7 @@ public:
 
 		if (!_data._root)
 			return NULL;
+
 		Element *res = _find(p_key);
 		return res;
 	}
@@ -552,6 +541,7 @@ public:
 
 		if (!_data._root)
 			return NULL;
+
 		const Element *res = _find_closest(p_key);
 		return res;
 	}
@@ -560,21 +550,28 @@ public:
 
 		if (!_data._root)
 			return NULL;
+
 		Element *res = _find_closest(p_key);
 		return res;
 	}
 
+	bool has(const K &p_key) const {
+
+		return find(p_key) != NULL;
+	}
+
 	Element *insert(const K &p_key, const V &p_value) {
 
 		if (!_data._root)
 			_data._create_root();
-		return _insert_rb(p_key, p_value);
+		return _insert(p_key, p_value);
 	}
 
 	void erase(Element *p_element) {
 
-		if (!_data._root)
+		if (!_data._root || !p_element)
 			return;
+
 		_erase(p_element);
 		if (_data.size_cache == 0 && _data._root)
 			_data._free_root();
@@ -584,20 +581,17 @@ public:
 
 		if (!_data._root)
 			return false;
+
 		Element *e = find(p_key);
 		if (!e)
 			return false;
+
 		_erase(e);
+		if (_data.size_cache == 0 && _data._root)
+			_data._free_root();
 		return true;
 	}
 
-	bool has(const K &p_key) const {
-
-		if (!_data._root)
-			return false;
-		return find(p_key) != NULL;
-	}
-
 	const V &operator[](const K &p_key) const {
 
 		CRASH_COND(!_data._root);
@@ -605,6 +599,7 @@ public:
 		CRASH_COND(!e);
 		return e->_value;
 	}
+
 	V &operator[](const K &p_key) {
 
 		if (!_data._root)
@@ -614,7 +609,6 @@ public:
 		if (!e)
 			e = insert(p_key, V());
 
-		CRASH_COND(!e);
 		return e->_value;
 	}
 
@@ -637,6 +631,7 @@ public:
 
 		if (!_data._root)
 			return NULL;
+
 		Element *e = _data._root->left;
 		if (e == _data._nil)
 			return NULL;
@@ -649,10 +644,12 @@ public:
 
 	inline bool empty() const { return _data.size_cache == 0; }
 	inline int size() const { return _data.size_cache; }
+
 	int calculate_depth() const {
 		// used for debug mostly
 		if (!_data._root)
 			return 0;
+
 		int max_d = 0;
 		_calculate_depth(_data._root->left, max_d, 0);
 		return max_d;
@@ -662,10 +659,10 @@ public:
 
 		if (!_data._root)
 			return;
+
 		_cleanup_tree(_data._root->left);
 		_data._root->left = _data._nil;
 		_data.size_cache = 0;
-		_data._nil->parent = _data._nil;
 		_data._free_root();
 	}