FIR IDE: invalidate sessions on exception

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

@@ -14,11 +14,13 @@ import org.jetbrains.kotlin.fir.resolve.transformers.FirPhaseManager
 import org.jetbrains.kotlin.fir.symbols.AbstractFirBasedSymbol
 import org.jetbrains.kotlin.idea.fir.low.level.api.file.builder.ModuleFileCache
 import org.jetbrains.kotlin.idea.fir.low.level.api.lazy.resolve.FirLazyDeclarationResolver
+import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.FirSessionInvalidator
 
 @ThreadSafeMutableState
 internal class IdeFirPhaseManager(
     private val lazyDeclarationResolver: FirLazyDeclarationResolver,
-    private val cache: ModuleFileCache
+    private val cache: ModuleFileCache,
+    private val sessionInvalidator: FirSessionInvalidator,
 ) : FirPhaseManager() {
     override fun ensureResolved(
         symbol: AbstractFirBasedSymbol<*>,
@@ -34,6 +36,11 @@ internal class IdeFirPhaseManager(
             "Incorrect resolvePhase: actual: $availablePhase, expected: $requiredPhase\n For: ${symbol.fir.render()}"
         }
 
-        lazyDeclarationResolver.lazyResolveDeclaration(result, cache, requiredPhase, checkPCE = false)
+        try {
+            lazyDeclarationResolver.lazyResolveDeclaration(result, cache, requiredPhase, checkPCE = true)
+        } catch (e: Throwable) {
+            sessionInvalidator.invalidate((symbol.fir as FirDeclaration).session)
+            throw e
+        }
     }
 }

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

@@ -49,6 +49,7 @@ internal object FirIdeSessionFactory {
         moduleInfo: ModuleSourceInfo,
         builtinsAndCloneableSession: FirIdeBuiltinsAndCloneableSession,
         firPhaseRunner: FirPhaseRunner,
+        sessionInvalidator: FirSessionInvalidator,
         builtinTypes: BuiltinTypes,
         sessionsCache: MutableMap<ModuleSourceInfo, FirIdeSourcesSession>,
         isRootModule: Boolean,
@@ -65,7 +66,7 @@ internal object FirIdeSessionFactory {
 
         return session.apply {
             val cache = ModuleFileCacheImpl(this)
-            val firPhaseManager = IdeFirPhaseManager(FirLazyDeclarationResolver(firFileBuilder), cache)
+            val firPhaseManager = IdeFirPhaseManager(FirLazyDeclarationResolver(firFileBuilder), cache, sessionInvalidator)
 
             registerCommonComponents()
             registerResolveComponents()
@@ -104,6 +105,7 @@ internal object FirIdeSessionFactory {
                                     it,
                                     builtinsAndCloneableSession,
                                     firPhaseRunner,
+                                    sessionInvalidator,
                                     builtinTypes,
                                     sessionsCache,
                                     isRootModule = false

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

@@ -32,6 +32,7 @@ internal class FirIdeSessionProviderStorage(private val project: Project) {
                     rootModule,
                     builtinsAndCloneableSession,
                     firPhaseRunner,
+                    cache.sessionInvalidator,
                     builtinTypes,
                     sessions,
                     isRootModule = true
@@ -45,11 +46,16 @@ internal class FirIdeSessionProviderStorage(private val project: Project) {
 }
 
 private class FromModuleViewSessionCache(
-    val root: ModuleSourceInfo
+    val root: ModuleSourceInfo,
 ) {
     @Volatile
     private var mappings: Map<ModuleSourceInfo, FirSessionWithModificationTracker> = emptyMap()
 
+    val sessionInvalidator: FirSessionInvalidator = FirSessionInvalidator { session ->
+        mappings[session.moduleInfo]?.invalidate()
+    }
+
+
     inline fun <R> withMappings(
         action: (Map<ModuleSourceInfo, FirIdeSourcesSession>) -> Pair<Map<ModuleSourceInfo, FirIdeSourcesSession>, R>
     ): Pair<Map<ModuleSourceInfo, FirIdeSourcesSession>, R> {
@@ -60,7 +66,36 @@ private class FromModuleViewSessionCache(
 
     @OptIn(ExperimentalStdlibApi::class)
     private fun getSessions(): Map<ModuleSourceInfo, FirIdeSourcesSession> = buildMap {
-        mappings.forEach { (module, session) -> if (session.isValid) put(module, session.firSession) }
+        val sessions = mappings.values
+        val sessionToValidity = hashMapOf<FirSessionWithModificationTracker, Boolean>()
+
+        var isValid = true
+        fun dfs(session: FirSessionWithModificationTracker) {
+            sessionToValidity[session]?.let { valid ->
+                if (!valid) isValid = false
+                return
+            }
+            sessionToValidity[session] = session.isValid
+            session.firSession.dependencies.forEach { dependency ->
+                mappings[dependency]?.let(::dfs)
+            }
+            if (!session.isValid) {
+                isValid = false
+            }
+            if (!isValid) {
+                sessionToValidity[session] = false
+            }
+        }
+
+        for (session in sessions) {
+            if (session !in sessionToValidity) {
+                isValid = true
+                dfs(session)
+            }
+        }
+        return sessionToValidity.entries
+            .mapNotNull { (session, valid) -> session.takeIf { valid } }
+            .associate { session -> session.firSession.moduleInfo to session.firSession }
     }
 }
 
@@ -70,5 +105,12 @@ private class FirSessionWithModificationTracker(
     private val modificationTracker = KotlinModuleOutOfCodeBlockModificationTracker(firSession.moduleInfo.module)
     private val timeStamp = modificationTracker.modificationCount
 
-    val isValid: Boolean get() = modificationTracker.modificationCount == timeStamp
+    @Volatile
+    private var isInvalidated = false
+
+    fun invalidate() {
+        isInvalidated = true
+    }
+
+    val isValid: Boolean get() = isInvalidated || modificationTracker.modificationCount == timeStamp
 }

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

@@ -0,0 +1,15 @@
+/*
+ * Copyright 2010-2020 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.idea.fir.low.level.api.sessions
+
+import org.jetbrains.kotlin.fir.FirSession
+
+internal class FirSessionInvalidator(private val invalidateSourcesSession: (FirIdeSourcesSession) -> Unit) {
+    fun invalidate(session: FirSession) {
+        require(session is FirIdeSourcesSession)
+        invalidateSourcesSession(session)
+    }
+}