Improve Streams implementation for performance and validity (based on JMH research)

2014-09-01 20:10:51 +04:00
parent 5a5c4145f2
commit c5c1ecf6c0
3 changed files with 275 additions and 118 deletions
@@ -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)
 }

 /**
@@ -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
+        }
+
    }
 }