referenceUtil: fix matchesTarget for import alias

#KT-30012 Fixed
#KT-30432 Fixed
#KT-30736 Fixed

idea/idea-analysis/src/org/jetbrains/kotlin/idea/references/referenceUtil.kt

@@ -16,7 +16,9 @@
 
 package org.jetbrains.kotlin.idea.references
 
+import com.intellij.openapi.project.Project
 import com.intellij.psi.*
+import com.intellij.psi.search.GlobalSearchScope
 import org.jetbrains.kotlin.asJava.unwrapped
 import org.jetbrains.kotlin.builtins.isExtensionFunctionType
 import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
@@ -26,8 +28,13 @@ import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.caches.resolve.resolveImportReference
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToCall
 import org.jetbrains.kotlin.idea.imports.canBeReferencedViaImport
+import org.jetbrains.kotlin.idea.imports.importableFqName
 import org.jetbrains.kotlin.idea.intentions.OperatorToFunctionIntention
 import org.jetbrains.kotlin.idea.kdoc.KDocReference
+import org.jetbrains.kotlin.idea.stubindex.KotlinFullClassNameIndex
+import org.jetbrains.kotlin.idea.stubindex.KotlinFunctionShortNameIndex
+import org.jetbrains.kotlin.idea.stubindex.KotlinPropertyShortNameIndex
+import org.jetbrains.kotlin.idea.stubindex.KotlinTypeAliasShortNameIndex
 import org.jetbrains.kotlin.idea.util.CallTypeAndReceiver
 import org.jetbrains.kotlin.idea.util.ProjectRootsUtil
 import org.jetbrains.kotlin.kdoc.psi.impl.KDocName
@@ -35,12 +42,18 @@ import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.*
 import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.FunctionImportedFromObject
+import org.jetbrains.kotlin.resolve.PropertyImportedFromObject
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.VariableAsFunctionResolvedCall
 import org.jetbrains.kotlin.resolve.calls.model.isReallySuccess
 import org.jetbrains.kotlin.resolve.descriptorUtil.isExtension
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.resolve.source.getPsi
+import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedClassDescriptor
+import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedPropertyDescriptor
+import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedSimpleFunctionDescriptor
+import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedTypeAliasDescriptor
 import org.jetbrains.kotlin.types.expressions.OperatorConventions
 import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.addToStdlib.constant
@@ -70,6 +83,20 @@ fun PsiReference.canBeReferenceTo(candidateTarget: PsiElement): Boolean {
             || ProjectRootsUtil.isInProjectOrLibSource(element, includeScriptsOutsideSourceRoots = true)
 }
 
