From 846db641be9e661c60d98ae7fb3af10af96b6a2d Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Mon, 18 May 2020 15:23:17 +0300
Subject: [PATCH] [FIR] Make ArrayMap iterable

---
 .../jetbrains/kotlin/fir/utils/ArrayMap.kt    | 60 ++++++++++++++++++-
 1 file changed, 59 insertions(+), 1 deletion(-)

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 5e4488fef84..02f4650354d 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
@@ -5,7 +5,7 @@
 
 package org.jetbrains.kotlin.fir.utils
 
-sealed class ArrayMap<T : Any> {
+sealed class ArrayMap<T : Any> : Iterable<T> {
     abstract val size: Int
 
     abstract operator fun set(index: Int, value: T)
@@ -23,6 +23,14 @@ internal object EmptyArrayMap : ArrayMap<Nothing>() {
     override fun get(index: Int): Nothing? {
         return null
     }
+
+    override fun iterator(): Iterator<Nothing> {
+        return object : Iterator<Nothing> {
+            override fun hasNext(): Boolean = false
+
+            override fun next(): Nothing = throw NoSuchElementException()
+        }
+    }
 }
 
 internal class OneElementArrayMap<T : Any>(val value: T, val index: Int) : ArrayMap<T>() {
@@ -36,6 +44,25 @@ internal class OneElementArrayMap<T : Any>(val value: T, val index: Int) : Array
     override fun get(index: Int): T? {
         return if (index == this.index) value else null
     }
+
+    override fun iterator(): Iterator<T> {
+        return object : Iterator<T> {
+            private var notVisited = true
+
+            override fun hasNext(): Boolean {
+                return notVisited
+            }
+
+            override fun next(): T {
+                if (notVisited) {
+                    notVisited = false
+                    return value
+                } else {
+                    throw NoSuchElementException()
+                }
+            }
+        }
+    }
 }
 
 internal class ArrayMapImpl<T : Any> : ArrayMap<T>() {
@@ -67,6 +94,37 @@ internal class ArrayMapImpl<T : Any> : ArrayMap<T>() {
         return data.getOrNull(index) as T?
     }
 
+    override fun iterator(): Iterator<T> {
+        return object : Iterator<T> {
+            private var currentIndex = -1
+            private var nextIndex = 0
+
+            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
+                } else {
+                    true
+                }
+            }
+
+            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
+            }
+        }
+    }
+
     fun remove(index: Int) {
         if (data[index] != null) {
             size--