From 45c0fa8d23f389d29d3dff0347cbe185433849e1 Mon Sep 17 00:00:00 2001
From: Yan Zhulanow <yan.zhulanow@jetbrains.com>
Date: Thu, 15 Feb 2024 16:45:13 +0900
Subject: [PATCH] [PSI] Refactor import handling in code fragments

Now there is a single place that performs import modification, as well
as reports events related to that modification.
---
 .../jetbrains/kotlin/psi/KtCodeFragment.kt    | 76 +++++++++++--------
 1 file changed, 46 insertions(+), 30 deletions(-)

diff --git a/compiler/psi/src/org/jetbrains/kotlin/psi/KtCodeFragment.kt b/compiler/psi/src/org/jetbrains/kotlin/psi/KtCodeFragment.kt
index 24e07a5924e..040f2f8ce25 100644
--- a/compiler/psi/src/org/jetbrains/kotlin/psi/KtCodeFragment.kt
+++ b/compiler/psi/src/org/jetbrains/kotlin/psi/KtCodeFragment.kt
@@ -28,7 +28,6 @@ import com.intellij.psi.tree.IElementType
 import com.intellij.testFramework.LightVirtualFile
 import org.jetbrains.kotlin.idea.KotlinFileType
 import org.jetbrains.kotlin.psi.psiUtil.getElementTextWithContext
-import java.util.*
 
 abstract class KtCodeFragment(
     private val myProject: Project,
@@ -44,7 +43,7 @@ abstract class KtCodeFragment(
     }, false
 ), KtCodeFragmentBase {
     private var viewProvider = super.getViewProvider() as SingleRootFileViewProvider
-    private var imports = LinkedHashSet<String>()
+    private val importDirectiveStrings = LinkedHashSet<String>()
 
     private val fakeContextForJavaFile: PsiElement? by lazy {
         this.getCopyableUserData(FAKE_CONTEXT_FOR_JAVA_FILE)?.invoke()
@@ -54,8 +53,9 @@ abstract class KtCodeFragment(
         @Suppress("LeakingThis")
         getViewProvider().forceCachedPsi(this)
         init(TokenType.CODE_FRAGMENT, elementType)
-        if (context != null) {
-            initImports(imports)
+
+        if (imports != null) {
+            appendImports(imports)
         }
     }
 
@@ -100,7 +100,8 @@ abstract class KtCodeFragment(
         return (cloneImpl(elementClone) as KtCodeFragment).apply {
             isPhysical = false
             myOriginalFile = this@KtCodeFragment
-            imports = this@KtCodeFragment.imports
+            importDirectiveStrings.clear()
+            importDirectiveStrings.addAll(this@KtCodeFragment.importDirectiveStrings)
             viewProvider = SingleRootFileViewProvider(
                 PsiManager.getInstance(myProject),
                 LightVirtualFile(name, KotlinFileType.INSTANCE, text),
@@ -125,40 +126,64 @@ abstract class KtCodeFragment(
     }
 
     override fun importsToString(): String {
-        return imports.joinToString(IMPORT_SEPARATOR)
+        return importDirectiveStrings.joinToString(IMPORT_SEPARATOR)
     }
 
     override fun addImportsFromString(imports: String?) {
-        if (imports == null || imports.isEmpty()) return
+        if (imports != null && appendImports(imports)) {
+            // For K1: This forces the code fragment to be re-highlighted.
+            add(KtPsiFactory(project).createColon()).delete()
 
-        imports.split(IMPORT_SEPARATOR).forEach {
-            addImport(it)
+            // Increment the modification stamp
+            clearCaches()
         }
-
-        // we need this code to force re-highlighting, otherwise it does not work by some reason
-        val tempElement = KtPsiFactory(project).createColon()
-        add(tempElement).delete()
     }
 
+    @Deprecated("Use 'addImportsFromString()' instead", ReplaceWith("addImportsFromString(import)"), level = DeprecationLevel.WARNING)
     fun addImport(import: String) {
-        val contextFile = getContextContainingFile()
-        if (contextFile != null) {
-            if (contextFile.importDirectives.find { it.text == import } == null) {
-                imports.add(import)
-                clearCaches() // Increment the modification stamp
+        addImportsFromString(import)
+    }
+
+    /**
+     * Parses raw [rawImports] and appends them to the list of code fragment imports.
+     *
+     * Import strings must be separated by the [IMPORT_SEPARATOR].
+     * Each import must be either a qualified name to import (e.g., 'foo.bar'), or a complete text representation of an import directive
+     * (e.g., 'import foo.bar as baz').
+     *
+     * Note that already present import directives will be ignored.
+     *
+     * @return `true` if new import directives were added.
+     */
+    private fun appendImports(rawImports: String): Boolean {
+        if (rawImports.isEmpty()) {
+            return false
+        }
+
+        var hasNewImports = false
+
+        for (rawImport in rawImports.split(IMPORT_SEPARATOR)) {
+            val importDirectiveString = if (rawImport.startsWith("import ")) rawImport else "import $rawImport"
+            if (importDirectiveStrings.add(importDirectiveString) && !hasNewImports) {
+                hasNewImports = true
             }
         }
+
+        return hasNewImports
     }
 
     fun importsAsImportList(): KtImportList? {
-        if (imports.isNotEmpty() && context != null) {
-            return KtPsiFactory.contextual(context).createFile("imports_for_codeFragment.kt", imports.joinToString("\n")).importList
+        if (importDirectiveStrings.isNotEmpty() && context != null) {
+            val ktPsiFactory = KtPsiFactory.contextual(context)
+            val fileText = importDirectiveStrings.joinToString("\n")
+            return ktPsiFactory.createFile("imports_for_codeFragment.kt", fileText).importList
         }
+
         return null
     }
 
     override val importLists: List<KtImportList>
-        get() = importsAsImportList().let(::listOfNotNull)
+        get() = listOfNotNull(importsAsImportList())
 
     override val importDirectives: List<KtImportDirective>
         get() = importsAsImportList()?.imports ?: emptyList()
@@ -186,15 +211,6 @@ abstract class KtCodeFragment(
         return contextElement
     }
 
-    private fun initImports(imports: String?) {
-        if (imports != null && imports.isNotEmpty()) {
-            val importsWithPrefix = imports.split(IMPORT_SEPARATOR).map { it.takeIf { it.startsWith("import ") } ?: "import ${it.trim()}" }
-            importsWithPrefix.forEach {
-                addImport(it)
-            }
-        }
-    }
-
     companion object {
         const val IMPORT_SEPARATOR: String = ","