KT-58588 Optimize Sequence::distinct

Use integer constants instead of enum to track AbstractIterator's state:
- for switches with a few cases, integer constants are
faster as there's no need to load ordinal value from
enum and map it into tableswitch index;
- on JVM, JIT-compiler can optimize away state-tracking code if it uses
integer constants, but when it comes to enums and tableswitches,
it can't apply a same set of optimizations to it.
This commit is contained in:
Filipp Zhinkin
2023-09-18 10:29:30 +02:00
committed by Space Team
parent 1062ca9646
commit 7ca506870a
@@ -5,12 +5,18 @@
package kotlin.collections package kotlin.collections
// Use integer constants instead of enum to improve performance:
private enum class State { // 1) switch over enum values require loading ordinal value and mapping it to a correct tableswitch offset
Ready, // using SwitchMap array, which is slightly slower compared to comparing integers (assuming that "switch" statement
NotReady, // has only a few cases);
Done, // 2) at least on JVM, usage of integer states may enable more optimizations; in particular, for the code looping over the iterator,
Failed // the JIT can fully eliminate state manipulation (assuming that iterator's methods were inlined and
// an iterator's allocation was eliminated) and use only a logic from a computeNext implementation to control iteration.
private object State {
const val NOT_READY: Int = 0
const val READY: Int = 1
const val DONE: Int = 2
const val FAILED: Int = 3
} }
/** /**
@@ -18,29 +24,36 @@ private enum class State {
* to implement the iterator, calling [done] when the iteration is complete. * to implement the iterator, calling [done] when the iteration is complete.
*/ */
public abstract class AbstractIterator<T> : Iterator<T> { public abstract class AbstractIterator<T> : Iterator<T> {
private var state = State.NotReady private var state = State.NOT_READY
private var nextValue: T? = null private var nextValue: T? = null
override fun hasNext(): Boolean { override fun hasNext(): Boolean {
require(state != State.Failed)
return when (state) { return when (state) {
State.Done -> false State.DONE -> false
State.Ready -> true State.READY -> true
else -> tryToComputeNext() State.NOT_READY -> tryToComputeNext()
else -> throw IllegalArgumentException("hasNext called when the iterator is in the FAILED state.")
} }
} }
override fun next(): T { override fun next(): T {
if (!hasNext()) throw NoSuchElementException() if (state == State.READY) {
state = State.NotReady state = State.NOT_READY
@Suppress("UNCHECKED_CAST")
return nextValue as T
}
if (state == State.DONE || !tryToComputeNext()) {
throw NoSuchElementException()
}
state = State.NOT_READY
@Suppress("UNCHECKED_CAST") @Suppress("UNCHECKED_CAST")
return nextValue as T return nextValue as T
} }
private fun tryToComputeNext(): Boolean { private fun tryToComputeNext(): Boolean {
state = State.Failed state = State.FAILED
computeNext() computeNext()
return state == State.Ready return state == State.READY
} }
/** /**
@@ -53,21 +66,21 @@ public abstract class AbstractIterator<T> : Iterator<T> {
* *
* Failure to call either method will result in the iteration terminating with a failed state * Failure to call either method will result in the iteration terminating with a failed state
*/ */
abstract protected fun computeNext(): Unit protected abstract fun computeNext(): Unit
/** /**
* Sets the next value in the iteration, called from the [computeNext] function * Sets the next value in the iteration, called from the [computeNext] function
*/ */
protected fun setNext(value: T): Unit { protected fun setNext(value: T): Unit {
nextValue = value nextValue = value
state = State.Ready state = State.READY
} }
/** /**
* Sets the state to done so that the iteration terminates. * Sets the state to done so that the iteration terminates.
*/ */
protected fun done() { protected fun done() {
state = State.Done state = State.DONE
} }
} }