Added a proper implementation of TreeSet, based on a Persistent set implementation.

classpath/java/util/Cell.java

@@ -0,0 +1,24 @@
+package java.util;
+
+public class Cell <T> {
+  public T value;
+  public Cell<T> next;
+  
+  public Cell(T value, Cell<T> next) {
+    this.value = value;
+    this.next = next;
+  }
+
+  public String toString() {
+    StringBuilder sb = new StringBuilder();
+    sb.append("(");
+    for (Cell c = this; c != null; c = c.next) {
+      sb.append(value);
+      if (c.next != null) {
+        sb.append(" ");
+      }
+    }
+    sb.append(")");
+    return sb.toString();
+  }
+}

classpath/java/util/PersistentSet.java

@@ -0,0 +1,590 @@
+package java.util;
+
+import java.util.Comparator;
+import java.lang.Iterable;
+
+public class PersistentSet <T> implements Iterable <T> {
+  private static final Node NullNode = new Node(null);
+
+  static {
+    NullNode.left = NullNode;
+    NullNode.right = NullNode;
+  }
+
+  private final Node<T> root;
+  private final Comparator<T> comparator;
+  private final int size;
+
+  public PersistentSet() {
+    this(NullNode, new Comparator<T>() {
+        public int compare(T a, T b) {
+          return ((Comparable<T>) a).compareTo(b);
+        }
+      }, 0);
+  }
+
+  public PersistentSet(Comparator<T> comparator) {
+    this(NullNode, comparator, 0);
+  }
+
+  private PersistentSet(Node<T> root, Comparator<T> comparator, int size) {
+    this.root = root;
+    this.comparator = comparator;
+    this.size = size;
+  }
+
+  public Comparator<T> comparator() {
+    return comparator;
+  }
+
+  public PersistentSet<T> add(T value) {
+    return add(value, false);
+  }
+
+  public int size() {
+    return size;
+  }
+
+  public PersistentSet<T> add(T value, boolean replaceExisting) {
+    Path<T> p = find(value);
+    if (! p.fresh) {
+      if (replaceExisting) {
+        return p.replaceWith(value);
+      } else {
+        return this;
+      }
+    }
+
+    return add(p);
+  }
+
+  private PersistentSet<T> add(Path<T> p) {
+    if (! p.fresh) throw new IllegalArgumentException();
+
+    Node<T> new_ = p.node;
+    Node<T> newRoot = p.root.root;
+    Cell<Node<T>> ancestors = p.ancestors;
+
+    // rebalance
+    new_.red = true;
+    while (ancestors != null && ancestors.value.red) {
+      if (ancestors.value == ancestors.next.value.left) {
+        if (ancestors.next.value.right.red) {
+          ancestors.value.red = false;
+          ancestors.next.value.right = new Node(ancestors.next.value.right);
+          ancestors.next.value.right.red = false;
+          ancestors.next.value.red = true;
+          new_ = ancestors.next.value;
+          ancestors = ancestors.next.next;
+        } else {
+          if (new_ == ancestors.value.right) {
+            new_ = ancestors.value;
+            ancestors = ancestors.next;
+
+            Node<T> n = leftRotate(new_);
+            if (ancestors.value.right == new_) {
+              ancestors.value.right = n;
+            } else {
+              ancestors.value.left = n;
+            }
+            ancestors = new Cell(n, ancestors);
+          }
+          ancestors.value.red = false;
+          ancestors.next.value.red = true;
+
+          Node<T> n = rightRotate(ancestors.next.value);
+          if (ancestors.next.next == null) {
+            newRoot = n;
+          } else if (ancestors.next.next.value.right == ancestors.next.value) {
+            ancestors.next.next.value.right = n;
+          } else {
+            ancestors.next.next.value.left = n;
+          }
+          // done
+        }
+      } else {
+        if (ancestors.next.value.left.red) {
+          ancestors.value.red = false;
+          ancestors.next.value.left = new Node(ancestors.next.value.left);
+          ancestors.next.value.left.red = false;
+          ancestors.next.value.red = true;
+          new_ = ancestors.next.value;
+          ancestors = ancestors.next.next;
+        } else {
+          if (new_ == ancestors.value.left) {
+            new_ = ancestors.value;
+            ancestors = ancestors.next;
+
+            Node<T> n = rightRotate(new_);
+            if (ancestors.value.right == new_) {
+              ancestors.value.right = n;
+            } else {
+              ancestors.value.left = n;
+            }
+            ancestors = new Cell(n, ancestors);
+          }
+          ancestors.value.red = false;
+          ancestors.next.value.red = true;
+
+          Node<T> n = leftRotate(ancestors.next.value);
+          if (ancestors.next.next == null) {
+            newRoot = n;
+          } else if (ancestors.next.next.value.right == ancestors.next.value) {
+            ancestors.next.next.value.right = n;
+          } else {
+            ancestors.next.next.value.left = n;
+          }
+          // done
+        }
+      }
+    }
+
+    newRoot.red = false;
+
+    return new PersistentSet(newRoot, comparator, size + 1);
+  }
+
+  private static <T> Node<T> leftRotate(Node<T> n) {
+    Node<T> child = new Node(n.right);
+    n.right = child.left;
+    child.left = n;
+    return child;
+  }
+
+  private static <T> Node<T> rightRotate(Node<T> n) {
+    Node<T> child = new Node(n.left);
+    n.left = child.right;
+    child.right = n;
+    return child;
+  }
+
+  public PersistentSet<T> remove(T value) {
+    Path<T> p = find(value);
+    if (! p.fresh) {
+      return remove(p);
+    }
+
+    return this;
+  }
+
+  private PersistentSet<T> remove(Path<T> p) {
+    Node<T> new_ = p.node;
+    Node<T> newRoot = p.root.root;
+    Cell<Node<T>> ancestors = p.ancestors;
+
+    Node<T> dead;
+    if (new_.left == NullNode || new_.right == NullNode) {
+      dead = new_;
+    } else {
+      Cell<Node<T>> path = successor(new_, ancestors);
+      dead = path.value;
+      ancestors = path.next;
+    }
+    
+    Node<T> child;
+    if (dead.left != NullNode) {
+      child = dead.left;
+    } else {
+      child = dead.right;
+    }
+
+    if (ancestors == null) {
+      child.red = false;
+      return new PersistentSet(child, comparator, 1);
+    } else if (dead == ancestors.value.left) {
+      ancestors.value.left = child;
+    } else {
+      ancestors.value.right = child;
+    }
+
+    if (dead != new_) {
+      new_.value = dead.value;
+    }
+
+    if (! dead.red) {
+      // rebalance
+      while (ancestors != null && ! child.red) {
+        if (child == ancestors.value.left) {
+          Node<T> sibling = ancestors.value.right
+            = new Node(ancestors.value.right);
+          if (sibling.red) {
+            sibling.red = false;
+            ancestors.value.red = true;
+            
+            Node<T> n = leftRotate(ancestors.value);
+            if (ancestors.next == null) {
+              newRoot = n;
+            } else if (ancestors.next.value.right == ancestors.value) {
+              ancestors.next.value.right = n;
+            } else {
+              ancestors.next.value.left = n;
+            }
+            ancestors.next = new Cell(n, ancestors.next);
+
+            sibling = ancestors.value.right;
+          }
+
+          if (! (sibling.left.red || sibling.right.red)) {
+            sibling.red = true;
+            child = ancestors.value;
+            ancestors = ancestors.next;
+          } else {
+            if (! sibling.right.red) {
+              sibling.left = new Node(sibling.left);
+              sibling.left.red = false;
+
+              sibling.red = true;
+              sibling = ancestors.value.right = rightRotate(sibling);
+            }
+
+            sibling.red = ancestors.value.red;
+            ancestors.value.red = false;
+
+            sibling.right = new Node(sibling.right);
+            sibling.right.red = false;
+            
+            Node<T> n = leftRotate(ancestors.value);
+            if (ancestors.next == null) {
+              newRoot = n;
+            } else if (ancestors.next.value.right == ancestors.value) {
+              ancestors.next.value.right = n;
+            } else {
+              ancestors.next.value.left = n;
+            }
+
+            child = newRoot;
+            ancestors = null;
+          }
+        } else {
+          Node<T> sibling = ancestors.value.left
+            = new Node(ancestors.value.left);
+          if (sibling.red) {
+            sibling.red = false;
+            ancestors.value.red = true;
+            
+            Node<T> n = rightRotate(ancestors.value);
+            if (ancestors.next == null) {
+              newRoot = n;
+            } else if (ancestors.next.value.left == ancestors.value) {
+              ancestors.next.value.left = n;
+            } else {
+              ancestors.next.value.right = n;
+            }
+            ancestors.next = new Cell(n, ancestors.next);
+
+            sibling = ancestors.value.left;
+          }
+
+          if (! (sibling.right.red || sibling.left.red)) {
+            sibling.red = true;
+            child = ancestors.value;
+            ancestors = ancestors.next;
+          } else {
+            if (! sibling.left.red) {
+              sibling.right = new Node(sibling.right);
+              sibling.right.red = false;
+
+              sibling.red = true;
+              sibling = ancestors.value.left = leftRotate(sibling);
+            }
+
+            sibling.red = ancestors.value.red;
+            ancestors.value.red = false;
+
+            sibling.left = new Node(sibling.left);
+            sibling.left.red = false;
+            
+            Node<T> n = rightRotate(ancestors.value);
+            if (ancestors.next == null) {
+              newRoot = n;
+            } else if (ancestors.next.value.left == ancestors.value) {
+              ancestors.next.value.left = n;
+            } else {
+              ancestors.next.value.right = n;
+            }
+
+            child = newRoot;
+            ancestors = null;
+          }
+        }
+      }
+
+      child.red = false;
+    }
+
+    return new PersistentSet(newRoot, comparator, size - 1);
+  }
+
+  private static <T> Cell<Node<T>> minimum(Node<T> n,
+                                           Cell<Node<T>> ancestors)
+  {
+    while (n.left != NullNode) {
+      n.left = new Node(n.left);
+      ancestors = new Cell(n, ancestors);
+      n = n.left;
+    }
+
+    return new Cell(n, ancestors);
+  }
+
+  private static <T> Cell<Node<T>> successor(Node<T> n,
+                                             Cell<Node<T>> ancestors)
+  {
+    if (n.right != NullNode) {
+      n.right = new Node(n.right);
+      return minimum(n.right, new Cell(n, ancestors));
+    }
+
+    while (ancestors != null && n == ancestors.value.right) {
+      n = ancestors.value;
+      ancestors = ancestors.next;
+    }
+
+    return ancestors;
+  }
+
+  public Path<T> find(T value) {
+    Node<T> newRoot = new Node(root);
+    Cell<Node<T>> ancestors = null;
+
+    Node<T> old = root;
+    Node<T> new_ = newRoot;
+    while (old != NullNode) {
+      ancestors = new Cell(new_, ancestors);
+
+      int difference = comparator.compare(value, old.value);
+      if (difference < 0) {
+        old = old.left;
+        new_ = new_.left = new Node(old);
+      } else if (difference > 0) {
+        old = old.right;
+        new_ = new_.right = new Node(old);
+      } else {
+        return new Path(false, new_,
+                        new PersistentSet(newRoot, comparator, size),
+                        ancestors.next);
+      }
+    }
+
+    new_.value = value;
+    return new Path(true, new_,
+                    new PersistentSet(newRoot, comparator, size),
+                    ancestors);
+  }
+
+  public Path<T> first() {
+    if (root == NullNode) return null;
+
+    Node<T> newRoot = new Node(root);
+    Cell<Node<T>> ancestors = null;
+
+    Node<T> old = root;
+    Node<T> new_ = newRoot;
+    while (old.left != NullNode) {
+      ancestors = new Cell(new_, ancestors);
+
+      old = old.left;
+      new_ = new_.left = new Node(old);
+    }
+
+    return new Path(true, new_,
+                    new PersistentSet(newRoot, comparator, size),
+                    ancestors);
+  }
+
+  public Path<T> last() {
+    if (root == NullNode) return null;
+
+    Node<T> newRoot = new Node(root);
+    Cell<Node<T>> ancestors = null;
+
+    Node<T> old = root;
+    Node<T> new_ = newRoot;
+    while (old.right != NullNode) {
+      ancestors = new Cell(new_, ancestors);
+
+      old = old.right;
+      new_ = new_.right = new Node(old);
+    }
+
+    return new Path(true, new_,
+                    new PersistentSet(newRoot, comparator, size),
+                    ancestors);
+  }
+
+  public java.util.Iterator<T> iterator() {
+    return new Iterator(first());
+  }
+
+  private Path<T> successor(Path<T> p) {
+    Cell<Node<T>> s = successor(p.node, p.ancestors);
+    if (s == null) {
+      return null;
+    } else {
+      return new Path(false, s.value, p.root, s.next);
+    }
+  }
+
+//   public void dump(java.io.PrintStream out) {
+//     dump(root, out, 0);
+//   }
+
+//   private static void indent(java.io.PrintStream out, int level) {
+//     for (int i = 0; i < level; ++i) out.print("  ");
+//   }
+
+//   private static <T> void dump(Node<T> n, java.io.PrintStream out, int level) {
+//     indent(out, level);
+//     out.print(n == NullNode ? null : n.value);
+//     out.println(n == NullNode ? "" : n.red ? " (red)" : " (black)");
+//     if (n.left != NullNode || n.right != NullNode) {
+//       dump(n.left, out, level + 1);
+//       dump(n.right, out, level + 1);
+//     }
+//   }
+
+//   private static int[] randomSet(java.util.Random r, int size) {
+//     int[] data = new int[size];
+//     for (int i = size - 1; i >= 0; --i) {
+//       data[i] = i + 1;
+//     }
+
+//     for (int i = size - 1; i >= 0; --i) {
+//       int n = r.nextInt(size);
+//       int tmp = data[i];
+//       data[i] = data[n];
+//       data[n] = tmp;
+//     }
+
+//     return data;
+//   }
+
+//   public static void main(String[] args) {
+//     java.util.Random r = new java.util.Random(Integer.parseInt(args[0]));
+//     int size = 18;
+//     PersistentSet<Integer>[] sets = new PersistentSet[size];
+//     PersistentSet<Integer> s = new PersistentSet();
+
+//     int[] data = randomSet(r, size);
+
+//     for (int i = 0; i < size; ++i) {
+//       System.out.println("-- add " + data[i] + " -- ");
+//       sets[i] = s = s.add(data[i]);
+//       dump(s.root, System.out, 0);
+//     }
+
+//     System.out.println("\npersistence:\n");
+//     for (int i = 0; i < size; ++i) {
+//       dump(sets[i].root, System.out, 0);
+//       System.out.println("--");
+//     }
+
+//     data = randomSet(r, size);
+
+//     System.out.println("\nremoval:\n");
+//     for (int i = 0; i < size; ++i) {
+//       System.out.println("-- remove " + data[i] + " -- ");
+//       sets[i] = s = s.remove(data[i]);
+//       dump(s.root, System.out, 0);
+//     }
+
+//     System.out.println("\npersistence:\n");
+//     for (int i = 0; i < size; ++i) {
+//       dump(sets[i].root, System.out, 0);
+//       System.out.println("--");
+//     }
+//   }
+
+  private static class Node <T> {
+    public T value;
+    public Node left;
+    public Node right;
+    public boolean red;
+    
+    public Node(Node<T> basis) {
+      if (basis != null) {
+        value = basis.value;
+        left = basis.left;
+        right = basis.right;
+        red = basis.red;
+      }
+    }
+  }
+
+  public static class Path <T> {
+    private final boolean fresh;
+    private final Node<T> node;
+    private final PersistentSet<T> root;
+    private final Cell<Node<T>> ancestors;
+    
+    public Path(boolean fresh, Node<T> node, PersistentSet<T> root,
+                Cell<Node<T>> ancestors)
+    {
+      this.fresh = fresh;
+      this.node = node;
+      this.root = root;
+      this.ancestors = ancestors;
+    }
+
+    public T value() {
+      return node.value;
+    }
+
+    public boolean fresh() {
+      return fresh;
+    }
+
+    public PersistentSet<T> root() {
+      return root;
+    }
+
+    public Path<T> successor() {
+      return root.successor(this);
+    }
+
+    public PersistentSet<T> remove() {
+      return root.remove(this);
+    }
+
+    public PersistentSet<T> add() {
+      if (! fresh) throw new IllegalStateException();
+
+      return root.add(this);
+    }
+
+    public PersistentSet<T> replaceWith(T value) {
+      if (fresh) throw new IllegalStateException();
+      if (root.comparator.compare(node.value, value) != 0)
+        throw new IllegalArgumentException();
+
+      node.value = value;
+      return root;
+    }
+  }
+  
+  public class Iterator <T> implements java.util.Iterator <T> {
+    private PersistentSet.Path<T> path;
+
+    private Iterator(PersistentSet.Path<T> path) {
+      this.path = path;
+    }
+
+    private Iterator(Iterator<T> start) {
+      path = start.path;
+    }
+
+    public boolean hasNext() {
+      return path != null;
+    }
+
+    public T next() {
+      PersistentSet.Path<T> p = path;
+      path = path.successor();
+      return p.value();
+    }
+
+    public void remove() {
+      throw new UnsupportedOperationException();
+    }
+  }
+}

