SI-6642 Adds iteratorFrom, keysIteratorFrom, and valuesIteratorFrom

Adds the ability to efficiently create an iterator that starts at a
given key or element of a sorted set or map. Similar work is done for
key and value only iterators on maps.

The bulk of the work is in RedBlackTree. Most of the rest is pushing
the new api methods throughout the appropriate spots in the collection
API.

This commit leaves undone some similar work possible on mutable TreeSets

11 files changed, 177 insertions, 10 deletions

diff --git a/src/library/scala/Enumeration.scala b/src/library/scala/Enumeration.scala
index 21f0c8fd3e..e7ce21b229 100644
--- a/src/library/scala/Enumeration.scala
+++ b/src/library/scala/Enumeration.scala
@@ -255,6 +255,7 @@ abstract class Enumeration (initial: Int) extends Serializable {
     def + (value: Value) = new ValueSet(nnIds + (value.id - bottomId))
     def - (value: Value) = new ValueSet(nnIds - (value.id - bottomId))
     def iterator = nnIds.iterator map (id => thisenum.apply(id + bottomId))
+    override def keysIteratorFrom(start: Value) = nnIds keysIteratorFrom start.id  map (id => thisenum.apply(id + bottomId))
     override def stringPrefix = thisenum + ".ValueSet"
     /** Creates a bit mask for the zero-adjusted ids in this set as a
      *  new array of longs */
diff --git a/src/library/scala/collection/BitSetLike.scala b/src/library/scala/collection/BitSetLike.scala
index d0f4e323c7..bf05331cb1 100644
--- a/src/library/scala/collection/BitSetLike.scala
+++ b/src/library/scala/collection/BitSetLike.scala
@@ -98,8 +98,10 @@ trait BitSetLike[+This <: BitSetLike[This] with SortedSet[Int]] extends SortedSe
     fromBitMaskNoCopy(a)
   }
 
-  def iterator: Iterator[Int] = new AbstractIterator[Int] {
-    private var current = 0
+  def iterator: Iterator[Int] = iteratorFrom(0)
+  
+  override def keysIteratorFrom(start: Int) = new AbstractIterator[Int] {
+    private var current = start
     private val end = nwords * WordLength
     def hasNext: Boolean = {
       while (current < end && !self.contains(current)) current += 1
diff --git a/src/library/scala/collection/SortedMapLike.scala b/src/library/scala/collection/SortedMapLike.scala
index 57ad3497c7..3c3e6095df 100644
--- a/src/library/scala/collection/SortedMapLike.scala
+++ b/src/library/scala/collection/SortedMapLike.scala
@@ -42,6 +42,7 @@ self =>
       val map = self.rangeImpl(from, until)
       new map.DefaultKeySortedSet
     }
+    override def keysIteratorFrom(start: A) = self.keysIteratorFrom(start)
   }
 
   /** Add a key/value pair to this map.
@@ -76,11 +77,17 @@ self =>
   override def filterKeys(p: A => Boolean): SortedMap[A, B] = new FilteredKeys(p) with SortedMap.Default[A, B] {
     implicit def ordering: Ordering[A] = self.ordering
     override def rangeImpl(from : Option[A], until : Option[A]): SortedMap[A, B] = self.rangeImpl(from, until).filterKeys(p)
+    override def iteratorFrom(start: A) = self iteratorFrom start filter {case (k, _) => p(k)}
+    override def keysIteratorFrom(start: A) = self keysIteratorFrom start filter p
+    override def valuesIteratorFrom(start: A) = self iteratorFrom start collect {case (k,v) if p(k) => v}
   }
   
   override def mapValues[C](f: B => C): SortedMap[A, C] = new MappedValues(f) with SortedMap.Default[A, C] {
     implicit def ordering: Ordering[A] = self.ordering
     override def rangeImpl(from : Option[A], until : Option[A]): SortedMap[A, C] = self.rangeImpl(from, until).mapValues(f)
+    override def iteratorFrom(start: A) = (self iteratorFrom start) map {case (k,v) => (k, f(v))}
+    override def keysIteratorFrom(start: A) = self keysIteratorFrom start
+    override def valuesIteratorFrom(start: A) = self valuesIteratorFrom start map f
   }
   
   /** Adds a number of elements provided by a traversable object
@@ -91,6 +98,28 @@ self =>
   override def ++[B1 >: B](xs: GenTraversableOnce[(A, B1)]): SortedMap[A, B1] =
     ((repr: SortedMap[A, B1]) /: xs.seq) (_ + _)
   
+  /**
+   * Creates an iterator over all the key/value pairs
+   * contained in this map having a key greater than or
+   * equal to `start` according to the ordering of
+   * this map. x.iteratorFrom(y) is equivalent
+   * to but often more efficient than x.from(y).iterator.
+   * 
+   * @param start The lower bound (inclusive)
+   * on the keys to be returned
+   */
+  def iteratorFrom(start: A): Iterator[(A, B)]
+  /**
+   * Creates an iterator over all the values contained in this
+   * map that are associated with a key greater than or equal to `start`
+   * according to the ordering of this map. x.valuesIteratorFrom(y) is
+   * equivalent to but often more efficient than 
+   * x.from(y).valuesIterator.
+   * 
+   * @param start The lower bound (inclusive)
+   * on the keys to be returned
+   */
+  def valuesIteratorFrom(start: A): Iterator[B]
 }
 
 
