diff options
-rw-r--r-- | src/library/scala/collection/Iterator.scala | 221 |
1 files changed, 162 insertions, 59 deletions
diff --git a/src/library/scala/collection/Iterator.scala b/src/library/scala/collection/Iterator.scala index c9037eb3e3..b69e51fdf5 100644 --- a/src/library/scala/collection/Iterator.scala +++ b/src/library/scala/collection/Iterator.scala @@ -392,8 +392,16 @@ trait Iterator[+A] extends TraversableOnce[A] { */ def flatMap[B](f: A => GenTraversableOnce[B]): Iterator[B] = new AbstractIterator[B] { private var cur: Iterator[B] = empty - def hasNext: Boolean = - cur.hasNext || self.hasNext && { cur = f(self.next()).toIterator; hasNext } + private def nextCur() { cur = f(self.next()).toIterator } + def hasNext: Boolean = { + // Equivalent to cur.hasNext || self.hasNext && { nextCur(); hasNext } + // but slightly shorter bytecode (better JVM inlining!) + while (!cur.hasNext) { + if (!self.hasNext) return false + nextCur() + } + true + } def next(): B = (if (hasNext) cur else empty).next() } @@ -405,6 +413,7 @@ trait Iterator[+A] extends TraversableOnce[A] { * @note Reuse: $consumesAndProducesIterator */ def filter(p: A => Boolean): Iterator[A] = new AbstractIterator[A] { + // TODO 2.12 - Make a full-fledged FilterImpl that will reverse sense of p private var hd: A = _ private var hdDefined: Boolean = false @@ -470,13 +479,27 @@ trait Iterator[+A] extends TraversableOnce[A] { * @note Reuse: $consumesAndProducesIterator */ @migration("`collect` has changed. The previous behavior can be reproduced with `toSeq`.", "2.8.0") - def collect[B](pf: PartialFunction[A, B]): Iterator[B] = { - val self = buffered - new AbstractIterator[B] { - private def skip() = while (self.hasNext && !pf.isDefinedAt(self.head)) self.next() - def hasNext = { skip(); self.hasNext } - def next() = { skip(); pf(self.next()) } + def collect[B](pf: PartialFunction[A, B]): Iterator[B] = new AbstractIterator[B] { + // Manually buffer to avoid extra layer of wrapping with buffered + private[this] var hd: A = _ + + // Little state machine to keep track of where we are + // Seek = 0; Found = 1; Empty = -1 + // Not in vals because scalac won't make them static (@inline def only works with -optimize) + // BE REALLY CAREFUL TO KEEP COMMENTS AND NUMBERS IN SYNC! + private[this] var status = 0/*Seek*/ + + def hasNext = { + while (status == 0/*Seek*/) { + if (self.hasNext) { + hd = self.next() + if (pf.isDefinedAt(hd)) status = 1/*Found*/ + } + else status = -1/*Empty*/ + } + status == 1/*Found*/ } + def next() = if (hasNext) { status = 0/*Seek*/; pf(hd) } else Iterator.empty.next() } /** Produces a collection containing cumulative results of applying the @@ -578,33 +601,105 @@ trait Iterator[+A] extends TraversableOnce[A] { * @note Reuse: $consumesOneAndProducesTwoIterators */ def span(p: A => Boolean): (Iterator[A], Iterator[A]) = { - val self = buffered - - // Must be a named class to avoid structural call to finish from trailing iterator + /* + * Giving a name to following iterator (as opposed to trailing) because + * anonymous class is represented as a structural type that trailing + * iterator is referring (the finish() method) and thus triggering + * handling of structural calls. It's not what's intended here. + */ class Leading extends AbstractIterator[A] { - private val drained = new mutable.Queue[A] - private var finished = false - def finish(): Unit = { - require(!finished) - finished = true - while (selfish) drained += self.next + var lookahead: mutable.Queue[A] = null + var hd: A = _ + /* Status is kept with magic numbers + * 1 means next element is in hd and we're still reading into this iterator + * 0 means we're still reading but haven't found a next element + * -1 means we are done reading into the iterator, so we must rely on lookahead + * -2 means we are done but have saved hd for the other iterator to use as its first element + */ + var status = 0 + private def store(a: A) { + if (lookahead == null) lookahead = new mutable.Queue[A] + lookahead += a + } + def hasNext = { + if (status < 0) (lookahead ne null) && lookahead.nonEmpty + else if (status > 0) true + else { + if (self.hasNext) { + hd = self.next() + status = if (p(hd)) 1 else -2 + } + else status = -1 + status > 0 + } } - private def selfish = self.hasNext && p(self.head) - def hasNext = if (finished) drained.nonEmpty else selfish def next() = { - if (finished) drained.dequeue() - else if (selfish) self.next() + if (hasNext) { + if (status == 1) { status = 0; hd } + else lookahead.dequeue() + } else empty.next() } + def finish(): Boolean = { + if (status == -1) false + else if (status == -2) { + status = -1 + true + } + else { + if (status == 1) store(hd) + while (self.hasNext) { + val a = self.next() + if (p(a)) store(a) + else { + hd = a + status = -1 + return true + } + } + false + } + } } + val leading = new Leading + val trailing = new AbstractIterator[A] { - private lazy val it = { - leading.finish() - self + private[this] var myLeading = leading + /* Status flags meanings: + * -1 not yet accesssed + * 0 single element waiting in leading + * 1 defer to self + */ + private[this] var status = -1 + def hasNext = { + if (status > 0) self.hasNext + else { + if (status == 0) true + else if (myLeading.finish()) { + status = 0 + true + } + else { + status = 1 + myLeading = null + self.hasNext + } + } } - def hasNext = it.hasNext - def next() = it.next() + def next() = { + if (hasNext) { + if (status > 0) self.next() + else { + status = 1 + val ans = myLeading.hd + myLeading = null + ans + } + } + else Iterator.empty.next() + } + override def toString = "unknown-if-empty iterator" } @@ -618,18 +713,35 @@ trait Iterator[+A] extends TraversableOnce[A] { * @return an iterator consisting of the remaining elements * @note Reuse: $consumesAndProducesIterator */ - def dropWhile(p: A => Boolean): Iterator[A] = { - val self = buffered - new AbstractIterator[A] { - var dropped = false - private def skip() = - if (!dropped) { - while (self.hasNext && p(self.head)) self.next() - dropped = true + def dropWhile(p: A => Boolean): Iterator[A] = new AbstractIterator[A] { + // Magic value: -1 = hasn't dropped, 0 = found first, 1 = defer to parent iterator + private[this] var status = -1 + // Local buffering to avoid double-wrap with .buffered + private[this] var fst: A = _ + def hasNext: Boolean = + if (status == 1) self.hasNext + else if (status == 0) true + else { + while (self.hasNext) { + val a = self.next() + if (!p(a)) { + fst = a + status = 0 + return true + } } - def hasNext = { skip(); self.hasNext } - def next() = { skip(); self.next() } - } + status = 1 + false + } + def next() = + if (hasNext) { + if (status == 1) self.next() + else { + status = 1 + fst + } + } + else Iterator.empty.next() } /** Creates an iterator formed from this iterator and another iterator @@ -777,7 +889,7 @@ trait Iterator[+A] extends TraversableOnce[A] { * is equal (as determined by `==`) to `elem`, `false` otherwise. * @note Reuse: $consumesIterator */ - def contains(elem: Any): Boolean = exists(_ == elem) + def contains(elem: Any): Boolean = exists(_ == elem) // Note--this seems faster than manual inlining! /** Finds the first value produced by the iterator satisfying a * predicate, if any. @@ -789,12 +901,11 @@ trait Iterator[+A] extends TraversableOnce[A] { * @note Reuse: $consumesIterator */ def find(p: A => Boolean): Option[A] = { - var res: Option[A] = None - while (res.isEmpty && hasNext) { - val e = next() - if (p(e)) res = Some(e) + while (hasNext) { + val a = next() + if (p(a)) return Some(a) } - res + None } /** Returns the index of the first produced value satisfying a predicate, or -1. @@ -807,15 +918,11 @@ trait Iterator[+A] extends TraversableOnce[A] { */ def indexWhere(p: A => Boolean): Int = { var i = 0 - var found = false - while (!found && hasNext) { - if (p(next())) { - found = true - } else { - i += 1 - } + while (hasNext) { + if (p(next())) return i + i += 1 } - if (found) i else -1 + -1 } /** Returns the index of the first occurrence of the specified @@ -829,15 +936,11 @@ trait Iterator[+A] extends TraversableOnce[A] { */ def indexOf[B >: A](elem: B): Int = { var i = 0 - var found = false - while (!found && hasNext) { - if (next() == elem) { - found = true - } else { - i += 1 - } + while (hasNext) { + if (next() == elem) return i + i += 1 } - if (found) i else -1 + -1 } /** Creates a buffered iterator from this iterator. |