/** Persistent binary search tree.
 * @author Andreas van Cranenburgh
 *
 * Uses node-copying (as opposed to path-copying, so has O(1) amortized space
 * complexity for each update operation instead of O(log n)).
 *
 * Each node is augmented with a mutation field, containing either an updated
 * left or right pointer. This way, nodes along the path of a node to be
 * updated can be re-used many times.
 *
 * For more details, see: Driscoll et al., 1989, "Making Data Structures
 * Persistent", journal of computer and system sciences, vol. 38, no. 1,
 * february 1989.
 **/

public class NodeCopyingBST
	<K extends Comparable<? super K>, V> implements BST<K,V> {
	/** Constants to signify whether a mutation is for the left or
 	  * right child. */
	final boolean LEFT = java.lang.Boolean.TRUE;
	final boolean RIGHT = java.lang.Boolean.FALSE;
	/** Array to contain root node for each version of the tree */
	private java.util.ArrayList<ANode<K,V>> versions; 
	/** Array to contain mutations, in order for rollback() to be able
 	  * to erase them. */
	private java.util.ArrayList<ANode<K,V>> mutations;
	/** Reference to the current (i.e., writable) root of the tree. */
	private ANode<K,V> root;
	/** Version number of the tree. */
	private int livever = 0;

	/** Default constructor */
	public NodeCopyingBST() {
		root = null;
		versions = new java.util.ArrayList<ANode<K,V>>();
		mutations = new java.util.ArrayList<ANode<K,V>>();
	}
	
	/** find key and return corresponding value
	 * @param key the key to find
	 * @return value associated with key, or null if not found. */
	public V search (K key) {
		if (key == null)
			throw new NullPointerException(
				"tried to search with null key");
		else
			return search(root, key);
	}
	
	/** find key and return corresponding value from a specific version
	 * @param key the key to find
	 * @param version the version number of the tree.
	 * @return value associated with key, or null if not found /
	 * 		non-existing version. */
	public V search (K key, int version) {
		if (key == null)
			throw new NullPointerException(
				"tried to search with null key");
		if (version > livever)
			return null;
		
		/* temporarily set current version to requested version,
		 * and restore it after performing search. */
		int temp = livever;
		livever = version;
		V value = search(root, key);
		livever = temp;
		return value;
	}
	
	/** find key in a specific subtree, return value
	 * pass a root node from an earlier insert or delete operation to 
	 * search in that version of the tree.
	 * @param node the (sub)tree where lookup starts
	 * @param key the key to find
	 * @return the value associated with key, in the specified tree */
	public V search (ANode<K,V> n, K key) {
		/* key not found */
		if (n == null)
			return null;

		int comp = key.compareTo(n.key);
		if (comp < 0)
			if (n.ver != 0 && n.ver <= livever
				&& n.lr == LEFT)
				return search(n.mutation, key);
			else
				return search(n.left, key);
		else if (comp > 0)
			if (n.ver != 0 && n.ver <= livever 
				&& n.lr == RIGHT)
				return search(n.mutation, key);
			else
				return search(n.right, key);
		else
			/* key found */
			return n.value;
	}

	/** associate key with value
	 * @param key the key to insert
	 * @param value the value to associate with key
	 * @return the version spawned by this operation. */
	public int insert (K key, V value) {
		versions.add(root);	
		livever++;
		if (key == null)
			throw new NullPointerException(
				"tried to insert with null key");
		
		root = insert(root, key, value);
		return livever;
	}
	
	/** insert key in a specific subtree */
	private ANode<K,V> insert(ANode<K,V> node, K key, V value) {
		/* place new node here, end recursion */
		if (node == null)
			return new ANode<K,V>(key, value);
		
		/* recurse until we reach the desired node */	
		ANode<K,V> next, result;
		int comp = key.compareTo(node.key);
		if (comp < 0) {
			/* traverse newest version of the tree */
			if (node.ver != 0 && node.ver <= livever 
				&& node.lr == LEFT)
				next = node.mutation;
			else
				next = node.left;
			
			/* recurse */
			result = insert(next, key, value);
			
			/* update reference in case new node was spawned */
			if (result != next) {
				if (node.ver != 0) {
					node = new ANode<K,V>(node);
					if (node.lr == RIGHT)
						node.right = node.mutation;
					node.ver = 0;
					node.left = result;
				} else {
					node.mutation = result;
					node.lr = LEFT;
					mutations.add(node);
					node.ver = livever;
				}
			}

		} else if (comp > 0) {
			/* traverse newest version of the tree */
			if (node.ver != 0 && node.ver <= livever 
				&& node.lr == RIGHT)
				next = node.mutation;
			else
				next = node.right;
			
			/* recurse */
			result = insert(next, key, value);
			
			/* update reference in case new node was spawned */
			if (result != next) {
				if (node.ver != 0) {
					node = new ANode<K,V>(node);
					if (node.lr == LEFT)
						node.left = node.mutation;
					node.ver = 0;
					node.right = result;
				} else {
					node.mutation = result;
					node.lr = RIGHT;
					mutations.add(node);
					node.ver = livever;
				}
			}

		} else {  /* old key found, add new value */ 
			node = new ANode<K,V>(node);
			node.key = key;
			node.value = value;
		}

		return node;
	}

	/** find key and remove corresponding element
	 * @param key the key to delete
	 * @return the version spawned by this operation. */
	public int delete (K key) {
		versions.add(root);
		livever++;
		if (key == null)
			throw new NullPointerException(
				"tried to delete with null key");
		root = delete(root, key);
		return livever;
	}

	/** find key and remove corresponding element from specified (sub)tree.
	 **/
	private ANode<K,V> delete (ANode<K,V> node, K key) {
		/* key not found, end recursion */
		if (node == null)
			return null;
		
		/* recurse until we reach the desired node */	
		ANode<K,V> next, result;
		int comp = key.compareTo(node.key);
		if (comp < 0) {
			/* traverse newest version of the tree */
			if (node.ver != 0 && node.ver <= livever 
				&& node.lr == LEFT)
				next = node.mutation;
			else
				next = node.left;
			
			/* recurse */
			result = delete(next, key);
			
			/* update reference in case a new node was spawned */
			if (result != next) {
				if (node.ver != 0) {
					node = new ANode<K,V>(node);
					if (node.lr == RIGHT)
						node.right = node.mutation;
					node.ver = 0;
					node.left = result;
				} else {
					node.mutation = result;
					node.lr = LEFT;
					mutations.add(node);
					node.ver = livever;
				}
			}
			return node;

		} else if (comp > 0) {
			/* traverse newest version of the tree */
			if (node.ver != 0 && node.ver <= livever 
				&& node.lr == RIGHT)
				next = node.mutation;
			else
				next = node.right;
			
			/* recurse */
			result = delete(next, key);
			
			/* update reference in case a new node was spawned */
			if (result != next) {
				if (node.ver != 0) {
					node = new ANode<K,V>(node);
					if (node.lr == LEFT)
						node.left = node.mutation;
					node.ver = 0;
					node.right = result;
				} else {
					node.mutation = result;
					node.lr = RIGHT;
					mutations.add(node);
					node.ver = livever;
				}
			}
			return node;
		}

		ANode<K,V> left = node.left, right = node.right;
		if (node.ver != 0 && node.ver <= livever) { 
			if (node.lr == LEFT)
				left = node.mutation;
			else
				right = node.mutation;
		}
		
		/* recursion reached the node to be deleted */
		/* simple cases, 0 or 1 child: */
		if (left == null && right == null)
			return null;
		else if (left == null)
			return right;
		else if (right == null)
			return left;
		/* 2 children, swap node with in-order predecessor 
		 * and delete it. to simplify matters the path to the
		 * predecessor is walked twice, first looking it up and
		 * then copying its path minus the node to be deleted. */
		else {
			/* find in-order predecessor */
			ANode<K,V> pred = predNode(left);
			/* delete it from this subtree */
			ANode<K,V> temp = delete(left, pred.key);
			/* return predecessor with children of old node */
			return new ANode<K,V>(temp, pred.key, pred.value, 
				right);
		}
	}
	
	/** return in-order predecessor of specified node */
	private ANode<K,V> predNode(ANode <K,V> node) {
		ANode<K,V> next;
		/* traverse newest version of the tree */
		if (node.ver != 0 && node.ver <= livever 
			&& node.lr == RIGHT)
			next = node.mutation;
		else
			next = node.right;
		
		if (next == null)
			return node;
		else

			return predNode(next);
	}

	/** undo last mutation. 
	 * undoing itself can not be undone, this operation is destructive,
	 * since the versions are stored in a list which can not keep track 
	 * of multiple branches. */
	public void rollback() {
		/* get last version from stack and make it current */
		root = versions.remove(versions.size() - 1);
		
		/* pop mutations of current version and erase them */
		while (mutations.get(mutations.size() - 1).ver == livever)
			mutations.remove(mutations.size() - 1).ver = 0;
		
		livever--;
	}

	/** return an in-order traversal of the tree */
	public String toString() {
		if (root == null)
			return "empty tree";
		else
			return inOrder(root);
	}
	
	/** recursive in-order helper function */
	private String inOrder (ANode<K,V> node) {
		if (node == null)
			return "";

		ANode<K,V> left = node.left, right = node.right;

		/* traverse newest version of tree */
		if (node.ver != 0 && node.ver <= livever) 
			if (node.lr == LEFT)
				left = node.mutation;
			else
				right = node.mutation;

		return inOrder(left) + " " + node.key + "=" + node.value
			+ inOrder(right);
	}
}

/** Augmented tree node with extra mutation pointer, version and left/right
 * field. */
class ANode<K extends Comparable<? super K>, V> extends Node<K,V> {
	protected K key;
	protected V value;
	protected ANode<K,V> left, right, mutation;
	/* version stamp of mutation (0 means no mutation) */
	protected int ver = 0;
	/* whether the mutation is for the left or right child */
	protected boolean lr;

	public ANode (ANode<K,V> left, K key, V value, ANode<K,V> right) {
		this(key, value);
		this.left = left;
		this.right = right;
	}

	public ANode (K key, V value) {
		this.key = key;
		this.value = value;
	}
	
	public ANode (ANode<K,V> node) {
		this(node.left, node.key, node.value, node.right);
		if (node.ver != 0) {
			this.mutation = node.mutation;
			this.lr = node.lr;
			this.ver = node.ver;
		}
	}
}