godot/core/set.h

/*************************************************************************/
/*  set.h                                                                */
/*************************************************************************/
/*                       This file is part of:                           */
/*                           GODOT ENGINE                                */
/*                    http://www.godotengine.org                         */
/*************************************************************************/
/* Copyright (c) 2007-2016 Juan Linietsky, Ariel Manzur.                 */
/*                                                                       */
/* Permission is hereby granted, free of charge, to any person obtaining */
/* a copy of this software and associated documentation files (the       */
/* "Software"), to deal in the Software without restriction, including   */
/* without limitation the rights to use, copy, modify, merge, publish,   */
/* distribute, sublicense, and/or sell copies of the Software, and to    */
/* permit persons to whom the Software is furnished to do so, subject to */
/* the following conditions:                                             */
/*                                                                       */
/* The above copyright notice and this permission notice shall be        */
/* included in all copies or substantial portions of the Software.       */
/*                                                                       */
/* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,       */
/* EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF    */
/* MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.*/
/* IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY  */
/* CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,  */
/* TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE     */
/* SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.                */
/*************************************************************************/
#ifndef SET_H
#define SET_H

#include "typedefs.h"
#include "os/memory.h"

/**
	@author Juan Linietsky <reduzio@gmail.com>
*/


// based on the very nice implementation of rb-trees by:
// http://web.mit.edu/~emin/www/source_code/red_black_tree/index.html

template <class T,class C=Comparator<T>,class A=DefaultAllocator >
class Set {

	enum Color {
		RED,
		BLACK
	};
	struct _Data;
public:



	class Element {

	private:
		friend class Set<T,C,A>;
		int color;
		Element* right;
		Element* left;
		Element* parent;
		Element* _next;
		Element* _prev;
		T value;
		//_Data *data;

	public:

		const Element *next() const {

			return _next;
		}
		Element *next() {

			return _next;
		}
		const Element *prev() const {

			return _prev;
		}
		Element *prev() {

			return _prev;
		}
		const T& get() const {
			return value;
		};
		Element() {
			color=RED;
			right=NULL;
			left=NULL;
			parent=NULL;
			_next=NULL;
			_prev=NULL;
		};
	};


private:

	struct _Data {

		Element* _root;
		Element* _nil;
		int size_cache;


		_Data() {
#ifdef GLOBALNIL_DISABLED
			_nil = memnew_allocator( Element,A );
			_nil->parent=_nil->left=_nil->right=_nil;
			_nil->color=BLACK;
#else

			_nil=(Element*)&_GlobalNilClass::_nil;
#endif
			_root = NULL;

			size_cache=0;
		}

		void _create_root() {

			_root = memnew_allocator( Element,A );
			_root->parent=_root->left=_root->right=_nil;
			_root->color=BLACK;
		}

		void _free_root() {

			if (_root) {
				memdelete_allocator<Element,A>(_root);
				_root=NULL;
			}
		}

		~_Data() {

			_free_root();
#ifdef GLOBALNIL_DISABLED
			memdelete_allocator<Element,A>(_nil);
#endif
//			memdelete_allocator<Element,A>(_root);
		}
	};

	_Data _data;

	inline void _set_color(Element *p_node, int p_color) {

		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;
		p_node->right=r->left;
		if (r->left != _data._nil )
			r->left->parent=p_node;
		r->parent=p_node->parent;
		if (p_node==p_node->parent->left)
			p_node->parent->left=r;
		else
			p_node->parent->right=r;

		r->left=p_node;
		p_node->parent=r;

	}

	inline void _rotate_right(Element *p_node) {

		Element *l=p_node->left;
		p_node->left=l->right;
		if (l->right != _data._nil)
			l->right->parent=p_node;
		l->parent=p_node->parent;
		if (p_node==p_node->parent->right)
			p_node->parent->right=l;
		else
			p_node->parent->left=l;

		l->right=p_node;
		p_node->parent=l;

	}

	inline Element* _successor(Element *p_node) const {

		Element *node=p_node;

		if (node->right != _data._nil) {

			node=node->right;
			while(node->left != _data._nil) { /* returns the minium of the right subtree of node */
				node=node->left;
			}
			return node;
		} else {

			while(node == node->parent->right) {
				node=node->parent;
			}
			if (node->parent == _data._root)
				return NULL;
			return node->parent;
		}
	}

	inline Element* _predecessor(Element *p_node) const {
		Element *node=p_node;

		if (node->left != _data._nil) {

			node=node->left;
			while(node->right != _data._nil) { /* returns the minium of the left subtree of node */
				node=node->right;
			}
			return node;
		} else {

			while(node == node->parent->left) {
				if (node->parent == _data._root)
					return NULL;

				node=node->parent;
			}
			return node->parent;
		}
	}


