From 0781156b42517f0f21283bf451a5f7c330e50dd6 Mon Sep 17 00:00:00 2001
From: Igor Yakovlev <igor.yakovlev@jetbrains.com>
Date: Tue, 10 Aug 2021 17:55:13 +0200
Subject: [PATCH] [FIR IDE] Fix invalid reference resolve into companion with
 custom name

---
 .../idea/references/FirReferenceResolveHelper.kt      |  6 +++---
 .../ResolveNamedCompanionInCompanionType.kt           | 11 +++++++++++
 .../ResolveNamedCompanionInCompanionType.txt          |  2 ++
 .../frontend/api/ReferenceResolveTestGenerated.java   |  6 ++++++
 4 files changed, 22 insertions(+), 3 deletions(-)
 create mode 100644 idea/idea-frontend-fir/testData/referenceResolve/nestedTypes/ResolveNamedCompanionInCompanionType.kt
 create mode 100644 idea/idea-frontend-fir/testData/referenceResolve/nestedTypes/ResolveNamedCompanionInCompanionType.txt

diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirReferenceResolveHelper.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirReferenceResolveHelper.kt
index e56381b4d5a..f0a475116d8 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirReferenceResolveHelper.kt
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirReferenceResolveHelper.kt
@@ -412,10 +412,10 @@ internal object FirReferenceResolveHelper {
             var qualifiedAccess: KtDotQualifiedExpression = firSourcePsi
             val referencedClassId =
                 if ((referencedClass as? FirRegularClass)?.isCompanion == true &&
-                    (qualifiedAccess.selectorExpression as? KtNameReferenceExpression)?.getReferencedName() != "Companion"
+                    (qualifiedAccess.selectorExpression as? KtNameReferenceExpression)?.getReferencedName() != referencedClass.classId.shortClassName.asString()
                 ) {
-                    // Remove the last "Companion" part if the qualified access does not contain it. This is needed because the "Companion"
-                    // part is optional.
+                    // Remove the last companion name part if the qualified access does not contain it.
+                    // This is needed because the companion name part is optional.
                     referencedClass.classId.outerClassId ?: return referencedSymbolsByFir
                 } else {
                     referencedClass.classId
diff --git a/idea/idea-frontend-fir/testData/referenceResolve/nestedTypes/ResolveNamedCompanionInCompanionType.kt b/idea/idea-frontend-fir/testData/referenceResolve/nestedTypes/ResolveNamedCompanionInCompanionType.kt
new file mode 100644
index 00000000000..8cb7e9bb9a6
--- /dev/null
+++ b/idea/idea-frontend-fir/testData/referenceResolve/nestedTypes/ResolveNamedCompanionInCompanionType.kt
@@ -0,0 +1,11 @@
+package foo.bar.baz
+
+class X {
+    companion object COMP {
+        fun lol() {}
+    }
+}
+
+class E {
+    val x = foo.bar.baz.X.CO<caret>MP
+}
\ No newline at end of file
diff --git a/idea/idea-frontend-fir/testData/referenceResolve/nestedTypes/ResolveNamedCompanionInCompanionType.txt b/idea/idea-frontend-fir/testData/referenceResolve/nestedTypes/ResolveNamedCompanionInCompanionType.txt
new file mode 100644
index 00000000000..37327649210
--- /dev/null
+++ b/idea/idea-frontend-fir/testData/referenceResolve/nestedTypes/ResolveNamedCompanionInCompanionType.txt
@@ -0,0 +1,2 @@
+Resolved to:
+0: (in foo.bar.baz.X) companion object COMP
\ No newline at end of file
diff --git a/idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/fir/frontend/api/ReferenceResolveTestGenerated.java b/idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/fir/frontend/api/ReferenceResolveTestGenerated.java
index 13a9c0f7bfc..1822f852c1c 100644
--- a/idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/fir/frontend/api/ReferenceResolveTestGenerated.java
+++ b/idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/fir/frontend/api/ReferenceResolveTestGenerated.java
@@ -971,6 +971,12 @@ public class ReferenceResolveTestGenerated extends AbstractReferenceResolveTest
             runTest("idea/idea-frontend-fir/testData/referenceResolve/nestedTypes/ResolveMiddleOfPackageInType.kt");
         }
 
+        @Test
+        @TestMetadata("ResolveNamedCompanionInCompanionType.kt")
+        public void testResolveNamedCompanionInCompanionType() throws Exception {
+            runTest("idea/idea-frontend-fir/testData/referenceResolve/nestedTypes/ResolveNamedCompanionInCompanionType.kt");
+        }
+
         @Test
         @TestMetadata("ResolveStartOfPackageInType.kt")
         public void testResolveStartOfPackageInType() throws Exception {