diff options
-rw-r--r-- | src/library/scala/Enumeration.scala | 1 | ||||
-rw-r--r-- | src/library/scala/collection/BitSetLike.scala | 6 | ||||
-rw-r--r-- | src/library/scala/collection/SortedMapLike.scala | 29 | ||||
-rw-r--r-- | src/library/scala/collection/SortedSetLike.scala | 10 | ||||
-rw-r--r-- | src/library/scala/collection/generic/Sorted.scala | 12 | ||||
-rw-r--r-- | src/library/scala/collection/immutable/RedBlackTree.scala | 32 | ||||
-rw-r--r-- | src/library/scala/collection/immutable/SortedMap.scala | 6 | ||||
-rw-r--r-- | src/library/scala/collection/immutable/TreeMap.scala | 4 | ||||
-rw-r--r-- | src/library/scala/collection/immutable/TreeSet.scala | 1 | ||||
-rw-r--r-- | src/library/scala/collection/mutable/TreeSet.scala | 17 | ||||
-rw-r--r-- | test/files/run/iterator-from.scala | 69 |
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)}) + } +} |