FIR LC: fix light facade caching

analysis/symbol-light-classes/build.gradle.kts

@@ -10,6 +10,7 @@ dependencies {
     implementation(project(":compiler:light-classes"))
     implementation(project(":analysis:analysis-api-providers"))
     implementation(project(":analysis:analysis-api"))
+    implementation(project(":analysis:analysis-internal-utils"))
     implementation(project(":analysis:project-structure"))
     implementation(intellijCoreDep()) { includeJars("intellij-core", rootProject = rootProject) }
 }

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/IDEKotlinAsJavaFirSupport.kt

@@ -15,8 +15,8 @@ import org.jetbrains.kotlin.asJava.KotlinAsJavaSupport
 import org.jetbrains.kotlin.asJava.classes.KtFakeLightClass
 import org.jetbrains.kotlin.asJava.classes.KtLightClass
 import org.jetbrains.kotlin.fileClasses.javaFileFacadeFqName
+import org.jetbrains.kotlin.light.classes.symbol.caches.SymbolLightClassFacadeCache
 import org.jetbrains.kotlin.light.classes.symbol.classes.getOrCreateFirLightClass
-import org.jetbrains.kotlin.light.classes.symbol.classes.getOrCreateFirLightFacade
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.parentOrNull
@@ -90,7 +90,9 @@ class IDEKotlinAsJavaFirSupport(private val project: Project) : KotlinAsJavaSupp
     override fun getFacadeClasses(facadeFqName: FqName, scope: GlobalSearchScope): Collection<PsiClass> =
         //TODO Split by modules
         findFilesForFacade(facadeFqName, scope).ifNotEmpty {
-            listOfNotNull(getOrCreateFirLightFacade(this.toList(), facadeFqName))
+            listOfNotNull(
+                project.getService(SymbolLightClassFacadeCache::class.java).getOrCreateSymbolLightFacade(this.toList(), facadeFqName)
+            )
         } ?: emptyList()
 
     override fun getScriptClasses(scriptFqName: FqName, scope: GlobalSearchScope): Collection<PsiClass> =
@@ -105,7 +107,10 @@ class IDEKotlinAsJavaFirSupport(private val project: Project) : KotlinAsJavaSupp
             .asSequence()
             .filter { it.isFromSource() }
             .groupBy { it.javaFileFacadeFqName }
-            .mapNotNull { getOrCreateFirLightFacade(it.value, it.key) }
+            .mapNotNull {
+                project.getService(SymbolLightClassFacadeCache::class.java)
+                    .getOrCreateSymbolLightFacade(it.value, it.key)
+            }
 
     override fun getFacadeNames(packageFqName: FqName, scope: GlobalSearchScope): Collection<String> =
         project.createDeclarationProvider(scope)

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/caches/SymbolLightClassFacadeCache.kt

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2010-2021 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.light.classes.symbol.caches
+
+import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.analysis.providers.createProjectWideOutOfBlockModificationTracker
+import org.jetbrains.kotlin.light.classes.symbol.FirLightClassForFacade
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.analysis.utils.caches.*
+import java.util.concurrent.ConcurrentHashMap
+
+class SymbolLightClassFacadeCache(project: Project) {
+    private val cache by softCachedValue(
+        project,
+        project.createProjectWideOutOfBlockModificationTracker()
+    ) {
+        ConcurrentHashMap<FacadeKey, FirLightClassForFacade>()
+    }
+
+    fun getOrCreateSymbolLightFacade(
+        ktFiles: List<KtFile>,
+        facadeClassFqName: FqName,
+    ): FirLightClassForFacade? {
+        if (ktFiles.isEmpty()) return null
+        val key = FacadeKey(facadeClassFqName, ktFiles.toSet())
+        return cache.computeIfAbsent(key) {
+            getOrCreateFirLightFacadeNoCache(ktFiles, facadeClassFqName)
+        }
+    }
+
+    private fun getOrCreateFirLightFacadeNoCache(
+        ktFiles: List<KtFile>,
+        facadeClassFqName: FqName,
+    ): FirLightClassForFacade {
+        val firstFile = ktFiles.first()
+        return FirLightClassForFacade(firstFile.manager, facadeClassFqName, ktFiles)
+    }
+
+    private data class FacadeKey(val fqName: FqName, val files: Set<KtFile>)
+}

analysis/symbol-light-classes/src/org/jetbrains/kotlin/light/classes/symbol/classes/firLightClassUtils.kt

@@ -92,28 +92,8 @@ internal fun KtClassOrObjectSymbol.createLightClassNoCache(manager: PsiManager):
     }
 }
 
-fun getOrCreateFirLightFacade(
-    ktFiles: List<KtFile>,
-    facadeClassFqName: FqName,
-): FirLightClassForFacade? {
-    val firstFile = ktFiles.firstOrNull() ?: return null
-    //TODO Make caching keyed by all files
-    return CachedValuesManager.getCachedValue(firstFile) {
-        CachedValueProvider.Result
-            .create(
-                getOrCreateFirLightFacadeNoCache(ktFiles, facadeClassFqName),
-                firstFile.project.createProjectWideOutOfBlockModificationTracker()
-            )
-    }
-}
 
-fun getOrCreateFirLightFacadeNoCache(
-    ktFiles: List<KtFile>,
-    facadeClassFqName: FqName,
-): FirLightClassForFacade? {
-    val firstFile = ktFiles.firstOrNull() ?: return null
-    return FirLightClassForFacade(firstFile.manager, facadeClassFqName, ktFiles)
-}
+
 
 
 private fun lightClassForEnumEntry(ktEnumEntry: KtEnumEntry): KtLightClass? {