FIR IDE: reuse up-to-date fir sessions

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/sessions/FirIdeLibrarySession.kt

@@ -16,8 +16,8 @@ import org.jetbrains.kotlin.fir.PrivateSessionConstructor
  */
 @OptIn(PrivateSessionConstructor::class)
 internal class FirIdeLibrariesSession @PrivateSessionConstructor constructor(
-    moduleInfo: ModuleInfo,
+    override val moduleInfo: ModuleInfo,
     override val project: Project,
     override val scope: GlobalSearchScope,
     builtinTypes: BuiltinTypes,
-) : FirIdeModuleSession(moduleInfo, builtinTypes)
+) : FirIdeModuleSession(builtinTypes)

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/sessions/FirIdeSession.kt

@@ -19,6 +19,7 @@ abstract class FirIdeSession(override val builtinTypes: BuiltinTypes) : FirSessi
 }
 
 @OptIn(PrivateSessionConstructor::class)
-abstract class FirIdeModuleSession(override val moduleInfo: ModuleInfo, builtinTypes: BuiltinTypes) : FirIdeSession(builtinTypes) {
+abstract class FirIdeModuleSession(builtinTypes: BuiltinTypes) : FirIdeSession(builtinTypes) {
+    abstract override val moduleInfo: ModuleInfo
     abstract val scope: GlobalSearchScope
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/sessions/FirIdeSessionFactory.kt

@@ -57,10 +57,11 @@ internal object FirIdeSessionFactory {
         val scopeProvider = KotlinScopeProvider(::wrapScopeWithJvmMapped)
         val firBuilder = FirFileBuilder(scopeProvider, firPhaseRunner)
         val searchScope = ModuleProductionSourceScope(moduleInfo.module)
-        val session = FirIdeSourcesSession(moduleInfo, project, searchScope, firBuilder, builtinTypes)
-        sessionsCache[moduleInfo] = session
         val dependentModules = moduleInfo.dependenciesWithoutSelf()
             .filterIsInstanceTo<ModuleSourceInfo, MutableList<ModuleSourceInfo>>(mutableListOf())
+        val session = FirIdeSourcesSession(moduleInfo, dependentModules, project, searchScope, firBuilder, builtinTypes)
+        sessionsCache[moduleInfo] = session
+
 
         return session.apply {
             val cache = ModuleFileCacheImpl(this)

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

@@ -8,29 +8,67 @@ package org.jetbrains.kotlin.idea.fir.low.level.api.sessions
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.fir.BuiltinTypes
 import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
+import org.jetbrains.kotlin.idea.caches.trackers.KotlinModuleOutOfCodeBlockModificationTracker
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirPhaseRunner
 import org.jetbrains.kotlin.idea.fir.low.level.api.FirTransformerProvider
 import org.jetbrains.kotlin.idea.fir.low.level.api.util.executeWithoutPCE
+import java.util.concurrent.ConcurrentHashMap
 
 internal class FirIdeSessionProviderStorage(private val project: Project) {
+    private val sessionsCache = ConcurrentHashMap<ModuleSourceInfo, FromModuleViewSessionCache>()
+
     fun getSessionProvider(rootModule: ModuleSourceInfo): FirIdeSessionProvider {
         val transformerProvider = FirTransformerProvider()
         val firPhaseRunner = FirPhaseRunner(transformerProvider)
+
         val builtinTypes = BuiltinTypes()
         val builtinsAndCloneableSession = FirIdeSessionFactory.createBuiltinsAndCloneableSession(project, builtinTypes)
-        val sessions = mutableMapOf<ModuleSourceInfo, FirIdeSourcesSession>()
-        val session = executeWithoutPCE {
-            FirIdeSessionFactory.createSourcesSession(
-                project,
-                rootModule,
-                builtinsAndCloneableSession,
-                firPhaseRunner,
-                builtinTypes,
-                sessions,
-                isRootModule = true
-            )
+        val cache = sessionsCache.getOrPut(rootModule) { FromModuleViewSessionCache(rootModule) }
+        val (sessions, session) = cache.withMappings { mappings ->
+            val sessions = mutableMapOf<ModuleSourceInfo, FirIdeSourcesSession>().apply { putAll(mappings) }
+            val session = executeWithoutPCE {
+                FirIdeSessionFactory.createSourcesSession(
+                    project,
+                    rootModule,
+                    builtinsAndCloneableSession,
+                    firPhaseRunner,
+                    builtinTypes,
+                    sessions,
+                    isRootModule = true
+                )
+            }
+            sessions to session
         }
 
         return FirIdeSessionProvider(project, session, sessions)
     }
+}
+
+private class FromModuleViewSessionCache(
+    val root: ModuleSourceInfo
+) {
+    @Volatile
+    private var mappings: Map<ModuleSourceInfo, FirSessionWithModificationTracker> = emptyMap()
+
+    inline fun <R> withMappings(
+        action: (Map<ModuleSourceInfo, FirIdeSourcesSession>) -> Pair<Map<ModuleSourceInfo, FirIdeSourcesSession>, R>
+    ): Pair<Map<ModuleSourceInfo, FirIdeSourcesSession>, R> {
+        val (newMappings, result) = action(getSessions().mapValues { it.value })
+        mappings = newMappings.mapValues { FirSessionWithModificationTracker(it.value) }
+        return newMappings to result
+    }
+
+    @OptIn(ExperimentalStdlibApi::class)
+    private fun getSessions(): Map<ModuleSourceInfo, FirIdeSourcesSession> = buildMap {
+        mappings.forEach { (module, session) -> if (session.isValid) put(module, session.firSession) }
+    }
+}
+
+private class FirSessionWithModificationTracker(
+    val firSession: FirIdeSourcesSession,
+) {
+    private val modificationTracker = KotlinModuleOutOfCodeBlockModificationTracker(firSession.moduleInfo.module)
+    private val timeStamp = modificationTracker.modificationCount
+
+    val isValid: Boolean get() = modificationTracker.modificationCount == timeStamp
 }

idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/sessions/FirIdeSourcesSession.kt

@@ -10,17 +10,20 @@ import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.analyzer.ModuleInfo
 import org.jetbrains.kotlin.fir.BuiltinTypes
 import org.jetbrains.kotlin.fir.PrivateSessionConstructor
+import org.jetbrains.kotlin.idea.caches.project.IdeaModuleInfo
+import org.jetbrains.kotlin.idea.caches.project.ModuleSourceInfo
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.FirFileBuilder
 import org.jetbrains.kotlin.idea.fir.low.level.api.providers.firIdeProvider
 
 @OptIn(PrivateSessionConstructor::class)
 internal class FirIdeSourcesSession @PrivateSessionConstructor constructor(
-    moduleInfo: ModuleInfo,
+    override val moduleInfo: ModuleSourceInfo,
+    val dependencies: List<ModuleSourceInfo>,
     override val project: Project,
     override val scope: GlobalSearchScope,
     val firFileBuilder: FirFileBuilder,
     builtinTypes: BuiltinTypes,
-) : FirIdeModuleSession(moduleInfo, builtinTypes) {
+) : FirIdeModuleSession( builtinTypes) {
     val cache get() = firIdeProvider.cache
 }