+fun DeclarationDescriptor.findPsiDeclarations(project: Project, resolveScope: GlobalSearchScope): Collection<PsiElement> {
+    val fqName = importableFqName ?: return emptyList()
+
+    fun Collection<KtNamedDeclaration>.fqNameFilter() = filter { it.fqName == fqName }
+    return when (this) {
+        is DeserializedClassDescriptor -> KotlinFullClassNameIndex.getInstance()[fqName.asString(), project, resolveScope]
+        is DeserializedTypeAliasDescriptor -> KotlinTypeAliasShortNameIndex.getInstance()[fqName.shortName().asString(), project, resolveScope].fqNameFilter()
+        is DeserializedSimpleFunctionDescriptor, is FunctionImportedFromObject -> KotlinFunctionShortNameIndex.getInstance()[fqName.shortName().asString(), project, resolveScope].fqNameFilter()
+        is DeserializedPropertyDescriptor, is PropertyImportedFromObject -> KotlinPropertyShortNameIndex.getInstance()[fqName.shortName().asString(), project, resolveScope].fqNameFilter()
+        is DeclarationDescriptorWithSource -> listOfNotNull(source.getPsi())
+        else -> emptyList()
+    }
+}
+
 fun PsiReference.matchesTarget(candidateTarget: PsiElement): Boolean {
     if (!canBeReferenceTo(candidateTarget)) return false
 
@@ -100,7 +127,14 @@ fun PsiReference.matchesTarget(candidateTarget: PsiElement): Boolean {
         val importableTargets = unwrappedTargets.mapNotNull {
             if (it is KtConstructor<*>) it.containingClassOrObject else it
         }
-        return importedDescriptors.any { (it as? DeclarationDescriptorWithSource)?.source?.getPsi() in importableTargets }
+
+        val project = element.project
+        val resolveScope = element.resolveScope
+        return importedDescriptors.any {
+            it.findPsiDeclarations(project, resolveScope).any { declaration ->
+                declaration in importableTargets
+            }
+        }
     }
 
     if (element is KtLabelReferenceExpression) {

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceImportAlias/KotlinIntroduceImportAliasHandler.kt

@@ -13,32 +13,27 @@ import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiNamedElement
-import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.search.searches.ReferencesSearch
 import com.intellij.refactoring.RefactoringActionHandler
 import com.intellij.refactoring.rename.inplace.VariableInplaceRenamer
-import org.jetbrains.kotlin.descriptors.DeclarationDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.caches.resolve.resolveImportReference
 import org.jetbrains.kotlin.idea.core.KotlinNameSuggester
-import org.jetbrains.kotlin.idea.core.util.CodeInsightUtils
 import org.jetbrains.kotlin.idea.core.moveCaret
+import org.jetbrains.kotlin.idea.core.util.CodeInsightUtils
 import org.jetbrains.kotlin.idea.imports.importableFqName
 import org.jetbrains.kotlin.idea.refactoring.selectElement
 import org.jetbrains.kotlin.idea.references.KtSimpleNameReference
+import org.jetbrains.kotlin.idea.references.findPsiDeclarations
 import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.idea.references.resolveMainReferenceToDescriptors
 import org.jetbrains.kotlin.idea.search.fileScope
 import org.jetbrains.kotlin.idea.search.usagesSearch.isImportUsage
-import org.jetbrains.kotlin.idea.stubindex.KotlinFullClassNameIndex
-import org.jetbrains.kotlin.idea.stubindex.KotlinFunctionShortNameIndex
-import org.jetbrains.kotlin.idea.stubindex.KotlinPropertyShortNameIndex
 import org.jetbrains.kotlin.idea.util.ImportInsertHelperImpl
 import org.jetbrains.kotlin.idea.util.getAllAccessibleFunctions
 import org.jetbrains.kotlin.idea.util.getAllAccessibleVariables
 import org.jetbrains.kotlin.idea.util.getResolutionScope
 import org.jetbrains.kotlin.incremental.components.NoLookupLocation
-import org.jetbrains.kotlin.js.resolve.diagnostics.findPsi
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtCallExpression
@@ -46,13 +41,10 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtNameReferenceExpression
 import org.jetbrains.kotlin.psi.KtUserType
 import org.jetbrains.kotlin.psi.psiUtil.*
-import org.jetbrains.kotlin.resolve.PropertyImportedFromObject
 import org.jetbrains.kotlin.resolve.descriptorUtil.isExtension
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.resolve.scopes.utils.findClassifier
 import org.jetbrains.kotlin.resolve.scopes.utils.findPackage
-import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedClassDescriptor
-import org.jetbrains.kotlin.serialization.deserialization.descriptors.DeserializedSimpleFunctionDescriptor
 
 object KotlinIntroduceImportAliasHandler : RefactoringActionHandler {
     const val REFACTORING_NAME = "Introduce Import Alias"
@@ -66,7 +58,7 @@ object KotlinIntroduceImportAliasHandler : RefactoringActionHandler {
         val resolveScope = file.resolveScope
         val usages = declarationDescriptors.flatMap { descriptor ->
             val isExtension = descriptor.isExtension
-            findPsiElements(project, resolveScope, descriptor).flatMap {
+            descriptor.findPsiDeclarations(project, resolveScope).flatMap {
                 ReferencesSearch.search(it, fileSearchScope)
                     .findAll()
                     .map { reference -> UsageContext(reference as KtSimpleNameReference, isExtension = isExtension) }
@@ -125,17 +117,6 @@ private fun cleanImport(file: KtFile, fqName: FqName) {
     file.importDirectives.find { it.alias == null && fqName == it.importedFqName }?.delete()
 }
 
-private fun findPsiElements(project: Project, resolveScope: GlobalSearchScope, descriptor: DeclarationDescriptor): Collection<PsiElement> {
-    descriptor.findPsi()?.let { return listOf(it) }
-    val fqName = descriptor.importableFqName ?: return emptyList()
-    return when (descriptor) {
-        is DeserializedClassDescriptor -> KotlinFullClassNameIndex.getInstance()[fqName.asString(), project, resolveScope]
-        is DeserializedSimpleFunctionDescriptor -> KotlinFunctionShortNameIndex.getInstance()[fqName.shortName().asString(), project, resolveScope]
-        is PropertyImportedFromObject -> KotlinPropertyShortNameIndex.getInstance()[fqName.shortName().asString(), project, resolveScope]
-        else -> emptyList()
-    }.filter { fqName == it.fqName }
-}
-
 private fun invokeRename(
     project: Project,
     editor: Editor,

idea/testData/usageHighlighter/importAliasFromStdLibClass.kt

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+import kotlin.collections.List as <info descr="null">~List1</info>
+
+fun main() {
+    val a: <info descr="null">List1</info><Int>
+    val b = <info descr="null">List1</info>(5) { 42 }
+}

idea/testData/usageHighlighter/importAliasFromStdLibFunction.kt

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+import kotlin.let as <info descr="null">~myLet</info>
+
+fun main() {
+    val a = 42
+    a.<info descr="null">myLet</info>{ println(it) }
+}

idea/testData/usageHighlighter/importAliasFromStdLibFunctionFromObject.kt

@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+import kotlin.Result.Companion.success as <info descr="null">~mySuccess</info>
+
+fun main() {
+    <info descr="null">mySuccess</info>(42)
+}

idea/testData/usageHighlighter/importAliasFromStdLibPropertyFromObject.kt

@@ -0,0 +1,6 @@
+// WITH_RUNTIME
+import kotlin.Int.Companion.MAX_VALUE as <info descr="null">~maxValue</info>
+
+fun main() {
+    <info descr="null">maxValue</info>
+}

idea/tests/org/jetbrains/kotlin/idea/highlighter/AbstractUsageHighlightingTest.kt

@@ -13,9 +13,7 @@ import com.intellij.openapi.editor.colors.EditorColorsManager
 import com.intellij.openapi.editor.markup.RangeHighlighter
 import com.intellij.openapi.editor.markup.TextAttributes
 import com.intellij.testFramework.ExpectedHighlightingData
-import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
-import org.jetbrains.kotlin.idea.test.KotlinLightProjectDescriptor
-import org.jetbrains.kotlin.idea.test.extractMarkerOffset
+import org.jetbrains.kotlin.idea.test.*
 
 abstract class AbstractUsageHighlightingTest : KotlinLightCodeInsightFixtureTestCase() {
     companion object {
@@ -63,6 +61,4 @@ abstract class AbstractUsageHighlightingTest : KotlinLightCodeInsightFixtureTest
 
         return info.textAttributes == readAttributes || info.textAttributes == writeAttributes
     }
-
-    override fun getProjectDescriptor(): KotlinLightProjectDescriptor = KotlinLightProjectDescriptor.INSTANCE
 }

idea/tests/org/jetbrains/kotlin/idea/highlighter/UsageHighlightingTestGenerated.java

@@ -49,6 +49,26 @@ public class UsageHighlightingTestGenerated extends AbstractUsageHighlightingTes
         runTest("idea/testData/usageHighlighter/importAlias.kt");
     }
 
+    @TestMetadata("importAliasFromStdLibClass.kt")
+    public void testImportAliasFromStdLibClass() throws Exception {
+        runTest("idea/testData/usageHighlighter/importAliasFromStdLibClass.kt");
+    }
+
+    @TestMetadata("importAliasFromStdLibFunction.kt")
+    public void testImportAliasFromStdLibFunction() throws Exception {
+        runTest("idea/testData/usageHighlighter/importAliasFromStdLibFunction.kt");
+    }
+
+    @TestMetadata("importAliasFromStdLibFunctionFromObject.kt")
+    public void testImportAliasFromStdLibFunctionFromObject() throws Exception {
+        runTest("idea/testData/usageHighlighter/importAliasFromStdLibFunctionFromObject.kt");
+    }
+
+    @TestMetadata("importAliasFromStdLibPropertyFromObject.kt")
+    public void testImportAliasFromStdLibPropertyFromObject() throws Exception {
+        runTest("idea/testData/usageHighlighter/importAliasFromStdLibPropertyFromObject.kt");
+    }
+
     @TestMetadata("labeledAnonymousFun.kt")
     public void testLabeledAnonymousFun() throws Exception {
         runTest("idea/testData/usageHighlighter/labeledAnonymousFun.kt");