//
// Cforall Version 1.0.0 Copyright (C) 2015 University of Waterloo
//
// The contents of this file are covered under the licence agreement in the
// file "LICENCE" distributed with Cforall.
//
// PersistentDisjointSet.h --
//
// Author           : Aaron B. Moss
// Created On       : Wed Jun 13 16:31:00 2018
// Last Modified By : Aaron B. Moss
// Last Modified On : Wed Jun 13 16:31:00 2018
// Update Count     : 1
//

#pragma once

#include <cassert>
#include <functional>
#include <unordered_map>
#include <utility>

#include "GC.h"

/// Persistent disjoint-set data structure based on the persistent array in 
/// Conchon & Filliatre "A Persistent Union-Find Data Structure". Path 
/// compression is not performed (to lower cost of persistent rollback). 
/// Auxilliary lists are kept for efficient retrieval of class members.
/// Find root should still operate in O(log k), for k the size of an 
/// equivalence class.

template<typename Elm, typename Hash = std::hash<Elm>, typename Eq = std::equal_to<Elm>>
class PersistentDisjointSet : public GC_Object {
public:
	/// Type of this class
	using Self = PersistentDisjointSet<Elm, Hash, Eq>;

	/// Types of version nodes
	enum Mode { 
		BASE,    ///< Root node of version tree
		ADD,     ///< Add key to set
		REM,     ///< Reverse add operation
		ADDTO,   ///< Merge one class root under another
		REMFROM  ///< Reverse addTo operation
	};

private:
	/// Type of node height
	using Height = unsigned char;

	/// Disjoint-set node
	struct Node {
		Elm parent;     ///< Parent node in equivalence class
		Elm next;       ///< Next node in equivalence class
		Height height;  ///< Tree height of the node

		template<typename E>
		Node(E&& e) : parent(e), next(std::forward<E>(e)), height(0) {}

		template<typename E, typename F>
		Node(E&& p, F&& n, Height h) 
			: parent(std::forward<E>(p)), next(std::forward<F>(n)), height(h) {}
	};

	/// Type of class map
	using Base = std::unordered_map<Elm, Node, Hash, Eq>;

	/// Node inserted into underlying map as new equivalence class
	struct Add {
		Self* base;  ///< Modified map
		Elm root;    ///< Element added

		template<typename E>
		Add(Self* b, E&& r) : base(b), root(std::forward<E>(r)) {}
	};

	/// Two classes merged
	struct AddTo {
		Self* base;       ///< Modified map
		Elm root;         ///< Root node
		Elm child;        ///< Child node, formerly root of own class
		bool new_height;  ///< Did the root node's height change?

		template<typename R, typename C>
		AddTo(Self* b, R&& r, C&& c, bool h) 
			: base(b), root(std::forward<R>(r)), child(std::forward<C>(c)), new_height(h) {}
	};

	/// Underlying storage
	union Data {
		char none;
		Base base;
		Add add;
		AddTo add_to;

		Data() : none('\0') {}
		~Data() {}
	} data;

	/// Type of version node
	mutable Mode mode;

	/// get mutable reference as T
	template<typename T>
	T& as() { return reinterpret_cast<T&>(data); }

	/// get const reference as T
	template<typename T>
	const T& as() const { return reinterpret_cast<const T&>(data); }

	/// get rvalue reference as T
	template<typename T>
	T&& take_as() { return std::move(as<T>()); }

	/// initialize as T
	template<typename T, typename... Args>
	void init( Args&&... args ) {
		new( &as<T>() ) T { std::forward<Args>(args)... };
	}

	/// reset according to current mode
	void reset() {
		switch( mode ) {
			case BASE:                as<Base>().~Base();   break;
			case ADD: case REM:       as<Add>().~Add();     break;
			case ADDTO: case REMFROM: as<AddTo>().~AddTo(); break;
			default: assertf(false, "invalid mode");
		}
	}

	/// reset as base
	void reset_as_base() {
		assertf( mode == BASE, "can only reset_as_base() on BASE" );
		as<Base>().~Base();
	}

