[LL FIR] move part of in-block modification logic into service

^KT-60611

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

@@ -11,13 +11,19 @@ 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.api.getFirResolveSession
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
+import org.jetbrains.kotlin.analysis.low.level.api.fir.api.resolveToFirSymbol
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.LLFirDeclarationModificationService.ModificationType
+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.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
+import org.jetbrains.kotlin.psi.KtDeclaration
 
 /**
  * This service is responsible for processing incoming [PsiElement] changes to reflect them on FIR tree.
@@ -56,14 +62,31 @@ class LLFirDeclarationModificationService(val project: Project) {
 
         when (val changeType = calculateChangeType(element, modificationType)) {
             is ChangeType.Invisible -> {}
-            is ChangeType.InBlock -> invalidateAfterInBlockModification(changeType.blockOwner)
+            is ChangeType.InBlock -> inBlockModification(changeType.blockOwner)
-            is ChangeType.OutOfBlock -> {
+            is ChangeType.OutOfBlock -> outOfBlockModification(element)
-                val ktModule = ProjectStructureProvider.getModule(project, element, contextualModule = null)
-                project.analysisMessageBus.syncPublisher(MODULE_OUT_OF_BLOCK_MODIFICATION).onModification(ktModule)
-            }
         }
     }
 
+    private fun inBlockModification(declaration: PsiElement) {
+        val ktModule = ProjectStructureProvider.getModule(project, declaration, contextualModule = null)
+        val resolveSession = ktModule.getFirResolveSession(project)
+        val firDeclaration = when (declaration) {
+            is KtCodeFragment -> declaration.getOrBuildFirFile(resolveSession).codeFragment
+            is KtDeclaration -> declaration.resolveToFirSymbol(resolveSession).fir
+            else -> errorWithFirSpecificEntries(
+                "Unexpected declaration kind: ${declaration::class.simpleName}",
+                psi = declaration,
+            )
+        }
+
+        invalidateAfterInBlockModification(firDeclaration)
+    }
+
+    private fun outOfBlockModification(element: PsiElement) {
+        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
      */

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

@@ -5,13 +5,7 @@
 
 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.api.getFirResolveSession
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFirFile
-import org.jetbrains.kotlin.analysis.low.level.api.fir.api.resolveToFirSymbol
-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.ProjectStructureProvider
 import org.jetbrains.kotlin.fir.contracts.impl.FirEmptyContractDescription
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirBlock
@@ -21,34 +15,18 @@ import org.jetbrains.kotlin.fir.expressions.FirLazyExpression
 import org.jetbrains.kotlin.fir.expressions.builder.buildLazyBlock
 import org.jetbrains.kotlin.fir.expressions.builder.buildLazyExpression
 import org.jetbrains.kotlin.fir.expressions.impl.FirContractCallBlock
-import org.jetbrains.kotlin.psi.KtAnnotated
+import org.jetbrains.kotlin.fir.psi
-import org.jetbrains.kotlin.psi.KtCodeFragment
-import org.jetbrains.kotlin.psi.KtDeclaration
 
 /**
  * Must be called in a write action.
  * @return **false** if it is not in-block modification
  */
-internal fun invalidateAfterInBlockModification(declaration: KtAnnotated): Boolean {
+internal fun invalidateAfterInBlockModification(declaration: FirDeclaration): Boolean = when (declaration) {
-    ApplicationManager.getApplication().assertIsWriteThread()
+    is FirSimpleFunction -> declaration.inBodyInvalidation()
-
+    is FirPropertyAccessor -> declaration.inBodyInvalidation()
-    val project = declaration.project
+    is FirProperty -> declaration.inBodyInvalidation()
-    val ktModule = ProjectStructureProvider.getModule(project, declaration, contextualModule = null)
+    is FirCodeFragment -> declaration.inBodyInvalidation()
-    val resolveSession = ktModule.getFirResolveSession(project)
+    else -> errorWithFirSpecificEntries("Unknown declaration with body", fir = declaration, psi = declaration.psi)
-
-    val firDeclaration = when (declaration) {
-        is KtCodeFragment -> declaration.getOrBuildFirFile(resolveSession).codeFragment
-        is KtDeclaration -> declaration.resolveToFirSymbol(resolveSession).fir
-        else -> errorWithFirSpecificEntries("Unexpected declaration kind: ${declaration.javaClass.simpleName}", psi = declaration)
-    }
-
-    return when (firDeclaration) {
-        is FirSimpleFunction -> firDeclaration.inBodyInvalidation()
-        is FirPropertyAccessor -> firDeclaration.inBodyInvalidation()
-        is FirProperty -> firDeclaration.inBodyInvalidation()
-        is FirCodeFragment -> firDeclaration.inBodyInvalidation()
-        else -> errorWithFirSpecificEntries("Unknown declaration with body", fir = firDeclaration, psi = declaration)
-    }
 }
 
 /**