Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with
or
.
Download ZIP
Browse files

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
  • Loading branch information...
commit 62bc99d3b20a7b37a977b19a6202cdac474eb5f6 1 parent a0b1db4
@JamesIry JamesIry authored
View
1  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 */
View
6 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
View
29 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]
}
View
10 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)
}
View
12 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
View
32 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
}
}
View
6 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
}
}
View
4 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)
View
1  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)
View
17 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
}
View
69 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)})
+ }
+}
Please sign in to comment.
Something went wrong with that request. Please try again.