FIR LC: fix light facade caching

analysis/analysis-internal-utils/build.gradle.kts

@@ -0,0 +1,18 @@
+plugins {
+    kotlin("jvm")
+    id("jps-compatible")
+}
+
+dependencies {
+    implementation(project(":compiler:psi"))
+    implementation(intellijCoreDep()) { includeJars("intellij-core", rootProject = rootProject) }
+}
+
+sourceSets {
+    "main" { projectDefault() }
+    "test" { none() }
+}
+
+kotlin {
+    explicitApi()
+}

analysis/analysis-internal-utils/src/org/jetbrains/kotlin/analysis/utils/caches/softCachedValue.kt

@@ -0,0 +1,23 @@
+package org.jetbrains.kotlin.analysis.utils.caches
+
+import com.intellij.openapi.project.Project
+import com.intellij.psi.util.CachedValue
+import com.intellij.psi.util.CachedValueProvider
+import com.intellij.psi.util.CachedValuesManager
+import kotlin.reflect.KProperty
+
+
+@Suppress("NOTHING_TO_INLINE")
+public inline operator fun <T> CachedValue<T>.getValue(thisRef: Any?, property: KProperty<*>): T = value
+
+public inline fun <T> softCachedValue(
+    project: Project,
+    vararg dependencies: Any,
+    crossinline createValue: () -> T
+): CachedValue<T> =
+    CachedValuesManager.getManager(project).createCachedValue {
+        CachedValueProvider.Result(
+            createValue(),
+            dependencies
+        )
+    }

analysis/analysis-internal-utils/src/org/jetbrains/kotlin/analysis/utils/caches/strongCachedValue.kt

@@ -0,0 +1,49 @@
+/*
+ * 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.analysis.utils.caches
+
+import com.intellij.openapi.util.ModificationTracker
+import java.util.concurrent.atomic.AtomicReference
+import kotlin.reflect.KProperty
+
+public class StrongRefModificationTrackerBasedCache<T> internal constructor(
+    private val dependencies: List<ModificationTracker>,
+    private val compute: () -> T,
+) {
+    private val cached = AtomicReference<CachedValue<T>?>(null)
+
+    public operator fun getValue(thisRef: Any?, property: KProperty<*>): T = cached.updateAndGet { value ->
+        when {
+            value == null -> createNewCachedValue()
+            value.isUpToDate(dependencies) -> value
+            else -> createNewCachedValue()
+        }
+    }!!.value
+
+    private fun createNewCachedValue() = CachedValue(compute(), dependencies.map { it.modificationCount })
+}
+
+private class CachedValue<T>(val value: T, val timestamps: List<Long>) {
+    fun isUpToDate(dependencies: List<ModificationTracker>): Boolean {
+        check(timestamps.size == dependencies.size)
+        for (i in timestamps.indices) {
+            if (dependencies[i].modificationCount != timestamps[i]) {
+                return false
+            }
+        }
+        return true
+    }
+}
+
+/**
+ * Create modification tracker which will be invalidated when dependencies change.
+ * The cached value is hold on the strong reference.
+ * So, the value will not be garbage collected until modification tracker changes.
+ */
+public fun <T> strongCachedValue(
+    vararg dependencies: ModificationTracker,
+    compute: () -> T,
+): StrongRefModificationTrackerBasedCache<T> = StrongRefModificationTrackerBasedCache(dependencies.toList(), compute)

analysis/low-level-api-fir/build.gradle.kts

