FIR: Optimize access to FirSymbolProvider via session

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

@@ -26,18 +26,23 @@ class FirJavaModuleBasedSession(
     override val sessionProvider: FirProjectSessionProvider,
     scope: GlobalSearchScope,
     dependenciesProvider: FirSymbolProvider? = null
-) : FirModuleBasedSession(moduleInfo) {
+) : FirModuleBasedSession(moduleInfo), FirSymbolProviderAwareSession {
+
+    override val firSymbolProvider: FirSymbolProvider
+
     init {
         sessionProvider.sessionCache[moduleInfo] = this
+        firSymbolProvider = FirCompositeSymbolProvider(
+            listOf(
+                service<FirProvider>(),
+                JavaSymbolProvider(this, sessionProvider.project, scope),
+                dependenciesProvider ?: FirDependenciesSymbolProviderImpl(this)
+            )
+        )
+
         registerComponent(
             FirSymbolProvider::class,
-            FirCompositeSymbolProvider(
+            firSymbolProvider
-                listOf(
-                    service<FirProvider>(),
-                    JavaSymbolProvider(this, sessionProvider.project, scope),
-                    dependenciesProvider ?: FirDependenciesSymbolProviderImpl(this)
-                )
-            )
         )
     }
 }
@@ -49,11 +54,13 @@ class FirLibrarySession private constructor(
     packagePartProvider: PackagePartProvider,
     kotlinClassFinder: KotlinClassFinder,
     javaClassFinder: JavaClassFinder
-) : FirSessionBase() {
+) : FirSessionBase(), FirSymbolProviderAwareSession {
+
+    override val firSymbolProvider: FirSymbolProvider
+
     init {
         sessionProvider.sessionCache[moduleInfo] = this
-        registerComponent(
+        firSymbolProvider =
-            FirSymbolProvider::class,
             FirCompositeSymbolProvider(
                 listOf(
                     KotlinDeserializedJvmSymbolsProvider(
@@ -66,6 +73,10 @@ class FirLibrarySession private constructor(
                     FirDependenciesSymbolProviderImpl(this)
                 )
             )
+
+        registerComponent(
+            FirSymbolProvider::class,
+            firSymbolProvider
         )
     }
 

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

@@ -0,0 +1,12 @@
+/*
+ * Copyright 2010-2019 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
+
+import org.jetbrains.kotlin.fir.resolve.FirSymbolProvider
+
+interface FirSymbolProviderAwareSession : FirSession {
+    val firSymbolProvider: FirSymbolProvider
+}

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve
 
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.FirSymbolProviderAwareSession
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.service
@@ -31,8 +32,13 @@ inline fun <K, V, VA : V> MutableMap<K, V>.getOrPut(key: K, defaultValue: (K) ->
     }
 }
 
-fun ConeClassLikeLookupTag.toSymbol(useSiteSession: FirSession): ConeClassifierSymbol? =
+fun ConeClassLikeLookupTag.toSymbol(useSiteSession: FirSession): ConeClassifierSymbol? {
-    useSiteSession.getService(FirSymbolProvider::class).getSymbolByLookupTag(this)
+    val firSymbolProvider =
+        (useSiteSession as? FirSymbolProviderAwareSession)?.firSymbolProvider
+            ?: useSiteSession.getService(FirSymbolProvider::class)
+
+    return firSymbolProvider.getSymbolByLookupTag(this)
+}
 
 fun ConeAbbreviatedType.directExpansionType(useSiteSession: FirSession): ConeClassLikeType? =
     abbreviationLookupTag