[LL FIR] remove retryOnInvalidSession as now we have correct FIR tree restore after PCE

Previously, `retryOnInvalidSession` was needed to try resolving declaration again if the `FirSession` is invalidated and PCE was thrown.
Currently, we do not invalidate `FirSession` on PCE as we can restore the FIR tree state after the PCE was thrown.

Now `FirSession` might be invalidated only on OOBM or module content root modification, in this case, we should not get access to the invalidated `FirSession` at all

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/builder/LLFirLockProvider.kt

@@ -30,7 +30,6 @@ internal class LLFirLockProvider(private val checker: LLFirLazyResolveContractCh
     private val implicitTypesLock = ReentrantLock()
 
     inline fun <R> withGlobalLock(
-        key: FirFile,
         lockingIntervalMs: Long = DEFAULT_LOCKING_INTERVAL,
         action: () -> R,
     ): R {
@@ -38,14 +37,11 @@ internal class LLFirLockProvider(private val checker: LLFirLazyResolveContractCh
 
         return globalLock.lockWithPCECheck(lockingIntervalMs) {
             val session = key.llFirSession
-            if (!session.isValid && shouldRetryFlag.get()) {
+            if (!session.isValid) {
                 val description = session.ktModule.moduleDescription
                 throw InvalidSessionException("Session '$description' is invalid", description)
             }
-
+            action()
-            // Normally, analysis should not be allowed on an invalid session.
-            // However, there isn't an easy way to cancel or redo it in general case, as it must then be supported on use-site.
-            withRetryFlag(false, action)
         }
     }
 
@@ -236,46 +232,3 @@ private val globalLockEnabled: Boolean by lazy(LazyThreadSafetyMode.PUBLICATION)
 private const val DEFAULT_LOCKING_INTERVAL = 50L
 
 internal class InvalidSessionException(message: String, val moduleDescription: String) : RuntimeException(message)
-
-/*
-    The flag specifies whether the analysis action should be repeated in case if it was originally started on an invalid session.
-
-    Possible values:
-        - `true` – throw the marker [InvalidSessionException] to trigger the retry.
-        - `false` – process analysis as usual (default).
- */
-private val shouldRetryFlag: ThreadLocal<Boolean> = ThreadLocal.withInitial { false }
-
-private val LOG = Logger.getInstance(LLFirLockProvider::class.java)
-
-/**
- * Retry the `action` calculation with a new FIR session if session passed to [LLFirLockProvider.withWriteLock] turns to be invalid.
- * This is a temporary solution to fix inconsistent analysis state in common cases of idempotent analysis.
- * The right solution would be to modify the FIR tree after the analysis is done, so the tree will always be in consistent state.
- */
-internal inline fun <R> retryOnInvalidSession(action: () -> R): R {
-    withRetryFlag(true) {
-        while (true) {
-            try {
-                return action()
-            } catch (e: InvalidSessionException) {
-                LOG.warn("Processing with invalid module '${e.moduleDescription}'")
-            }
-        }
-    }
-}
-
-@OptIn(ExperimentalContracts::class)
-private inline fun <R> withRetryFlag(value: Boolean, action: () -> R): R {
-    contract {
-        callsInPlace(action, InvocationKind.EXACTLY_ONCE)
-    }
-
-    val oldValue = shouldRetryFlag.get()
-    shouldRetryFlag.set(value)
-    try {
-        return action()
-    } finally {
-        shouldRetryFlag.set(oldValue)
-    }
-}

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

@@ -14,7 +14,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.LLFirResolveSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getModule
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.FirTowerContextProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingDeclaration
-import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.retryOnInvalidSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirResolvableModuleSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionCache
@@ -80,17 +79,13 @@ internal abstract class LLFirResolvableResolveSession(
     }
 
     override fun getOrBuildFirFor(element: KtElement): FirElement? {
-        retryOnInvalidSession {
+        val moduleComponents = getModuleComponentsForElement(element)
-            val moduleComponents = getModuleComponentsForElement(element)
+        return moduleComponents.elementsBuilder.getOrBuildFirFor(element, this)
-            return moduleComponents.elementsBuilder.getOrBuildFirFor(element, this)
-        }
     }
 
     override fun getOrBuildFirFile(ktFile: KtFile): FirFile {
-        retryOnInvalidSession {
+        val moduleComponents = getModuleComponentsForElement(ktFile)
-            val moduleComponents = getModuleComponentsForElement(ktFile)
+        return moduleComponents.firFileBuilder.buildRawFirFileWithCaching(ktFile)
-            return moduleComponents.firFileBuilder.buildRawFirFileWithCaching(ktFile)
-        }
     }
 
     protected fun getModuleComponentsForElement(element: KtElement): LLFirModuleResolveComponents {
@@ -100,15 +95,13 @@ internal abstract class LLFirResolvableResolveSession(
 
     override fun resolveToFirSymbol(
         ktDeclaration: KtDeclaration,
-        phase: FirResolvePhase
+        phase: FirResolvePhase,
     ): FirBasedSymbol<*> {
         val containingKtFile = ktDeclaration.containingKtFile
         val module = getModule(containingKtFile.originalKtFile ?: containingKtFile)
-        retryOnInvalidSession {
+        return when (getModuleKind(module)) {
-            return when (getModuleKind(module)) {
+            ModuleKind.RESOLVABLE_MODULE -> findSourceFirSymbol(ktDeclaration, module).also { resolveFirToPhase(it.fir, phase) }
-                ModuleKind.RESOLVABLE_MODULE -> findSourceFirSymbol(ktDeclaration, module).also { resolveFirToPhase(it.fir, phase) }
+            ModuleKind.BINARY_MODULE -> findFirCompiledSymbol(ktDeclaration, module)
-                ModuleKind.BINARY_MODULE -> findFirCompiledSymbol(ktDeclaration, module)
-            }
         }
     }
 

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

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.state
 
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilter
-import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.retryOnInvalidSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.project.structure.*
 import org.jetbrains.kotlin.analysis.utils.errors.unexpectedElementError
@@ -16,20 +15,16 @@ import org.jetbrains.kotlin.psi.KtFile
 
 internal class LLFirScriptResolveSession(
     useSiteKtModule: KtModule,
-    useSiteSessionFactory: (KtModule) -> LLFirSession
+    useSiteSessionFactory: (KtModule) -> LLFirSession,
 ) : LLFirResolvableResolveSession(useSiteKtModule, useSiteSessionFactory) {
     override fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic> {
-        retryOnInvalidSession {
+        val moduleComponents = getModuleComponentsForElement(element)
-            val moduleComponents = getModuleComponentsForElement(element)
+        return moduleComponents.diagnosticsCollector.getDiagnosticsFor(element, filter)
-            return moduleComponents.diagnosticsCollector.getDiagnosticsFor(element, filter)
-        }
     }
 
     override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic> {
-        retryOnInvalidSession {
+        val moduleComponents = getModuleComponentsForElement(ktFile)
-            val moduleComponents = getModuleComponentsForElement(ktFile)
+        return moduleComponents.diagnosticsCollector.collectDiagnosticsForFile(ktFile, filter)
-            return moduleComponents.diagnosticsCollector.collectDiagnosticsForFile(ktFile, filter)
-        }
     }
 
     override fun getModuleKind(module: KtModule): ModuleKind {

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

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.state
 
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.DiagnosticCheckerFilter
-import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.retryOnInvalidSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.project.structure.KtBuiltinsModule
 import org.jetbrains.kotlin.analysis.project.structure.KtLibraryModule
@@ -19,20 +18,16 @@ import org.jetbrains.kotlin.psi.KtFile
 
 internal class LLFirSourceResolveSession(
     useSiteKtModule: KtModule,
-    useSiteSessionFactory: (KtModule) -> LLFirSession
+    useSiteSessionFactory: (KtModule) -> LLFirSession,
 ) : LLFirResolvableResolveSession(useSiteKtModule, useSiteSessionFactory) {
     override fun getDiagnostics(element: KtElement, filter: DiagnosticCheckerFilter): List<KtPsiDiagnostic> {
-        retryOnInvalidSession {
+        val moduleComponents = getModuleComponentsForElement(element)
-            val moduleComponents = getModuleComponentsForElement(element)
+        return moduleComponents.diagnosticsCollector.getDiagnosticsFor(element, filter)
-            return moduleComponents.diagnosticsCollector.getDiagnosticsFor(element, filter)
-        }
     }
 
     override fun collectDiagnosticsForFile(ktFile: KtFile, filter: DiagnosticCheckerFilter): Collection<KtPsiDiagnostic> {
-        retryOnInvalidSession {
+        val moduleComponents = getModuleComponentsForElement(ktFile)
-            val moduleComponents = getModuleComponentsForElement(ktFile)
+        return moduleComponents.diagnosticsCollector.collectDiagnosticsForFile(ktFile, filter)
-            return moduleComponents.diagnosticsCollector.collectDiagnosticsForFile(ktFile, filter)
-        }
     }
 
     override fun getModuleKind(module: KtModule): ModuleKind {

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

@@ -22,7 +22,7 @@ internal object LLFirLazyResolverRunner {
         val lazyResolver = LLFirLazyPhaseResolverByPhase.getByPhase(phase)
         val firFile = target.firFile
         val session = firFile.moduleData.session
-        lockProvider.withGlobalLock(firFile) {
+        lockProvider.withGlobalLock {
             lockProvider.withGlobalPhaseLock(phase) {
                 lazyResolver.resolve(target, lockProvider, session, scopeSession, towerDataContextCollector)
             }