From e803977924c5e0eea043297d9b1736693e3cda4d Mon Sep 17 00:00:00 2001
From: Filipp Zhinkin <filipp.zhinkin@jetbrains.com>
Date: Wed, 20 Sep 2023 10:36:04 +0200
Subject: [PATCH] KT-58588 Optimize Sequence::flatten

Rewrite FlatteningSequence to use
the same optimization technique as AbstractIterator.
Didn't extend AbstractIterator as storing and reloading a value
from a field is slightly slower compared to the implemented approach.
---
 .../src/kotlin/collections/Sequences.kt       | 46 +++++++++++++------
 1 file changed, 32 insertions(+), 14 deletions(-)

diff --git a/libraries/stdlib/src/kotlin/collections/Sequences.kt b/libraries/stdlib/src/kotlin/collections/Sequences.kt
index 0b50d93818a..7cdcf47af62 100644
--- a/libraries/stdlib/src/kotlin/collections/Sequences.kt
+++ b/libraries/stdlib/src/kotlin/collections/Sequences.kt
@@ -289,37 +289,55 @@ constructor(
     private val transformer: (T) -> R,
     private val iterator: (R) -> Iterator<E>
 ) : Sequence<E> {
+    private object State {
+        const val UNDEFINED = 0
+        const val READY = 1
+        const val DONE = 2
+    }
+
     override fun iterator(): Iterator<E> = object : Iterator<E> {
         val iterator = sequence.iterator()
         var itemIterator: Iterator<E>? = null
 
+        // Use state to avoid excessive ensureItemIterator calls.
+        // The state is represented by the integer to avoid an overhead associated with the enum.
+        var state = State.UNDEFINED
+
         override fun next(): E {
-            if (!ensureItemIterator())
+            if (state == State.DONE) throw NoSuchElementException()
+            if (state == State.UNDEFINED && !ensureItemIterator()) {
                 throw NoSuchElementException()
+            }
+            state = State.UNDEFINED
             return itemIterator!!.next()
         }
 
         override fun hasNext(): Boolean {
+            if (state == State.READY) return true
+            if (state == State.DONE) return false
             return ensureItemIterator()
         }
 
         private fun ensureItemIterator(): Boolean {
-            if (itemIterator?.hasNext() == false)
-                itemIterator = null
+            val itemIterator = itemIterator
+            if (itemIterator != null && itemIterator.hasNext()) {
+                state = State.READY
+                return true
+            }
 
-            while (itemIterator == null) {
-                if (!iterator.hasNext()) {
-                    return false
-                } else {
-                    val element = iterator.next()
-                    val nextItemIterator = iterator(transformer(element))
-                    if (nextItemIterator.hasNext()) {
-                        itemIterator = nextItemIterator
-                        return true
-                    }
+            while (iterator.hasNext()) {
+                val element = iterator.next()
+                val nextItemIterator = iterator(transformer(element))
+                if (nextItemIterator.hasNext()) {
+                    this.itemIterator = nextItemIterator
+                    state = State.READY
+                    return true
                 }
             }
-            return true
+
+            state = State.DONE
+            this.itemIterator = null
+            return false
         }
     }
 }