From c8bd6d3203755d3b5a30efafea0b3eacec1d2d37 Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Thu, 2 Sep 2021 14:58:47 +0300
Subject: [PATCH] [FIR] Store SAM constructors in global session component

---
 .../kotlin/fir/resolve/SamResolution.kt         | 17 +++++++----------
 .../kotlin/fir/scopes/FirFakeOverrideStorage.kt |  4 ++++
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/SamResolution.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/SamResolution.kt
index c6e6ba31b4b..db305bece74 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/SamResolution.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/SamResolution.kt
@@ -22,12 +22,10 @@ import org.jetbrains.kotlin.fir.diagnostics.ConeIntermediateDiagnostic
 import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.fir.resolve.calls.FirSyntheticFunctionSymbol
 import org.jetbrains.kotlin.fir.resolve.substitution.substitutorByMap
+import org.jetbrains.kotlin.fir.scopes.fakeOverrideStorage
 import org.jetbrains.kotlin.fir.scopes.impl.hasTypeOf
 import org.jetbrains.kotlin.fir.scopes.unsubstitutedScope
-import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
-import org.jetbrains.kotlin.fir.symbols.impl.FirValueParameterSymbol
+import org.jetbrains.kotlin.fir.symbols.impl.*
 import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.types.builder.buildResolvedTypeRef
@@ -52,7 +50,7 @@ class FirSamResolverImpl(
     private val outerClassManager: FirOuterClassManager? = null,
 ) : FirSamResolver() {
     private val resolvedFunctionType: NullableMap<FirRegularClass, ConeLookupTagBasedType?> = NullableMap()
-    private val samConstructor: NullableMap<FirRegularClass, FirSimpleFunction?> = NullableMap()
+    private val samConstructorsCache = firSession.fakeOverrideStorage.samConstructorCache
 
     override fun getFunctionTypeForPossibleSamType(type: ConeKotlinType): ConeKotlinType? {
         return when (type) {
@@ -113,12 +111,11 @@ class FirSamResolverImpl(
     }
 
     override fun getSamConstructor(firRegularClass: FirRegularClass): FirSimpleFunction? {
-        return samConstructor.getOrPut(firRegularClass) {
-            buildSamConstructor(firRegularClass)
-        }
+        return samConstructorsCache.getValue(firRegularClass.symbol, this)?.fir
     }
 
-    private fun buildSamConstructor(firRegularClass: FirRegularClass): FirSimpleFunction? {
+    fun buildSamConstructor(classSymbol: FirRegularClassSymbol): FirNamedFunctionSymbol? {
+        val firRegularClass = classSymbol.fir
         val functionType = resolveFunctionTypeIfSamInterface(firRegularClass) ?: return null
 
         val classId = firRegularClass.classId
@@ -218,7 +215,7 @@ class FirSamResolverImpl(
             resolvePhase = FirResolvePhase.BODY_RESOLVE
         }.apply {
             containingClassForStaticMemberAttr = outerClassManager?.outerClass(firRegularClass.symbol)?.toLookupTag()
-        }
+        }.symbol
     }
 
     private fun resolveFunctionTypeIfSamInterface(firRegularClass: FirRegularClass): ConeLookupTagBasedType? {
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/FirFakeOverrideStorage.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/FirFakeOverrideStorage.kt
index 5dc2aa896e3..5a81fe6a92d 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/FirFakeOverrideStorage.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/scopes/FirFakeOverrideStorage.kt
@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.caches.FirCache
 import org.jetbrains.kotlin.fir.caches.FirCachesFactory
 import org.jetbrains.kotlin.fir.caches.createCache
 import org.jetbrains.kotlin.fir.caches.firCachesFactory
+import org.jetbrains.kotlin.fir.resolve.FirSamResolverImpl
 import org.jetbrains.kotlin.fir.resolve.ScopeSessionKey
 import org.jetbrains.kotlin.fir.scopes.impl.FirClassSubstitutionScope
 import org.jetbrains.kotlin.fir.scopes.impl.FirTypeIntersectionScope
@@ -56,6 +57,9 @@ class FirFakeOverrideStorage(val session: FirSession) : FirSessionComponent {
 
     val intersectionOverrideCacheByScope: FirCache<ConeKotlinType, IntersectionOverrideCache, Nothing?> =
         cachesFactory.createCache { _ -> IntersectionOverrideCache(cachesFactory) }
+
+    val samConstructorCache: FirCache<FirRegularClassSymbol, FirNamedFunctionSymbol?, FirSamResolverImpl> =
+        cachesFactory.createCache { classSymbol, samResolver -> samResolver.buildSamConstructor(classSymbol) }
 }
 
 val FirSession.fakeOverrideStorage: FirFakeOverrideStorage by FirSession.sessionComponentAccessor()