diff options
-rw-r--r-- | src/library/scala/collection/Iterator.scala | 165 | ||||
-rw-r--r-- | test/files/run/unittest_iterator.scala | 37 |
2 files changed, 154 insertions, 48 deletions
diff --git a/src/library/scala/collection/Iterator.scala b/src/library/scala/collection/Iterator.scala index 6c184bdaec..ffcb6eda70 100644 --- a/src/library/scala/collection/Iterator.scala +++ b/src/library/scala/collection/Iterator.scala @@ -727,64 +727,133 @@ trait Iterator[+A] { self => buf } - /** Returns an iterator which groups this iterator into fixed size - * blocks, possibly except for the final block. For instance: - * <code>(1 to 8).iterator grouped 3</code> - * will return an iterator equivalent to - * Iterator(Seq(1,2,3), Seq(4,5,6), Seq(7,8)) - */ - def grouped(size: Int): Iterator[Sequence[A]] = new Iterator[Sequence[A]] { - def hasNext = self.hasNext - def next = self takeDestructively size toList - } - /** Returns an iterator which presents a "sliding window" of the - * given size across this iterator. For instance: - * <code>(1 to 5).iterator sliding 3</code> - * will return an iterator equivalent to - * Iterator(Seq(1,2,3), Seq(2,3,4), Seq(3,4,5)) + /** A flexible iterator for transforming an Iterator[A] into an + * Iterator[Sequence[A]], with configurable sequence size, step, and + * strategy for dealing with elements which don't fit evenly. * - * The optional 'step' parameter if given advances the window - * by the given number of positions. - * - * Note: if your parameters don't perfectly "fit" the sequence, - * the last result may be of a size less than windowSize, and the - * last step may be less than the given step, or both. What is - * guaranteed is that each element of the original iterator will - * appear in at least one sequence in the sliding iterator, UNLESS - * the step size is larger than the windowSize. - */ - def sliding(windowSize: Int, step: Int = 1): Iterator[Sequence[A]] = { - require(windowSize >= 1 && step >= 1) - - val buf = takeDestructively(windowSize) - - if (!self.hasNext) Iterator single buf.toList - else new Iterator[Sequence[A]] { - private[this] var filled = true - private[this] var buffer = ArrayBuffer(buf: _*) - def fill() = { - val xs = self takeDestructively step - val len = xs.length min buf.size - - (len > 0) && { - buffer trimStart len - buffer ++= (xs takeRight len) + * Typical uses can be achieved via methods `grouped' and `sliding'. + */ + class GroupedIterator[B >: A](self: Iterator[A], size: Int, step: Int) extends Iterator[Sequence[B]] { + require(size >= 1 && size >= 1) + + private[this] var buffer: ArrayBuffer[B] = ArrayBuffer() // the buffer + private[this] var filled = false // whether the buffer is "hot" + private[this] var _partial = true // whether we deliver short sequences + private[this] var pad: Option[() => B] = None // what to pad short sequences with + + /** Public functions which can be used to configure the iterator before use. */ + def withPadding(x: => B): this.type = { + pad = Some(() => x) + this + } + def withPartial(x: Boolean): this.type = { + _partial = x + if (_partial == true) // reset pad since otherwise it will take precedence + pad = None + + this + } + + private def padding(x: Int) = List.fill(x)(pad.get()) + private def gap = (step - size) max 0 + + private def go(count: Int) = { + val prevSize = buffer.size + def isFirst = prevSize == 0 + // If there is padding defined we insert it immediately + // so the rest of the code can be oblivious + val xs: Seq[B] = { + val res = self takeDestructively count + // extra checks so we don't calculate length unless there's reason + if (pad.isDefined && !self.hasNext) { + val shortBy = count - res.length + if (shortBy > 0) res ++ padding(shortBy) else res + } + else res + } + lazy val len = xs.length + lazy val incomplete = len < count + + // if 0 elements are requested, or if the number of newly obtained + // elements is less than the gap between sequences, we are done. + def deliver(howMany: Int) = { + (howMany > 0 && len > gap) && { + if (!isFirst) + buffer trimStart (step min prevSize) + + val available = + if (isFirst) len + else howMany min (len - gap) + + buffer ++= (xs takeRight available) filled = true true } } - def hasNext = filled || fill() - def next = { - if (!filled) - fill() + if (xs.isEmpty) false // self ran out of elements + else if (_partial) deliver(len min size) // if _partial is true, we deliver regardless + else if (incomplete) false // !_partial && incomplete means no more seqs + else if (isFirst) deliver(len) // first element + else deliver(step min size) // the typical case + } - filled = false - buffer.toList - } + // fill() returns false if no more sequences can be produced + private def fill(): Boolean = { + if (!self.hasNext) false + // the first time we grab size, but after that we grab step + else if (buffer.isEmpty) go(size) + else go(step) + } + + def hasNext = filled || fill() + def next = { + if (!filled) + fill() + + filled = false + buffer.toList } } + /** Returns an iterator which groups this iterator into fixed size + * blocks. Example usages: + * + * <pre> + * // Returns List(List(1, 2, 3), List(4, 5, 6), List(7))) + * (1 to 7).iterator grouped 3 toList + * // Returns List(List(1, 2, 3), List(4, 5, 6)) + * (1 to 7).iterator grouped 3 withPartial false toList + * // Returns List(List(1, 2, 3), List(4, 5, 6), List(7, 20, 25) + * // Illustrating that withPadding's argument is by-name. + * val it2 = Iterator.iterate(20)(_ + 5) + * (1 to 7).iterator grouped 3 withPadding it2.next toList + * </pre> + */ + def grouped[B >: A](size: Int): GroupedIterator[B] = + new GroupedIterator[B](self, size, size) + + /** Returns an iterator which presents a "sliding window" view of + * another iterator. The first argument is the window size, and + * the second is how far to advance the window on each iteration; + * defaults to 1. Example usages: + * + * <pre> + * // returns List(List(1, 2, 3), List(2, 3, 4), List(3, 4, 5)) + * (1 to 5).iterator.sliding(3).toList + * // returns List(List(1, 2, 3, 4), List(4, 5)) + * (1 to 5).iterator.sliding(4, 3).toList + * // returns List(List(1, 2, 3, 4)) + * (1 to 5).iterator.sliding(4, 3).withPartial(false).toList + * // List(List(1, 2, 3, 4), List(4, 5, 30, 35)) + * // Illustrating that withPadding's argument is by-name. + * val it2 = Iterator.iterate(20)(_ + 5) + * (1 to 5).iterator.sliding(4, 3).withPadding(it2.next).toList + * </pre> + */ + def sliding[B >: A](size: Int, step: Int = 1): GroupedIterator[B] = + new GroupedIterator[B](self, size, step) + /** Returns the number of elements in this iterator. * @note The iterator is at its end after this method returns. */ diff --git a/test/files/run/unittest_iterator.scala b/test/files/run/unittest_iterator.scala new file mode 100644 index 0000000000..93aaa4a834 --- /dev/null +++ b/test/files/run/unittest_iterator.scala @@ -0,0 +1,37 @@ +// Some iterator grouped/sliding unit tests +object Test +{ + def it = (1 to 10).iterator + def assertThat[T](expectedLength: Int, expectedLast: Seq[T])(it: Iterator[Seq[T]]) { + val xs = it.toList + def fail(msg: String) = "assertion failed on %s: %s".format(xs, msg) + assert(xs.size == expectedLength, fail("expected length " + expectedLength)) + assert(xs.last == expectedLast, fail("expected last " + expectedLast)) + } + + def main(args: Array[String]): Unit = { + val itSum = it.toStream.sum + for (i <- it) { + // sum of the groups == sum of the original + val thisSum = ((it grouped i) map (_.sum)).toStream.sum + assert(thisSum == itSum, thisSum + " != " + itSum) + } + + // grouped + assertThat(4, List(10)) { it grouped 3 } + assertThat(3, List(7, 8, 9)) { it grouped 3 withPartial false } + assertThat(4, List(10, -1, -1)) { it grouped 3 withPadding -1 } + + // testing by-name padding + val padIt = it + assertThat(4, List(10, 1, 2)) { it grouped 3 withPadding padIt.next } + + // sliding + assertThat(8, List(8, 9, 10)) { it sliding 3 } + assertThat(3, (3 to 10).toList) { it sliding 8 } + assertThat(2, List(9, 10)) { it.sliding(8, 8) } + assertThat(1, (1 to 8).toList) { it.sliding(8, 8) withPartial false } + assertThat(2, List(9, 10, -1, -1, -1)) { it.sliding(5, 8) withPadding -1 } + assertThat(1, (1 to 5).toList) { it.sliding(5, 8) withPartial false } + } +} |