Optimized implementations of head/headOption/last/lastOption for

TreeMap/TreeSet.

3 files changed, 19 insertions, 0 deletions

diff --git a/src/library/scala/collection/immutable/RedBlack.scala b/src/library/scala/collection/immutable/RedBlack.scala
index 097df54af2..2d3d839851 100644
--- a/src/library/scala/collection/immutable/RedBlack.scala
+++ b/src/library/scala/collection/immutable/RedBlack.scala
@@ -40,6 +40,7 @@ abstract class RedBlack[A] extends Serializable {
     def upd[B1 >: B](k: A, v: B1): Tree[B1]
     def del(k: A): Tree[B]
     def smallest: NonEmpty[B]
+    def greatest: NonEmpty[B]
     def rng(from: Option[A], until: Option[A]): Tree[B]
     def first : A
     def last : A
@@ -148,6 +149,7 @@ abstract class RedBlack[A] extends Serializable {
     }
 
     def smallest: NonEmpty[B] = if (left.isEmpty) this else left.smallest
+    def greatest: NonEmpty[B] = if (right.isEmpty) this else right.greatest
 
     def toStream: Stream[(A,B)] = iterator.toStream
 
@@ -262,6 +264,7 @@ abstract class RedBlack[A] extends Serializable {
     def upd[B](k: A, v: B): Tree[B] = RedTree(k, v, Empty, Empty)
     def del(k: A): Tree[Nothing] = this
     def smallest: NonEmpty[Nothing] = throw new NoSuchElementException("empty map")
+    def greatest: NonEmpty[Nothing] = throw new NoSuchElementException("empty map")
     def iterator: Iterator[(A, Nothing)] = Iterator.empty
     def toStream: Stream[(A,Nothing)] = Stream.empty
 
diff --git a/src/library/scala/collection/immutable/TreeMap.scala b/src/library/scala/collection/immutable/TreeMap.scala
index 2fd5208991..be67a45b1e 100644
--- a/src/library/scala/collection/immutable/TreeMap.scala
+++ b/src/library/scala/collection/immutable/TreeMap.scala
@@ -71,6 +71,17 @@ class TreeMap[A, +B](override val size: Int, t: RedBlack[A]#Tree[B])(implicit va
   override def lastKey = t.last
   override def compare(k0: A, k1: A): Int = ordering.compare(k0, k1)
 
+  override def head = {
+    val smallest = t.smallest
+    (smallest.key, smallest.value)
+  }
+  override def headOption = if (t.isEmpty) None else Some(head)
+  override def last = {
+    val greatest = t.greatest
+    (greatest.key, greatest.value)
+  }
+  override def lastOption = if (t.isEmpty) None else Some(last)
+
   /** A factory to create empty maps of the same type of keys.
    */
   override def empty: TreeMap[A, B] = TreeMap.empty[A, B](ordering)
diff --git a/src/library/scala/collection/immutable/TreeSet.scala b/src/library/scala/collection/immutable/TreeSet.scala
index 05f27d0d93..85f4ae4b0e 100644
--- a/src/library/scala/collection/immutable/TreeSet.scala
+++ b/src/library/scala/collection/immutable/TreeSet.scala
@@ -53,6 +53,11 @@ class TreeSet[A](override val size: Int, t: RedBlack[A]#Tree[Unit])
 
   override def stringPrefix = "TreeSet"
 
+  override def head = t.smallest.key
+  override def headOption = if (t.isEmpty) None else Some(head)
+  override def last = t.greatest.key
+  override def lastOption = if (t.isEmpty) None else Some(last)
+
   def isSmaller(x: A, y: A) = compare(x,y) < 0
 
   def this()(implicit ordering: Ordering[A]) = this(0, null)(ordering)