[AA] support `#isReferenceToImportAlias for all fir references

- reference on constructor through import alias:
  reference in import directive resolves to class,
     thus we have to check constructor explicitly
^KT-68272 fixed

- references in kdoc, etc. have to support import aliases explicitly
^KT-68273 fixed

Author: akozlova

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/FirSimpleNameReference.kt

@@ -16,7 +16,6 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFir
 import org.jetbrains.kotlin.fir.expressions.FirLoopJump
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.containingClassOrObject
 
 internal class KaFirSimpleNameReference(
     expression: KtSimpleNameExpression,
@@ -42,7 +41,7 @@ internal class KaFirSimpleNameReference(
     }
 
     override fun isReferenceToImportAlias(alias: KtImportAlias): Boolean {
-        return getImportAlias(alias.importDirective) != null
+        return super<KaFirReference>.isReferenceToImportAlias(alias)
     }
 
     override fun KaSession.resolveToSymbols(): Collection<KaSymbol> {
@@ -85,31 +84,6 @@ internal class KaFirSimpleNameReference(
         return false
     }
 
-    private fun getImportAlias(importDirective: KtImportDirective?): KtImportAlias? {
-        val fqName = importDirective?.importedFqName ?: return null
-        val codeFragment = KtPsiFactory(element.project).createExpressionCodeFragment(fqName.asString(), element)
-        val contentElement = codeFragment.getContentElement()
-        val importResults =
-            when (contentElement) {
-                is KtDotQualifiedExpression -> contentElement.selectorExpression?.mainReference?.multiResolve(false)
-                is KtSimpleNameExpression -> contentElement.mainReference.multiResolve(false)
-                else -> null
-            } ?: return null
-        val resolveResults = multiResolve(false)
-        if (resolveResults.any { it in importResults }) {
-            return importDirective.alias
-        }
-        val targets = resolveResults.mapNotNull { it.element }
-        val adjustedImportTargets = importResults.mapNotNull {
-            val e = it.element
-            if (e is KtObjectDeclaration && e.isCompanion()) e.containingClassOrObject else null
-        }
-        if (adjustedImportTargets.any { it in targets }) {
-            return importDirective.alias
-        }
-        return null
-    }
-
     override fun getImportAlias(): KtImportAlias? {
         val name = element.getReferencedName()
         val file = element.containingKtFile

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KtFirArrayAccessReference.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.getOrBuildFir
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.psi.KtArrayAccessExpression
+import org.jetbrains.kotlin.psi.KtImportAlias
 
 class KaFirArrayAccessReference(
     expression: KtArrayAccessExpression
@@ -27,4 +28,9 @@ class KaFirArrayAccessReference(
             else -> emptyList()
         }
     }
+
+    override fun isReferenceToImportAlias(alias: KtImportAlias): Boolean {
+        return super<KaFirReference>.isReferenceToImportAlias(alias)
+    }
+
 }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KtFirCollectionLiteralReference.kt

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.fir.types.ConeClassLikeType
 import org.jetbrains.kotlin.fir.types.resolvedType
 
 import org.jetbrains.kotlin.psi.KtCollectionLiteralExpression
+import org.jetbrains.kotlin.psi.KtImportAlias
 
 class KaFirCollectionLiteralReference(
     expression: KtCollectionLiteralExpression
@@ -29,4 +30,8 @@ class KaFirCollectionLiteralReference(
         val call = arrayTypeToArrayOfCall[type.lookupTag.classId] ?: arrayOf
         return listOfNotNull(arrayOfSymbol(call))
     }
+
+    override fun isReferenceToImportAlias(alias: KtImportAlias): Boolean {
+        return super<KaFirReference>.isReferenceToImportAlias(alias)
+    }
 }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KtFirConstructorDelegationReference.kt

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.analysis.api.KaSession
 import org.jetbrains.kotlin.analysis.api.fir.KaFirSession
 import org.jetbrains.kotlin.analysis.api.symbols.KaSymbol
 import org.jetbrains.kotlin.psi.KtConstructorDelegationReferenceExpression
+import org.jetbrains.kotlin.psi.KtImportAlias
 
 class KaFirConstructorDelegationReference(
     expression: KtConstructorDelegationReferenceExpression
@@ -22,4 +23,8 @@ class KaFirConstructorDelegationReference(
         val fir = expression.getOrBuildFirSafe<FirDelegatedConstructorCall>(firResolveSession) ?: return emptyList()
         return listOfNotNull(fir.calleeReference.getResolvedKtSymbolOfNameReference(firSymbolBuilder))
     }
+
+    override fun isReferenceToImportAlias(alias: KtImportAlias): Boolean {
+        return super<KaFirReference>.isReferenceToImportAlias(alias)
+    }
 }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KtFirDestructuringDeclarationReference.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.analysis.api.fir.KaFirSession
 import org.jetbrains.kotlin.analysis.api.fir.buildSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KaSymbol
 import org.jetbrains.kotlin.psi.KtDestructuringDeclarationEntry
+import org.jetbrains.kotlin.psi.KtImportAlias
 
 class KaFirDestructuringDeclarationReference(
     element: KtDestructuringDeclarationEntry
@@ -33,4 +34,9 @@ class KaFirDestructuringDeclarationReference(
         val componentFunctionSymbol = (fir.initializer as? FirComponentCall)?.getCalleeSymbol() ?: return null
         return componentFunctionSymbol.fir.buildSymbol(firSymbolBuilder)
     }
+
+    override fun isReferenceToImportAlias(alias: KtImportAlias): Boolean {
+        return super<KaFirReference>.isReferenceToImportAlias(alias)
+    }
+
 }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KtFirForLoopInReference.kt

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.analysis.api.fir.KaFirSession
 import org.jetbrains.kotlin.analysis.api.fir.buildSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KaSymbol
 import org.jetbrains.kotlin.psi.KtForExpression
+import org.jetbrains.kotlin.psi.KtImportAlias
 
 open class KaFirForLoopInReference(expression: KtForExpression) : KtForLoopInReference(expression), KaFirReference {
     override fun KaSession.resolveToSymbols(): Collection<KaSymbol> {
@@ -37,4 +38,9 @@ open class KaFirForLoopInReference(expression: KtForExpression) : KtForLoopInRef
 
     private fun FirProperty.getInitializerFunctionCall() =
         (initializer as? FirFunctionCall)?.calleeReference?.getResolvedSymbolOfNameReference()
+
+    override fun isReferenceToImportAlias(alias: KtImportAlias): Boolean {
+        return super<KaFirReference>.isReferenceToImportAlias(alias)
+    }
+
 }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KtFirInvokeFunctionReference.kt

@@ -12,6 +12,7 @@ import org.jetbrains.kotlin.analysis.api.calls.symbol
 import org.jetbrains.kotlin.analysis.api.symbols.KaFunctionSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KaSymbol
 import org.jetbrains.kotlin.psi.KtCallExpression
+import org.jetbrains.kotlin.psi.KtImportAlias
 
 class KaFirInvokeFunctionReference(expression: KtCallExpression) : KtInvokeFunctionReference(expression), KaFirReference {
     override fun KaSession.resolveToSymbols(): Collection<KaSymbol> {
@@ -23,4 +24,8 @@ class KaFirInvokeFunctionReference(expression: KtCallExpression) : KtInvokeFunct
                 ?.takeUnless { it is KaFunctionSymbol && it.isBuiltinFunctionInvoke }
         }
     }
+
+    override fun isReferenceToImportAlias(alias: KtImportAlias): Boolean {
+        return super<KaFirReference>.isReferenceToImportAlias(alias)
+    }
 }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KtFirKDocReference.kt

@@ -14,6 +14,7 @@ import org.jetbrains.kotlin.idea.references.KDocReference
 import org.jetbrains.kotlin.idea.references.KaFirReference
 import org.jetbrains.kotlin.idea.references.getPsiDeclarations
 import org.jetbrains.kotlin.kdoc.psi.impl.KDocName
+import org.jetbrains.kotlin.psi.KtImportAlias
 
 internal class KaFirKDocReference(element: KDocName) : KDocReference(element), KaFirReference {
     override fun KaSession.resolveToSymbols(): Collection<KaSymbol> {
@@ -34,4 +35,9 @@ internal class KaFirKDocReference(element: KDocName) : KDocReference(element), K
             }
         }
     }
+
+    override fun isReferenceToImportAlias(alias: KtImportAlias): Boolean {
+        return super<KaFirReference>.isReferenceToImportAlias(alias)
+    }
+
 }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KtFirPropertyDelegationMethodsReference.kt

@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.analysis.api.fir.KaFirSession
 import org.jetbrains.kotlin.analysis.api.fir.buildSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KaSymbol
 import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtImportAlias
 import org.jetbrains.kotlin.psi.KtPropertyDelegate
 
 class KtFirPropertyDelegationMethodsReference(
@@ -36,4 +37,9 @@ class KtFirPropertyDelegationMethodsReference(
 
     private inline fun <reified S : FirStatement> FirPropertyAccessor.singleStatementOfType(): S? =
         body?.statements?.singleOrNull() as? S
+
+    override fun isReferenceToImportAlias(alias: KtImportAlias): Boolean {
+        return super<KaFirReference>.isReferenceToImportAlias(alias)
+    }
+
 }

analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/references/KtFirReference.kt

@@ -9,6 +9,13 @@ import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.analysis.api.KaSession
 import org.jetbrains.kotlin.analysis.api.KaSymbolBasedReference
 import org.jetbrains.kotlin.analysis.api.fir.findReferencePsi
+import org.jetbrains.kotlin.psi.KtClass
+import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
+import org.jetbrains.kotlin.psi.KtImportAlias
+import org.jetbrains.kotlin.psi.KtImportDirective
+import org.jetbrains.kotlin.psi.KtObjectDeclaration
+import org.jetbrains.kotlin.psi.KtSimpleNameExpression
+import org.jetbrains.kotlin.psi.psiUtil.getNonStrictParentOfType
 import org.jetbrains.kotlin.analysis.api.fir.symbols.KaFirSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KaCallableSymbol
 import org.jetbrains.kotlin.analysis.api.symbols.KaSymbol
@@ -31,6 +38,33 @@ interface KaFirReference : KtReference, KaSymbolBasedReference {
         }
 
     override val resolver get() = KaFirReferenceResolver
+
+    fun isReferenceToImportAlias(alias: KtImportAlias): Boolean {
+        return getImportAlias(alias.importDirective) != null
+    }
+
+    fun getImportAlias(importDirective: KtImportDirective?): KtImportAlias? {
+        val importedReference = importDirective?.importedReference ?: return null
+        val importResults =
+            when (importedReference) {
+                is KtDotQualifiedExpression -> importedReference.selectorExpression?.mainReference?.multiResolve(false)
+                is KtSimpleNameExpression -> importedReference.mainReference.multiResolve(false)
+                else -> null
+            } ?: return null
+        val targets = multiResolve(false).mapNotNull { it.element }
+        val adjustedImportTargets = importResults.mapNotNull { it.element }
+        val manager = importDirective.manager
+        if (adjustedImportTargets.any { importTarget ->
+                targets.any { target ->
+                    manager.areElementsEquivalent(target, importTarget) ||
+                            target.isConstructorOf(importTarget) ||
+                            importTarget is KtObjectDeclaration && importTarget.isCompanion() && importTarget.getNonStrictParentOfType<KtClass>() == target
+                }
+            }) {
+            return importDirective.alias
+        }
+        return null
+    }
 }
 
 internal fun KaSession.getPsiDeclarations(symbol: KaFirSymbol<*>): Collection<PsiElement> {

analysis/analysis-api-fir/tests-gen/org/jetbrains/kotlin/analysis/api/fir/test/cases/generated/cases/references/FirIdeNormalAnalysisSourceModuleReferenceImportAliasTestGenerated.java

@@ -8,13 +8,13 @@ package org.jetbrains.kotlin.analysis.api.impl.base.test.cases.references
 import org.jetbrains.kotlin.analysis.test.framework.base.AbstractAnalysisApiBasedTest
 import org.jetbrains.kotlin.analysis.test.framework.project.structure.KtTestModule
 import org.jetbrains.kotlin.analysis.test.framework.services.expressionMarkerProvider
-import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.psi.KtFile
-import org.jetbrains.kotlin.psi.KtSimpleNameExpression
 import org.jetbrains.kotlin.test.builders.TestConfigurationBuilder
 import org.jetbrains.kotlin.test.directives.model.SimpleDirectivesContainer
+import org.jetbrains.kotlin.test.directives.model.singleValue
 import org.jetbrains.kotlin.test.services.TestServices
 import org.jetbrains.kotlin.test.services.assertions
+import kotlin.test.assertNotNull
 
 abstract class AbstractReferenceImportAliasTest : AbstractAnalysisApiBasedTest() {
     override fun configureTest(builder: TestConfigurationBuilder) {
@@ -25,10 +25,15 @@ abstract class AbstractReferenceImportAliasTest : AbstractAnalysisApiBasedTest()
     }
 
     override fun doTestByMainFile(mainFile: KtFile, mainModule: KtTestModule, testServices: TestServices) {
-        val element = testServices.expressionMarkerProvider.getElementOfTypeAtCaret<KtSimpleNameExpression>(mainFile)
-        val alias = element.mainReference.getImportAlias()
-        val expectedAlias = mainModule.testModule.directives[Directives.TYPE_ALIAS].joinToString(" ")
-        testServices.assertions.assertEquals(expectedAlias, alias?.text)
+        val position = testServices.expressionMarkerProvider.getCaretPosition(mainFile)
+        val reference = mainFile.findReferenceAt(position)
+        assertNotNull(reference)
+        val expectedAlias = mainModule.testModule.directives.singleValue(Directives.TYPE_ALIAS)
+        val importDirective = mainFile.importDirectives.find { it.aliasName == expectedAlias }
+        assertNotNull(importDirective)
+        val importAlias = importDirective.alias
+        assertNotNull(importAlias)
+        testServices.assertions.assertTrue(reference.isReferenceTo(importAlias))
     }
 
     private object Directives : SimpleDirectivesContainer() {

analysis/analysis-api-impl-base/tests/org/jetbrains/kotlin/analysis/api/impl/base/test/cases/references/AbstractReferenceImportAliasTest.kt

@@ -1,4 +1,4 @@
-// TYPE_ALIAS: as bar
+// TYPE_ALIAS: bar
 import Foo as bar
 
 object Foo {

analysis/analysis-api-standalone/tests-gen/org/jetbrains/kotlin/analysis/api/standalone/fir/test/cases/generated/cases/references/FirStandaloneNormalAnalysisSourceModuleReferenceImportAliasTestGenerated.java

@@ -1,4 +1,4 @@
-// TYPE_ALIAS: as bar
+// TYPE_ALIAS: bar
 import Foo as bar
 
 class Foo {

analysis/analysis-api/testData/imports/importAliases/classAlias.kt

@@ -0,0 +1,10 @@
+// TYPE_ALIAS: bar
+import Foo as bar
+
+class Foo() {
+    constructor(i: Int): this()
+}
+
+fun test() {
+    val foo = ba<caret>r()
+}

analysis/analysis-api/testData/imports/importAliases/classAliasWithCompanionObject.kt

@@ -1,4 +1,4 @@
-// TYPE_ALIAS: as bar
+// TYPE_ALIAS: bar
 import Foo.foo as bar
 
 object Foo {

analysis/analysis-api/testData/imports/importAliases/constructorAlias.kt

@@ -0,0 +1,8 @@
+// TYPE_ALIAS: bar
+import Foo as bar
+
+/**
+ * [ba<caret>r]
+ */
+object Foo {
+}

analysis/analysis-api/testData/imports/importAliases/functionAlias.kt

@@ -1,4 +1,4 @@
-// TYPE_ALIAS: as bar
+// TYPE_ALIAS: bar
 import Foo.foo as bar
 
 object Foo {

analysis/analysis-api/testData/imports/importAliases/kdocReferenceAlias.kt

@@ -129,15 +129,6 @@ abstract class AbstractKtReference<T : KtElement>(element: T) : PsiPolyVariantRe
                 && element.parent !is KtCallExpression
                 && !element.isInImportDirective()
     }
-
-    private fun PsiElement.isConstructorOf(unwrappedCandidate: PsiElement) =
-        when {
-            // call to Java constructor
-            this is PsiMethod && isConstructor && containingClass == unwrappedCandidate -> true
-            // call to Kotlin constructor
-            this is KtConstructor<*> && getContainingClassOrObject().isEquivalentTo(unwrappedCandidate) -> true
-            else -> false
-        }
 }
 
 abstract class KtSimpleReference<T : KtReferenceExpression>(expression: T) : AbstractKtReference<T>(expression)