@@ -22,6 +22,7 @@ dependencies {
     implementation(project(":compiler:ir.psi2ir"))
     implementation(project(":compiler:fir:entrypoint"))
     implementation(project(":analysis:analysis-api-providers"))
+    implementation(project(":analysis:analysis-internal-utils"))
 
     api(intellijCoreDep()) { includeJars("intellij-core", "guava", rootProject = rootProject) }
 

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/FirIdeResolveStateService.kt

@@ -13,18 +13,16 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirModuleResolveState
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.FirIdeSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.FirIdeSessionProviderStorage
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.cachedValue
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.getValue
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.analysis.providers.createProjectWideOutOfBlockModificationTracker
+import org.jetbrains.kotlin.analysis.utils.caches.strongCachedValue
 import java.util.concurrent.ConcurrentHashMap
 
 internal class FirIdeResolveStateService(project: Project) {
     private val sessionProviderStorage = FirIdeSessionProviderStorage(project)
 
-    private val stateCache by cachedValue(
+    private val stateCache by strongCachedValue(
-        project,
         project.createProjectWideOutOfBlockModificationTracker(),
         ProjectRootModificationTracker.getInstance(project),
     ) {

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/sessions/FirIdeSessionProviderStorage.kt

@@ -16,16 +16,16 @@ import org.jetbrains.kotlin.fir.FirModuleData
 import org.jetbrains.kotlin.fir.moduleData
 import org.jetbrains.kotlin.analysis.low.level.api.fir.FirPhaseRunner
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.addValueFor
-import org.jetbrains.kotlin.analysis.low.level.api.fir.util.*
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.executeWithoutPCE
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
+import org.jetbrains.kotlin.analysis.utils.caches.*
 import java.util.concurrent.ConcurrentHashMap
 
 class FirIdeSessionProviderStorage(val project: Project) {
     private val sessionsCache = ConcurrentHashMap<KtSourceModule, FromModuleViewSessionCache>()
 
-    private val librariesCache by cachedValue(project, project.createLibrariesModificationTracker()) { LibrariesCache() }
+    private val librariesCache by softCachedValue(project, project.createLibrariesModificationTracker()) { LibrariesCache() }
 
     fun getSessionProvider(
         rootModule: KtSourceModule,

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/util/utils.kt

@@ -95,25 +95,6 @@ internal fun KtDeclaration.isNonAnonymousClassOrObject() =
             && this !is KtEnumEntry
 
 
-@Suppress("NOTHING_TO_INLINE")
-internal inline operator fun <T> CachedValue<T>.getValue(thisRef: Any?, property: KProperty<*>): T = value
-
-internal inline fun <T> cachedValue(project: Project, vararg dependencies: Any, crossinline createValue: () -> T) =
-    CachedValuesManager.getManager(project).createCachedValue {
-        CachedValueProvider.Result(
-            createValue(),
-            dependencies
-        )
-    }
-
-/**
- * Creates a value which will be cached until until any physical PSI change happens
- *
- * @see com.intellij.psi.util.CachedValue
- * @see com.intellij.psi.util.PsiModificationTracker.MODIFICATION_COUNT
- */
-internal fun <T> psiModificationTrackerBasedCachedValue(project: Project, createValue: () -> T) =
-    cachedValue(project, PsiModificationTracker.MODIFICATION_COUNT, createValue = createValue)
 
 fun KtElement.getElementTextInContext(): String {
     val context = parentOfType<KtImportDirective>()

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/compiler/based/testServices.kt

@@ -32,6 +32,7 @@ import org.jetbrains.kotlin.asJava.KotlinAsJavaSupport
 import org.jetbrains.kotlin.asJava.finder.JavaElementFinder
 import org.jetbrains.kotlin.fir.declarations.SealedClassInheritorsProvider
 import org.jetbrains.kotlin.light.classes.symbol.IDEKotlinAsJavaFirSupport
+import org.jetbrains.kotlin.light.classes.symbol.caches.SymbolLightClassFacadeCache
 import org.jetbrains.kotlin.load.kotlin.PackagePartProvider
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.test.model.TestModule
@@ -60,6 +61,7 @@ fun MockProject.registerTestServices(
     registerService(KotlinDeclarationProviderFactory::class.java, KotlinStaticDeclarationProviderFactory(allKtFiles))
     registerService(KotlinPackageProviderFactory::class.java, KotlinStaticPackageProviderFactory(allKtFiles))
     registerService(ProjectStructureProvider::class.java, KotlinProjectStructureProviderTestImpl(testServices))
+    registerService(SymbolLightClassFacadeCache::class.java)
     reRegisterJavaElementFinder(this)
 }
 

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? {

prepare/ide-plugin-dependencies/high-level-api-impl-base-for-ide/build.gradle.kts

@@ -2,4 +2,10 @@ plugins {
     kotlin("jvm")
 }
 
-publishJarsForIde(listOf(":analysis:analysis-api-impl-barebone", ":analysis:analysis-api-impl-base"))
+publishJarsForIde(
+    listOf(
+        ":analysis:analysis-api-impl-barebone",
+        ":analysis:analysis-api-impl-base",
+        ":analysis:analysis-internal-utils"
+    )
+)

settings.gradle

@@ -482,6 +482,7 @@ include ":generators:analysis-api-generator",
         ":analysis:analysis-api-impl-barebone",
         ":analysis:analysis-api-impl-base",
         ":analysis:analysis-api-providers",
+        ":analysis:analysis-internal-utils",
         ":analysis:symbol-light-classes",
         ":analysis:project-structure",
         ":analysis:analysis-api-fe10"