diff --git a/src/library/scala/collection/SortedSetLike.scala b/src/library/scala/collection/SortedSetLike.scala
index 71b45c72ff..6d1d1ac111 100644
--- a/src/library/scala/collection/SortedSetLike.scala
+++ b/src/library/scala/collection/SortedSetLike.scala
@@ -40,4 +40,14 @@ self =>
     case that: SortedSet[_] if that.ordering == ordering => that.hasAll(this.iterator)
     case that => super.subsetOf(that)
   }
+
+  /**
+   * Creates an iterator that contains all values from this collection
+   * greater than or equal to `start` according to the ordering of
+   * this collection. x.iteratorFrom(y) is equivalent to but will usually
+   * be more efficient than x.from(y).iterator
+   * 
+   * @param start The lower-bound (inclusive) of the iterator
+   */
+  def iteratorFrom(start: A): Iterator[A] = keysIteratorFrom(start)
 }
diff --git a/src/library/scala/collection/generic/Sorted.scala b/src/library/scala/collection/generic/Sorted.scala
index f962b26bd3..b3847fffc9 100644
--- a/src/library/scala/collection/generic/Sorted.scala
+++ b/src/library/scala/collection/generic/Sorted.scala
@@ -78,6 +78,18 @@ trait Sorted[K, +This <: Sorted[K, This]] {
     else
       until(next)
   }
