From 1c1da6bcedf244c605e0fec4ee61b2b284cd34e7 Mon Sep 17 00:00:00 2001
From: Yan Zhulanow <yan.zhulanow@jetbrains.com>
Date: Thu, 15 Feb 2024 20:13:54 +0900
Subject: [PATCH] [Analysis API] Cause OOB when imports change in code
 fragments

In K1, code fragment analysis was completely invalidated on any PSI
change. Because imports are not a part of the PSI tree of
'KtCodeFragment's, a colon tremble happened on 'addImportsFromString()'.

In K2, changes inside code fragments are always considered in-body
modifications. So, even with the colon trembling, the 'FirFile',
together with its 'FirImport's was not recreated.

^KT-65600 Fixed
---
 .../LLFirDeclarationModificationService.kt          |  6 ++++++
 .../src/org/jetbrains/kotlin/psi/KtCodeFragment.kt  | 13 +++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/LLFirDeclarationModificationService.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/LLFirDeclarationModificationService.kt
index 9319495e061..9f14088b722 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/LLFirDeclarationModificationService.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/file/structure/LLFirDeclarationModificationService.kt
@@ -35,6 +35,7 @@ import org.jetbrains.kotlin.fir.declarations.FirResolvePhase
 import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
 import org.jetbrains.kotlin.idea.KotlinLanguage
 import org.jetbrains.kotlin.psi
+import org.jetbrains.kotlin.psi.KotlinCodeFragmentImportModificationListener
 import org.jetbrains.kotlin.psi.KtAnnotated
 import org.jetbrains.kotlin.psi.KtBlockExpression
 import org.jetbrains.kotlin.psi.KtCodeFragment
@@ -69,6 +70,11 @@ class LLFirDeclarationModificationService(val project: Project) : Disposable {
             },
             this,
         )
+
+        project.messageBus.connect(this).subscribe(
+            KtCodeFragment.IMPORT_MODIFICATION,
+            KotlinCodeFragmentImportModificationListener { codeFragment -> outOfBlockModification(codeFragment) }
+        )
     }
 
     private var inBlockModificationQueue: MutableSet<ChangeType.InBlock>? = null
diff --git a/compiler/psi/src/org/jetbrains/kotlin/psi/KtCodeFragment.kt b/compiler/psi/src/org/jetbrains/kotlin/psi/KtCodeFragment.kt
index 040f2f8ce25..6162fd33f17 100644
--- a/compiler/psi/src/org/jetbrains/kotlin/psi/KtCodeFragment.kt
+++ b/compiler/psi/src/org/jetbrains/kotlin/psi/KtCodeFragment.kt
@@ -26,6 +26,7 @@ import com.intellij.psi.impl.source.tree.FileElement
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.tree.IElementType
 import com.intellij.testFramework.LightVirtualFile
+import com.intellij.util.messages.Topic
 import org.jetbrains.kotlin.idea.KotlinFileType
 import org.jetbrains.kotlin.psi.psiUtil.getElementTextWithContext
 
@@ -136,6 +137,10 @@ abstract class KtCodeFragment(
 
             // Increment the modification stamp
             clearCaches()
+
+            if (viewProvider.isEventSystemEnabled) {
+                project.messageBus.syncPublisher(IMPORT_MODIFICATION).onCodeFragmentImportsModification(this)
+            }
         }
     }
 
@@ -214,8 +219,16 @@ abstract class KtCodeFragment(
     companion object {
         const val IMPORT_SEPARATOR: String = ","
 
+        @Suppress("UnstableApiUsage")
+        val IMPORT_MODIFICATION: Topic<KotlinCodeFragmentImportModificationListener> =
+            Topic(KotlinCodeFragmentImportModificationListener::class.java, Topic.BroadcastDirection.TO_CHILDREN, true)
+
         val FAKE_CONTEXT_FOR_JAVA_FILE: Key<Function0<KtElement>> = Key.create("FAKE_CONTEXT_FOR_JAVA_FILE")
 
         private val LOG = Logger.getInstance(KtCodeFragment::class.java)
     }
 }
+
+fun interface KotlinCodeFragmentImportModificationListener {
+    fun onCodeFragmentImportsModification(codeFragment: KtCodeFragment)
+}
\ No newline at end of file