[LL FIR] invalidate caches on exceptions in more places where we modify some fir elements

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

@@ -5,30 +5,24 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir
 
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirResolvableModuleSession
 import org.jetbrains.kotlin.fir.ThreadSafeMutableState
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
-import org.jetbrains.kotlin.fir.symbols.FirLazyDeclarationResolver
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
-import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirResolvableModuleSession
+import org.jetbrains.kotlin.fir.symbols.FirLazyDeclarationResolver
-import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionInvalidator
 
 @ThreadSafeMutableState
-internal class LLFirLazyDeclarationResolver(private val sessionInvalidator: LLFirSessionInvalidator) : FirLazyDeclarationResolver() {
+internal class LLFirLazyDeclarationResolver : FirLazyDeclarationResolver() {
     override fun lazyResolveToPhase(symbol: FirBasedSymbol<*>, toPhase: FirResolvePhase) {
         val fir = symbol.fir
         val session = fir.moduleData.session
         if (session !is LLFirResolvableModuleSession) return
         val moduleComponents = session.moduleComponents
-        try {
+        moduleComponents.firModuleLazyDeclarationResolver.lazyResolveDeclaration(
-            moduleComponents.firModuleLazyDeclarationResolver.lazyResolveDeclaration(
+            firDeclarationToResolve = fir,
-                firDeclarationToResolve = fir,
+            scopeSession = moduleComponents.scopeSessionProvider.getScopeSession(),
-                scopeSession = moduleComponents.scopeSessionProvider.getScopeSession(),
+            toPhase = toPhase,
-                toPhase = toPhase,
+            checkPCE = true,
-                checkPCE = true,
+        )
-            )
-        } catch (e: Throwable) {
-            sessionInvalidator.invalidate(session)
-            throw e
-        }
     }
 }

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

@@ -15,6 +15,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.ModuleFileCa
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.FileStructureCache
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.LLFirModuleLazyDeclarationResolver
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirResolvableModuleSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionInvalidator
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.LLFirScopeSessionProvider
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.fir.scopes.FirScopeProvider
@@ -24,6 +25,7 @@ internal class LLFirModuleResolveComponents(
     val module: KtModule,
     val globalResolveComponents: LLFirGlobalResolveComponents,
     val scopeProvider: FirScopeProvider,
+    val sessionInvalidator: LLFirSessionInvalidator,
 ) {
     val cache: ModuleFileCache = ModuleFileCacheImpl(this)
     val firFileBuilder: LLFirFileBuilder = LLFirFileBuilder(this)

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/FileStructureElement.kt

@@ -117,24 +117,27 @@ internal class ReanalyzableFunctionStructureElement(
 
         return moduleComponents.globalResolveComponents.lockProvider.withWriteLock(firFile) {
             val upgradedPhase = minOf(originalFunction.resolvePhase, FirResolvePhase.DECLARATIONS)
-            with(originalFunction) {
-                replaceBody(temporaryFunction.body)
-                replaceContractDescription(temporaryFunction.contractDescription)
-                replaceResolvePhase(upgradedPhase)
-            }
-            designation.toSequence(includeTarget = true).forEach {
-                it.replaceResolvePhase(minOf(it.resolvePhase, upgradedPhase))
-            }
 
-            originalFunction.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
+            moduleComponents.sessionInvalidator.withInvalidationOnException(moduleComponents.session) {
+                with(originalFunction) {
+                    replaceBody(temporaryFunction.body)
+                    replaceContractDescription(temporaryFunction.contractDescription)
+                    replaceResolvePhase(upgradedPhase)
+                }
+                designation.toSequence(includeTarget = true).forEach {
+                    it.replaceResolvePhase(minOf(it.resolvePhase, upgradedPhase))
+                }
 
-            ReanalyzableFunctionStructureElement(
+                originalFunction.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
-                firFile,
+
-                newKtDeclaration,
+                ReanalyzableFunctionStructureElement(
-                originalFunction.symbol,
+                    firFile,
-                newKtDeclaration.modificationStamp,
+                    newKtDeclaration,
-                moduleComponents,
+                    originalFunction.symbol,
-            )
+                    newKtDeclaration.modificationStamp,
+                    moduleComponents,
+                )
+            }
         }
     }
 }
@@ -160,30 +163,31 @@ internal class ReanalyzablePropertyStructureElement(
         ) as FirProperty
 
         return moduleComponents.globalResolveComponents.lockProvider.withWriteLock(firFile) {
-
             val getterPhase = originalProperty.getter?.resolvePhase ?: originalProperty.resolvePhase
             val setterPhase = originalProperty.setter?.resolvePhase ?: originalProperty.resolvePhase
             val upgradedPhase = minOf(originalProperty.resolvePhase, getterPhase, setterPhase, FirResolvePhase.DECLARATIONS)
 
-            with(originalProperty) {
+            moduleComponents.sessionInvalidator.withInvalidationOnException(moduleComponents.session) {
-                getter?.replaceBody(temporaryProperty.getter?.body)
+                with(originalProperty) {
-                setter?.replaceBody(temporaryProperty.setter?.body)
+                    getter?.replaceBody(temporaryProperty.getter?.body)
-                replaceInitializer(temporaryProperty.initializer)
+                    setter?.replaceBody(temporaryProperty.setter?.body)
-                getter?.replaceResolvePhase(upgradedPhase)
+                    replaceInitializer(temporaryProperty.initializer)
-                setter?.replaceResolvePhase(upgradedPhase)
+                    getter?.replaceResolvePhase(upgradedPhase)
-                replaceResolvePhase(upgradedPhase)
+                    setter?.replaceResolvePhase(upgradedPhase)
-                replaceBodyResolveState(FirPropertyBodyResolveState.NOTHING_RESOLVED)
+                    replaceResolvePhase(upgradedPhase)
+                    replaceBodyResolveState(FirPropertyBodyResolveState.NOTHING_RESOLVED)
+                }
+
+                originalProperty.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
+
+                ReanalyzablePropertyStructureElement(
+                    firFile,
+                    newKtDeclaration,
+                    originalProperty.symbol,
+                    newKtDeclaration.modificationStamp,
+                    moduleComponents,
+                )
             }
-
-            originalProperty.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE)
-
-            ReanalyzablePropertyStructureElement(
-                firFile,
-                newKtDeclaration,
-                originalProperty.symbol,
-                newKtDeclaration.modificationStamp,
-                moduleComponents,
-            )
         }
     }
 }

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/lazy/resolve/LLFirModuleLazyDeclarationResolver.kt

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.tryCollectDesignation
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.runCustomResolveUnderLock
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.llFirModuleData
+import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionInvalidator
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.LLFirFileAnnotationsResolveTransformer
 import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.LLFirFirProviderInterceptor
