Merge pull request scala#4333 from Ichoran/opt-Iterator-2.11.x

Performance optimization - Iterator

Author: SethTisue

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.