Merge pull request #4086 from som-snytt/issue/8835-2.12.x

SI-8835 Lazier slice for Iterator

3 files changed, 72 insertions, 22 deletions

diff --git a/src/library/scala/collection/Iterator.scala b/src/library/scala/collection/Iterator.scala
index 660cc5a42a..adb97e27e3 100644
--- a/src/library/scala/collection/Iterator.scala
+++ b/src/library/scala/collection/Iterator.scala
@@ -182,7 +182,7 @@ object Iterator {
       }
     }
     def hasNext = (current ne null) && (current.hasNext || advance())
-    def next()  = if (hasNext) current.next else Iterator.empty.next
+    def next()  = if (hasNext) current.next() else Iterator.empty.next()
 
     override def ++[B >: A](that: => GenTraversableOnce[B]): Iterator[B] =
       new ConcatIterator(current, queue :+ (() => that.toIterator))
@@ -191,11 +191,55 @@ object Iterator {
   private[scala] final class JoinIterator[+A](lhs: Iterator[A], that: => GenTraversableOnce[A]) extends Iterator[A] {
     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 next()  = if (lhs.hasNext) lhs.next() else rhs.next()
 
     override def ++[B >: A](that: => GenTraversableOnce[B]) =
       new ConcatIterator(this, Vector(() => that.toIterator))
   }
+
+  /** Creates a delegating iterator capped by a limit count. Negative limit means unbounded.
+   *  Lazily skip to start on first evaluation.  Avoids daisy-chained iterators due to slicing.
+   */
+  private[scala] final class SliceIterator[A](val underlying: Iterator[A], start: Int, limit: Int) extends AbstractIterator[A] {
+    private var remaining = limit
+    private var dropping  = start
+    @inline private def unbounded = remaining < 0
+    private def skip(): Unit =
+      while (dropping > 0) {
+        if (underlying.hasNext) {
+          underlying.next()
+          dropping -= 1
+        } else
+          dropping = 0
+      }
+    def hasNext = { skip(); remaining != 0 && underlying.hasNext }
+    def next()  = {
+      skip()
+      if (remaining > 0) {
+        remaining -= 1
+        underlying.next()
+      }
+      else if (unbounded) underlying.next()
+      else empty.next()
+    }
+    override protected def sliceIterator(from: Int, until: Int): Iterator[A] = {
+      val lo = from max 0
+      def adjustedBound =
+        if (unbounded) -1
+        else 0 max (remaining - lo)
+      val rest =
+        if (until < 0) adjustedBound          // respect current bound, if any
+        else if (until <= lo) 0               // empty
+        else if (unbounded) until - lo        // now finite
+        else adjustedBound min (until - lo)   // keep lesser bound
+      if (rest == 0) empty
+      else {
+        dropping += lo
+        remaining = rest
+        this
+      }
+    }
+  }
 }
 
 import Iterator.empty
