From 1ffec2b28714d062646f582a405db174a4024e0a Mon Sep 17 00:00:00 2001
From: Marco Pennekamp <marco.pennekamp@jetbrains.com>
Date: Wed, 26 Jun 2024 19:17:16 +0200
Subject: [PATCH] [kotlin] Replace `LLFirDeclarationModificationService` with
 `KaSourceModificationService`

^KT-69436

GitOrigin-RevId: 5612f8082f322e8f65f5e235e2c83356879a5109
---
 ...rIdeOutOfBlockPsiTreeChangePreprocessor.kt | 35 +++++++++----------
 .../KotlinChangeLocalityDetector.kt           |  7 ++--
 2 files changed, 19 insertions(+), 23 deletions(-)

diff --git a/plugins/kotlin/base/fir/analysis-api-platform/src/org/jetbrains/kotlin/idea/base/fir/analysisApiPlatform/FirIdeOutOfBlockPsiTreeChangePreprocessor.kt b/plugins/kotlin/base/fir/analysis-api-platform/src/org/jetbrains/kotlin/idea/base/fir/analysisApiPlatform/FirIdeOutOfBlockPsiTreeChangePreprocessor.kt
index dd2e2782b111..cc99d48bac49 100644
--- a/plugins/kotlin/base/fir/analysis-api-platform/src/org/jetbrains/kotlin/idea/base/fir/analysisApiPlatform/FirIdeOutOfBlockPsiTreeChangePreprocessor.kt
+++ b/plugins/kotlin/base/fir/analysis-api-platform/src/org/jetbrains/kotlin/idea/base/fir/analysisApiPlatform/FirIdeOutOfBlockPsiTreeChangePreprocessor.kt
@@ -12,11 +12,10 @@ import com.intellij.psi.impl.PsiTreeChangePreprocessor
 import com.intellij.psi.impl.source.tree.injected.InjectedLanguageUtilBase
 import com.intellij.psi.util.parentOfType
 import org.jetbrains.kotlin.analysis.api.platform.analysisMessageBus
+import org.jetbrains.kotlin.analysis.api.platform.modification.KaElementModificationType
+import org.jetbrains.kotlin.analysis.api.platform.modification.KaSourceModificationService
 import org.jetbrains.kotlin.analysis.api.platform.modification.KotlinModificationTopics
 import org.jetbrains.kotlin.analysis.api.platform.projectStructure.KotlinProjectStructureProvider
-import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirInternals
-import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.LLFirDeclarationModificationService
-import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.LLFirDeclarationModificationService.ModificationType
 import org.jetbrains.kotlin.idea.util.publishGlobalSourceOutOfBlockModification
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
@@ -56,22 +55,20 @@ internal class FirIdeOutOfBlockPsiTreeChangePreprocessor(private val project: Pr
             return
         }
 
-        @OptIn(LLFirInternals::class)
-        LLFirDeclarationModificationService.getInstance(project).elementModified(
-            element = child ?: rootElement,
-            modificationType = when (event.code) {
-                PsiEventType.CHILD_ADDED -> ModificationType.ElementAdded
-                PsiEventType.CHILD_REMOVED -> {
-                    val removedElement = event.child ?:
-                        errorWithAttachment("A ${PsiEventType.CHILD_REMOVED} PSI tree change event should have a child element") {
-                            withEntry("psiTreeChangeEvent", event.toString())
-                            withPsiEntry("rootElement", rootElement)
-                        }
-                    ModificationType.ElementRemoved(removedElement)
-                }
-                else -> ModificationType.Unknown
-            },
-        )
+        val modificationType = when (event.code) {
+            PsiEventType.CHILD_ADDED -> KaElementModificationType.ElementAdded
+            PsiEventType.CHILD_REMOVED -> {
+                val removedElement = event.child ?:
+                    errorWithAttachment("A ${PsiEventType.CHILD_REMOVED} PSI tree change event should have a child element") {
+                        withEntry("psiTreeChangeEvent", event.toString())
+                        withPsiEntry("rootElement", rootElement)
+                    }
+                KaElementModificationType.ElementRemoved(removedElement)
+            }
+            else -> KaElementModificationType.Unknown
+        }
+
+        KaSourceModificationService.getInstance(project).handleElementModification(child ?: rootElement, modificationType)
     }
 
     private fun invalidateCachesInInjectedDocuments(rootElement: PsiElement) {
diff --git a/plugins/kotlin/base/fir/analysis-api-platform/src/org/jetbrains/kotlin/idea/base/fir/analysisApiPlatform/KotlinChangeLocalityDetector.kt b/plugins/kotlin/base/fir/analysis-api-platform/src/org/jetbrains/kotlin/idea/base/fir/analysisApiPlatform/KotlinChangeLocalityDetector.kt
index ddea45f5bd1f..1c4ca463e1bb 100644
--- a/plugins/kotlin/base/fir/analysis-api-platform/src/org/jetbrains/kotlin/idea/base/fir/analysisApiPlatform/KotlinChangeLocalityDetector.kt
+++ b/plugins/kotlin/base/fir/analysis-api-platform/src/org/jetbrains/kotlin/idea/base/fir/analysisApiPlatform/KotlinChangeLocalityDetector.kt
@@ -3,14 +3,13 @@ package org.jetbrains.kotlin.idea.base.fir.analysisApiPlatform
 
 import com.intellij.codeInsight.daemon.ChangeLocalityDetector
 import com.intellij.psi.PsiElement
-import org.jetbrains.kotlin.analysis.low.level.api.fir.LLFirInternals
-import org.jetbrains.kotlin.analysis.low.level.api.fir.file.structure.LLFirDeclarationModificationService
+import org.jetbrains.kotlin.analysis.api.platform.modification.KaSourceModificationService
 
+// Note: The K2 change locality detector is currently unused due to missing logic in `KotlinDiagnosticHighlightingVisitor`. See KTIJ-26691.
 internal class KotlinChangeLocalityDetector : ChangeLocalityDetector {
     override fun getChangeHighlightingDirtyScopeFor(changedElement: PsiElement): PsiElement? {
         // we shouldn't process comments here because the default detector will do that for us
 
-        @OptIn(LLFirInternals::class)
-        return LLFirDeclarationModificationService.getInstance(changedElement.project).elementToRehighlight(changedElement)
+        return KaSourceModificationService.getInstance(changedElement.project).ancestorAffectedByInBlockModification(changedElement)
     }
 }