	/// Non-initializing constructor; should call init() before use
	PersistentDisjointSet( Mode m ) : data(), mode(m) {}

	PersistentDisjointSet( Mode m, Base&& b ) : data(), mode(m) {
		assertf(m == BASE, "invalid mode");
		init<Base>(std::move(b));
	}

	template<typename R>
	PersistentDisjointSet( Mode m, const Self* b, R&& r ) : data(), mode(m) {
		assertf(m == ADD || m == REM, "invalid mode");
		init<Add>(b, std::forward<R>(r));
	}

	template<typename R, typename C>
	PersistentDisjointSet( Mode m, const Self* b, R&& r, C&& c, bool h ) : data(), mode(m) {
		assertf(m == ADDTO || m == REMFROM, "invalid mode");
		init<AddTo>(b, std::forward<R>(r), std::forward<C>(c), h);
	}

	/// Adds (also removes) graph edges.
	/// * `from.parent` updated to `new_root`, 
	/// * `from.next` and `to.next` swapped (splices or un-splices class lists)
	/// * `to.height` adjusted by change
	template<typename R>
	static void addEdge( Node& from, Node& to, R&& new_root, Height change ) {
		from.parent = std::forward<R>(new_root);
		std::swap(from.next, to.next);
		to.height += change;
	}

protected:
	void trace( const GC& gc ) const {
		switch( mode ) {
			case BASE: {
				for (const auto& entry : as<Base>()) {
					gc.maybe_trace( entry.first );
				}
				return;
			}
			case ADD: case REM: {
				const Add& self = as<Add>();
				gc << self.base; 
				gc.maybe_trace( self.root );
				return;
			}
			case ADDTO: case REMFROM: {
				const AddTo& self = as<AddTo>();
				gc << self.base;
				gc.maybe_trace( self.root, self.child );
				return;
			}
			default: assertf(false, "invalid mode");
		}
	}

public:
	using size_type = std::size_t;

	using iterator = typename Base::const_iterator;

	PersistentDisjointSet() : data(), mode(BASE) { init<Base>(); }

	PersistentDisjointSet( const Self& o ) = delete;

	Self& operator= ( const Self& o ) = delete;

	~PersistentDisjointSet() { reset(); }

	/// reroot persistent data structure at current node
	void reroot() const {
		if ( mode == BASE ) return;

		// reroot base
		Self* mut_this = const_cast<Self*>(this);
		Self* base = ( mode == ADD || mode == REM ) ? 
			mut_this->as<Add>().base : 
			mut_this->as<AddTo>().base;
		base->reroot();
		assertf(base->mode == BASE, "reroot results in base");

		// take map out of base
		Base base_map = base->take_as<Base>();
		base->reset_as_base();

		// switch base to inverse of self and mutate base map
		switch ( mode ) {
			case ADD: {
				Add& self = mut_this->as<Add>();

				base->init<Add>( mut_this, self.root );
				base->mode = REM;

				base_map.emplace( self.root, Node{ std::move(self.root) } );
			} break;
			case REM: {
				Add& self = mut_this->as<Add>();

				base->init<Add>( mut_this, self.root );
				base->mode = ADD;

				base_map.erase( self.root );
			} break;
			case ADDTO: {
				AddTo& self = mut_this->as<AddTo>();

				base->init<AddTo>( mut_this, self.root, self.child, self.new_height );
				base->mode = REMFROM;

				auto child_it = base_map.find( self.child );
				auto root_it = base_map.find( self.root );
				assertf(child_it != base_map.end() && root_it != base_map.end(), 
					"nodes must exist in base");
				Node& child = child_it->second;
				Node& root = root_it->second;

				addEdge( child, root, std::move(self.root), Height(self.new_height) );
			} break;
			case REMFROM: {
				AddTo& self = mut_this->as<AddTo>();

				base->init<AddTo>( mut_this, self.root, self.child, self.new_height );
				base->mode = ADDTO;

				auto child_it = base_map.find( self.child );
				auto root_it = base_map.find( self.root );
				assertf(child_it != base_map.end() && root_it != base_map.end(), 
					"nodes must exist in base");
				Node& child = child_it->second;
				Node& root = root_it->second;

				addEdge( child, root, std::move(self.child), Height(-1 * self.new_height) );
			} break;
			default: assertf(false, "invalid mode");
		}

		// set base map into self
		mut_this->reset();
		mut_this->init<Base>( std::move(base_map) );
		mode = BASE;
	}

private:
	/// Gets the base after rerooting at the current node 
	const Base& rerooted() const {
		reroot();
		assertf(mode == BASE, "reroot results in base");
		return as<Base>();
	}

public:
	/// true if the set of sets is empty
	bool empty() const { return rerooted().empty(); }

