[FIR] Introduce different implementations of ArrayMapOwner

- ComponentArrayOwner for services
- AttributeArrayOwner for attributes like userdata or type attributes

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/extensions/FirExtensionsService.kt

@@ -38,7 +38,7 @@ class FirRegisteredExtension<P : FirExtension>(
 class FirExtensionsService(
     val session: FirSession
 ) : ComponentArrayOwner<FirExtension, FirRegisteredExtension<*>>(), FirSessionComponent {
-    companion object : ComponentTypeRegistry<FirExtension, FirRegisteredExtension<*>>() {
+    companion object : TypeRegistry<FirExtension, FirRegisteredExtension<*>>() {
         inline fun <reified P : FirExtension, V : FirRegisteredExtension<P>> registeredExtensions(): ReadOnlyProperty<FirExtensionsService, ExtensionsAccessor<P>> {
             val accessor = generateAccessor<V, P>(P::class)
             return object : ReadOnlyProperty<FirExtensionsService, ExtensionsAccessor<P>> {
@@ -51,7 +51,7 @@ class FirExtensionsService(
         private fun <K, V> createMultimap(): Multimap<K, V> = LinkedHashMultimap.create()
     }
 
-    override val typeRegistry: ComponentTypeRegistry<FirExtension, FirRegisteredExtension<*>>
+    override val typeRegistry: TypeRegistry<FirExtension, FirRegisteredExtension<*>>
         get() = Companion
 
     fun <P : FirExtension> registerExtensions(extensionClass: KClass<P>, extensionFactories: List<FirExtension.Factory<P>>) {

compiler/fir/tree/src/org/jetbrains/kotlin/fir/FirSession.kt

@@ -8,16 +8,16 @@ package org.jetbrains.kotlin.fir
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.analyzer.ModuleInfo
 import org.jetbrains.kotlin.fir.types.impl.*
-import org.jetbrains.kotlin.fir.utils.ComponentArrayAccessor
+import org.jetbrains.kotlin.fir.utils.ArrayMapAccessor
 import org.jetbrains.kotlin.fir.utils.ComponentArrayOwner
-import org.jetbrains.kotlin.fir.utils.ComponentTypeRegistry
+import org.jetbrains.kotlin.fir.utils.TypeRegistry
 import org.jetbrains.kotlin.utils.Jsr305State
 
 interface FirSessionComponent
 
 abstract class FirSession(val sessionProvider: FirSessionProvider?) : ComponentArrayOwner<FirSessionComponent, FirSessionComponent>() {
-    companion object : ComponentTypeRegistry<FirSessionComponent, FirSessionComponent>() {
+    companion object : TypeRegistry<FirSessionComponent, FirSessionComponent>() {
-        inline fun <reified T : FirSessionComponent> sessionComponentAccessor(): ComponentArrayAccessor<FirSessionComponent, FirSessionComponent, T> {
+        inline fun <reified T : FirSessionComponent> sessionComponentAccessor(): ArrayMapAccessor<FirSessionComponent, FirSessionComponent, T> {
             return generateAccessor(T::class)
         }
     }
@@ -28,7 +28,7 @@ abstract class FirSession(val sessionProvider: FirSessionProvider?) : ComponentA
 
     val builtinTypes: BuiltinTypes = BuiltinTypes()
 
-    final override val typeRegistry: ComponentTypeRegistry<FirSessionComponent, FirSessionComponent> = Companion
+    final override val typeRegistry: TypeRegistry<FirSessionComponent, FirSessionComponent> = Companion
 }
 
 interface FirSessionProvider {

compiler/fir/tree/src/org/jetbrains/kotlin/fir/utils/ArrayMap.kt

@@ -0,0 +1,68 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.utils
+
+sealed class ArrayMap<T : Any> {
+    abstract val size: Int
+
+    abstract operator fun set(index: Int, value: T)
+    abstract operator fun get(index: Int): T?
+}
+
+internal object EmptyArrayMap : ArrayMap<Nothing>() {
+    override val size: Int
+        get() = 0
+
+    override fun set(index: Int, value: Nothing) {
+        throw IllegalStateException()
+    }
+
+    override fun get(index: Int): Nothing? {
+        return null
+    }
+}
+
+internal class OneElementArrayMap<T : Any>(val value: T, val index: Int) : ArrayMap<T>() {
+    override val size: Int
+        get() = 1
+
+    override fun set(index: Int, value: T) {
+        throw IllegalStateException()
+    }
+
+    override fun get(index: Int): T? {
+        return if (index == this.index) value else null
+    }
+}
+
+internal  class ArrayMapImpl<T : Any> : ArrayMap<T>() {
+    companion object {
+        private const val DEFAULT_SIZE = 20
+        private const val INCREASE_K = 2
+    }
+
+    override var size: Int = 0
+        private set
+
+    private var data = arrayOfNulls<Any>(DEFAULT_SIZE)
+    private fun ensureCapacity(index: Int) {
+        if (data.size < index) {
+            data = data.copyOf(data.size * INCREASE_K)
+        }
+    }
+
+    override operator fun set(index: Int, value: T) {
+        ensureCapacity(index)
+        data[index] = value
+
+        size++
+    }
+
+    override operator fun get(index: Int): T? {
+        @Suppress("UNCHECKED_CAST")
+        return data.getOrNull(index) as T?
+    }
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/utils/ArrayMapOwner.kt

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.utils
+
+import java.util.*
+import kotlin.properties.ReadOnlyProperty
+import kotlin.reflect.KClass
+import kotlin.reflect.KProperty
+
+abstract class AbstractArrayMapOwner<K : Any, V : Any> {
+    protected abstract val arrayMap: ArrayMap<V>
+    protected abstract val typeRegistry: TypeRegistry<K, V>
+
+    abstract class AbstractArrayMapAccessor<K : Any, V : Any, T : V>(
+        protected val key: KClass<out K>,
+        protected val id: Int
+    ) {
+        protected fun extractValue(thisRef: AbstractArrayMapOwner<K, V>): T? {
+            @Suppress("UNCHECKED_CAST")
+            return thisRef.arrayMap[id] as T?
+        }
+    }
+
+    protected abstract fun registerComponent(tClass: KClass<out K>, value: V)
+}
+
+class ArrayMapAccessor<K : Any, V : Any, T : V>(
+    key: KClass<out K>,
+    id: Int
+) : AbstractArrayMapOwner.AbstractArrayMapAccessor<K, V, T>(key, id), ReadOnlyProperty<AbstractArrayMapOwner<K, V>, V> {
+    override fun getValue(thisRef: AbstractArrayMapOwner<K, V>, property: KProperty<*>): T {
+        return extractValue(thisRef) ?: error("No '$key'($id) in array owner: $thisRef")
+    }
+}
+
+class NullableArrayMapAccessor<K : Any, V : Any, T : V>(
+    key: KClass<out K>,
+    id: Int
+) : AbstractArrayMapOwner.AbstractArrayMapAccessor<K, V, T>(key, id), ReadOnlyProperty<AbstractArrayMapOwner<K, V>, V?> {
+    override fun getValue(thisRef: AbstractArrayMapOwner<K, V>, property: KProperty<*>): T? {
+        return extractValue(thisRef)
+    }
+}
+
+abstract class TypeRegistry<K : Any, V : Any> {
+    private val idPerType = HashMap<KClass<out K>, Int>()
+
+    fun <T : V, KK : K> generateAccessor(kClass: KClass<KK>): ArrayMapAccessor<K, V, T> {
+        return ArrayMapAccessor(kClass, getId(kClass))
+    }
+
+    fun <T : V, KK : K> generateNullableAccessor(kClass: KClass<KK>): NullableArrayMapAccessor<K, V, T> {
+        return NullableArrayMapAccessor(kClass, getId(kClass))
+    }
+
+    fun <T : K> getId(kClass: KClass<T>): Int {
+        return idPerType.getOrPut(kClass) { idPerType.size }
+    }
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/utils/AttributeArrayOwner.kt

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.utils
+
+import kotlin.reflect.KClass
+
+/**
+ * [AttributeArrayOwner] based on different implementations of [ArrayMap] and switches them
+ *   depending on array map fullness
+ * [AttributeArrayOwner] can be used in classes with many instances,
+ *   like user data for Fir elements or attributes for cone types
+ */
+abstract class AttributeArrayOwner<K : Any, T : Any> : AbstractArrayMapOwner<K, T>() {
+    @Suppress("UNCHECKED_CAST")
+    final override var arrayMap: ArrayMap<T> = EmptyArrayMap as ArrayMap<T>
+        private set
+
+    final override fun registerComponent(tClass: KClass<out K>, value: T) {
+        val id = typeRegistry.getId(tClass)
+        when (arrayMap.size) {
+            0 -> {
+                arrayMap = OneElementArrayMap(value, id)
+                return
+            }
+
+            1 -> {
+                arrayMap = ArrayMapImpl<T>().apply {
+                    val map = arrayMap as OneElementArrayMap<T>
+                    this[map.index] = map.value
+                }
+            }
+        }
+
+        arrayMap[id] = value
+    }
+}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/utils/ComponentArray.kt

@@ -1,71 +0,0 @@
-/*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
- * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.fir.utils
-
-import kotlin.properties.ReadOnlyProperty
-import kotlin.reflect.KClass
-import kotlin.reflect.KProperty
-
-abstract class ComponentArrayOwner<K : Any, V : Any> {
-    internal val componentArray: ComponentArray<V> = ComponentArray()
-    protected abstract val typeRegistry: ComponentTypeRegistry<K, V>
-
-    protected fun registerComponent(tClass: KClass<out K>, value: V) {
-        componentArray[(typeRegistry.getId(tClass))] = value
-    }
-
-    protected operator fun get(tClass: KClass<out K>): V {
-        return componentArray[typeRegistry.getId(tClass)]
-    }
-}
-
-
-abstract class ComponentTypeRegistry<K : Any, V : Any> {
-    private val idPerType = mutableMapOf<KClass<out K>, Int>()
-
-    fun <T : V, KK : K> generateAccessor(kClass: KClass<KK>): ComponentArrayAccessor<K, V, T> {
-        return ComponentArrayAccessor(kClass, getId(kClass))
-    }
-
-    fun <T : K> getId(kClass: KClass<T>): Int {
-        return idPerType.getOrPut(kClass) { idPerType.size }
-    }
-}
-
-
-class ComponentArrayAccessor<K : Any, V : Any, T : V>(
-    private val key: KClass<out K>,
-    private val id: Int
-) : ReadOnlyProperty<ComponentArrayOwner<K, V>, V> {
-    override fun getValue(thisRef: ComponentArrayOwner<K, V>, property: KProperty<*>): T {
-        @Suppress("UNCHECKED_CAST")
-        return thisRef.componentArray[id] as T? ?: error("No '$key'($id) component in session: $thisRef")
-    }
-}
-
-class ComponentArray<T : Any> {
-    companion object {
-        private const val DEFAULT_SIZE = 20
-        private const val INCREASE_K = 2
-    }
-
-    private var data = arrayOfNulls<Any>(DEFAULT_SIZE)
-    private fun ensureCapacity(index: Int) {
-        if (data.size < index) {
-            data = data.copyOf(data.size * INCREASE_K)
-        }
-    }
-
-    operator fun set(index: Int, value: T) {
-        ensureCapacity(index)
-        data[index] = value
-    }
-
-    operator fun get(index: Int): T {
-        @Suppress("UNCHECKED_CAST")
-        return data.getOrNull(index) as T
-    }
-}

compiler/fir/tree/src/org/jetbrains/kotlin/fir/utils/ComponentArrayOwner.kt

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.utils
+
+import kotlin.reflect.KClass
+
+/**
+ * [ComponentArrayOwner] based on [ArrayMap] with flexible size and should be used for
+ *   storing services in entities with limited number of instances, like FirSession
+ */
+abstract class ComponentArrayOwner<K : Any, V : Any> : AbstractArrayMapOwner<K, V>() {
+    final override val arrayMap: ArrayMap<V> =
+        ArrayMapImpl()
+
+    final override fun registerComponent(tClass: KClass<out K>, value: V) {
+        arrayMap[typeRegistry.getId(tClass)] = value
+    }
+
+    protected operator fun get(key: KClass<out K>): V {
+        val id = typeRegistry.getId(key)
+        return arrayMap[id] ?: error("No '$key'($id) component in array: $this")
+    }
+}