1 files changed, 126 insertions, 119 deletions

diff --git a/core/set.h b/core/set.h
index f68d78cea1..d91dd9b3ea 100644
--- a/core/set.h
+++ b/core/set.h
@@ -135,7 +135,6 @@ private:
 #ifdef GLOBALNIL_DISABLED
 			memdelete_allocator<Element, A>(_nil);
 #endif
-			//memdelete_allocator<Element,A>(_root);
 		}
 	};
 
@@ -146,6 +145,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;
@@ -194,8 +194,9 @@ private:
 			while (node == node->parent->right) {
 				node = node->parent;
 			}
+
 			if (node->parent == _data._root)
-				return NULL;
+				return NULL; // No successor, as p_node is the last node.
 			return node->parent;
 		}
 	}
@@ -213,11 +214,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 is the first node.
 			return node->parent;
 		}
 	}
@@ -228,16 +229,15 @@ private:
 		C less;
 
 		while (node != _data._nil) {
-
 			if (less(p_value, node->value))
 				node = node->left;
 			else if (less(node->value, p_value))
 				node = node->right;
 			else
-				break; // found
+				return node; // found
 		}
 
-		return (node != _data._nil) ? node : NULL;
+		return NULL;
 	}
 
 	Element *_lower_bound(const T &p_value) const {
@@ -254,21 +254,16 @@ private:
 			else if (less(node->value, p_value))
 				node = node->right;
 			else
-				break; // found
+				return node; // found
 		}
 
-		if (node == _data._nil) {
-			if (prev == NULL)
-				return NULL;
-			if (less(prev->value, p_value)) {
-
-				prev = prev->_next;
-			}
+		if (prev == NULL)
+			return NULL; // tree empty
 
-			return prev;
+		if (less(prev->value, p_value))
+			prev = prev->_next;
 
-		} else
-			return node;
+		return prev;
 	}
 
 	Element *_insert(const T &p_value, bool &r_exists) {
@@ -291,22 +286,21 @@ private:
 			}
 		}
 
-		Element *new_node = memnew_allocator(Element, A);
+		r_exists = false;
 
+		Element *new_node = memnew_allocator(Element, A);
 		new_node->parent = new_parent;
 		new_node->right = _data._nil;
 		new_node->left = _data._nil;
 		new_node->value = p_value;
 		//new_node->data=_data;
-		if (new_parent == _data._root || less(p_value, new_parent->value)) {
 
+		if (new_parent == _data._root || less(p_value, new_parent->value)) {
 			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)
@@ -324,154 +318,159 @@ private:
 		if (exists)
 			return new_node;
 
-		Element *node = new_node;
 		_data.size_cache++;
+		Element *node = new_node;
+		Element *nparent = node->parent;
+		Element *ngrand_parent;
 
-		while (node->parent->color == RED) {
-
-			if (node->parent == node->parent->parent->left) {
+		while (nparent->color == RED) {
 
-				Element *aux = node->parent->parent->right;
+			ngrand_parent = nparent->parent;
 
-				if (aux->color == RED) {
-					_set_color(node->parent, BLACK);
-					_set_color(aux, BLACK);
-					_set_color(node->parent->parent, RED);
-					node = node->parent->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 == node->parent->right) {
-						node = node->parent;
-						_rotate_left(node);
+					if (node == nparent->right) {
+						_rotate_left(nparent);
+						node = nparent;
+						nparent = node->parent;
 					}
-					_set_color(node->parent, BLACK);
-					_set_color(node->parent->parent, RED);
-					_rotate_right(node->parent->parent);
+					_set_color(nparent, BLACK);
+					_set_color(ngrand_parent, RED);
+					_rotate_right(ngrand_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;
+				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 == node->parent->left) {
-						node = node->parent;
-						_rotate_right(node);
+					if (node == nparent->left) {
+						_rotate_right(nparent);
+						node = nparent;
+						nparent = node->parent;
 					}
-					_set_color(node->parent, BLACK);
-					_set_color(node->parent->parent, RED);
-					_rotate_left(node->parent->parent);
+					_set_color(nparent, BLACK);
+					_set_color(ngrand_parent, RED);
+					_rotate_left(ngrand_parent);
 				}
 			}
 		}
+
 		_set_color(_data._root->left, BLACK);
+
 		return new_node;
 	}
 
 	void _erase_fix(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(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)
@@ -529,6 +528,7 @@ public:
 
 		if (!_data._root)
 			return NULL;
+
 		Element *res = _find(p_value);
 		return res;
 	}
@@ -549,8 +549,9 @@ public:
 
 	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();
@@ -560,9 +561,11 @@ public:
 
 		if (!_data._root)
 			return false;
+
 		Element *e = find(p_value);
 		if (!e)
 			return false;
+
 		_erase(e);
 		if (_data.size_cache == 0 && _data._root)
 			_data._free_root();
@@ -573,6 +576,7 @@ public:
 
 		if (!_data._root)
 			return NULL;
+
 		Element *e = _data._root->left;
 		if (e == _data._nil)
 			return NULL;
@@ -587,6 +591,7 @@ public:
 
 		if (!_data._root)
 			return NULL;
+
 		Element *e = _data._root->left;
 		if (e == _data._nil)
 			return NULL;
@@ -603,10 +608,12 @@ public:
 	}
 
 	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;
@@ -620,7 +627,6 @@ public:
 		_cleanup_tree(_data._root->left);
 		_data._root->left = _data._nil;
 		_data.size_cache = 0;
-		_data._nil->parent = _data._nil;
 		_data._free_root();
 	}
 
@@ -633,6 +639,7 @@ public:
 
 		_copy_from(p_set);
 	}
+
 	_FORCE_INLINE_ Set() {
 	}