[LL API] Clean invalidated sessions in 'KtFirAnalysisSessionProvider'

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/KtFirAnalysisSessionProvider.kt

@@ -8,64 +8,65 @@ package org.jetbrains.kotlin.analysis.api.fir
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.ProjectRootModificationTracker
 import com.intellij.openapi.util.LowMemoryWatcher
-import com.intellij.psi.util.PsiModificationTracker
+import com.intellij.psi.util.CachedValue
-import org.jetbrains.annotations.TestOnly
+import com.intellij.psi.util.CachedValueProvider
+import com.intellij.psi.util.CachedValuesManager
+import com.intellij.util.CachedValueBase
+import com.intellij.util.containers.CollectionFactory
 import org.jetbrains.kotlin.analysis.api.KtAnalysisApiInternals
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeTokenFactory
 import org.jetbrains.kotlin.analysis.api.session.KtAnalysisSessionProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getFirResolveSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.FirSessionValidityStamp
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.getKtModule
 import org.jetbrains.kotlin.analysis.providers.createProjectWideOutOfBlockModificationTracker
-import org.jetbrains.kotlin.analysis.utils.caches.SoftCachedMap
 import org.jetbrains.kotlin.psi.KtElement
+import java.util.concurrent.ConcurrentMap
 import kotlin.reflect.KClass
 
 @OptIn(KtAnalysisApiInternals::class)
 class KtFirAnalysisSessionProvider(project: Project) : KtAnalysisSessionProvider(project) {
-    private val cache = KtAnalysisSessionCache(project)
+    private val cache: ConcurrentMap<Pair<KtModule, KClass<out KtLifetimeToken>>, CachedValue<KtAnalysisSession>> =
+        CollectionFactory.createConcurrentWeakValueMap()
+
+    init {
+        LowMemoryWatcher.register(::clearCaches, project)
+    }
 
     override fun getAnalysisSession(useSiteKtElement: KtElement, factory: KtLifetimeTokenFactory): KtAnalysisSession {
         return getAnalysisSessionByUseSiteKtModule(useSiteKtElement.getKtModule(project), factory)
     }
 
     override fun getAnalysisSessionByUseSiteKtModule(useSiteKtModule: KtModule, factory: KtLifetimeTokenFactory): KtAnalysisSession {
-        return cache.getAnalysisSession(useSiteKtModule to factory.identifier) {
+        val key = Pair(useSiteKtModule, factory.identifier)
-            val firResolveSession = useSiteKtModule.getFirResolveSession(project)
+        return cache.computeIfAbsent(key) {
-            val validityToken = factory.create(project)
+            CachedValuesManager.getManager(project).createCachedValue {
-            KtFirAnalysisSession.createAnalysisSessionByFirResolveSession(firResolveSession, validityToken)
+                val firResolveSession = useSiteKtModule.getFirResolveSession(project)
+                val validityToken = factory.create(project)
+
+                val validityStamp = FirSessionValidityStamp(firResolveSession.useSiteFirSession)
+                require(validityStamp.isValid) { "Got an invalid session for module $useSiteKtModule" }
+
+                CachedValueProvider.Result(
+                    KtFirAnalysisSession.createAnalysisSessionByFirResolveSession(firResolveSession, validityToken),
+                    validityStamp,
+                    ProjectRootModificationTracker.getInstance(project),
+                    project.createProjectWideOutOfBlockModificationTracker()
+                )
+            }
+        }.value
+    }
+
+    override fun clearCaches() {
+        for (cachedValue in cache.values) {
+            check(cachedValue is CachedValueBase<*>) {
+                "Unsupported 'CachedValue' of type ${cachedValue.javaClass}'"
+            }
+
+            cachedValue.clear()
         }
     }
-
-    @TestOnly
-    override fun clearCaches() {
-        cache.clear()
-    }
-}
-
-private class KtAnalysisSessionCache(project: Project) {
-    private val cache = SoftCachedMap.create<Pair<KtModule, KClass<out KtLifetimeToken>>, KtAnalysisSession>(
-        project,
-        SoftCachedMap.Kind.STRONG_KEYS_SOFT_VALUES,
-        listOf(
-            PsiModificationTracker.MODIFICATION_COUNT,
-            ProjectRootModificationTracker.getInstance(project),
-            project.createProjectWideOutOfBlockModificationTracker()
-        )
-    )
-
-    init {
-        LowMemoryWatcher.register({ cache.clearCachedValues() }, project)
-    }
-
-    @TestOnly
-    fun clear() {
-        cache.clear()
-    }
-
-    fun getAnalysisSession(key: Pair<KtModule, KClass<out KtLifetimeToken>>, create: () -> KtAnalysisSession): KtAnalysisSession {
-        return cache.getOrPut(key) { create() }
-    }
 }

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

@@ -27,7 +27,7 @@ import org.jetbrains.kotlin.psi.KtFile
 abstract class LLFirResolveSession {
     abstract val project: Project
 
-    abstract val useSiteFirSession: FirSession
+    abstract val useSiteFirSession: LLFirSession
 
     abstract val useSiteKtModule: KtModule
 

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

@@ -0,0 +1,21 @@
+/*
+ * Copyright 2010-2023 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.low.level.api.fir.sessions
+
+import com.intellij.openapi.util.ModificationTracker
+import java.lang.ref.WeakReference
+
+class FirSessionValidityStamp(session: LLFirSession) : ModificationTracker {
+    private val sessionRef = WeakReference(session)
+
+    val isValid: Boolean
+        get() = modificationCount == 0L
+
+    override fun getModificationCount(): Long {
+        val session = sessionRef.get() ?: return 1
+        return if (session.isValid) 0 else 2
+    }
+}