Optimize FirSession component storage

compiler/fir/java/src/org/jetbrains/kotlin/fir/java/FirJavaModuleBasedSession.kt

@@ -34,23 +34,20 @@ class FirJavaModuleBasedSession(
     init {
         sessionProvider.sessionCache[moduleInfo] = this
 
-        _firSymbolProvider = FirCompositeSymbolProvider(
-            listOf(
-                service<FirProvider>(),
-                JavaSymbolProvider(this, sessionProvider.project, scope),
-                dependenciesProvider ?: FirDependenciesSymbolProviderImpl(this)
-            )
-        )
-
         registerComponent(
             FirSymbolProvider::class,
-            _firSymbolProvider as FirSymbolProvider
+            FirCompositeSymbolProvider(
+                listOf(
+                    service<FirProvider>(),
+                    JavaSymbolProvider(this, sessionProvider.project, scope),
+                    dependenciesProvider ?: FirDependenciesSymbolProviderImpl(this)
+                )
+            ) as FirSymbolProvider
         )
 
-        _correspondingSupertypesCache = FirCorrespondingSupertypesCache(this)
         registerComponent(
             FirCorrespondingSupertypesCache::class,
-            _correspondingSupertypesCache as FirCorrespondingSupertypesCache
+            FirCorrespondingSupertypesCache(this)
         )
     }
 }
