From c5c1ecf6c0cf86ffa28a6d6ff0b2117c9c9ecd35 Mon Sep 17 00:00:00 2001
From: Ilya Ryzhenkov <orangy@jetbrains.com>
Date: Mon, 1 Sep 2014 20:10:51 +0400
Subject: [PATCH] Improve Streams implementation for performance and validity
 (based on JMH research)

---
 libraries/stdlib/src/generated/_Filtering.kt  |  19 +-
 .../stdlib/src/kotlin/collections/Stream.kt   | 355 +++++++++++++-----
 .../src/templates/Filtering.kt                |  19 +-
 3 files changed, 275 insertions(+), 118 deletions(-)

diff --git a/libraries/stdlib/src/generated/_Filtering.kt b/libraries/stdlib/src/generated/_Filtering.kt
index c90b3373d98..b7fd46ff924 100644
--- a/libraries/stdlib/src/generated/_Filtering.kt
+++ b/libraries/stdlib/src/generated/_Filtering.kt
@@ -163,8 +163,7 @@ public fun <T> Iterable<T>.drop(n: Int): List<T> {
  * Returns a stream containing all elements except first *n* elements
  */
 public fun <T> Stream<T>.drop(n: Int): Stream<T> {
-    var count = 0;
-    return FilteringStream(this) { count++ >= n }
+    return DropStream(this, n)
 }
 
 /**
@@ -338,16 +337,7 @@ public inline fun <T> Iterable<T>.dropWhile(predicate: (T) -> Boolean): List<T>
  * Returns a stream containing all elements except first elements that satisfy the given *predicate*
  */
 public fun <T> Stream<T>.dropWhile(predicate: (T) -> Boolean): Stream<T> {
-    var yielding = false
-    return FilteringStream(this) {
-        if (yielding)
-            true
-        else if (!predicate(it)) {
-            yielding = true
-            true
-        } else
-            false
-    }
+    return DropWhileStream(this, predicate)
 }
 
 /**
@@ -1058,8 +1048,7 @@ public fun <T> Iterable<T>.take(n: Int): List<T> {
  * Returns a stream containing first *n* elements
  */
 public fun <T> Stream<T>.take(n: Int): Stream<T> {
-    var count = 0
-    return LimitedStream(this) { count++ == n }
+    return TakeStream(this, n)
 }
 
 /**
@@ -1203,7 +1192,7 @@ public inline fun <T> Iterable<T>.takeWhile(predicate: (T) -> Boolean): List<T>
  * Returns a stream containing first elements satisfying the given *predicate*
  */
 public fun <T> Stream<T>.takeWhile(predicate: (T) -> Boolean): Stream<T> {
-    return LimitedStream(this, false, predicate)
+    return TakeWhileStream(this, predicate)
 }
 
 /**
diff --git a/libraries/stdlib/src/kotlin/collections/Stream.kt b/libraries/stdlib/src/kotlin/collections/Stream.kt
index 5ce7c097d01..fbeecbe32a0 100644
--- a/libraries/stdlib/src/kotlin/collections/Stream.kt
+++ b/libraries/stdlib/src/kotlin/collections/Stream.kt
@@ -1,6 +1,5 @@
 package kotlin
 
-import kotlin.support.AbstractIterator
 import java.util.*
 
 public trait Stream<out T> {
@@ -9,144 +8,324 @@ public trait Stream<out T> {
 
 public fun <T> streamOf(vararg elements: T): Stream<T> = elements.stream()
 
-public class FilteringStream<T>(
-        private val stream: Stream<T>, private val sendWhen: Boolean = true, private val predicate: (T) -> Boolean
-) : Stream<T> {
-    override fun iterator(): Iterator<T> = object : AbstractIterator<T>() {
-        val iterator = stream.iterator()
-        override fun computeNext() {
+public fun <T> streamOf(progression: Progression<T>): Stream<T> = object : Stream<T> {
+    override fun iterator(): Iterator<T> = progression.iterator()
+}
+
+public class FilteringStream<T>(private val stream: Stream<T>,
+                                private val sendWhen: Boolean = true,
+                                private val predicate: (T) -> Boolean
+                               ) : Stream<T> {
+
+    override fun iterator(): Iterator<T> = object : Iterator<T> {
+        val iterator = stream.iterator();
+        var nextState: Int = -1 // -1 for unknown, 0 for done, 1 for continue
+        var nextItem: T? = null
+
+        private fun calcNext() {
             while (iterator.hasNext()) {
                 val item = iterator.next()
                 if (predicate(item) == sendWhen) {
-                    setNext(item)
+                    nextItem = item
+                    nextState = 1
                     return
                 }
             }
-            done()
+            nextState = 0
+        }
+
+        override fun next(): T {
+            if (nextState == -1)
+                calcNext()
+            if (nextState == 0)
+                throw NoSuchElementException()
+            val result = nextItem
+            nextItem = null
+            nextState = -1
+            return result as T
+        }
+
+        override fun hasNext(): Boolean {
+            if (nextState == -1)
+                calcNext()
+            return nextState == 1
         }
     }
 }
 
 public class TransformingStream<T, R>(private val stream: Stream<T>, private val transformer: (T) -> R) : Stream<R> {
-    override fun iterator(): Iterator<R> = object : AbstractIterator<R>() {
+    override fun iterator(): Iterator<R> = object : Iterator<R> {
         val iterator = stream.iterator()
-        override fun computeNext() {
-            if (iterator.hasNext()) {
-                setNext(transformer(iterator.next()))
-            } else {
-                done()
-            }
+        override fun next(): R {
+            return transformer(iterator.next())
+        }
+        override fun hasNext(): Boolean {
+            return iterator.hasNext()
         }
     }
 }
 
-public class MergingStream<T1, T2, V>(
-        private val stream1: Stream<T1>, private val stream2: Stream<T2>, private val transform: (T1, T2) -> V
-) : Stream<V> {
-    override fun iterator(): Iterator<V> = object : AbstractIterator<V>() {
+public class MergingStream<T1, T2, V>(private val stream1: Stream<T1>,
+                                      private val stream2: Stream<T2>,
+                                      private val transform: (T1, T2) -> V
+                                     ) : Stream<V> {
+    override fun iterator(): Iterator<V> = object : Iterator<V> {
         val iterator1 = stream1.iterator()
         val iterator2 = stream2.iterator()
-        override fun computeNext() {
-            if (iterator1.hasNext() && iterator2.hasNext()) {
-                setNext(transform(iterator1.next(), iterator2.next()))
-            } else {
-                done()
-            }
+        override fun next(): V {
+            return transform(iterator1.next(), iterator2.next())
+        }
+        override fun hasNext(): Boolean {
+            return iterator1.hasNext() && iterator2.hasNext()
         }
     }
 }
 
-public class FlatteningStream<T, R>(
-        private val stream: Stream<T>, private val transformer: (T) -> Stream<R>
-) : Stream<R> {
-    override fun iterator(): Iterator<R> = object : AbstractIterator<R>() {
+public class FlatteningStream<T, R>(private val stream: Stream<T>,
+                                    private val transformer: (T) -> Stream<R>
+                                   ) : Stream<R> {
+    override fun iterator(): Iterator<R> = object : Iterator<R> {
         val iterator = stream.iterator()
         var itemIterator: Iterator<R>? = null
-        override fun computeNext() {
+
+        override fun next(): R {
+            if (!ensureItemIterator())
+                throw NoSuchElementException()
+            return itemIterator!!.next()
+        }
+
+        override fun hasNext(): Boolean {
+            return ensureItemIterator()
+        }
+
+        private fun ensureItemIterator(): Boolean {
+            if (itemIterator?.hasNext() == false)
+                itemIterator = null
+
             while (itemIterator == null) {
                 if (!iterator.hasNext()) {
-                    done()
-                    break;
+                    return false
                 } else {
                     val element = iterator.next()
                     val nextItemIterator = transformer(element).iterator()
-                    if (nextItemIterator.hasNext())
+                    if (nextItemIterator.hasNext()) {
                         itemIterator = nextItemIterator
+                        return true
+                    }
                 }
             }
-
-            val currentItemIterator = itemIterator
-            if (currentItemIterator == null) {
-                done()
-            } else {
-                setNext(currentItemIterator.next())
-                if (!currentItemIterator.hasNext())
-                    itemIterator = null
-            }
+            return true
         }
     }
 }
 
-public class Multistream<T>(private val streams: Stream<Stream<T>>) : Stream<T> {
-    override fun iterator(): Iterator<T> = object : AbstractIterator<T>() {
-        val iterator = streams.iterator()
-        var streamIterator: Iterator<T>? = null
-        override fun computeNext() {
-            while (streamIterator == null) {
-                if (!iterator.hasNext()) {
-                    done()
-                    break;
-                } else {
-                    val stream = iterator.next()
-                    val nextStreamIterator = stream.iterator()
-                    if (nextStreamIterator.hasNext())
-                        streamIterator = nextStreamIterator
-                }
-            }
-
-            val currentStreamIterator = streamIterator
-            if (currentStreamIterator == null) {
-                done()
-            } else {
-                setNext(currentStreamIterator.next())
-                if (!currentStreamIterator.hasNext())
-                    streamIterator = null
-            }
-        }
-    }
-}
-
-public class LimitedStream<T>(
-        private val stream: Stream<T>, private val stopWhen: Boolean = true, private val predicate: (T) -> Boolean
-) : Stream<T> {
-    override fun iterator(): Iterator<T> = object : AbstractIterator<T>() {
+public class Multistream<T>(private val stream: Stream<Stream<T>>) : Stream<T> {
+    override fun iterator(): Iterator<T> = object : Iterator<T> {
         val iterator = stream.iterator()
-        override fun computeNext() {
-            if (!iterator.hasNext()) {
-                done()
-            } else {
-                val item = iterator.next()
-                if (predicate(item) == stopWhen) {
-                    done()
+        var itemIterator: Iterator<T>? = null
+
+        override fun next(): T {
+            if (!ensureItemIterator())
+                throw NoSuchElementException()
+            return itemIterator!!.next()
+        }
+
+        override fun hasNext(): Boolean {
+            return ensureItemIterator()
+        }
+
+        private fun ensureItemIterator(): Boolean {
+            if (itemIterator?.hasNext() == false)
+                itemIterator = null
+
+            while (itemIterator == null) {
+                if (!iterator.hasNext()) {
+                    return false
                 } else {
-                    setNext(item)
+                    val element = iterator.next()
+                    val nextItemIterator = element.iterator()
+                    if (nextItemIterator.hasNext()) {
+                        itemIterator = nextItemIterator
+                        return true
+                    }
                 }
             }
+            return true
+        }
+    }
+}
+
+public class TakeStream<T>(private val stream: Stream<T>,
+                           private var count: Int
+                          ) : Stream<T> {
+    {
+        if (count < 0)
+            throw IllegalArgumentException("count should be non-negative, but is $count")
+    }
+
+    override fun iterator(): Iterator<T> = object : Iterator<T> {
+        val iterator = stream.iterator();
+
+        override fun next(): T {
+            if (count == 0)
+                throw NoSuchElementException()
+            count--
+            return iterator.next()
+        }
+
+        override fun hasNext(): Boolean {
+            return count > 0 && iterator.hasNext()
+        }
+    }
+}
+
+public class TakeWhileStream<T>(private val stream: Stream<T>,
+                                private val predicate: (T) -> Boolean
+                               ) : Stream<T> {
+    override fun iterator(): Iterator<T> = object : Iterator<T> {
+        val iterator = stream.iterator();
+        var nextState: Int = -1 // -1 for unknown, 0 for done, 1 for continue
+        var nextItem: T? = null
+
+        private fun calcNext() {
+            if (iterator.hasNext()) {
+                val item = iterator.next()
+                if (predicate(item)) {
+                    nextState = 1
+                    nextItem = item
+                    return
+                }
+            }
+            nextState = 0
+        }
+
+        override fun next(): T {
+            if (nextState == -1)
+                calcNext() // will change nextState
+            if (nextState == 0)
+                throw NoSuchElementException()
+            val result = nextItem as T
+
+            // Clean next to avoid keeping reference on yielded instance
+            nextItem = null
+            nextState = -1
+            return result
+        }
+
+        override fun hasNext(): Boolean {
+            if (nextState == -1)
+                calcNext() // will change nextState
+            return nextState == 1
+        }
+    }
+}
+
+public class DropStream<T>(private val stream: Stream<T>,
+                           private var count: Int
+                          ) : Stream<T> {
+    {
+        if (count < 0)
+            throw IllegalArgumentException("count should be non-negative, but is $count")
+    }
+
+    override fun iterator(): Iterator<T> = object : Iterator<T> {
+        val iterator = stream.iterator();
+
+        // Shouldn't be called from constructor to avoid premature iteration
+        private fun drop() {
+            while (count > 0 && iterator.hasNext()) {
+                iterator.next()
+                count--
+            }
+        }
+
+        override fun next(): T {
+            drop()
+            return iterator.next()
+        }
+
+        override fun hasNext(): Boolean {
+            drop()
+            return iterator.hasNext()
+        }
+    }
+}
+
+public class DropWhileStream<T>(private val stream: Stream<T>,
+                                private val predicate: (T) -> Boolean
+                               ) : Stream<T> {
+
+    override fun iterator(): Iterator<T> = object : Iterator<T> {
+        val iterator = stream.iterator();
+        var dropState: Int = -1 // -1 for not dropping, 1 for nextItem, 0 for normal iteration
+        var nextItem: T? = null
+
+        private fun drop() {
+            while (iterator.hasNext()) {
+                val item = iterator.next()
+                if (!predicate(item)) {
+                    nextItem = item
+                    dropState = 1
+                    return
+                }
+            }
+            dropState = 0
+        }
+
+        override fun next(): T {
+            if (dropState == -1)
+                drop()
+
+            if (dropState == 1) {
+                val result = nextItem as T
+                nextItem = null
+                dropState = 0
+                return result
+            }
+            return iterator.next()
+        }
+
+        override fun hasNext(): Boolean {
+            if (dropState == -1)
+                drop()
+            return dropState == 1 || iterator.hasNext()
         }
     }
 }
 
 public class FunctionStream<T : Any>(private val producer: () -> T?) : Stream<T> {
-    override fun iterator(): Iterator<T> = object : AbstractIterator<T>() {
+    override fun iterator(): Iterator<T> = object : Iterator<T> {
+        var nextState: Int = -1 // -1 for unknown, 0 for done, 1 for continue
+        var nextItem: T? = null
 
-        override fun computeNext() {
+        private fun calcNext() {
             val item = producer()
             if (item == null) {
-                done()
+                nextState = 0
             } else {
-                setNext(item)
+                nextState = 1
+                nextItem = item
             }
         }
+
+        override fun next(): T {
+            if (nextState == -1)
+                calcNext()
+            if (nextState == 0)
+                throw NoSuchElementException()
+            val result = nextItem as T
+            // Clean next to avoid keeping reference on yielded instance
+            nextItem = null
+            nextState = -1
+            return result
+        }
+
+        override fun hasNext(): Boolean {
+            if (nextState == -1)
+                calcNext()
+            return nextState == 1
+        }
+
     }
 }
 
diff --git a/libraries/tools/kotlin-stdlib-gen/src/templates/Filtering.kt b/libraries/tools/kotlin-stdlib-gen/src/templates/Filtering.kt
index 4e041edd93c..61e65b31e11 100644
--- a/libraries/tools/kotlin-stdlib-gen/src/templates/Filtering.kt
+++ b/libraries/tools/kotlin-stdlib-gen/src/templates/Filtering.kt
@@ -23,8 +23,7 @@ fun filtering(): List<GenericFunction> {
         returns(Streams) { "Stream<T>" }
         body(Streams) {
             """
-            var count = 0;
-            return FilteringStream(this) { count++ >= n }
+            return DropStream(this, n)
             """
         }
 
@@ -69,8 +68,7 @@ fun filtering(): List<GenericFunction> {
         returns(Streams) { "Stream<T>" }
         body(Streams) {
             """
-            var count = 0
-            return LimitedStream(this) { count++ == n }
+            return TakeStream(this, n)
             """
         }
 
@@ -126,16 +124,7 @@ fun filtering(): List<GenericFunction> {
         returns(Streams) { "Stream<T>" }
         body(Streams) {
             """
-            var yielding = false
-            return FilteringStream(this) {
-                if (yielding)
-                    true
-                else if (!predicate(it)) {
-                    yielding = true
-                    true
-                } else
-                    false
-            }
+            return DropWhileStream(this, predicate)
             """
         }
 
@@ -174,7 +163,7 @@ fun filtering(): List<GenericFunction> {
         returns(Streams) { "Stream<T>" }
         body(Streams) {
             """
-            return LimitedStream(this, false, predicate)
+            return TakeWhileStream(this, predicate)
             """
         }
     }