@@ -20,7 +21,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkCanceled
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.findSourceNonLocalFirDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.withFirEntry
 import org.jetbrains.kotlin.analysis.utils.errors.buildErrorWithAttachment
-import org.jetbrains.kotlin.analysis.utils.errors.buildErrorWithAttachment
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
@@ -61,7 +61,7 @@ internal class LLFirModuleLazyDeclarationResolver(val moduleComponents: LLFirMod
                 )
             }
         } catch (e: Throwable) {
-            rethrowWithDetails(e, firFile, fromPhase, toPhase = null)
+            handleExceptionFromResolve(e, moduleComponents.sessionInvalidator, firFile, fromPhase, toPhase = null)
         }
     }
 
@@ -142,7 +142,7 @@ internal class LLFirModuleLazyDeclarationResolver(val moduleComponents: LLFirMod
                 )
             }
         } catch (e: Throwable) {
-            rethrowWithDetails(e, firFile, fromPhase, toPhase)
+            handleExceptionFromResolve(e,moduleComponents.sessionInvalidator, firFile, fromPhase, toPhase)
         }
     }
 
@@ -243,7 +243,7 @@ internal class LLFirModuleLazyDeclarationResolver(val moduleComponents: LLFirMod
         try {
             doLazyResolveDeclaration(firDeclarationToResolve, scopeSession, toPhase, checkPCE)
         } catch (e: Throwable) {
-            rethrowWithDetails(e, firDeclarationToResolve, fromPhase, toPhase)
+            handleExceptionFromResolve(e, moduleComponents.sessionInvalidator, firDeclarationToResolve, fromPhase, toPhase)
         }
     }
 
@@ -410,12 +410,14 @@ internal class LLFirModuleLazyDeclarationResolver(val moduleComponents: LLFirMod
     }
 }
 
