From 019635189a83bf9255f2daf85b7b7f79cb7960e1 Mon Sep 17 00:00:00 2001
From: Igor Yakovlev <igor.yakovlev@jetbrains.com>
Date: Fri, 22 Jan 2021 00:41:12 +0300
Subject: [PATCH] [FIR IDE] Fix annotation call reference resolve to ctor

---
 .../level/api/file/structure/FileStructure.kt |  2 +-
 .../idea/references/FirSimpleNameReference.kt | 25 ++++++++++++++++---
 .../fir/lazyResolve/import/foo/main.txt       |  2 +-
 .../resolve/references/AnnotationOnFile.kt    |  2 --
 .../references/AnnotationOnFileWithImport.kt  |  2 --
 5 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructure.kt b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructure.kt
index a0595684c3f..fe01b3a6557 100644
--- a/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructure.kt
+++ b/idea/idea-frontend-fir/idea-fir-low-level-api/src/org/jetbrains/kotlin/idea/fir/low/level/api/file/structure/FileStructure.kt
@@ -89,7 +89,7 @@ internal class FileStructure(
             val firFile = firFileBuilder.getFirFileResolvedToPhaseWithCaching(
                 container,
                 moduleFileCache,
-                FirResolvePhase.IMPORTS,
+                FirResolvePhase.TYPES,
                 checkPCE = true
             )
             RootStructureElement(
diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirSimpleNameReference.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirSimpleNameReference.kt
index 4d60ef2b400..48c020c1936 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirSimpleNameReference.kt
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirSimpleNameReference.kt
@@ -8,17 +8,36 @@ package org.jetbrains.kotlin.idea.references
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.idea.frontend.api.KtAnalysisSession
 import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.fir.symbols.KtFirClassOrObjectSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtClassKind
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
 import org.jetbrains.kotlin.name.FqName
-import org.jetbrains.kotlin.psi.KtImportAlias
-import org.jetbrains.kotlin.psi.KtSimpleNameExpression
+import org.jetbrains.kotlin.psi.*
 
 internal class KtFirSimpleNameReference(
     expression: KtSimpleNameExpression
 ) : KtSimpleNameReference(expression), KtFirReference {
+
+    private fun KtAnalysisSession.fixUpAnnotationCallResolveToCtor(resultsToFix: Collection<KtSymbol>): Collection<KtSymbol> {
+        if (resultsToFix.isEmpty()) return resultsToFix
+
+        val isAnnotationCall = (((expression.parent as? KtUserType)
+            ?.parent as? KtTypeReference)
+            ?.parent as? KtConstructorCalleeExpression)
+            ?.parent is KtAnnotationEntry
+        if (!isAnnotationCall) return resultsToFix
+
+        return resultsToFix.map { targetSymbol ->
+            if (targetSymbol is KtFirClassOrObjectSymbol && targetSymbol.classKind == KtClassKind.ANNOTATION_CLASS) {
+                targetSymbol.getMemberScope().getConstructors().firstOrNull() ?: targetSymbol
+            } else targetSymbol
+        }
+    }
+
     override fun KtAnalysisSession.resolveToSymbols(): Collection<KtSymbol> {
         check(this is KtFirAnalysisSession)
-        return FirReferenceResolveHelper.resolveSimpleNameReference(this@KtFirSimpleNameReference, this)
+        val results = FirReferenceResolveHelper.resolveSimpleNameReference(this@KtFirSimpleNameReference, this)
+        return fixUpAnnotationCallResolveToCtor(results)
     }
 
     override fun doCanBeReferenceTo(candidateTarget: PsiElement): Boolean {
diff --git a/idea/testData/fir/lazyResolve/import/foo/main.txt b/idea/testData/fir/lazyResolve/import/foo/main.txt
index 74b0e15ee36..386af2d81b2 100644
--- a/idea/testData/fir/lazyResolve/import/foo/main.txt
+++ b/idea/testData/fir/lazyResolve/import/foo/main.txt
@@ -1,4 +1,4 @@
 FILE: main.kt
-    public? final? [RAW_FIR] fun foofoo(): Int {
+    public? final? [TYPES] fun foofoo(): R|kotlin/Int| {
         ^foofoo barbar#()
     }
diff --git a/idea/testData/resolve/references/AnnotationOnFile.kt b/idea/testData/resolve/references/AnnotationOnFile.kt
index 85f6ce26681..491f48d1035 100644
--- a/idea/testData/resolve/references/AnnotationOnFile.kt
+++ b/idea/testData/resolve/references/AnnotationOnFile.kt
@@ -1,5 +1,3 @@
-// IGNORE_FIR
-
 @file:[kotlin.<caret>Deprecated("message")]
 
 package foo
diff --git a/idea/testData/resolve/references/AnnotationOnFileWithImport.kt b/idea/testData/resolve/references/AnnotationOnFileWithImport.kt
index 21e0ccdd293..8497c5dc2e5 100644
--- a/idea/testData/resolve/references/AnnotationOnFileWithImport.kt
+++ b/idea/testData/resolve/references/AnnotationOnFileWithImport.kt
@@ -1,5 +1,3 @@
-// IGNORE_FIR
-
 @file:[<caret>D("message")]
 
 package foo