path: root/src/org/clojure/runtime/PersistentTree.java

/**
 *   Copyright (c) Rich Hickey. All rights reserved.
 *   The use and distribution terms for this software are covered by the
 *   Common Public License 1.0 (http://opensource.org/licenses/cpl.php)
 *   which can be found in the file CPL.TXT at the root of this distribution.
 *   By using this software in any fashion, you are agreeing to be bound by
 * 	 the terms of this license.
 *   You must not remove this notice, or any other, from this software.
 **/

/* rich May 20, 2006 */

package org.clojure.runtime;

import java.util.*;

/**
 * Persistent Red Black Tree
 * Note that instances of this class are constant values
 * i.e. add/remove etc return new values
 * <p/>
 * See Okasaki, Kahrs, Larsen et al
 */

public class PersistentTree implements IPersistentMap, ISequential {

public final Comparator comp;
public final Node tree;
public final int _count;

public PersistentTree(){
	this(null);
}

public PersistentTree(Comparator comp){
	this.comp = comp;
	tree = null;
	_count = 0;
}

public boolean contains(Object key){
	return find(key) != null;
}

public PersistentTree add(Object key){
	return put(key, null);
}

public PersistentTree put(Object key, Object val){
	Box found = new Box(null);
	Node t = add(tree, key, val, found);
	if(t == null)   //null == already contains key
		{
		Node foundNode = (Node) found.val;
		if(foundNode.val() == val)  //note only get same collection on identity of val, not equals()
			return this;
		return new PersistentTree(comp, replace(tree, key, val), _count);
		}
	return new PersistentTree(comp, t.blacken(), _count + 1);
}


public PersistentTree remove(Object key){
	Box found = new Box(null);
	Node t = remove(tree, key, found);
	if(t == null)
		{
		if(found.val == null)//null == doesn't contain key
			return this;
		//empty
		return new PersistentTree(comp);
		}
	return new PersistentTree(comp, t.blacken(), _count - 1);
}

public ISeq seq() {
    if(_count > 0)
        return Seq.create(tree, true);
    return null;
}

public ISeq rseq() {
    if(_count > 0)
        return Seq.create(tree, false);
    return null;
}


public NodeIterator iterator(){
	return new NodeIterator(tree, true);
}

public NodeIterator reverseIterator(){
	return new NodeIterator(tree, false);
}

public Iterator keys(){
	return keys(iterator());
}

public Iterator vals(){
	return vals(iterator());
}

public Iterator keys(NodeIterator it){
	return new KeyIterator(it);
}

public Iterator vals(NodeIterator it){
	return new ValIterator(it);
}

public Object minKey(){
	Node t = min();
	return t!=null?t.key:null;
}

public Node min(){
	Node t = tree;
	if(t != null)
		{
		while(t.left() != null)
			t = t.left();
		}
	return t;
}

public Object maxKey(){
	Node t = max();
	return t!=null?t.key:null;
}

public Node max(){
	Node t = tree;
	if(t != null)
		{
		while(t.right() != null)
			t = t.right();
		}
	return t;
}

public int depth(){
	return depth(tree);
}

int depth(Node t){
	if(t == null)
		return 0;
	return 1 + Math.max(depth(t.left()), depth(t.right()));
}

public Object get(Object key){
	Node n = find(key);
	return (n != null) ? n.val() : null;
}

public int capacity() {
    return _count;
}

public int count() {
    return _count;
}

public Node find(Object key){
    Node t = tree;
    while(t != null)
        {
        int c = compare(key, t.key);
        if(c == 0)
            return t;
        else if(c < 0)
            t = t.left();
        else
            t = t.right();
        }
    return t;
}

int compare(Object k1, Object k2){
	if(comp != null)
		return comp.compare(k1, k2);
	return ((Comparable) k1).compareTo(k2);
}

Node add(Node t, Object key, Object val, Box found){
	if(t == null)
		{
		if(val == null)
			return new Red(key);
		return new RedVal(key, val);
		}
	int