[Low Level FIR] introduce service for incoming psi modifications processing

This service encapsulates all logic and simplifies future evolution

^KT-60610

analysis/analysis-api-standalone/analysis-api-fir-standalone-base/src/org/jetbrains/kotlin/analysis/api/standalone/base/project/structure/FirStandaloneServiceRegistrar.kt

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.analysis.api.session.KtAnalysisSessionProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirGlobalResolveComponents
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirInternals
 import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirResolveSessionService
+import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.LLFirDeclarationModificationService
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.JvmFirDeserializedSymbolProviderFactory
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.LLFirBuiltinsSessionFactory
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSessionCache
@@ -60,6 +61,8 @@ object FirStandaloneServiceRegistrar : AnalysisApiStandaloneServiceRegistrar {
 
             registerService(LLFirSessionInvalidationService::class.java)
             LLFirSessionInvalidationService.getInstance(project).subscribeToModificationEvents()
+
+            registerService(LLFirDeclarationModificationService::class.java)
         }
     }
 

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

@@ -117,9 +117,7 @@ internal object FileElementFactory {
  *
  * @return The declaration in which a change of the passed receiver parameter can be treated as in-block modification
  */
-@LLFirInternals
+internal fun PsiElement.getNonLocalReanalyzableContainingDeclaration(): KtDeclaration? {
-@Suppress("unused") // Used in the IDE plugin
-fun PsiElement.getNonLocalReanalyzableContainingDeclaration(): KtDeclaration? {
     return when (val declaration = getNonLocalContainingOrThisDeclaration()) {
         is KtNamedFunction -> declaration.takeIf { function ->
             function.isReanalyzableContainer() && isElementInsideBody(declaration = function, child = this)

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

@@ -0,0 +1,104 @@
+/*
+ * 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.low.level.api.fir.file.structure
+
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiComment
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiWhiteSpace
+import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirInternals
+import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.LLFirDeclarationModificationService.ModificationType
+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
+import org.jetbrains.kotlin.idea.KotlinLanguage
+import org.jetbrains.kotlin.psi.KtAnnotated
+import org.jetbrains.kotlin.psi.KtCodeFragment
+
+/**
+ * This service is responsible for processing incoming [PsiElement] changes to reflect them on FIR tree.
+ *
+ * For local changes (in-block modification), this service will do all required work.
+ *
+ * In case of non-local changes (out-of-block modification), this service will publish event to [MODULE_OUT_OF_BLOCK_MODIFICATION].
+ *
+ * @see MODULE_OUT_OF_BLOCK_MODIFICATION
+ * @see org.jetbrains.kotlin.analysis.providers.topics.KotlinModuleOutOfBlockModificationListener
+ */
+@LLFirInternals
+class LLFirDeclarationModificationService(val project: Project) {
+    sealed class ModificationType {
+        object NewElement : ModificationType()
+        object Unknown : ModificationType()
+    }
+
+    /**
+     * This method should be called during some [PsiElement] modification.
+     * This method must be called from write action.
+     *
+     * Will publish event to [MODULE_OUT_OF_BLOCK_MODIFICATION] in case of out-of-block modification.
+     *
+     * @param element is an element that we want to modify, remove or add.
+     * Some examples:
+     * * [element] is [KtNamedFunction][org.jetbrains.kotlin.psi.KtNamedFunction] if we
+     * dropped body ([KtBlockExpression][org.jetbrains.kotlin.psi.KtBlockExpression]) of this function
+     * * [element] is [KtBlockExpression][org.jetbrains.kotlin.psi.KtBlockExpression] if we replaced one body-expression with another one
+     * * [element] is [KtBlockExpression][org.jetbrains.kotlin.psi.KtBlockExpression] if added a body to the function without body
+     *
+     * @param modificationType additional information to make more accurate decisions
+     */
+    fun elementModified(element: PsiElement, modificationType: ModificationType = ModificationType.Unknown) {
+        ApplicationManager.getApplication().assertIsWriteThread()
+
+        when (val changeType = calculateChangeType(element, modificationType)) {
+            is ChangeType.Invisible -> {}
+            is ChangeType.InBlock -> invalidateAfterInBlockModification(changeType.blockOwner)
+            is ChangeType.OutOfBlock -> {
+                val ktModule = ProjectStructureProvider.getModule(project, element, contextualModule = null)
+                project.analysisMessageBus.syncPublisher(MODULE_OUT_OF_BLOCK_MODIFICATION).onModification(ktModule)
+            }
+        }
+    }
+
+    /**
+     * @return the psi element (ancestor of the changedElement) which should be re-highlighted in case of in-block changes or null if unsure
+     */
+    fun elementToRehighlight(changedElement: PsiElement): PsiElement? {
+        return nonLocalDeclarationForLocalChange(changedElement)
+    }
+
+    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
+}
+
+private sealed class ChangeType {
+    object OutOfBlock : ChangeType()
+    object Invisible : ChangeType()
+    class InBlock(val blockOwner: KtAnnotated) : ChangeType()
+}

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

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure
 
 import com.intellij.openapi.application.ApplicationManager
-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.api.getOrBuildFirFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.resolveToFirSymbol
@@ -26,23 +25,11 @@ import org.jetbrains.kotlin.psi.KtAnnotated
 import org.jetbrains.kotlin.psi.KtCodeFragment
 import org.jetbrains.kotlin.psi.KtDeclaration
 
-@Deprecated(
-    "Temporarily left for binary compatibility. Use invalidateAfterInBlockModification(KtElement) instead.",
-    replaceWith = ReplaceWith("invalidateAfterInBlockModification(declaration)", "org.jetbrains.kotlin.psi.KtElement"),
-    level = DeprecationLevel.HIDDEN,
-)
-@LLFirInternals
-@Suppress("unused")
-fun invalidateAfterInBlockModification(declaration: KtDeclaration): Boolean {
-    return invalidateAfterInBlockModification(declaration as KtAnnotated)
-}
-
 /**
  * Must be called in a write action.
  * @return **false** if it is not in-block modification
  */
-@LLFirInternals
+internal fun invalidateAfterInBlockModification(declaration: KtAnnotated): Boolean {
-fun invalidateAfterInBlockModification(declaration: KtAnnotated): Boolean {
     ApplicationManager.getApplication().assertIsWriteThread()
 
     val project = declaration.project

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

@@ -21,8 +21,7 @@ abstract class AbstractCodeFragmentInBlockModificationTest : AbstractLowLevelApi
 
         assertNull(targetElement.getNonLocalReanalyzableContainingDeclaration())
 
-        val (before, after) = testInBlockModification(ktCodeFragment, ktCodeFragment, testServices, dumpFirFile = false)
+        val actualText = testInBlockModification(ktCodeFragment, ktCodeFragment, testServices, dumpFirFile = false)
-        val actualText = "BEFORE MODIFICATION:\n$before\nAFTER MODIFICATION:\n$after"
 
         testServices.assertions.assertEqualsToTestDataFileSibling(actualText)
     }

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

@@ -6,20 +6,25 @@
 package org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure
 
 import com.intellij.extapi.psi.ASTDelegatePsiElement
+import com.intellij.openapi.util.Disposer
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
 import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirOfType
+import org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.getNonLocalContainingOrThisDeclaration
 import org.jetbrains.kotlin.analysis.low.level.api.fir.lazyResolveRenderer
 import org.jetbrains.kotlin.analysis.low.level.api.fir.resolveWithCaches
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.base.AbstractLowLevelApiSingleFileTest
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.AnalysisApiFirOutOfContentRootTestConfigurator
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.AnalysisApiFirScriptTestConfigurator
 import org.jetbrains.kotlin.analysis.low.level.api.fir.test.configurators.AnalysisApiFirSourceTestConfigurator
+import org.jetbrains.kotlin.analysis.providers.analysisMessageBus
+import org.jetbrains.kotlin.analysis.providers.topics.KotlinModuleOutOfBlockModificationListener
+import org.jetbrains.kotlin.analysis.providers.topics.KotlinTopics
 import org.jetbrains.kotlin.analysis.test.framework.services.expressionMarkerProvider
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.declarations.FirPropertyAccessor
 import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.symbols.lazyResolveToPhase
-import org.jetbrains.kotlin.psi.KtAnnotated
 import org.jetbrains.kotlin.psi.KtCodeFragment
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
@@ -41,19 +46,12 @@ abstract class AbstractInBlockModificationTest : AbstractLowLevelApiSingleFileTe
             module = moduleStructure.modules.last(),
         )
 
-        val declaration = selectedElement.getNonLocalReanalyzableContainingDeclaration()
+        val actual = testInBlockModification(
-        val actual = if (declaration != null) {
+            file = ktFile,
-            val (before, after) = testInBlockModification(
+            elementToModify = selectedElement,
-                file = ktFile,
+            testServices = testServices,
-                declaration = declaration,
+            dumpFirFile = Directives.DUMP_FILE in moduleStructure.allDirectives,
-                testServices = testServices,
+        )
-                dumpFirFile = Directives.DUMP_FILE in moduleStructure.allDirectives,
-            )
-
-            "BEFORE MODIFICATION:\n$before\nAFTER MODIFICATION:\n$after"
-        } else {
-            "IN-BLOCK MODIFICATION IS NOT APPLICABLE FOR THIS PLACE"
-        }
 
         testServices.assertions.assertEqualsToTestDataFileSibling(actual)
     }
@@ -65,10 +63,11 @@ abstract class AbstractInBlockModificationTest : AbstractLowLevelApiSingleFileTe
 
 internal fun testInBlockModification(
     file: KtFile,
-    declaration: KtAnnotated,
+    elementToModify: PsiElement,
     testServices: TestServices,
     dumpFirFile: Boolean,
-): Pair<String, String> = resolveWithCaches(file) { firSession ->
+): String = resolveWithCaches(file) { firSession ->
+    val declaration = elementToModify.getNonLocalContainingOrThisDeclaration() ?: file
     val firDeclarationBefore = declaration.getOrBuildFirOfType<FirDeclaration>(firSession)
     val declarationToRender = if (dumpFirFile) {
         file.getOrBuildFirFile(firSession).also { it.lazyResolveToPhase(FirResolvePhase.BODY_RESOLVE) }
@@ -78,8 +77,12 @@ internal fun testInBlockModification(
 
     val textBefore = declarationToRender.render()
 
-    declaration.modifyBody()
+    val isOutOfBlock = LLFirDeclarationModificationService.getInstance(elementToModify.project).modifyElement(elementToModify)
-    invalidateAfterInBlockModification(declaration)
+    if (isOutOfBlock) {
+        return "IN-BLOCK MODIFICATION IS NOT APPLICABLE FOR THIS PLACE"
+    }
+
+    elementToModify.modify()
 
     val textAfterModification = declarationToRender.render()
     testServices.assertions.assertNotEquals(textBefore, textAfterModification) {
@@ -105,13 +108,33 @@ internal fun testInBlockModification(
         "The declaration must have the same in the resolved state"
     }
 
-    Pair(textBefore, textAfterModification)
+    "BEFORE MODIFICATION:\n$textBefore\nAFTER MODIFICATION:\n$textAfterModification"
+}
+
+/**
+ * @return **true** if out-of-block happens
+ */
+private fun LLFirDeclarationModificationService.modifyElement(element: PsiElement): Boolean {
+    val disposable = Disposer.newDisposable()
+    var isOutOfBlock = false
+    try {
+        project.analysisMessageBus.connect(disposable).subscribe(
+            KotlinTopics.MODULE_OUT_OF_BLOCK_MODIFICATION,
+            KotlinModuleOutOfBlockModificationListener { isOutOfBlock = true },
+        )
+
+        elementModified(element)
+    } finally {
+        Disposer.dispose(disposable)
+    }
+
+    return isOutOfBlock
 }
 
 /**
  * Emulate modification inside the body
  */
-private fun KtAnnotated.modifyBody() {
+private fun PsiElement.modify() {
     for (parent in parentsWithSelf) {
         when (parent) {
             is ASTDelegatePsiElement -> parent.subtreeChanged()