	/// Get number of entries in the map
	size_type size() const { return rerooted().size(); }

	/// Get begin iterator for map; may be invalidated by calls to non-iteration functions 
	/// or functions on other maps in the same chain
	iterator begin() const { return rerooted().begin(); }

	/// Get end iterator for map; may be invalidated by calls to non-iteration functions 
	/// or functions on other maps in the same chain
	iterator end() const { return rerooted().end(); }

	/// Check if value is present
	size_type count(Elm i) const { return rerooted().count( i ); }

	/// Finds root for element i, undefined behaviour if i is not present
	Elm find(Elm i) const {
		const Base& self = rerooted();
		
		auto it = self.find( i );
		while (true) {
			assertf(it != self.end(), "find target not present");

			if ( it->first == it->second.parent ) return it->first;

			it = self.find( it->second.parent );
		}
	}

	/// Finds root for element i, or default if i is not present
	template<typename E>
	Elm find_or_default(Elm i, E&& d) const {
		const Base& self = rerooted();

		auto it = self.find( i );
		if ( it == self.end() ) return d;

		while ( it->first != it->second.parent ) {
			it = self.find( it->second.parent );

			assertf(it != self.end(), "find target not present");
		}
		return it->first;
	}

	/// Adds fresh class including only one item; returns updated map (or self if no change)
	template<typename E>
	Self* add(E&& i) {
		reroot();

		// add new element to node
		Base base_map = take_as<Base>();
		bool added = base_map.emplace( i, Node{ i } ).second;

		// fail early on node already present
		if ( ! added ) {
			as<Base>() = std::move(base_map);
			return this;
		}

		// make new return node and reset self as REM node
		Self* ret = new Self{ BASE, std::move(base_map) };
		reset_as_base();
		init<Add>( ret, i );
		mode = REM;

		return ret;
	}

