- Added HashTreeMap and HashTreeSet

- Updated red-black trees to carry their size,
  thus speeding up TreeMap and TreeSet
- Added "updatef" to various tree classes

7 files changed, 234 insertions, 55 deletions

diff --git a/src/library/scala/collection/immutable/HashTreeMap.scala b/src/library/scala/collection/immutable/HashTreeMap.scala
new file mode 100644
index 0000000000..db334310b9
--- /dev/null
+++ b/src/library/scala/collection/immutable/HashTreeMap.scala
@@ -0,0 +1,103 @@
+/*                     __                                               *\
+**     ________ ___   / /  ___     Scala API                            **
+**    / __/ __// _ | / /  / _ |    (c) 2003-2007, LAMP/EPFL             **
+**  __\ \/ /__/ __ |/ /__/ __ |                                         **
+** /____/\___/_/ |_/____/_/ | |                                         **
+**                          |/                                          **
+\*                                                                      */
+
+// $Id$
+
+
+
+package scala.collection.immutable
+
+
+object HashTreeMap {
+  /** The empty map of this type */
+  def empty[A, B] = new HashTreeMap[A, B]
+
+  /** The canonical factory for this type
+   */
+  def apply[A, B](elems: Pair[A, B]*) = empty[A, B] ++ elems
+}
+
+/** An immutable tree-based map that uses hashes to order
+ * the elements.  While it is not as efficient as
+ * a regular TreeMap, it has the advantage of not requiring
+ * an Ordered view being available for the key data.
+ *
+ * The implementation a TreeMap mapping hash codes into
+ * ListMap's, where each ListMap is a bucket of entries
+ * whose keys have the same hash code.
+ */
+class HashTreeMap[Key, +Value](treemap: TreeMap[int, Map[Key, Value]],  val size: int)
+extends Map[Key, Value]
+{
+  /** Create an empty HashTreeMap */
+  def this() = this(new TreeMap, 0)
+
+  /** A central, convenenient updating function.  Given a key,
+    * it finds the bucket for that key, allows the specified
+    * routine to update the bucket, and then creates a new
+    * HashTreeMap using the updated bucket.
+    */
+  private def updateBucket[Value1 >: Value]
+    (key: Key)
+    (upd: Map[Key, Value1] => Map[Key, Value1])
+    : Map[Key, Value1] =
+  {
+    val hash: int = key.hashCode
+    var sizechange = 0
+
+    val newmap =
+      treemap.updatef(
+          hash,
+          maybeBucket => {
+            val oldBucket = maybeBucket match {
+              case None => new ListMap[Key, Value]
+              case Some(oldBucket) => oldBucket
+            }
+            val newBucket = upd(oldBucket)
+            sizechange = newBucket.size - oldBucket.size
+            newBucket
+          })
+
+    new HashTreeMap(newmap, size + sizechange)
+  }
+
+  def update[Value1 >: Value](key: Key, value: Value1) =
+    updateBucket[Value1](key)(b => b + {key, value})
+
+  def get(key: Key) = {
+    treemap.get(key.hashCode) match {
+      case None => None
+      case Some(bucket) => bucket.get(key)
+    }
+  }
+
+  def - (key: Key) = updateBucket(key)(b => b - key)
+
+  def empty[C] = new HashTreeMap[Key, C]
+
+  def elements = new scala.Iterator[Pair[Key,Value]] {
+    val topElements = treemap.elements
+    var bucket: Iterator[Pair[Key, Value]] = null
+
+    def moveToNext = {
+      while(topElements.hasNext &&
+            (bucket == null || !bucket.hasNext))
+        bucket = topElements.next._2.elements
+    }
+
+    def hasNext = {
+      moveToNext
+      (bucket != null) && bucket.hasNext
+    }
+
+    def next = {
+      moveToNext
+      bucket.next
+    }
+  }
+}
\ No newline at end of file
diff --git a/src/library/scala/collection/immutable/HashTreeSet.scala b/src/library/scala/collection/immutable/HashTreeSet.scala
new file mode 100644
index 0000000000..c7ceedb66f
--- /dev/null
+++ b/src/library/scala/collection/immutable/HashTreeSet.scala
@@ -0,0 +1,46 @@
+/*                     __                                               *\
+**     ________ ___   / /  ___     Scala API                            **
+**    / __/ __// _ | / /  / _ |    (c) 2003-2007, LAMP/EPFL             **
+**  __\ \/ /__/ __ |/ /__/ __ |                                         **
+** /____/\___/_/ |_/____/_/ | |                                         **
+**                          |/                                          **
+\*                                                                      */
+
+// $Id$
+
+
+
+package scala.collection.immutable
+
+
+/** An immutable tree-based set that uses hashes to order
+ * the elements.  While it is not as efficient as
+ * a regular TreeSet, it has the advantage of not requiring
+ * an Ordered view being available for the data.
+ *
+ * The implementation is a TreeSet mapping hash codes into
+ * ListSet's, where each ListSet is a bucket of entries
+ * whose elements all have the same hash code.
+ */
+class HashTreeSet[T](unitmap: Map[T,Unit])
+extends Set[T]
+{
+  /** Create an empty HashTreeMap */
+  def this() = this(new HashTreeMap)
+
+  def size = unitmap.size
+
+  def +(x: T) = new HashTreeSet(unitmap + {x,()})
+
+  def -(x: T) = new HashTreeSet(unitmap - x)
+
+  def contains(x: T) = unitmap.contains(x)
+
+  def empty[S]: Set[S] = new HashTreeSet[S]
+
+  def elements = new scala.Iterator[T] {
+    val mapElements = unitmap.elements
+    def hasNext = mapElements.hasNext
+    def next = mapElements.next._1
+  }
+}
\ No newline at end of file
diff --git a/src/library/scala/collection/immutable/Map.scala b/src/library/scala/collection/immutable/Map.scala
index 8c9b3b75b8..31c0de211b 100644
--- a/src/library/scala/collection/immutable/Map.scala
+++ b/src/library/scala/collection/immutable/Map.scala
@@ -35,12 +35,14 @@ package scala.collection.immutable
  */
 object Map {
 
-  /** The empty map of this type; this is implemented as a treemap */
-  def empty[A <% Ordered[A], B] = new TreeMap[A, B]
-
-  /** The canonical factory for this type
+  /** An empty immutable map.  This is implemented with a HashTreeMap.
+   * Note that for key types that have an Ordered view, it is more
+   * efficient to use a TreeMap, instead.
    */
-  def apply[A <% Ordered[A], B](elems: Pair[A, B]*) = empty[A, B] ++ elems
+  def empty[A, B] = new HashTreeMap[A, B]
+
+  /** Create an immutable map from a specified sequence of elements.  */
+  def apply[A, B](elems: Pair[A, B]*) = empty[A, B] ++ elems
 }
 
 trait Map[A, +B] extends collection.Map[A, B] {
@@ -63,6 +65,7 @@ trait Map[A, +B] extends collection.Map[A, B] {
    */
   def update [B1 >: B] (key: A, value: B1): Map[A, B1]
 
+
   /** Add a key/value pair to this map.
    *  @param    kv the key/value pair.
    *  @return   A new map with the new binding added to this map
@@ -143,6 +146,7 @@ trait Map[A, +B] extends collection.Map[A, B] {
     res
   }
 
+
   /** This method removes all the mappings for which the predicate
    *  <code>p</code> returns <code>false</code>.
    *
diff --git a/src/library/scala/collection/immutable/RedBlack.scala b/src/library/scala/collection/immutable/RedBlack.scala
index dd3c7ebc94..aac583bb2c 100755
--- a/src/library/scala/collection/immutable/RedBlack.scala
+++ b/src/library/scala/collection/immutable/RedBlack.scala
@@ -18,11 +18,14 @@ abstract class RedBlack[A] {
     def isBlack: boolean
     def lookup(x: A): Tree[B]
     def update[B1 >: B](k: A, v: B1): Tree[B1] = blacken(upd(k, v))
+    def updatef[B1 >: B](k: A, f: Option[B]=>B1): Tree[B1] = blacken(updf(k, f))
     def delete(k: A): Tree[B] = del(k)
     def elements: Iterator[Pair[A, B]]
     def upd[B1 >: B](k: A, v: B1): Tree[B1]
+    def updf[B1 >: B](k: A, f: Option[B]=>B1): Tree[B1]
     def del(k: A): Tree[B]
     def smallest: NonEmpty[B]
+    val size: Int
   }
   [serializable]
   abstract class NonEmpty[+B] extends Tree[B] {
@@ -35,27 +38,37 @@ abstract class RedBlack[A] {
       if (isSmaller(k, key)) left.lookup(k)
       else if (isSmaller(key, k)) right.lookup(k)
       else this
+
+    /** helper function for upd and updf */
+    def balanceLeft[B1 >: B](isBlack: boolean, z: A, zv: B1, l: Tree[B1], d: Tree[B1]) = l match {
+      case RedTree(y, yv, RedTree(x, xv, a, b), c) =>
+        RedTree(y, yv, BlackTree(x, xv, a, b), BlackTree(z, zv, c, d))
+      case RedTree(x, xv, a, RedTree(y, yv, b, c)) =>
+        RedTree(y, yv, BlackTree(x, xv, a, b), BlackTree(z, zv, c, d))
+      case _ =>
+        mkTree(isBlack, z, zv, l, d)
+    }
+
+    /** helper function for upd and updf */
+    def balanceRight[B1 >: B](isBlack: boolean, x: A, xv: B1, a: Tree[B1], r: Tree[B1]) = r match {
+      case RedTree(z, zv, RedTree(y, yv, b, c), d) =>
+        RedTree(y, yv, BlackTree(x, xv, a, b), BlackTree(z, zv, c, d))
+      case RedTree(y, yv, b, RedTree(z, zv, c, d)) =>
+        RedTree(y, yv, BlackTree(x, xv, a, b), BlackTree(z, zv, c, d))
+      case _ =>
+        mkTree(isBlack, x, xv, a, r)
+    }
+
     def upd[B1 >: B](k: A, v: B1): Tree[B1] = {
-      def balanceLeft(isBlack: boolean, z: A, zv: B, l: Tree[B1], d: Tree[B1]) = l match {
-        case RedTree(y, yv, RedTree(x, xv, a, b), c) =>
-          RedTree(y, yv, BlackTree(x, xv, a, b), BlackTree(z, zv, c, d))
-        case RedTree(x, xv, a, RedTree(y, yv, b, c)) =>
-          RedTree(y, yv, BlackTree(x, xv, a, b), BlackTree(z, zv, c, d))
-        case _ =>
-          mkTree(isBlack, z, zv, l, d)
-      }
-      def balanceRight(isBlack: boolean, x: A, xv: B, a: Tree[B1], r: Tree[B1]) = r match {
-        case RedTree(z, zv, RedTree(y, yv, b, c), d) =>
-          RedTree(y, yv, BlackTree(x, xv, a, b), BlackTree(z, zv, c, d))
-        case RedTree(y, yv, b, RedTree(z, zv, c, d)) =>
-          RedTree(y, yv, BlackTree(x, xv, a, b), BlackTree(z, zv, c, d))
-        case _ =>
-          mkTree(isBlack, x, xv, a, r)
-      }
       if (isSmaller(k, key)) balanceLeft(isBlack, key, value, left.upd(k, v), right)
       else if (isSmaller(key, k)) balanceRight(isBlack, key, value, left, right.upd(k, v))
       else mkTree(isBlack, k, v, left, right)
     }
+    def updf[B1 >: B](k: A, f: Option[B]=>B1): Tree[B1] = {
+      if (isSmaller(k, key)) balanceLeft(isBlack, key, value, left.updf(k, f), right)
+      else if (isSmaller(key, k)) balanceRight(isBlack, key, value, left, right.updf(k, f))
+      else mkTree(isBlack, k, f(Some(value)), left, right)
+    }
     def del(k: A): Tree[B] = {
       if (isSmaller(k, key)) mkTree(isBlack, key, value, left.del(k), right)
       else if (isSmaller(key, k)) mkTree(isBlack, key, value, left, right.del(k))
@@ -69,6 +82,8 @@ abstract class RedBlack[A] {
     def smallest: NonEmpty[B] = if (left.isEmpty) this else left.smallest
     def elements: Iterator[Pair[A, B]] =
       left.elements append Iterator.single(Pair(key, value)) append right.elements
+
+    val size = left.size + right.size
   }
   [serializable]
   case object Empty extends Tree[Nothing] {
@@ -76,9 +91,11 @@ abstract class RedBlack[A] {
     def isBlack = true
     def lookup(k: A): Tree[Nothing] = this
     def upd[B](k: A, v: B): Tree[B] = RedTree(k, v, Empty, Empty)
+    def updf[B](k: A, f: Option[Nothing]=>B) = upd(k, f(None))
     def del(k: A): Tree[Nothing] = this
     def smallest: NonEmpty[Nothing] = throw new NoSuchElementException("empty map")
     def elements: Iterator[Pair[A, Nothing]] = Iterator.empty
+    val size = 0
   }
   [serializable]
   case class RedTree[+B](override val key: A,
diff --git a/src/library/scala/collection/immutable/Set.scala b/src/library/scala/collection/immutable/Set.scala
index eadc656470..42c106eddd 100644
--- a/src/library/scala/collection/immutable/Set.scala
+++ b/src/library/scala/collection/immutable/Set.scala
@@ -25,13 +25,14 @@ package scala.collection.immutable
  *  @version 1.1, 03/05/2004
  */
 object Set {
-  /** The empty set of this type
+  /** An empty immutable set.  This is implemented with a HashTreeSet.
+   * Note that for element types that have an Ordered view, it is more
+   * efficient to use a TreeSet, instead.
    */
-  def empty[A <% Ordered[A]] = new TreeSet[A]
+  def empty[A] = new HashTreeSet[A]
 
-  /** The canonical factory for this type
-   */
-  def apply[A <% Ordered[A]](elems: A*) = empty[A] ++ elems
+  /** Create an immutable map from a specified sequence of elements.  */
+  def apply[A](elems: A*) = empty[A] ++ elems
 }
 
 trait Set[A] extends AnyRef with collection.Set[A] {
diff --git a/src/library/scala/collection/immutable/TreeMap.scala b/src/library/scala/collection/immutable/TreeMap.scala
index 9e43f092df..cd8210a524 100644
--- a/src/library/scala/collection/immutable/TreeMap.scala
+++ b/src/library/scala/collection/immutable/TreeMap.scala
@@ -35,20 +35,19 @@ object TreeMap {
  *  @version 1.1, 03/05/2004
  */
 [serializable]
-class TreeMap[A <% Ordered[A], +B](val size: int, t: RedBlack[A]#Tree[B])
+class TreeMap[A <% Ordered[A], +B](tree: RedBlack[A]#Tree[B])
 extends RedBlack[A] with Map[A, B] {
 
   def isSmaller(x: A, y: A) = x < y
 
-  def this() = this(0, null)
+  def this() = this(Empty)
+  def size = tree.size
 
-  protected val tree: RedBlack[A]#Tree[B] = if (size == 0) Empty else t
-
-  private def newMap[B](s: int, t: RedBlack[A]#Tree[B]) = new TreeMap[A, B](s, t)
+  private def newMap[B](t: RedBlack[A]#Tree[B]) = new TreeMap[A, B](t)
 
   /** A factory to create empty maps of the same type of keys.
    */
-  def empty[C] = ListMap.empty[A, C]
+  def empty[C] = new TreeMap[A,C]
 
   /** A new TreeMap with the entry added is returned,
    *  if key is <em>not</em> in the TreeMap, otherwise
@@ -59,8 +58,16 @@ extends RedBlack[A] with Map[A, B] {
    *  @return ...
    */
   def update [B1 >: B](key: A, value: B1): TreeMap[A, B1] = {
-    val newsize = if (tree.lookup(key).isEmpty) size + 1 else size
-    newMap(newsize, tree.update(key, value))
+     val newtree = tree.update(key, value)
+     newMap(newtree)
+  }
+
+  /** Update the item at a single key.  This method
+   * is functionally equivalent to: update(key, f(get(key))) .
+   */
+  def updatef [B1 >: B](key: A, f: Option[B]=>B1) = {
+    val newtree = tree.updatef(key, f)
+    newMap(newtree)
   }
 
   /** A new TreeMap with the entry added is returned,
@@ -68,12 +75,16 @@ extends RedBlack[A] with Map[A, B] {
    */
   def insert [B1 >: B](key: A, value: B1): TreeMap[A, B1] = {
     assert(tree.lookup(key).isEmpty)
-    newMap(size + 1, tree.update(key, value))
+    newMap(tree.update(key, value))
   }
 
-  def - (key:A): TreeMap[A, B] =
-    if (tree.lookup(key).isEmpty) this
-    else newMap(size - 1, tree.delete(key))
+  def - (key:A): TreeMap[A, B] = {
+    val newtree = tree.delete(key)
+    if(newtree eq tree)
+      this
+    else
+      newMap(newtree)
+  }
 
   /** Check if this map maps <code>key</code> to a value and return the
    *  value if it exists.
@@ -106,7 +117,3 @@ extends RedBlack[A] with Map[A, B] {
    */
   def elements: Iterator[Pair[A, B]] = tree.elements
 }
-
-
-
-
diff --git a/src/library/scala/collection/immutable/TreeSet.scala b/src/library/scala/collection/immutable/TreeSet.scala
index 1226923d66..e5f4ed312b 100644
--- a/src/library/scala/collection/immutable/TreeSet.scala
+++ b/src/library/scala/collection/immutable/TreeSet.scala
@@ -34,26 +34,23 @@ object TreeSet {
  */
 
 [serializable]
-class TreeSet[A <% Ordered[A]](val size: int, t: RedBlack[A]#Tree[Unit])
+class TreeSet[A <% Ordered[A]](tree: RedBlack[A]#Tree[Unit])
 extends RedBlack[A] with Set[A] {
-
   def isSmaller(x: A, y: A) = x < y
+  def this() = this(Empty)
 
-  def this() = this(0, null)
-
-  protected val tree: RedBlack[A]#Tree[Unit] = if (size == 0) Empty else t
+  def size = tree.size
 
-  private def newSet(s: int, t: RedBlack[A]#Tree[Unit]) = new TreeSet[A](s, t)
+  private def newSet(t: RedBlack[A]#Tree[Unit]) = new TreeSet[A](t)
 
-  /** A factory to create empty maps of the same type of keys.
+  /** @return an empty set of arbitrary element type
    */
   def empty[B]: Set[B] = ListSet.empty[B]
 
   /** A new TreeSet with the entry added is returned,
    */
   def + (elem: A): TreeSet[A] = {
-    val newsize = if (tree.lookup(elem).isEmpty) size + 1 else size
-    newSet(newsize, tree.update(elem, ()))
+    newSet(tree.update(elem, ()))
   }
 
   /** A new TreeSet with the entry added is returned,
@@ -61,12 +58,16 @@ extends RedBlack[A] with Set[A] {
    */
   def insert (elem: A): TreeSet[A] = {
     assert(tree.lookup(elem).isEmpty)
-    newSet(size + 1, tree.update(elem, ()))
+    newSet(tree.update(elem, ()))
   }
 
-  def - (elem:A): TreeSet[A] =
-    if (tree.lookup(elem).isEmpty) this
-    else newSet(size - 1, tree.delete(elem))
+  def - (elem:A): TreeSet[A] = {
+    val newtree = tree.delete(elem)
+    if(newtree eq tree)
+      this
+    else
+      newSet(newtree)
+  }
 
   /** Checks if this set contains element <code>elem</code>.
    *