+  
+  /**
+   * Creates an iterator over all the keys(or elements)  contained in this
+   * collection greater than or equal to `start`
+   * according to the ordering of this collection. x.keysIteratorFrom(y) 
+   * is equivalent to but often more efficient than 
+   * x.from(y).keysIterator.
+   * 
+   * @param start The lower bound (inclusive)
+   * on the keys to be returned
+   */
+  def keysIteratorFrom(start: K): Iterator[K]
 
   protected def hasAll(j: Iterator[K]): Boolean = {
     val i = keySet.iterator
diff --git a/src/library/scala/collection/immutable/RedBlackTree.scala b/src/library/scala/collection/immutable/RedBlackTree.scala
index 004c0ae8c6..c0d46765be 100644
--- a/src/library/scala/collection/immutable/RedBlackTree.scala
+++ b/src/library/scala/collection/immutable/RedBlackTree.scala
@@ -91,9 +91,9 @@ object RedBlackTree {
     if (tree.right ne null) _foreachKey(tree.right, f)
   }
 
-  def iterator[A, B](tree: Tree[A, B]): Iterator[(A, B)] = new EntriesIterator(tree)
-  def keysIterator[A, _](tree: Tree[A, _]): Iterator[A] = new KeysIterator(tree)
-  def valuesIterator[_, B](tree: Tree[_, B]): Iterator[B] = new ValuesIterator(tree)
+  def iterator[A, B](tree: Tree[A, B], start: Option[A] = None)(implicit ordering: Ordering[A]): Iterator[(A, B)] = new EntriesIterator(tree, start)
+  def keysIterator[A, _](tree: Tree[A, _], start: Option[A] = None)(implicit ordering: Ordering[A]): Iterator[A] = new KeysIterator(tree, start)
+  def valuesIterator[A, B](tree: Tree[A, B], start: Option[A] = None)(implicit ordering: Ordering[A]): Iterator[B] = new ValuesIterator(tree, start)
 
   @tailrec
   def nth[A, B](tree: Tree[A, B], n: Int): Tree[A, B] = {
@@ -425,7 +425,7 @@ object RedBlackTree {
     def unapply[A, B](t: BlackTree[A, B]) = Some((t.key, t.value, t.left, t.right))
   }
 
-  private[this] abstract class TreeIterator[A, B, R](root: Tree[A, B]) extends Iterator[R] {
+  private[this] abstract class TreeIterator[A, B, R](root: Tree[A, B], start: Option[A])(implicit ordering: Ordering[A]) extends Iterator[R] {
     protected[this] def nextResult(tree: Tree[A, B]): R
 
     override def hasNext: Boolean = lookahead ne null
@@ -481,7 +481,23 @@ object RedBlackTree {
       new Array[Tree[A, B]](maximumHeight)
     }
     private[this] var index = 0
-    private[this] var lookahead: Tree[A, B] = findLeftMostOrPopOnEmpty(root)
+    private[this] var lookahead: Tree[A, B] = start map startFrom getOrElse findLeftMostOrPopOnEmpty(root)
+    
+    /**
+     * Find the leftmost subtree whose key is equal to the given key, or if no such thing, 
+     * the leftmost subtree with the key that would be "next" after it according
+     * to the ordering. Along the way build up the iterator's path stack so that "next"
+     * functionality works.
+     */
+    private[this] def startFrom(key: A) : Tree[A,B] = if (root eq null) null else {
+      @tailrec def find(tree: Tree[A, B]): Tree[A, B] = 
+        if (tree == null) popNext
+        else find(
+          if (ordering.lteq(key, tree.key)) goLeft(tree)
+          else goRight(tree)
+        )      
+      find(root)
+    }
    
     private[this] def goLeft(tree: Tree[A, B]) = {
       pushNext(tree)
@@ -491,15 +507,15 @@ object RedBlackTree {
     private[this] def goRight(tree: Tree[A, B]) = tree.right
   }
 
-  private[this] class EntriesIterator[A, B](tree: Tree[A, B]) extends TreeIterator[A, B, (A, B)](tree) {
+  private[this] class EntriesIterator[A, B](tree: Tree[A, B], focus: Option[A])(implicit ordering: Ordering[A]) extends TreeIterator[A, B, (A, B)](tree, focus) {
     override def nextResult(tree: Tree[A, B]) = (tree.key, tree.value)
   }
 
-  private[this] class KeysIterator[A, B](tree: Tree[A, B]) extends TreeIterator[A, B, A](tree) {
+  private[this] class KeysIterator[A, B](tree: Tree[A, B], focus: Option[A])(implicit ordering: Ordering[A]) extends TreeIterator[A, B, A](tree, focus) {
     override def nextResult(tree: Tree[A, B]) = tree.key
   }
 
-  private[this] class ValuesIterator[A, B](tree: Tree[A, B]) extends TreeIterator[A, B, B](tree) {
+  private[this] class ValuesIterator[A, B](tree: Tree[A, B], focus: Option[A])(implicit ordering: Ordering[A]) extends TreeIterator[A, B, B](tree, focus) {
     override def nextResult(tree: Tree[A, B]) = tree.value
   }
 }
diff --git a/src/library/scala/collection/immutable/SortedMap.scala b/src/library/scala/collection/immutable/SortedMap.scala
index eb04231c55..5e833f87af 100644
--- a/src/library/scala/collection/immutable/SortedMap.scala
+++ b/src/library/scala/collection/immutable/SortedMap.scala
@@ -82,11 +82,17 @@ self =>
   override def filterKeys(p: A => Boolean): SortedMap[A, B] = new FilteredKeys(p) with SortedMap.Default[A, B] {
     implicit def ordering: Ordering[A] = self.ordering
     override def rangeImpl(from : Option[A], until : Option[A]): SortedMap[A, B] = self.rangeImpl(from, until).filterKeys(p)
+    override def iteratorFrom(start: A) = self iteratorFrom start filter {case (k, _) => p(k)} 
+    override def keysIteratorFrom(start : A) = self keysIteratorFrom start filter p
+    override def valuesIteratorFrom(start : A) = self iteratorFrom start collect {case (k,v) if p(k) => v}
   }
 
   override def mapValues[C](f: B => C): SortedMap[A, C] = new MappedValues(f) with SortedMap.Default[A, C] {
     implicit def ordering: Ordering[A] = self.ordering
     override def rangeImpl(from : Option[A], until : Option[A]): SortedMap[A, C] = self.rangeImpl(from, until).mapValues(f)
+    override def iteratorFrom(start: A) = self iteratorFrom start map {case (k, v) => (k, f(v))} 
+    override def keysIteratorFrom(start : A) = self keysIteratorFrom start
+    override def valuesIteratorFrom(start : A) = self valuesIteratorFrom start map f
   }
 
 }
diff --git a/src/library/scala/collection/immutable/TreeMap.scala b/src/library/scala/collection/immutable/TreeMap.scala
index 9a87d8636b..a6a6b75c32 100644
--- a/src/library/scala/collection/immutable/TreeMap.scala
+++ b/src/library/scala/collection/immutable/TreeMap.scala
@@ -189,9 +189,13 @@ class TreeMap[A, +B] private (tree: RB.Tree[A, B])(implicit val ordering: Orderi
    *  @return the new iterator
    */
   override def iterator: Iterator[(A, B)] = RB.iterator(tree)
+  override def iteratorFrom(start: A): Iterator[(A, B)] = RB.iterator(tree, Some(start))
 
   override def keysIterator: Iterator[A] = RB.keysIterator(tree)
+  override def keysIteratorFrom(start: A): Iterator[A] = RB.keysIterator(tree, Some(start))
+  
   override def valuesIterator: Iterator[B] = RB.valuesIterator(tree)
+  override def valuesIteratorFrom(start: A): Iterator[B] = RB.valuesIterator(tree, Some(start))
 
   override def contains(key: A): Boolean = RB.contains(tree, key)
   override def isDefinedAt(key: A): Boolean = RB.contains(tree, key)
diff --git a/src/library/scala/collection/immutable/TreeSet.scala b/src/library/scala/collection/immutable/TreeSet.scala
index 8bceb936aa..67668b3bef 100644
--- a/src/library/scala/collection/immutable/TreeSet.scala
+++ b/src/library/scala/collection/immutable/TreeSet.scala
@@ -144,6 +144,7 @@ class TreeSet[A] private (tree: RB.Tree[A, Unit])(implicit val ordering: Orderin
    *  @return the new iterator
    */
   def iterator: Iterator[A] = RB.keysIterator(tree)
+  override def keysIteratorFrom(start: A): Iterator[A] = RB.keysIterator(tree, Some(start))
 
   override def foreach[U](f: A =>  U) = RB.foreachKey(tree, f)
 
diff --git a/src/library/scala/collection/mutable/TreeSet.scala b/src/library/scala/collection/mutable/TreeSet.scala
index 5197af1b04..4fd35658fa 100644
--- a/src/library/scala/collection/mutable/TreeSet.scala
+++ b/src/library/scala/collection/mutable/TreeSet.scala
@@ -116,8 +116,25 @@ class TreeSet[A](implicit val ordering: Ordering[A]) extends SortedSet[A] with S
     resolve.avl.contains(elem, ordering)
   }
 
+  // TODO see the discussion on keysIteratorFrom
   override def iterator: Iterator[A] = resolve.avl.iterator
     .dropWhile(e => !isLeftAcceptable(from, ordering)(e))
       .takeWhile(e => isRightAcceptable(until, ordering)(e))
+  
+  // TODO because TreeSets are potentially ranged views into other TreeSets
+  // what this really needs to do is walk the whole stack of tree sets, find
+  // the highest "from", and then do a tree walk of the underlying avl tree
+  // to find that spot in max(O(stack depth), O(log tree.size)) time which
+  // should effectively be O(log size) since ranged views are rare and
+  // even more rarely deep. With the following implementation it's
+  // O(N log N) to get an iterator from a start point.  
+  // But before engaging that endeavor I think mutable.TreeSet should be
+  // based on the same immutable RedBlackTree that immutable.TreeSet is
+  // based on. There's no good reason to have these two collections based
+  // on two different balanced binary trees. That'll save
+  // having to duplicate logic for finding the starting point of a
+  // sorted binary tree iterator, logic that has already been
+  // baked into RedBlackTree.
+  override def keysIteratorFrom(start: A) = from(start).iterator
 
 }
diff --git a/test/files/run/iterator-from.scala b/test/files/run/iterator-from.scala
new file mode 100644
index 0000000000..8dc6ae4e51
--- /dev/null
+++ b/test/files/run/iterator-from.scala
@@ -0,0 +1,69 @@
+// This file tests iteratorFrom, keysIteratorFrom, and valueIteratorFrom on various sorted sets and maps
+
+import scala.util.{Random => R}
+import scala.collection._
+import scala.math.Ordered
+
+object Test extends App {
+  val maxLength = 25
+  val maxKey = 50
+  val maxValue = 50
+
+  def testSet[A <% Ordered[A]](set: SortedSet[A], list: List[A]) {
+    val distinctSorted = list.distinct.sorted
+    assertEquals("Set size wasn't the same as list sze", set.size, distinctSorted.size)
+
+    for(key <- distinctSorted) {
+      val clazz = set.getClass
+      val iteratorFrom = (set iteratorFrom key).toList
+      check(clazz, list, s"set iteratorFrom $key", s"(set from $key).iterator", iteratorFrom, (set from key).iterator.toList)
+      check(clazz, list, s"set.iteratorFrom $key", s"distinctSorted dropWhile (_ < $key)", iteratorFrom, distinctSorted dropWhile (_ < key))
+      check(clazz, list, s"set iteratorFrom $key", s"set keysIterator from $key", iteratorFrom, (set keysIteratorFrom key).toList)
+    }
+  }
+
+  def testMap[A <% Ordered[A], B](map: SortedMap[A, B], list: List[(A, B)]) {
+    val distinctSorted = distinctByKey(list).sortBy(_._1) 
+    assertEquals("Map size wasn't the same as list sze", map.size, distinctSorted.size)
+
+    for(keyValue <- distinctSorted) {
+      val key = keyValue._1
+      val clazz = map.getClass
+      val iteratorFrom = (map iteratorFrom key).toList
+      check(clazz, list, s"map iteratorFrom $key", s"(map from $key).iterator", iteratorFrom, (map from key).iterator.toList)
+      check(clazz, list, s"map iteratorFrom $key", s"distinctSorted dropWhile (_._1 < $key)", iteratorFrom, distinctSorted dropWhile (_._1 < key))
+      check(clazz, list, s"map iteratorFrom $key map (_._1)", s"map keysIteratorFrom $key", iteratorFrom map (_._1), (map keysIteratorFrom key).toList)
+      check(clazz, list, s"map iteratorFrom $key map (_._2)", s"map valuesIteratorFrom $key", iteratorFrom map (_._2), (map valuesIteratorFrom key).toList)     
+    }
+  }
+
+  def check[A](clazz: Class[_], list: List[_], m1: String, m2: String, l1: List[A], l2: List[A]) {
+    assertEquals(s"$clazz: `$m1` didn't match `$m2` on list $list", l1, l2)
+  }
+
+  def assertEquals[A](msg: String, x: A, y: A) {
+    assert(x == y, s"$msg\n1: $x\n2: $y")
+  }
+
+  def distinctByKey[A,B](list: List[(A, B)]) : List[(A,B)] = list.groupBy(_._1).map(_._2.last).toList
+
+  object Weekday extends Enumeration {
+    type Weekday = Value
+    val Mon, Tue, Wed, Thu, Fri, Sat, Sun = Value
+  }
+
+  0 until maxLength foreach {length => 
+    val keyValues = (0 until length map {_ => (R nextInt maxKey, R nextInt maxValue)}).toList
+    val keys = keyValues map (_._2)
+    testSet(immutable.BitSet(keys:_*), keys)
+    testSet(immutable.TreeSet(keys:_*), keys)
+    testSet(mutable.TreeSet(keys:_*), keys)
+    val days = keys map {n => Weekday(n % Weekday.values.size)}
+    testSet(Weekday.ValueSet(days:_*), days)
+
+    val treeMap = immutable.TreeMap(keyValues:_*)
+    testMap(treeMap, keyValues)
+    testMap(treeMap.filterKeys(_ % 2 == 0), keyValues  filter (_._1 % 2 == 0))
+    testMap(treeMap mapValues (_ + 1), keyValues map {case (k,v) => (k, v + 1)})
+  }
+}