classpath/java/util/TreeSet.java

@@ -1,81 +1,131 @@
 package java.util;
 
 public class TreeSet<T> implements Iterable<T> {
-  private final Comparator<T> comparator;
+  private PersistentSet<Cell<T>> set;
   private int size;
-  private Cell<T> root;
 
-  public TreeSet(Comparator<T> comparator) {
-    this.comparator = comparator;
-    size=0;
-    root=null;
+  public TreeSet(final Comparator<T> comparator) {
+    set = new PersistentSet(new Comparator<Cell<T>>() {
+      public int compare(Cell<T> a, Cell<T> b) {
+        return comparator.compare(a.value, b.value);
+      }
+    });
+    size = 0;
   }
 
   public Iterator<T> iterator() {
-    return walk().iterator();
+    return new MyIterator<T>(set.first());
   }
 
-  private ArrayList<T> walk() {
-    return walk(root, new ArrayList<T>(size));
-  }
-
-  private ArrayList<T> walk(Cell<T> cell, ArrayList<T> list) {
-    if (cell != null) {
-      walk(cell.left, list);
-      list.add(cell.value);
-      walk(cell.right, list);
-    }
-    return list;
-  }
-
-  public boolean add(T o) {
-    ++size;
-    if (root == null) {
-      root = new Cell<T>(o);
+  public boolean add(T value) {
+    PersistentSet.Path<Cell<T>> p = set.find(new Cell(value, null));
+    if (p.fresh()) {
+      set = p.add();
+      ++size;
       return true;
-    } else {
-      Cell<T> newElt = new Cell<T>(o);
-      Cell<T> cur = root;
-      do {
-        int result = comparator.compare(o, cur.value);
-        if (result == 0) return false;
-        if (result < 0) {
-          if (cur.left == null) {
-            newElt.parent = cur;
-            cur.left = newElt;
-            return false;
-          } else {
-            cur = cur.left;
-          }
-        } else {
-          if (cur.right == null) {
-            newElt.parent = cur;
-            cur.right = newElt;
-            return false;
-          } else {
-            cur = cur.right;
-          }
-        }
-      } while (cur != null);
-      throw new RuntimeException("Fell off end of TreeSet");
     }
+    return false;
   }
 
-  public boolean remove(T o) {
-    throw new UnsupportedOperationException();
+  // Used by hashMaps for replacement
+  public void addAndReplace(T value) {
+    PersistentSet.Path<Cell<T>> p = set.find(new Cell(value, null));
+    if (p.fresh()) {
+      set = p.add();
+      ++size;
+    } else {
+      set = p.replaceWith(new Cell(value, null));
+    }
+  }
+    
+  public boolean remove(T value) {
+    PersistentSet.Path<Cell<T>> p = set.find(new Cell(value, null));
+    if (p.fresh()) {
+      return false;
+    } else {
+      --size;
+
+      if (p.value().next != null) {
+        set = p.replaceWith(p.value().next);
+      } else {
+        set = p.remove();
+      }
+
+      return true;
+    }
   }
 
   public int size() {
     return size;
   }
 
-  private static class Cell<T> {
-    public final T value;
-    public Cell<T> parent;
-    public Cell<T> left;
-    public Cell<T> right;
-    public Cell(T val) {
-      value = val;
+  private class MyIterator<T> implements java.util.Iterator<T> {
+    private PersistentSet.Path<Cell<T>> path;
+    private PersistentSet.Path<Cell<T>> nextPath;
+    private Cell<T> cell;
+    private Cell<T> prevCell;
+    private Cell<T> prevPrevCell;
+    private boolean canRemove = false;
+
+    private MyIterator(PersistentSet.Path<Cell<T>> path) {
+      this.path = path;
+      if (path != null) {
+        cell = path.value();
+        nextPath = path.successor();
+      }
     }
-  }
+
+    private MyIterator(MyIterator<T> start) {
+      path = start.path;
+      nextPath = start.nextPath;
+      cell = start.cell;
+      prevCell = start.prevCell;
+      prevPrevCell = start.prevPrevCell;
+      canRemove = start.canRemove;
+    }
+
+    public boolean hasNext() {
+      return cell != null || nextPath != null;
+    }
+
+    public T next() {
+      if (cell == null) {
+        path = nextPath;
+        nextPath = path.successor();
+        cell = path.value();
+      }
+      prevPrevCell = prevCell;
+      prevCell = cell;
+      cell = cell.next;
+      canRemove = true;
+      return prevCell.value;
+    }
+
+    public void remove() {
+      if (! canRemove) throw new IllegalStateException();
+
+      --size;
+
+      if (prevPrevCell != null && prevPrevCell.next == prevCell) {
+        // cell to remove is not the first in the list.
+        prevPrevCell.next = prevCell.next;
+        prevCell = prevPrevCell;
+      } else if (prevCell.next == cell && cell != null) {
+        // cell to remove is the first in the list, but not the last.
+        set = (PersistentSet) path.replaceWith(cell);
+        prevCell = null;
+      } else {
+        // cell is alone in the list.
+        set = (PersistentSet) path.remove();
+        path = path.successor();
+        if (path != null) {
+          prevCell = null;
+          cell = path.value();
+          path = (PersistentSet.Path) set.find((Cell) cell);
+        }
+      }
+
+      canRemove = false;
+    }
+  }    
 }

test/Tree.java

@@ -6,11 +6,11 @@ public class Tree {
     if (! v) throw new RuntimeException();
   }
 
-  private static String printList(TreeSet<Integer> list) {
+  private static String printList(TreeSet<?> list) {
     StringBuilder sb = new StringBuilder();
 
-    for (Integer i : list) {
-      sb.append(i);
+    for (Object o : list) {
+      sb.append(o);
       sb.append(", ");
     }
     sb.setLength(sb.length()-2);
@@ -29,12 +29,29 @@ public class Tree {
   }
 
   public static void main(String args[]) {
-    TreeSet<Integer> l = new TreeSet<Integer>(new MyCompare());
-    l.add(5); l.add(2); l.add(1); l.add(8); l.add(3);
-    isEqual(printList(l), "1, 2, 3, 5, 8");
-    l.add(4);
-    isEqual(printList(l), "1, 2, 3, 4, 5, 8");
-    l.remove(3);
-    isEqual(printList(l), "1, 2, 4, 5, 8");
+    TreeSet<Integer> t1 = new TreeSet<Integer>(new MyCompare());
+    t1.add(5); t1.add(2); t1.add(1); t1.add(8); t1.add(3);
+    isEqual(printList(t1), "1, 2, 3, 5, 8");
+    t1.add(4);
+    isEqual(printList(t1), "1, 2, 3, 4, 5, 8");
+    t1.remove(3);
+    isEqual(printList(t1), "1, 2, 4, 5, 8");
+    TreeSet<String> t2 = new TreeSet<String>(new Comparator<String>() {
+      public int compare(String s1, String s2) {
+        return s1.compareTo(s2);
+      }
+    });
+    t2.add("one"); t2.add("two"); t2.add("three"); t2.add("four"); t2.add("five");
+    isEqual(printList(t2), "five, four, one, three, two");
+    for (int i=0; i < 1000; i++) {
+      t2.add(Integer.toString(i));
+    }
+    expect(t2.size() == 1005);
+    for (int i=0; i < 999; i++) {
+      t2.remove(Integer.toString(i));
+    }
+    expect(t2.size() == 6);
+    t2.add("kappa");
+    isEqual(printList(t2), "999, five, four, kappa, one, three, two");
   }
 }