FIR IDE: Add Add Import quickfix

This quickfix can import unresolved types and callables It currently does not support Java types and is not as advanced as in the old plugin Also, enable tests that now pass
2021-04-16 13:26:52 +05:00
parent 3ad9e4cb5c
commit 22db894076
13 changed files with 293 additions and 2 deletions
@@ -12,6 +12,8 @@ import org.jetbrains.kotlin.idea.fir.api.fixes.KtQuickFixesListBuilder
 import org.jetbrains.kotlin.idea.frontend.api.fir.diagnostics.KtFirDiagnostic
 import org.jetbrains.kotlin.idea.quickfix.fixes.*
 import org.jetbrains.kotlin.idea.quickfix.fixes.InitializePropertyQuickFixFactory
+import org.jetbrains.kotlin.idea.quickfix.fixes.ChangeTypeQuickFix
+import org.jetbrains.kotlin.idea.quickfix.fixes.ReplaceCallFixFactories

 class MainKtQuickFixRegistrar : KtQuickFixRegistrar() {
    private val modifiers = KtQuickFixesListBuilder.registerPsiQuickFix {
@@ -85,6 +87,10 @@ class MainKtQuickFixRegistrar : KtQuickFixRegistrar() {
        registerApplicator(MemberNotImplementedQuickfixFactories.manyImplMemberNotImplemented)
    }

+    private val imports = KtQuickFixesListBuilder.registerPsiQuickFix {
+        registerApplicator(ImportQuickFix.FACTORY)
+    }
+
    private val mutability = KtQuickFixesListBuilder.registerPsiQuickFix {
        registerPsiQuickFixes(KtFirDiagnostic.VarOverriddenByVal::class, ChangeVariableMutabilityFix.VAR_OVERRIDDEN_BY_VAL_FACTORY)
        registerPsiQuickFixes(KtFirDiagnostic.VarAnnotationParameter::class, ChangeVariableMutabilityFix.VAR_ANNOTATION_PARAMETER_FACTORY)
@@ -106,6 +112,7 @@ class MainKtQuickFixRegistrar : KtQuickFixRegistrar() {
        modifiers,
        propertyInitialization,
        overrides,
+        imports,
        mutability,
        expressions,
    )
@@ -0,0 +1,212 @@
+/*
+ * Copyright 2010-2021 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.idea.quickfix.fixes
+
+import com.intellij.codeInsight.daemon.QuickFixBundle
+import com.intellij.codeInsight.hint.HintManager
+import com.intellij.codeInsight.hint.QuestionAction
+import com.intellij.codeInspection.HintAction
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.ui.popup.JBPopup
+import com.intellij.openapi.ui.popup.JBPopupFactory
+import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
+import com.intellij.psi.search.GlobalSearchScope
+import org.jetbrains.kotlin.idea.KotlinBundle
+import org.jetbrains.kotlin.idea.caches.project.getModuleInfo
+import org.jetbrains.kotlin.idea.fir.api.fixes.diagnosticFixFactory
+import org.jetbrains.kotlin.idea.fir.low.level.api.IndexHelper
+import org.jetbrains.kotlin.idea.fir.low.level.api.util.createScopeForModuleLibraries
+import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.fir.diagnostics.KtFirDiagnostic
+import org.jetbrains.kotlin.idea.frontend.api.fir.utils.addImportToFile
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCallableSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtNamedClassOrObjectSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtSymbolWithVisibility
+import org.jetbrains.kotlin.idea.quickfix.QuickFixActionBase
+import org.jetbrains.kotlin.idea.util.application.executeWriteCommand
+import org.jetbrains.kotlin.idea.util.module
+import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.KtPsiUtil.isSelectorInQualified
+import org.jetbrains.kotlin.psi.psiUtil.unwrapNullability
+
+internal class ImportQuickFix(
+    element: KtElement,
+    private val importCandidates: List<FqName>
+) : QuickFixActionBase<KtElement>(element), HintAction {
+
+    init {
+        require(importCandidates.isNotEmpty())
+    }
+
+    override fun getText(): String = KotlinBundle.message("fix.import")
+
+    override fun getFamilyName(): String = KotlinBundle.message("fix.import")
+
+    override fun invoke(project: Project, editor: Editor, file: PsiFile) {
+        if (file !is KtFile) return
+
+        createAddImportAction(project, editor, file).execute()
+    }
+
+    private fun createAddImportAction(project: Project, editor: Editor, file: KtFile): QuestionAction {
+        return ImportQuestionAction(project, editor, file, importCandidates)
+    }
+
+    override fun showHint(editor: Editor): Boolean {
+        val element = element ?: return false
+        val file = element.containingKtFile
+        val project = file.project
+
+        val elementRange = element.textRange
+        val autoImportHintText = "Import ${importCandidates.first().asString()}?"
+
+        HintManager.getInstance().showQuestionHint(
+            editor,
+            autoImportHintText,
+            elementRange.startOffset,
+            elementRange.endOffset,
+            createAddImportAction(project, editor, file)
+        )
+
+        return true
+    }
+
+    private class ImportQuestionAction(
+        private val project: Project,
+        private val editor: Editor,
+        private val file: KtFile,
+        private val importCandidates: List<FqName>
+    ) : QuestionAction {
+
+        init {
+            require(importCandidates.isNotEmpty())
+        }
+
+        override fun execute(): Boolean {
+            val unambiguousImport = importCandidates.singleOrNull()
+            if (unambiguousImport != null) {
+                addImport(unambiguousImport)
+                return true
+            }
+
+            createImportSelectorPopup().showInBestPositionFor(editor)
+
+            return true
+        }
+
+        private fun createImportSelectorPopup(): JBPopup {
+            return JBPopupFactory.getInstance()
+                .createPopupChooserBuilder(importCandidates)
+                .setTitle(KotlinBundle.message("action.add.import.chooser.title"))
+                .setItemChosenCallback { selectedValue: FqName -> addImport(selectedValue) }
+                .createPopup()
+        }
+
+        private fun addImport(nameToImport: FqName) {
+            project.executeWriteCommand(QuickFixBundle.message("add.import")) {
+                addImportToFile(project, file, nameToImport)
+            }
+        }
+    }
+
+    internal companion object {
+        val FACTORY = diagnosticFixFactory<KtFirDiagnostic.UnresolvedReference> { diagnostic ->
+            val element = diagnostic.psi
+
+            val indexHelper = IndexHelper(element.project, createSearchScope(element))
+
+            val quickFix = when (element) {
+                is KtTypeReference -> createImportTypeFix(indexHelper, element)
+                is KtNameReferenceExpression -> createImportNameFix(indexHelper, element)
+                else -> null
+            }
+
+            listOfNotNull(quickFix)
+        }
+
+        private fun createSearchScope(element: PsiElement): GlobalSearchScope {
+            val contentScope = element.getModuleInfo().contentScope()
+            val librariesScope = element.module?.let(::createScopeForModuleLibraries)
+                ?: GlobalSearchScope.EMPTY_SCOPE
+
+            return contentScope.uniteWith(librariesScope)
+        }
+
+        private fun KtAnalysisSession.createImportNameFix(
+            indexHelper: IndexHelper,
+            element: KtNameReferenceExpression
+        ): ImportQuickFix? {
+            if (isSelectorInQualified(element)) return null
+
+            val firFile = element.containingKtFile.getFileSymbol()
+            val unresolvedName = element.getReferencedNameAsName()
+
+            val isVisible: (KtSymbol) -> Boolean =
+                { it !is KtSymbolWithVisibility || isVisible(it, firFile, null, element) }
+
+            val callableCandidates = collectCallableCandidates(indexHelper, unresolvedName, isVisible)
+            val typeCandidates = collectTypesCandidates(indexHelper, unresolvedName, isVisible)
+
+            val importCandidates = (callableCandidates + typeCandidates).distinct()
+            if (importCandidates.isEmpty()) return null
+
+            return ImportQuickFix(element, importCandidates)
+        }
+
+        private fun KtAnalysisSession.createImportTypeFix(indexHelper: IndexHelper, element: KtTypeReference): ImportQuickFix? {
+            val firFile = element.containingKtFile.getFileSymbol()
+            val unresolvedName = element.typeName ?: return null
+
+            val isVisible: (KtNamedClassOrObjectSymbol) -> Boolean =
+                { isVisible(it, firFile, null, element) }
+
+            val acceptableClasses = collectTypesCandidates(indexHelper, unresolvedName, isVisible)
+            if (acceptableClasses.isEmpty()) return null
+
+            return ImportQuickFix(element, acceptableClasses)
+        }
+
+        private fun KtAnalysisSession.collectCallableCandidates(
+            indexHelper: IndexHelper,
+            unresolvedName: Name,
+            isVisible: (KtCallableSymbol) -> Boolean
+        ): List<FqName> {
+            val callablesCandidates = indexHelper.getKotlinCallablesByName(unresolvedName)
+
+            return callablesCandidates
+                .asSequence()
+                .map { it.getSymbol() as KtCallableSymbol }
+                .filter(isVisible)
+                .mapNotNull { it.callableIdIfNonLocal?.asSingleFqName() }
+                .toList()
+        }
+
+        private fun KtAnalysisSession.collectTypesCandidates(
+            indexHelper: IndexHelper,
+            unresolvedName: Name,
+            isVisible: (KtNamedClassOrObjectSymbol) -> Boolean
+        ): List<FqName> {
+            val classesCandidates = indexHelper.getKotlinClassesByName(unresolvedName)
+
+            return classesCandidates.asSequence()
+                .mapNotNull { it.getNamedClassOrObjectSymbol() }
+                .filter(isVisible)
+                .mapNotNull { it.classIdIfNonLocal?.asSingleFqName() }
+                .toList()
+        }
+    }
+}
+
+private val KtTypeReference.typeName: Name?
+    get() {
+        val userType = typeElement?.unwrapNullability() as? KtUserType
+        return userType?.referencedName?.let(Name::identifier)
+    }
@@ -164,6 +164,11 @@ public class HighLevelQuickFixMultiFileTestGenerated extends AbstractHighLevelQu
        runTest("idea/testData/quickfix/autoImports/importInFirstPartInQualifiedExpression.before.Main.kt");
    }

+    @TestMetadata("importNullableTraitWithGenerics.before.Main.kt")
+    public void testImportNullableTraitWithGenerics() throws Exception {
+        runTest("idea/testData/quickfix/autoImports/importNullableTraitWithGenerics.before.Main.kt");
+    }
+
    @TestMetadata("ImportOperatorInvokeWithConvention.before.Main.kt")
    public void testImportOperatorInvokeWithConvention() throws Exception {
        runTest("idea/testData/quickfix/autoImports/ImportOperatorInvokeWithConvention.before.Main.kt");
@@ -219,6 +224,11 @@ public class HighLevelQuickFixMultiFileTestGenerated extends AbstractHighLevelQu
        runTest("idea/testData/quickfix/autoImports/noImportForPrivateClass.before.Main.kt");
    }

+    @TestMetadata("noImportForPrivateFunction.before.Main.kt")
+    public void testNoImportForPrivateFunction() throws Exception {
+        runTest("idea/testData/quickfix/autoImports/noImportForPrivateFunction.before.Main.kt");
+    }
+
    @TestMetadata("noImportInImports.before.Main.kt")
    public void testNoImportInImports() throws Exception {
        runTest("idea/testData/quickfix/autoImports/noImportInImports.before.Main.kt");
@@ -7,17 +7,22 @@ package org.jetbrains.kotlin.idea.fir.low.level.api

 import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiElement
 import com.intellij.psi.search.GlobalSearchScope
+import com.intellij.psi.search.PsiShortNamesCache
 import com.intellij.psi.stubs.StringStubIndexExtension
 import com.intellij.psi.stubs.StubIndex
 import com.intellij.psi.stubs.StubIndexKey
+import org.jetbrains.kotlin.descriptors.ClassDescriptor
+import org.jetbrains.kotlin.idea.caches.resolve.util.resolveToDescriptor
 import org.jetbrains.kotlin.name.CallableId
 import org.jetbrains.kotlin.idea.stubindex.*
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.contains
 import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult

 /*
@@ -103,6 +108,22 @@ public class IndexHelper(val project: Project, private val scope: GlobalSearchSc
        return (functions + properties).toList()
    }

+    fun getKotlinCallablesByName(name: Name): Collection<KtCallableDeclaration> {
+        val functions: Sequence<KtCallableDeclaration> =
+            KotlinFunctionShortNameIndex.getInstance().get(name.asString(), project, scope).asSequence()
+
+        val properties: Sequence<KtNamedDeclaration> =
+            KotlinPropertyShortNameIndex.getInstance().get(name.asString(), project, scope).asSequence()
+
+        return (functions + properties)
+            .filterIsInstance<KtCallableDeclaration>()
+            .toList()
+    }
+
+    fun getKotlinClassesByName(name: Name): Collection<KtClassOrObject> {
+        return KotlinClassShortNameIndex.getInstance().get(name.asString(), project, scope)
+    }
+
    fun getTopLevelExtensions(nameFilter: (Name) -> Boolean, receiverTypeNames: Set<String>): Collection<KtCallableDeclaration> {
        val index = KotlinTopLevelExtensionsByReceiverTypeIndex.INSTANCE

@@ -303,7 +303,6 @@ The Kotlin FIR plugin provides language support in IntelliJ IDEA and Android Stu
              fieldName="INSTANCE"
              extensions="kotlin_module"/>

-
-
+    <java.shortNamesCache implementation="org.jetbrains.kotlin.idea.caches.KotlinShortNamesCache"/>
  </extensions>
 </idea-plugin>
@@ -0,0 +1,9 @@
+import importTrait.data.TestTrait
+
+// "Import" "true"
+// FIR_IGNORE
+// ERROR: Unresolved reference: TestTrait
+
+fun test() {
+    val a: TestTrait<String, Int>? = null
+}
@@ -0,0 +1,3 @@
+package importTrait.data
+
+interface TestTrait<A, B>
@@ -0,0 +1,7 @@
+// "Import" "true"
+// FIR_IGNORE
+// ERROR: Unresolved reference: TestTrait
+
+fun test() {
+    val a: <caret>TestTrait<String, Int>? = null
+}
@@ -1,6 +1,7 @@
 import importTrait.data.TestTrait

 // "Import" "true"
+// IGNORE_FIR
 // ERROR: Unresolved reference: TestTrait

 fun test() {
@@ -1,4 +1,5 @@
 // "Import" "true"
+// IGNORE_FIR
 // ERROR: Unresolved reference: TestTrait

 fun test() {
@@ -0,0 +1,3 @@
+package sometest
+
+private fun privateFun() {}
@@ -0,0 +1,8 @@
+// "Import" "false"
+// ACTION: Create function 'privateFun'
+// ACTION: Rename reference
+// ERROR: Unresolved reference: privateFun
+
+fun test() {
+    <caret>privateFun()
+}
@@ -761,6 +761,11 @@ public class QuickFixMultiFileTestGenerated extends AbstractQuickFixMultiFileTes
            runTest("idea/testData/quickfix/autoImports/importKotlinStaticPropertyOverloadedSetterFromJava.test");
        }

+        @TestMetadata("importNullableTraitWithGenerics.before.Main.kt")
+        public void testImportNullableTraitWithGenerics() throws Exception {
+            runTest("idea/testData/quickfix/autoImports/importNullableTraitWithGenerics.before.Main.kt");
+        }
+
        @TestMetadata("ImportOperatorInvokeWithConvention.before.Main.kt")
        public void testImportOperatorInvokeWithConvention() throws Exception {
            runTest("idea/testData/quickfix/autoImports/ImportOperatorInvokeWithConvention.before.Main.kt");
@@ -951,6 +956,11 @@ public class QuickFixMultiFileTestGenerated extends AbstractQuickFixMultiFileTes
            runTest("idea/testData/quickfix/autoImports/noImportForPrivateClass.before.Main.kt");
        }

+        @TestMetadata("noImportForPrivateFunction.before.Main.kt")
+        public void testNoImportForPrivateFunction() throws Exception {
+            runTest("idea/testData/quickfix/autoImports/noImportForPrivateFunction.before.Main.kt");
+        }
+
        @TestMetadata("noImportInImports.before.Main.kt")
        public void testNoImportInImports() throws Exception {
            runTest("idea/testData/quickfix/autoImports/noImportInImports.before.Main.kt");