@@ -68,7 +65,9 @@ class FirLibrarySession private constructor(
     init {
         sessionProvider.sessionCache[moduleInfo] = this
         val javaSymbolProvider = JavaSymbolProvider(this, sessionProvider.project, scope)
-        _firSymbolProvider =
+
+        registerComponent(
+            FirSymbolProvider::class,
             FirCompositeSymbolProvider(
                 listOf(
                     FirLibrarySymbolProviderImpl(this),
@@ -82,18 +81,13 @@ class FirLibrarySession private constructor(
                     javaSymbolProvider,
                     FirDependenciesSymbolProviderImpl(this)
                 )
-            )
-
-        registerComponent(
-            FirSymbolProvider::class,
-            _firSymbolProvider as FirSymbolProvider
+            ) as FirSymbolProvider
         )
         registerComponent(FirClassDeclaredMemberScopeProvider::class, FirClassDeclaredMemberScopeProvider())
 
-        _correspondingSupertypesCache = FirCorrespondingSupertypesCache(this)
         registerComponent(
             FirCorrespondingSupertypesCache::class,
-            _correspondingSupertypesCache as FirCorrespondingSupertypesCache
+            FirCorrespondingSupertypesCache(this)
         )
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/FirSymbolProvider.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve
 
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSessionComponent
 import org.jetbrains.kotlin.fir.declarations.toFirClassLike
 import org.jetbrains.kotlin.fir.scopes.FirScope
 import org.jetbrains.kotlin.fir.scopes.ProcessorAction
@@ -17,7 +18,7 @@ import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 
-abstract class FirSymbolProvider {
+abstract class FirSymbolProvider : FirSessionComponent {
 
     abstract fun getClassLikeSymbolByFqName(classId: ClassId): FirClassLikeSymbol<*>?
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt

@@ -7,6 +7,8 @@ package org.jetbrains.kotlin.fir.resolve
 
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.*
+import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.componentArrayAccessor
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccess
 import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
@@ -36,9 +38,8 @@ inline fun <K, V, VA : V> MutableMap<K, V>.getOrPut(key: K, defaultValue: (K) ->
     }
 }
 
-val FirSession.firSymbolProvider: FirSymbolProvider get() = _firSymbolProvider as FirSymbolProvider? ?: service()
-val FirSession.correspondingSupertypesCache: FirCorrespondingSupertypesCache
-    get() = _correspondingSupertypesCache as FirCorrespondingSupertypesCache? ?: service()
+val FirSession.firSymbolProvider: FirSymbolProvider by componentArrayAccessor()
+val FirSession.correspondingSupertypesCache: FirCorrespondingSupertypesCache by componentArrayAccessor()
 
 fun ConeClassLikeLookupTag.toSymbol(useSiteSession: FirSession): ConeClassifierSymbol? {
     val firSymbolProvider = useSiteSession.firSymbolProvider

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/types/FirCorrespondingSupertypesCache.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.types
 
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSessionComponent
 import org.jetbrains.kotlin.fir.declarations.toFirClassLike
 import org.jetbrains.kotlin.fir.resolve.constructClassType
 import org.jetbrains.kotlin.fir.resolve.constructType
@@ -16,7 +17,7 @@ import org.jetbrains.kotlin.types.model.CaptureStatus
 import org.jetbrains.kotlin.types.model.SimpleTypeMarker
 import org.jetbrains.kotlin.types.model.TypeConstructorMarker
 
-class FirCorrespondingSupertypesCache(private val session: FirSession) {
+class FirCorrespondingSupertypesCache(private val session: FirSession) : FirSessionComponent {
     private val context = ConeTypeCheckerContext(false, session)
     private val cache = HashMap<ConeClassLikeSymbol, Map<ConeClassLikeSymbol, List<ConeClassLikeType>>?>(1000, 0.5f)
 

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

@@ -7,7 +7,9 @@ package org.jetbrains.kotlin.fir
 
 import org.jetbrains.kotlin.analyzer.ModuleInfo
 import org.jetbrains.kotlin.utils.Jsr305State
+import kotlin.properties.ReadOnlyProperty
 import kotlin.reflect.KClass
+import kotlin.reflect.KProperty
 
 abstract class FirSession(val sessionProvider: FirSessionProvider?) {
     open val moduleInfo: ModuleInfo? get() = null
@@ -17,16 +19,21 @@ abstract class FirSession(val sessionProvider: FirSessionProvider?) {
 
     val components: MutableMap<KClass<*>, Any> = mutableMapOf()
 
-    var _firSymbolProvider: Any? = null
-
-    var _correspondingSupertypesCache: Any? = null
+    internal val componentArray = ComponentArray()
 
+    @Suppress("UNCHECKED_CAST")
     fun <T : Any> getService(kclass: KClass<T>): T =
         components[kclass] as T
 
-    protected fun <T : Any> registerComponent(tClass: KClass<T>, t: T) {
+    protected fun <T : Any /* TODO: FirSessionComponent */> registerComponent(tClass: KClass<T>, t: T) {
         assert(tClass !in components) { "Already registered component" }
         components[tClass] = t
+
+        // TODO: Make t of FirSessionComponent
+        if (t is FirSessionComponent) {
+            @Suppress("UNCHECKED_CAST")
+            componentArray[(tClass as KClass<FirSessionComponent>).componentId()] = t
+        }
     }
 }
 
@@ -36,3 +43,51 @@ interface FirSessionProvider {
 
 inline fun <reified T : Any> FirSession.service(): T =
     getService(T::class)
+
+internal object ComponentTypeRegistry {
+    private val idPerType = mutableMapOf<KClass<out FirSessionComponent>, Int>()
+
+    fun <T : FirSessionComponent> id(kClass: KClass<T>): Int {
+        return idPerType.getOrPut(kClass) { idPerType.size }
+    }
+}
+
+
+private fun <T : FirSessionComponent> KClass<T>.componentId(): Int {
+    return ComponentTypeRegistry.id(this)
+}
+
+
+class ComponentArrayAccessor<T : FirSessionComponent>(val type: KClass<T>) : ReadOnlyProperty<FirSession, T> {
+    val id: Int = type.componentId()
+    override fun getValue(thisRef: FirSession, property: KProperty<*>): T {
+        @Suppress("UNCHECKED_CAST")
+        return thisRef.componentArray.getOrNull(id) as? T ?: error("No '$type'($id) component in session: $thisRef")
+    }
+}
+
+inline fun <reified T : FirSessionComponent> componentArrayAccessor(): ComponentArrayAccessor<T> {
+    return ComponentArrayAccessor(T::class)
+}
+
+interface FirSessionComponent
+
+internal class ComponentArray : AbstractList<FirSessionComponent?>() {
+    override val size: Int
+        get() = data.size
+    private var data = arrayOfNulls<FirSessionComponent>(20)
+    private fun ensureCapacity(index: Int) {
+        if (data.size < index) {
+            data = data.copyOf(data.size * 2)
+        }
+    }
+
+    operator fun set(index: Int, value: FirSessionComponent) {
+        ensureCapacity(index)
+        data[index] = value
+    }
+
+    override operator fun get(index: Int): FirSessionComponent? {
+        return data[index]
+    }
+}