[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()
+    }
+}

analysis/analysis-api/src/org/jetbrains/kotlin/analysis/api/lifetime/KtReadActionConfinementLifetimeToken.kt

@@ -111,6 +111,29 @@ public inline fun <T> allowAnalysisOnEdt(action: () -> T): T {
  * Analysis is not supposed to be called from write action.
  * Such actions can lead to IDE freezes and incorrect behavior in some cases.
  *
+ * There is no guarantee that PSI changes will be reflected in an Analysis API world inside
+ * one [analyze] session.
+ * Example:
+ * ```
+ * // code to be analyzed
+ * fun foo(): Int = 0
+ *
+ * // use case code
+ * fun useCase() {
+ *   analyse(function) {
+ *    // 'getConstantFromExpressionBody' is an imaginary function
+ *    val valueBefore = function.getConstantFromExpressionBody() // valueBefore is 0
+ *
+ *    changeExpressionBodyTo(1) // now function will looks like `fun foo(): Int = 1`
+ *    val valueAfter = function.getConstantFromExpressionBody() // Wrong way: valueAfter is not guarantied to be '1'
+ *   }
+ *
+ *   analyse(function) {
+ *    val valueAfter = function.getConstantFromExpressionBody() // OK: valueAfter is guarantied to be '1'
+ *   }
+ * }
+ * ```
+ *
  * @see KtAnalysisSession
  * @see KtReadActionConfinementLifetimeToken
  */

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

@@ -5,6 +5,8 @@
 
 package org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure
 
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.application.ApplicationListener
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiComment
@@ -19,6 +21,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirResolvableM
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.llFirResolvableSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.codeFragment
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecificEntries
+import org.jetbrains.kotlin.analysis.project.structure.KtModule
 import org.jetbrains.kotlin.analysis.project.structure.ProjectStructureProvider
 import org.jetbrains.kotlin.analysis.providers.analysisMessageBus
 import org.jetbrains.kotlin.analysis.providers.topics.KotlinTopics.MODULE_OUT_OF_BLOCK_MODIFICATION
@@ -39,7 +42,70 @@ import org.jetbrains.kotlin.psi.KtElement
  * @see org.jetbrains.kotlin.analysis.providers.topics.KotlinModuleOutOfBlockModificationListener
  */
 @LLFirInternals
-class LLFirDeclarationModificationService(val project: Project) {
+class LLFirDeclarationModificationService(val project: Project) : Disposable {
+    init {
+        ApplicationManager.getApplication().addApplicationListener(
+            object : ApplicationListener {
+                override fun writeActionFinished(action: Any) {
+                    flushModifications()
+                }
+            },
+            this,
+        )
+    }
+
+    private var inBlockModificationQueue: MutableSet<ChangeType.InBlock>? = null
+
+    private fun addModificationToQueue(modification: ChangeType.InBlock) {
+        val queue = inBlockModificationQueue ?: HashSet<ChangeType.InBlock>().also { inBlockModificationQueue = it }
+        queue += modification
+    }
+
+    /**
+     * We can avoid processing of in-block modification with the same [KtModule] because they
+     * will be invalidated anyway by OOBM
+     */
+    private fun dropOutdatedModifications(ktModuleWithOutOfBlockModification: KtModule) {
+        processQueue { value, iterator ->
+            if (value.ktModule == ktModuleWithOutOfBlockModification) iterator.remove()
+        }
+    }
+
+    /**
+     * Process valid elements in the current queue.
+     * Non-valid elements will be dropped from the queue during this iteration.
+     *
+     * @param action will be executed for each valid element in the queue;
+     * **value** is a current element;
+     * **iterator** is the corresponding iterator for this element.
+     */
+    private inline fun processQueue(action: (value: ChangeType.InBlock, iterator: MutableIterator<ChangeType.InBlock>) -> Unit) {
+        val queue = inBlockModificationQueue ?: return
+        val iterator = queue.iterator()
+        while (iterator.hasNext()) {
+            val element = iterator.next()
+            if (!element.blockOwner.isValid) {
+                iterator.remove()
+                continue
+            }
+
+            action(element, iterator)
+        }
+    }
+
+    /**
+     * Force the service to publish delayed modifications. This action is required to fix inconsistencies in FirFile tree.
+     */
+    fun flushModifications() {
+        ApplicationManager.getApplication().assertIsWriteThread()
+
+        processQueue { value, _ ->
+            inBlockModification(value.blockOwner, value.ktModule)
+        }
+
+        inBlockModificationQueue = null
+    }
+
     sealed class ModificationType {
         object NewElement : ModificationType()
         object Unknown : ModificationType()
@@ -65,13 +131,28 @@ class LLFirDeclarationModificationService(val project: Project) {
 
         when (val changeType = calculateChangeType(element, modificationType)) {
             is ChangeType.Invisible -> {}
-            is ChangeType.InBlock -> inBlockModification(changeType.blockOwner)
+            is ChangeType.InBlock -> addModificationToQueue(changeType)
             is ChangeType.OutOfBlock -> outOfBlockModification(element)
         }
     }
 
-    private fun inBlockModification(declaration: KtElement) {
+    private fun calculateChangeType(element: PsiElement, modificationType: ModificationType): ChangeType = when {
-        val ktModule = ProjectStructureProvider.getModule(project, declaration, contextualModule = null)
+        // If PSI is not valid, well something bad happened; OOBM won't hurt
+        !element.isValid -> ChangeType.OutOfBlock
+        element is PsiWhiteSpace || element is PsiComment -> ChangeType.Invisible
+        // TODO improve for Java KTIJ-21684
+        element.language !is KotlinLanguage -> ChangeType.OutOfBlock
+        else -> {
+            val inBlockModificationOwner = nonLocalDeclarationForLocalChange(element)
+            if (inBlockModificationOwner != null && (element.parent != inBlockModificationOwner || modificationType != ModificationType.NewElement)) {
+                ChangeType.InBlock(inBlockModificationOwner, project)
+            } else {
+                ChangeType.OutOfBlock
+            }
+        }
+    }
+
+    private fun inBlockModification(declaration: KtElement, ktModule: KtModule) {
         val resolveSession = ktModule.getFirResolveSession(project)
         val firDeclaration = when (declaration) {
             is KtCodeFragment -> declaration.getOrBuildFirFile(resolveSession).codeFragment
@@ -102,6 +183,9 @@ class LLFirDeclarationModificationService(val project: Project) {
 
     private fun outOfBlockModification(element: PsiElement) {
         val ktModule = ProjectStructureProvider.getModule(project, element, contextualModule = null)
+
+        // We should check outdated modifications before to avoid cache dropping (e.g., KtModule cache)
+        dropOutdatedModifications(ktModule)
         project.analysisMessageBus.syncPublisher(MODULE_OUT_OF_BLOCK_MODIFICATION).onModification(ktModule)
     }
 
@@ -112,28 +196,14 @@ class LLFirDeclarationModificationService(val project: Project) {
         return nonLocalDeclarationForLocalChange(changedElement)
     }
 
+    override fun dispose() {}
+
     companion object {
         fun getInstance(project: Project): LLFirDeclarationModificationService =
             project.getService(LLFirDeclarationModificationService::class.java)
     }
 }
 
-private fun calculateChangeType(element: PsiElement, modificationType: ModificationType): ChangeType = when {
-    // If PSI is not valid, well something bad happened, OOBM won't hurt
-    !element.isValid -> ChangeType.OutOfBlock
-    element is PsiWhiteSpace || element is PsiComment -> ChangeType.Invisible
-    // TODO improve for Java KTIJ-21684
-    element.language !is KotlinLanguage -> ChangeType.OutOfBlock
-    else -> {
-        val inBlockModificationOwner = nonLocalDeclarationForLocalChange(element)
-        if (inBlockModificationOwner != null && (element.parent != inBlockModificationOwner || modificationType != ModificationType.NewElement)) {
-            ChangeType.InBlock(inBlockModificationOwner)
-        } else {
-            ChangeType.OutOfBlock
-        }
-    }
-}
-
 private fun nonLocalDeclarationForLocalChange(psi: PsiElement): KtAnnotated? {
     return psi.getNonLocalReanalyzableContainingDeclaration() ?: psi.containingFile as? KtCodeFragment
 }
@@ -141,5 +211,13 @@ private fun nonLocalDeclarationForLocalChange(psi: PsiElement): KtAnnotated? {
 private sealed class ChangeType {
     object OutOfBlock : ChangeType()
     object Invisible : ChangeType()
-    class InBlock(val blockOwner: KtAnnotated) : ChangeType()
+
+    class InBlock(val blockOwner: KtAnnotated, val project: Project) : ChangeType() {
+        val ktModule: KtModule by lazy(LazyThreadSafetyMode.NONE) {
+            ProjectStructureProvider.getModule(project, blockOwner, contextualModule = null)
+        }
+
+        override fun equals(other: Any?): Boolean = other === this || other is InBlock && other.blockOwner == blockOwner
+        override fun hashCode(): Int = blockOwner.hashCode()
+    }
 }

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/AbstractInBlockModificationTest.kt

@@ -77,13 +77,21 @@ internal fun testInBlockModification(
 
     val textBefore = declarationToRender.render()
 
-    val isOutOfBlock = LLFirDeclarationModificationService.getInstance(elementToModify.project).modifyElement(elementToModify)
+    val modificationService = LLFirDeclarationModificationService.getInstance(elementToModify.project)
+    val isOutOfBlock = modificationService.modifyElement(elementToModify)
     if (isOutOfBlock) {
         return "IN-BLOCK MODIFICATION IS NOT APPLICABLE FOR THIS PLACE"
     }
 
     elementToModify.modify()
 
+    val textAfterPsiModification = declarationToRender.render()
+    testServices.assertions.assertEquals(textBefore, textAfterPsiModification) {
+        "The declaration before and after modification must be in the same state, because changes in not flushed yet"
+    }
+
+    modificationService.flushModifications()
+
     val textAfterModification = declarationToRender.render()
     testServices.assertions.assertNotEquals(textBefore, textAfterModification) {
         "The declaration before and after modification must be in different state"