	Element *_find(const T& p_value) const {

		Element *node = _data._root->left;
		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!=_data._nil)?node:NULL;
	}

	Element *_lower_bound(const T& p_value) const {

		Element *node = _data._root->left;
		Element *prev = NULL;
		C less;

		while(node!=_data._nil) {
			prev=node;

			if (less(p_value,node->value))
				node=node->left;
			else if (less(node->value,p_value))
				node=node->right;
			else
				break; // found
		}

		if (node==_data._nil) {
			if (prev==NULL)
				return NULL;
			if (less(prev->value,p_value)) {

				prev=prev->_next;
			}

			return prev;

		} else
			return node;
	}


	Element *_insert(const T& p_value, bool& r_exists) {

		Element *new_parent=_data._root;
		Element *node = _data._root->left;
		C less;

		while (node!=_data._nil) {

			new_parent=node;

			if (less(p_value,node->value))
				node=node->left;
			else if (less(node->value,p_value))
				node=node->right;
			else {
				r_exists=true;
				return node;
			}
		}

		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)) {

			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)
			new_node->_next->_prev=new_node;
		if (new_node->_prev)
			new_node->_prev->_next=new_node;


		return new_node;
	}

	Element * _insert_rb(const T& p_value) {

		bool exists=false;
		Element *new_node = _insert(p_value,exists);
		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);
		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;
				} 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 */
				}
			} 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 ( (aux->right->color==BLACK) && (aux->left->color==BLACK) ) {
					_set_color(aux,RED);
					node=node->parent;
				} else {
					if (aux->left->color==BLACK) {
						_set_color(aux->right,BLACK);
						_set_color(aux,RED);
						_rotate_left(aux);
						aux=node->parent->left;
					}
					_set_color(aux,node->parent->color);
					_set_color(node->parent,BLACK);
					_set_color(aux->left,BLACK);
					_rotate_right(node->parent);
					node=root;
				}
			}
		}

		_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 *node= (rp->left == _data._nil) ? rp->right : rp->left;

		if (_data._root == (node->parent=rp->parent) ) {
			_data._root->left=node;
		} else {
			if (rp == rp->parent->left) {
				rp->parent->left=node;
			} else {
				rp->parent->right=node;
			}
		}

		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 == 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)
			p_node->_next->_prev=p_node->_prev;
		if (p_node->_prev)
			p_node->_prev->_next=p_node->_next;

		memdelete_allocator<Element,A>(p_node);
		_data.size_cache--;
		ERR_FAIL_COND( _data._nil->color==RED );
	}


	void _calculate_depth(Element *p_element,int &max_d,int d) const {

		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;
	}

	void _cleanup_tree(Element *p_element) {

		if (p_element==_data._nil)
			return;

		_cleanup_tree(p_element->left);
		_cleanup_tree(p_element->right);
		memdelete_allocator<Element,A>( p_element );
	}

	void _copy_from( const Set& p_set) {

		clear();
		// not the fastest way, but safeset to write.
		for(Element *I=p_set.front();I;I=I->next()) {

			insert(I->get());
		}
	}
public:

	const Element *find(const T& p_value) const {

		if (!_data._root)
			return NULL;

		const Element *res=_find(p_value);
		return res;
	}

	Element *find(const T& p_value) {

		if (!_data._root)
			return NULL;
		Element *res=_find(p_value);
		return res;
	}

	bool has(const T& p_value) const {

		if (!_data._root)
			return false;
		return find(p_value)!=NULL;
	}

	Element *insert(const T& p_value) {

		if (!_data._root)
			_data._create_root();
		return _insert_rb(p_value);

	}

	void erase(Element* p_element) {

		if (!_data._root)
			return;
		_erase(p_element);
		if (_data.size_cache==0 && _data._root)
			_data._free_root();
	}

	bool erase(const T& p_value) {

		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();
		return true;
	}

	Element *front() const {

		if (!_data._root)
			return NULL;
		Element *e=_data._root->left;
		if (e==_data._nil)
			return NULL;

		while(e->left!=_data._nil)
			e=e->left;

		return e;
	}

	Element *back() const {

		if (!_data._root)
			return NULL;
		Element *e=_data._root->left;
		if (e==_data._nil)
			return NULL;

		while(e->right!=_data._nil)
			e=e->right;

		return e;
	}

	Element *lower_bound(const T& p_value) const {

		return _lower_bound(p_value);
	}


	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;
	}

	void clear() {

		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();
	}

	void operator=(const Set& p_set) {

		_copy_from( p_set );
	}

	Set(const Set& p_set) {

		_copy_from( p_set );
	}
	_FORCE_INLINE_ Set() {


	}


	~Set() {

		clear();
	}

};

#endif