	/// Merges two classes given by their roots; returns updated map.
	/// If two classes have same height, `i` is new root.
	Self* merge(Elm i, Elm j) {
		reroot();

		// transfer map to new node
		Self* ret = new Self{ BASE, take_as<Base>() };
		reset_as_base();

		// find set nodes
		Base& base_map = ret->as<Base>();
		auto it = base_map.find( i );
		auto jt = base_map.find( j );
		assertf(it != base_map.end() && jt != base_map.end(), "nodes must exist in base");
		Node& in = it->second;
		Node& jn = jt->second;

		// update returned map and set self to appropriate REMFROM node
		if ( in.height < jn.height ) {
			addEdge( in, jn, j, 0 );
			init<AddTo>( ret, j, i, false );
		} else if ( jn.height < in.height ) {
			addEdge( jn, in, i, 0 );
			init<AddTo>( ret, i, j, false );
		} else /* if ( jn.height == in.height ) */ {
			addEdge( jn, in, i, 1 );
			init<AddTo>( ret, i, j, true );
		}
		mode = REMFROM;

		return ret;
	}

private:
	/// Removes the children of the tree rooted at `root` as `root_node`, returning a new 
	/// uninitialized node and editing `next_edit` to point toward this new node.
	static Self* remove_children(Elm root, Node& root_node, Base& base_map, Self* next_edit) {
		// Invariant: root.next is *always* a pointer to a leaf of the most-recently added
		// subtree.
		//
		// Proof: By induction on height:
		// * height == 0: root.next == root, trivially true
		// * added.height < root.height: true by ind. hyp.
		// * added.height == root.height: no node may have a child the same height, ergo 
		//   property holds for added, therefore root.
		//
		// Corollary: next of most-recently-added subtree root is previously added subtree

		// remove all subtrees
		while ( root != root_node.next ) {
			// find added child
			auto it = base_map.find( root_node.next );
			while ( it->second.parent != root ) it = base_map.find( it->second.parent );
			Elm added = it->first;
			Node& added_node = it->second;

			// unlink subtree and set up previous map as ADDTO new map
			std::swap( root_node.next, added_node.next );
			Self* new_edit = new Self{ BASE };
			next_edit->init<AddTo>( new_edit, root, added, false );  // assume unchanged root height
			next_edit->mode = ADDTO;

			// remove subtree children from map
			next_edit = remove_children( added, added_node, base_map, new_edit );

			// remove subtree root from map and set up previous map as ADD to new map
			base_map.erase( added );
			new_edit = new Self{ BASE };
			next_edit->init<Add>( new_edit, added );
			next_edit->mode = ADD;

			next_edit = new_edit;
		}

		return next_edit;
	}

public:
	/// Removes all elements of a class given by its root; returns updated map.
	Self* remove_class(Elm root) {
		reroot();

		// remove map from self
		Base base_map = take_as<Base>();
		reset_as_base();

		// find root node and remove children
		auto it = base_map.find( root );
		assertf(it != base_map.end(), "root node must exist in map");
		Self* next_edit = remove_children( root, it->second, base_map, this);

		// root is alone in class, remove from map
		base_map.erase( root );
		Self* ret = new Self{ BASE, std::move(base_map) };
		// reinitialize previous node as ADD to new map
		next_edit->init<Add>( ret, root );
		next_edit->mode = ADD;

		return ret;
	}

	/// Applies `f` to all members of a class containing `i` (undefined behaviour if not present).
	/// `f` should take members by const reference or copy
	template<typename F>
	void for_class(Elm i, F&& f) const {
		const Base& self = rerooted();

		// exit early if class not present
		auto it = self.find( i );
		if ( it == self.end() ) return;

		// apply f to each member of class
		f( i );
		for ( Elm crnt = it->second.next; crnt != i; crnt = it->second.next ) {
			f( crnt );
			it = self.find( crnt );
			assertf( it != self.end(), "current node must exist in base" );
		}
	}

	/// Get version node type
	Mode get_mode() const { return mode; }

	/// Get next version up the revision tree (self if base node)
	const Self* get_base() const {
		switch ( mode ) {
			case BASE:                return this;
			case ADD: case REM:       return as<Add>().base;
			case ADDTO: case REMFROM: return as<AddTo>().base;
			default: assertf(false, "invalid mode");
		}
	}

	/// Get root of new class formed/removed/split (undefined if called on base)
	Elm get_root() const {
		switch ( mode ) {
			case ADD: case REM:       return as<Add>().root;
			case ADDTO: case REMFROM: return as<AddTo>().root;
			default: assertf(false, "invalid mode for get_root()");
		}
	}

	/// Get child root of new class formed/split (undefined if called on base or add/remove node)
	Elm get_child() const {
		switch ( mode ) {
			case ADDTO: case REMFROM: return as<AddTo>().child;
			default: assertf(false, "invalid mode for get_child()");
		}
	}

	/// Gets acted-upon key for new class (root unless for add/remove node, child for add to/remove 
	/// from node, undefined otherwise)
	Elm get_key() const {
		switch ( mode ) {
			case ADD: case REM:       return as<Add>().root;
			case ADDTO: case REMFROM: return as<AddTo>().child;
			default: assertf(false, "invalid mode for get_key()");
		}
	}
};

// Local Variables: //
// tab-width: 4 //
// mode: c++ //
// compile-command: "make install" //
// End: //