-private fun rethrowWithDetails(
+private fun handleExceptionFromResolve(
     e: Throwable,
+    sessionInvalidator: LLFirSessionInvalidator,
     firDeclarationToResolve: FirDeclaration,
     fromPhase: FirResolvePhase,
     toPhase: FirResolvePhase?
 ): Nothing {
+    sessionInvalidator.invalidate(firDeclarationToResolve.llFirSession)
     if (shouldIjPlatformExceptionBeRethrown(e)) throw e
     buildErrorWithAttachment(
         buildString {

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/project/structure/LLFirNonUnderContentRootSessionFactory.kt

@@ -7,8 +7,8 @@ package org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure
 
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirGlobalResolveComponents
-import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirLazyDeclarationResolver
+import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirDependentModuleProvidersBySessions
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirModuleWithDependenciesSymbolProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.LLFirProvider
@@ -54,7 +54,7 @@ internal class LLFirNonUnderContentRootSessionFactory(private val project: Proje
         val languageVersionSettings = LanguageVersionSettingsImpl.DEFAULT
         val scopeProvider = FirKotlinScopeProvider(::wrapScopeWithJvmMapped)
         val globalResolveComponents = LLFirGlobalResolveComponents(project)
-        val components = LLFirModuleResolveComponents(module, globalResolveComponents, scopeProvider)
+        val components = LLFirModuleResolveComponents(module, globalResolveComponents, scopeProvider, sessionInvalidator)
         val contentScope = module.contentScope
         val session = LLFirNonUnderContentRootSession(module, project, components, builtinsSession.builtinTypes)
         components.session = session
@@ -79,7 +79,7 @@ internal class LLFirNonUnderContentRootSessionFactory(private val project: Proje
             )
 
             register(FirProvider::class, provider)
-            register(FirLazyDeclarationResolver::class, LLFirLazyDeclarationResolver(sessionInvalidator))
+            register(FirLazyDeclarationResolver::class, LLFirLazyDeclarationResolver())
 
             val dependencyProvider = LLFirDependentModuleProvidersBySessions(this) {
                 add(builtinsSession)

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

@@ -9,11 +9,9 @@ import com.intellij.openapi.project.Project
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.search.ProjectScope
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirGlobalResolveComponents
-import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirLazyDeclarationResolver
+import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirModuleResolveComponents
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.*
-import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirLibrarySessionFactory
-import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.registerIdeComponents
 import org.jetbrains.kotlin.analysis.low.level.api.fir.providers.*
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkCanceled
 import org.jetbrains.kotlin.analysis.project.structure.*
@@ -70,7 +68,7 @@ internal object LLFirSessionFactory {
 
         val scopeProvider = FirKotlinScopeProvider(::wrapScopeWithJvmMapped)
 
-        val components = LLFirModuleResolveComponents(module, globalResolveComponents, scopeProvider)
+        val components = LLFirModuleResolveComponents(module, globalResolveComponents, scopeProvider, sessionInvalidator)
 
         val contentScope = module.contentScope
         val session = LLFirSourcesSession(
@@ -103,7 +101,7 @@ internal object LLFirSessionFactory {
             )
 
             register(FirProvider::class, provider)
-            register(FirLazyDeclarationResolver::class, LLFirLazyDeclarationResolver(sessionInvalidator))
+            register(FirLazyDeclarationResolver::class, LLFirLazyDeclarationResolver())
 
             registerCompilerPluginServices(contentScope, project, module)
             registerCompilerPluginExtensions(project, module)
@@ -181,7 +179,7 @@ internal object LLFirSessionFactory {
         }
 
         val scopeProvider = FirKotlinScopeProvider()
-        val components = LLFirModuleResolveComponents(module, globalComponents, scopeProvider)
+        val components = LLFirModuleResolveComponents(module, globalComponents, scopeProvider, sessionInvalidator)
 
         val contentScope = module.contentScope
         val session = LLFirLibraryOrLibrarySourceResolvableModuleSession(module, project, components, builtinSession.builtinTypes)
@@ -209,7 +207,7 @@ internal object LLFirSessionFactory {
 
             register(FirProvider::class, provider)
 
-            register(FirLazyDeclarationResolver::class, LLFirLazyDeclarationResolver(sessionInvalidator))
+            register(FirLazyDeclarationResolver::class, LLFirLazyDeclarationResolver())
 
             // We need FirRegisteredPluginAnnotations during extensions' registration process
             val annotationsResolver = project.createAnnotationResolver(contentScope)

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

@@ -12,4 +12,13 @@ internal class LLFirSessionInvalidator(private val invalidateSourcesSession: (LL
         require(session is LLFirResolvableModuleSession)
         invalidateSourcesSession(session)
     }
+
+    internal inline fun <R : Any> withInvalidationOnException(session: LLFirResolvableModuleSession, action: () -> R): R {
+        try {
+            return action()
+        } catch (e: Throwable) {
+            invalidate(session)
+            throw e
+        }
+    }
 }

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

@@ -11,8 +11,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.FirDeclarationDesigna
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.collectDesignation
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.runCustomResolveUnderLock
-import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.LLFirModuleLazyDeclarationResolver
-import org.jetbrains.kotlin.analysis.low.level.api.fir.lazy.resolve.ResolveTreeBuilder
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirResolvableSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.transformers.LLFirLazyTransformer.Companion.updatePhaseDeep
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkCanceled
@@ -120,7 +118,11 @@ internal class LLFirDesignatedSupertypeResolverTransformer(
         for (designationsPerFile in filesToDesignations) {
             if (checkPCE) checkCanceled()
             lockProvider.runCustomResolveUnderLock(designationsPerFile.key, checkPCE) {
-                applyToFileSymbols(designationsPerFile.value)
+                val session = designationsPerFile.key.llFirResolvableSession
+                    ?: error("When FirFile exists for the declaration, the session should be resolvevablable")
+                session.moduleComponents.sessionInvalidator.withInvalidationOnException(session) {
+                    applyToFileSymbols(designationsPerFile.value)
+                }
             }
         }
     }