SI-9623 Avoid unnecessary hasNext calls in JoinIterator & ConcatIterator

These iterator implementations are used to concatenate two (JoinIterator)
or more (ConcatIterator) other iterators with `++`. They used to perform
many unnecessary calls to the child iterators’ `hasNext` methods.

This improved state machine-based implementation reduces that number to
the bare minimum, i.e. iterating over concatenated iterators with
`foreach` calls the children's `hasNext` methods a total of (number of
children) + (number of elements) times, the same as when iterating over
all children separately.

2 files changed, 73 insertions, 6 deletions

diff --git a/src/library/scala/collection/Iterator.scala b/src/library/scala/collection/Iterator.scala
index ed536f10a8..8d88b1c6b1 100644
--- a/src/library/scala/collection/Iterator.scala
+++ b/src/library/scala/collection/Iterator.scala
@@ -10,7 +10,7 @@ package scala
 package collection
 
 import mutable.ArrayBuffer
-import scala.annotation.migration
+import scala.annotation.{tailrec, migration}
 import immutable.Stream
 import scala.collection.generic.CanBuildFrom
 import scala.annotation.unchecked.{ uncheckedVariance => uV }
@@ -168,8 +168,10 @@ object Iterator {
   private[scala] final class ConcatIterator[+A](private[this] var current: Iterator[A], initial: Vector[() => Iterator[A]]) extends Iterator[A] {
     @deprecated def this(initial: Vector[() => Iterator[A]]) = this(Iterator.empty, initial) // for binary compatibility
     private[this] var queue: Vector[() => Iterator[A]] = initial
+    private[this] var currentHasNextChecked = false
     // Advance current to the next non-empty iterator
     // current is set to null when all iterators are exhausted
+    @tailrec
     private[this] def advance(): Boolean = {
       if (queue.isEmpty) {
         current = null
@@ -178,20 +180,57 @@ object Iterator {
       else {
         current = queue.head()
         queue = queue.tail
-        current.hasNext || advance()
+        if (current.hasNext) {
+          currentHasNextChecked = true
+          true
+        } else advance()
       }
     }
-    def hasNext = (current ne null) && (current.hasNext || advance())
-    def next()  = if (hasNext) current.next else Iterator.empty.next
+    def hasNext =
+      if (currentHasNextChecked) true
+      else if (current eq null) false
+      else if (current.hasNext) {
+        currentHasNextChecked = true
+        true
+      } else advance()
+    def next()  =
+      if (hasNext) {
+        currentHasNextChecked = false
+        current.next()
+      } else Iterator.empty.next()
 
     override def ++[B >: A](that: => GenTraversableOnce[B]): Iterator[B] =
       new ConcatIterator(current, queue :+ (() => that.toIterator))
   }
 
   private[scala] final class JoinIterator[+A](lhs: Iterator[A], that: => GenTraversableOnce[A]) extends Iterator[A] {
+    private[this] var state = 0 // 0: lhs not checked, 1: lhs has next, 2: switched to rhs
     private[this] lazy val rhs: Iterator[A] = that.toIterator
-    def hasNext = lhs.hasNext || rhs.hasNext
-    def next    = if (lhs.hasNext) lhs.next else rhs.next
+    def hasNext = state match {
+      case 0 =>
+        if (lhs.hasNext) {
+          state = 1
+          true
+        } else {
+          state = 2
+          rhs.hasNext
+        }
+      case 1 => true
+      case _ => rhs.hasNext
+    }
+    def next() = state match {
+      case 0 =>
+        if (lhs.hasNext) lhs.next()
+        else {
+          state = 2
+          rhs.next()
+        }
+      case 1 =>
+        state = 0
+        lhs.next()
+      case _ =>
+        rhs.next()
+    }
 
     override def ++[B >: A](that: => GenTraversableOnce[B]) =
       new ConcatIterator(this, Vector(() => that.toIterator))
diff --git a/test/junit/scala/collection/IteratorTest.scala b/test/junit/scala/collection/IteratorTest.scala
index 1c1e50aed9..329c85127a 100644
--- a/test/junit/scala/collection/IteratorTest.scala
+++ b/test/junit/scala/collection/IteratorTest.scala
@@ -164,4 +164,32 @@ class IteratorTest {
     assertEquals(1, y.next)
     assertFalse(x.hasNext)   // was true, after advancing underlying iterator
   }
+  // SI-9623
+  @Test def noExcessiveHasNextInJoinIterator: Unit = {
+    var counter = 0
+    val exp = List(1,2,3,1,2,3)
+    def it: Iterator[Int] = new Iterator[Int] {
+      val parent = List(1,2,3).iterator
+      def next(): Int = parent.next
+      def hasNext: Boolean = { counter += 1; parent.hasNext }
+    }
+    // Iterate separately
+    val res = new mutable.ArrayBuffer[Int]
+    it.foreach(res += _)
+    it.foreach(res += _)
+    assertSameElements(exp, res)
+    assertEquals(8, counter)
+    // JoinIterator
+    counter = 0
+    res.clear
+    (it ++ it).foreach(res += _)
+    assertSameElements(exp, res)
+    assertEquals(8, counter) // was 17
+    // ConcatIterator
+    counter = 0
+    res.clear
+    (Iterator.empty ++ it ++ it).foreach(res += _)
+    assertSameElements(exp, res)
+    assertEquals(8, counter) // was 14
+  }
 }