@@ -307,11 +351,11 @@ trait Iterator[+A] extends TraversableOnce[A] {
   /** Selects first ''n'' values of this iterator.
    *
    *  @param  n    the number of values to take
-   *  @return an iterator producing only of the first `n` values of this iterator, or else the
+   *  @return an iterator producing only the first `n` values of this iterator, or else the
    *          whole iterator, if it produces fewer than `n` values.
    *  @note   Reuse: $consumesAndProducesIterator
    */
-  def take(n: Int): Iterator[A] = slice(0, n)
+  def take(n: Int): Iterator[A] = sliceIterator(0, n max 0)
 
   /** Advances this iterator past the first ''n'' elements, or the length of the iterator, whichever is smaller.
    *
@@ -320,34 +364,29 @@ trait Iterator[+A] extends TraversableOnce[A] {
    *           it omits the first `n` values.
    *  @note    Reuse: $consumesAndProducesIterator
    */
-  def drop(n: Int): Iterator[A] = slice(n, Int.MaxValue)
+  def drop(n: Int): Iterator[A] = sliceIterator(n, -1)
 
   /** Creates an iterator returning an interval of the values produced by this iterator.
    *
    *  @param from   the index of the first element in this iterator which forms part of the slice.
-   *  @param until  the index of the first element following the slice.
+   *                If negative, the slice starts at zero.
+   *  @param until  the index of the first element following the slice. If negative, the slice is empty.
    *  @return an iterator which advances this iterator past the first `from` elements using `drop`,
    *  and then takes `until - from` elements, using `take`.
    *  @note         Reuse: $consumesAndProducesIterator
    */
-  def slice(from: Int, until: Int): Iterator[A] = {
+  def slice(from: Int, until: Int): Iterator[A] = sliceIterator(from, until max 0)
+
+  /** Creates an optionally bounded slice, unbounded if `until` is negative. */
+  protected def sliceIterator(from: Int, until: Int): Iterator[A] = {
     val lo = from max 0
-    var toDrop = lo
-    while (toDrop > 0 && self.hasNext) {
-      self.next()
-      toDrop -= 1
-    }
+    val rest =
+      if (until < 0) -1            // unbounded
+      else if (until <= lo) 0      // empty
+      else until - lo              // finite
 
-    new AbstractIterator[A] {
-      private var remaining = until - lo
-      def hasNext = remaining > 0 && self.hasNext
-      def next(): A =
-        if (remaining > 0) {
-          remaining -= 1
-          self.next()
-        }
-        else empty.next()
-    }
+    if (rest == 0) empty
+    else new Iterator.SliceIterator(this, lo, rest)
   }
 
   /** Creates a new iterator that maps all produced values of this iterator
diff --git a/src/library/scala/collection/parallel/RemainsIterator.scala b/src/library/scala/collection/parallel/RemainsIterator.scala
index 5f2ceac0e0..7bb278b038 100644
--- a/src/library/scala/collection/parallel/RemainsIterator.scala
+++ b/src/library/scala/collection/parallel/RemainsIterator.scala
@@ -456,6 +456,15 @@ self =>
     }
     it
   }
+  /** Drop implemented as simple eager consumption. */
+  override def drop(n: Int): IterableSplitter[T] = {
+    var i = 0
+    while (i < n && hasNext) {
+      next()
+      i += 1
+    }
+    this
+  }
   override def take(n: Int): IterableSplitter[T] = newTaken(n)
   override def slice(from1: Int, until1: Int): IterableSplitter[T] = newSliceInternal(newTaken(until1), from1)
 
diff --git a/src/library/scala/runtime/ScalaRunTime.scala b/src/library/scala/runtime/ScalaRunTime.scala
index f50059ce54..6c69ebae9b 100644
--- a/src/library/scala/runtime/ScalaRunTime.scala
+++ b/src/library/scala/runtime/ScalaRunTime.scala
@@ -13,6 +13,7 @@ import scala.collection.{ Seq, IndexedSeq, TraversableView, AbstractIterator }
 import scala.collection.mutable.WrappedArray
 import scala.collection.immutable.{ StringLike, NumericRange, List, Stream, Nil, :: }
 import scala.collection.generic.{ Sorted, IsTraversableLike }
+import scala.collection.parallel.ParIterable
 import scala.reflect.{ ClassTag, classTag }
 import scala.util.control.ControlThrowable
 import java.lang.{ Class => jClass }
@@ -326,6 +327,7 @@ object ScalaRunTime {
       case x: AnyRef if isArray(x)      => arrayToString(x)
       case x: scala.collection.Map[_, _]      => x.iterator take maxElements map mapInner mkString (x.stringPrefix + "(", ", ", ")")
       case x: Iterable[_]               => x.iterator take maxElements map inner mkString (x.stringPrefix + "(", ", ", ")")
+      case x: ParIterable[_]            => x.iterator take maxElements map inner mkString (x.stringPrefix + "(", ", ", ")")
       case x: Traversable[_]            => x take maxElements map inner mkString (x.stringPrefix + "(", ", ", ")")
       case x: Product1[_] if isTuple(x) => "(" + inner(x._1) + ",)" // that special trailing comma
       case x: Product if isTuple(x)     => x.productIterator map inner mkString ("(", ",", ")")