Merge pull request scala#2119 from JamesIry/master_SI-6642

SI-6642 Adds iteratorFrom, keysIteratorFrom, and valuesIteratorFrom

Author: JamesIry

src/library/scala/Enumeration.scala

@@ -254,7 +254,8 @@ abstract class Enumeration (initial: Int) extends Serializable {
     def contains(v: Value) = nnIds contains (v.id - bottomId)
     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))
+    def iterator = nnIds.iterator map (id => thisenum.apply(bottomId + id))
+    override def keysIteratorFrom(start: Value) = nnIds keysIteratorFrom start.id  map (id => thisenum.apply(bottomId + id))
     override def stringPrefix = thisenum + ".ValueSet"
     /** Creates a bit mask for the zero-adjusted ids in this set as a
      *  new array of longs */

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

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]
 }
 
 

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)
 }

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

src/library/scala/collection/immutable/RedBlackTree.scala

@@ -24,13 +24,13 @@ import scala.annotation.meta.getter
  *
  *  @since 2.10
  */
-private[immutable]
+private[collection]
 object RedBlackTree {
 
   def isEmpty(tree: Tree[_, _]): Boolean = tree eq null
 
-  def contains[A](tree: Tree[A, _], x: A)(implicit ordering: Ordering[A]): Boolean = lookup(tree, x) ne null
-  def get[A, B](tree: Tree[A, B], x: A)(implicit ordering: Ordering[A]): Option[B] = lookup(tree, x) match {
+  def contains[A: Ordering](tree: Tree[A, _], x: A): Boolean = lookup(tree, x) ne null
+  def get[A: Ordering, B](tree: Tree[A, B], x: A): Option[B] = lookup(tree, x) match {
     case null => None
     case tree => Some(tree.value)
   }
@@ -44,8 +44,27 @@ object RedBlackTree {
   }
 
   def count(tree: Tree[_, _]) = if (tree eq null) 0 else tree.count
-  def update[A, B, B1 >: B](tree: Tree[A, B], k: A, v: B1, overwrite: Boolean)(implicit ordering: Ordering[A]): Tree[A, B1] = blacken(upd(tree, k, v, overwrite))
-  def delete[A, B](tree: Tree[A, B], k: A)(implicit ordering: Ordering[A]): Tree[A, B] = blacken(del(tree, k))
+  /**
+   * Count all the nodes with keys greater than or equal to the lower bound and less than the upper bound.
+   * The two bounds are optional.
+   */
+  def countInRange[A](tree: Tree[A, _], from: Option[A], to:Option[A])(implicit ordering: Ordering[A]) : Int = 
+    if (tree eq null) 0 else
+    (from, to) match {
+      // with no bounds use this node's count
+      case (None, None) => tree.count
+      // if node is less than the lower bound, try the tree on the right, it might be in range
+      case (Some(lb), _) if ordering.lt(tree.key, lb) => countInRange(tree.right, from, to)
+      // if node is greater than or equal to the upper bound, try the tree on the left, it might be in range
+      case (_, Some(ub)) if ordering.gteq(tree.key, ub) => countInRange(tree.left, from, to)
+      // node is in range so the tree on the left will all be less than the upper bound and the tree on the
+      // right will all be greater than or equal to the lower bound. So 1 for this node plus
+      // count the subtrees by stripping off the bounds that we don't need any more
+      case _ => 1 + countInRange(tree.left, from, None) + countInRange(tree.right, None, to)
+    
+    }
+  def update[A: Ordering, B, B1 >: B](tree: Tree[A, B], k: A, v: B1, overwrite: Boolean): Tree[A, B1] = blacken(upd(tree, k, v, overwrite))
+  def delete[A: Ordering, B](tree: Tree[A, B], k: A): Tree[A, B] = blacken(del(tree, k))
   def rangeImpl[A: Ordering, B](tree: Tree[A, B], from: Option[A], until: Option[A]): Tree[A, B] = (from, until) match {
     case (Some(from), Some(until)) => this.range(tree, from, until)
     case (Some(from), None)        => this.from(tree, from)
@@ -91,9 +110,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: Ordering, B](tree: Tree[A, B], start: Option[A] = None): Iterator[(A, B)] = new EntriesIterator(tree, start)
+  def keysIterator[A: Ordering](tree: Tree[A, _], start: Option[A] = None): Iterator[A] = new KeysIterator(tree, start)
+  def valuesIterator[A: Ordering, B](tree: Tree[A, B], start: Option[A] = None): Iterator[B] = new ValuesIterator(tree, start)
 
   @tailrec
   def nth[A, B](tree: Tree[A, B], n: Int): Tree[A, B] = {
@@ -425,32 +444,28 @@ 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](tree: 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 = next ne null
+    override def hasNext: Boolean = lookahead ne null
 
-    override def next: R = next match {
+    override def next: R = lookahead match {
       case null =>
         throw new NoSuchElementException("next on empty iterator")
       case tree =>
-        next = findNext(tree.right)
+        lookahead = findLeftMostOrPopOnEmpty(goRight(tree))
         nextResult(tree)
     }
 
     @tailrec
-    private[this] def findNext(tree: Tree[A, B]): Tree[A, B] = {
-      if (tree eq null) popPath()
+    private[this] def findLeftMostOrPopOnEmpty(tree: Tree[A, B]): Tree[A, B] =
+      if (tree eq null) popNext()
       else if (tree.left eq null) tree
-      else {
-        pushPath(tree)
-        findNext(tree.left)
-      }
-    }
+      else findLeftMostOrPopOnEmpty(goLeft(tree))
 
-    private[this] def pushPath(tree: Tree[A, B]) {
+    private[this] def pushNext(tree: Tree[A, B]) {
       try {
-        path(index) = tree
+        stackOfNexts(index) = tree
         index += 1
       } catch {
         case _: ArrayIndexOutOfBoundsException =>
@@ -462,17 +477,17 @@ object RedBlackTree {
            * An exception handler is used instead of an if-condition to optimize the normal path.
            * This makes a large difference in iteration speed!
            */
-          assert(index >= path.length)
-          path :+= null
-          pushPath(tree)
+          assert(index >= stackOfNexts.length)
+          stackOfNexts :+= null
+          pushNext(tree)
       }
     }
-    private[this] def popPath(): Tree[A, B] = if (index == 0) null else {
+    private[this] def popNext(): Tree[A, B] = if (index == 0) null else {
       index -= 1
-      path(index)
+      stackOfNexts(index)
     }
 
-    private[this] var path = if (tree eq null) null else {
+    private[this] var stackOfNexts = if (root eq null) null else {
       /*
        * According to "Ralf Hinze. Constructing red-black trees" [http://www.cs.ox.ac.uk/ralf.hinze/publications/#P5]
        * the maximum height of a red-black tree is 2*log_2(n + 2) - 2.
@@ -481,22 +496,45 @@ object RedBlackTree {
        *
        * We also don't store the deepest nodes in the path so the maximum path length is further reduced by one.
        */
-      val maximumHeight = 2 * (32 - Integer.numberOfLeadingZeros(tree.count + 2 - 1)) - 2 - 1
+      val maximumHeight = 2 * (32 - Integer.numberOfLeadingZeros(root.count + 2 - 1)) - 2 - 1
       new Array[Tree[A, B]](maximumHeight)
     }
     private[this] var index = 0
-    private[this] var next: Tree[A, B] = findNext(tree)
+    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 eq 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)
+      tree.left
+    }
+    
+    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: Ordering, B](tree: Tree[A, B], focus: Option[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: Ordering, B](tree: Tree[A, B], focus: Option[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: Ordering, B](tree: Tree[A, B], focus: Option[A]) extends TreeIterator[A, B, B](tree, focus) {
     override def nextResult(tree: Tree[A, B]) = tree.value
   }
 }

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
   }
 
 }

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)

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)
 

src/library/scala/collection/mutable/AVLTree.scala

@@ -15,7 +15,7 @@ package mutable
  * An immutable AVL Tree implementation used by mutable.TreeSet
  *
  * @author Lucien Pereira
- *
+ * @deprecated("AVLTree and its related classes are being removed from the standard library since they're not different enough from RedBlackTree to justify keeping them.", "2.11")
  */
 private[mutable] sealed trait AVLTree[+A] extends Serializable {
   def balance: Int
@@ -65,12 +65,18 @@ private[mutable] sealed trait AVLTree[+A] extends Serializable {
   def doubleRightRotation[B >: A]: Node[B] = sys.error("Should not happen.")
 }
 
+/**
+ * @deprecated("AVLTree and its related classes are being removed from the standard library since they're not different enough from RedBlackTree to justify keeping them.", "2.11")
+ */
 private case object Leaf extends AVLTree[Nothing] {
   override val balance: Int = 0
 
   override val depth: Int = -1
 }
 
+/**
+ * @deprecated("AVLTree and its related classes are being removed from the standard library since they're not different enough from RedBlackTree to justify keeping them.", "2.11")
+ */
 private case class Node[A](val data: A, val left: AVLTree[A], val right: AVLTree[A]) extends AVLTree[A] {
   override val balance: Int = right.depth - left.depth
 
@@ -205,6 +211,9 @@ private case class Node[A](val data: A, val left: AVLTree[A], val right: AVLTree
   }
 }
 
+/**
+ * @deprecated("AVLTree and its related classes are being removed from the standard library since they're not different enough from RedBlackTree to justify keeping them.", "2.11")
+ */
 private class AVLIterator[A](root: Node[A]) extends Iterator[A] {
   val stack = mutable.ArrayStack[Node[A]](root)
   diveLeft()