From 9fa4ff750bf5bcbc92ea88fc252b94ad7deffed7 Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Tue, 2 Jun 2020 13:00:42 +0300
Subject: [PATCH] [FIR] Fix ArrayMapImpl.iterator()

---
 .../jetbrains/kotlin/fir/utils/ArrayMap.kt    | 33 ++++++-------------
 1 file changed, 10 insertions(+), 23 deletions(-)

diff --git a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/utils/ArrayMap.kt b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/utils/ArrayMap.kt
index 02f4650354d..d1f8ae0fdd0 100644
--- a/compiler/fir/tree/src/org/jetbrains/kotlin/fir/utils/ArrayMap.kt
+++ b/compiler/fir/tree/src/org/jetbrains/kotlin/fir/utils/ArrayMap.kt
@@ -95,33 +95,20 @@ internal class ArrayMapImpl<T : Any> : ArrayMap<T>() {
     }
 
     override fun iterator(): Iterator<T> {
-        return object : Iterator<T> {
-            private var currentIndex = -1
-            private var nextIndex = 0
+        return object : AbstractIterator<T>() {
+            private var index = -1
 
-            override fun hasNext(): Boolean {
-                if (nextIndex < 0) return false
-                while (nextIndex < data.size && data[nextIndex] == null) {
-                    nextIndex++
-                }
-                return if (nextIndex >= data.size) {
-                    nextIndex = -1
-                    false
+            override fun computeNext() {
+                do {
+                    index++
+                } while (index < data.size && data[index] == null)
+                if (index >= data.size) {
+                    done()
                 } else {
-                    true
+                    @Suppress("UNCHECKED_CAST")
+                    setNext(data[index] as T)
                 }
             }
-
-            override fun next(): T {
-                if (!hasNext()) throw NoSuchElementException()
-                if (currentIndex < 0) {
-                    currentIndex = nextIndex
-                }
-                @Suppress("UNCHECKED_CAST")
-                val result = data[currentIndex] as T
-                currentIndex = nextIndex++
-                return result
-            }
         }
     }