[LL FIR] reduce number of in-block modifications

The idea of optimization is simple – delay in-block modification
until the end of a write action.

Q: Why the end of write action?
A: We have to publish the result of modification somewhere during
write action to be sure that the next read action will see the
updated result. We can't publish the result somewhere later
because we can't guarantee that the modification will be visible
to all customers.

In the case of out-of-block modification, we still have to publish
the result immediately because it is hard to evaluate the consequences
of the opposite decision yet.

If an out-of-block modification happens, we can drop all previous
in-block modifications from the queue because they don't make sense
due to invalidation of the entire KtModule.

A corner case is the analysis under write action. The delay means that
there is no longer any guarantee that all PSI changes will be reflected
into FIR tree immediately. So now we can guaranty only that at
the start of `analyze` block you will have the up-to-date FIR tree.

^KT-60611 Fixed

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

@@ -1,27 +1,30 @@
 /*
- * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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.api.fir
 
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.LowMemoryWatcher
 import com.intellij.psi.util.CachedValue
 import com.intellij.psi.util.CachedValueProvider
 import com.intellij.psi.util.CachedValuesManager
 import com.intellij.util.CachedValueBase
-import java.util.concurrent.ConcurrentHashMap
 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.lifetime.KtReadActionConfinementLifetimeToken
 import org.jetbrains.kotlin.analysis.api.session.KtAnalysisSessionProvider
+import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirInternals
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getFirResolveSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.LLFirDeclarationModificationService
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.analysis.providers.createProjectWideOutOfBlockModificationTracker
 import org.jetbrains.kotlin.psi.KtElement
+import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.ConcurrentMap
 import kotlin.reflect.KClass
 
@@ -39,7 +42,10 @@ class KtFirAnalysisSessionProvider(project: Project) : KtAnalysisSessionProvider
     }
 
     override fun getAnalysisSessionByUseSiteKtModule(useSiteKtModule: KtModule): KtAnalysisSession {
-        val key = Pair(useSiteKtModule, tokenFactory.identifier)
+        val identifier = tokenFactory.identifier
+        identifier.flushPendingChanges(project)
+
+        val key = Pair(useSiteKtModule, identifier)
         return cache.computeIfAbsent(key) {
             CachedValuesManager.getManager(project).createCachedValue {
                 val firResolveSession = useSiteKtModule.getFirResolveSession(project)
@@ -63,4 +69,15 @@ class KtFirAnalysisSessionProvider(project: Project) : KtAnalysisSessionProvider
             cachedValue.clear()
         }
     }
-}
+}
+
+private fun KClass<out KtLifetimeToken>.flushPendingChanges(project: Project) {
+    if (this == KtReadActionConfinementLifetimeToken::class &&
+        KtReadActionConfinementLifetimeToken.allowFromWriteAction.get() &&
+        ApplicationManager.getApplication().isWriteAccessAllowed
+    ) {
+        // We must flush modifications to publish local modifications into FIR tree
+        @OptIn(LLFirInternals::class)
+        LLFirDeclarationModificationService.getInstance(project).flushModifications()
+    }
+}