From 91d8f1fc88e05f29494bb5cf44d76076eb0ab7d9 Mon Sep 17 00:00:00 2001
From: Denis Zharkov <denis.zharkov@jetbrains.com>
Date: Wed, 29 May 2019 16:24:26 +0300
Subject: [PATCH] FIR: Optimize access to FirSymbolProvider via session

---
 .../fir/java/FirJavaModuleBasedSession.kt     | 33 ++++++++++++-------
 .../fir/FirSymbolProviderAwareSession.kt      | 12 +++++++
 .../kotlin/fir/resolve/ResolveUtils.kt        | 10 ++++--
 3 files changed, 42 insertions(+), 13 deletions(-)
 create mode 100644 compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirSymbolProviderAwareSession.kt

diff --git a/compiler/fir/java/src/org/jetbrains/kotlin/fir/java/FirJavaModuleBasedSession.kt b/compiler/fir/java/src/org/jetbrains/kotlin/fir/java/FirJavaModuleBasedSession.kt
index daa3650f6c3..a2d942a9287 100644
--- a/compiler/fir/java/src/org/jetbrains/kotlin/fir/java/FirJavaModuleBasedSession.kt
+++ b/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(
-                listOf(
-                    service<FirProvider>(),
-                    JavaSymbolProvider(this, sessionProvider.project, scope),
-                    dependenciesProvider ?: FirDependenciesSymbolProviderImpl(this)
-                )
-            )
+            firSymbolProvider
         )
     }
 }
@@ -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::class,
+        firSymbolProvider =
             FirCompositeSymbolProvider(
                 listOf(
                     KotlinDeserializedJvmSymbolsProvider(
@@ -66,6 +73,10 @@ class FirLibrarySession private constructor(
                     FirDependenciesSymbolProviderImpl(this)
                 )
             )
+
+        registerComponent(
+            FirSymbolProvider::class,
+            firSymbolProvider
         )
     }
 
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirSymbolProviderAwareSession.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirSymbolProviderAwareSession.kt
new file mode 100644
index 00000000000..f246593b2ee
--- /dev/null
+++ b/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
+}
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt
index 30752ac93d9..dee796af94d 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/ResolveUtils.kt
+++ b/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? =
-    useSiteSession.getService(FirSymbolProvider::class).getSymbolByLookupTag(this)
+fun ConeClassLikeLookupTag.toSymbol(useSiteSession: FirSession): ConeClassifierSymbol? {
+    val firSymbolProvider =
+        (useSiteSession as? FirSymbolProviderAwareSession)?.firSymbolProvider
+            ?: useSiteSession.getService(FirSymbolProvider::class)
+
+    return firSymbolProvider.getSymbolByLookupTag(this)
+}
 
 fun ConeAbbreviatedType.directExpansionType(useSiteSession: FirSession): ConeClassLikeType? =
     abbreviationLookupTag