FIR: Fix bug with incorrect source element for qualifiers

compiler/fir/resolve/build.gradle.kts

@@ -14,9 +14,10 @@ dependencies {
     api(project(":compiler:fir:tree"))
     api(kotlinxCollectionsImmutable())
     implementation(project(":core:util.runtime"))
+    implementation(project(":compiler:psi"))
 
     compileOnly(project(":kotlin-reflect-api"))
-    compileOnly(intellijCoreDep()) { includeJars("guava", rootProject = rootProject) }
+    compileOnly(intellijCoreDep()) { includeJars("intellij-core", "guava", rootProject = rootProject) }
 }
 
 sourceSets {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirQualifiedNameResolver.kt

@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.fir.resolve.typeForQualifier
 import org.jetbrains.kotlin.fir.types.FirTypeProjection
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
 
 const val ROOT_PREFIX_FOR_IDE_RESOLUTION_MODE = "_root_ide_package_"
 
@@ -91,7 +92,7 @@ class FirQualifiedNameResolver(private val components: BodyResolveComponents) {
             }
 
             return buildResolvedQualifier {
-                this.source = source
+                this.source = getWholeQualifierSource(source, qualifierPartsToDrop)
                 packageFqName = resolved.packageFqName
                 relativeClassFqName = resolved.relativeClassFqName
                 symbol = resolved.classSymbol
@@ -104,4 +105,11 @@ class FirQualifiedNameResolver(private val components: BodyResolveComponents) {
         return null
     }
 
+    private fun getWholeQualifierSource(qualifierStartSource: FirSourceElement?, stepsToWholeQualifier: Int): FirSourceElement? {
+        if (qualifierStartSource !is FirRealPsiSourceElement<*>) return qualifierStartSource
+
+        val qualifierStart = qualifierStartSource.psi
+        val wholeQualifier = qualifierStart.parentsWithSelf.drop(stepsToWholeQualifier).first()
+        return wholeQualifier.toFirPsiSourceElement()
+    }
 }

idea/idea-fir/tests/org/jetbrains/kotlin/idea/resolve/FirReferenceResolveTestGenerated.java

@@ -144,6 +144,11 @@ public class FirReferenceResolveTestGenerated extends AbstractFirReferenceResolv
         runTest("idea/testData/resolve/references/EnumValues.kt");
     }
 
+    @TestMetadata("ExternalCompanionObject.kt")
+    public void testExternalCompanionObject() throws Exception {
+        runTest("idea/testData/resolve/references/ExternalCompanionObject.kt");
+    }
+
     @TestMetadata("FakeJavaLang1.kt")
     public void testFakeJavaLang1() throws Exception {
         runTest("idea/testData/resolve/references/FakeJavaLang1.kt");

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/references/FirReferenceResolveHelper.kt

@@ -7,6 +7,7 @@ package org.jetbrains.kotlin.idea.references
 
 import com.intellij.psi.tree.TokenSet
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.ROOT_PREFIX_FOR_IDE_RESOLUTION_MODE
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
 import org.jetbrains.kotlin.fir.expressions.*
@@ -121,6 +122,7 @@ internal object FirReferenceResolveHelper {
             else -> {
                 qualified
                     .collectDescendantsOfType<KtSimpleNameExpression>()
+                    .dropWhile { it.getReferencedName() == ROOT_PREFIX_FOR_IDE_RESOLUTION_MODE }
                     .joinToString(separator = ".") { it.getReferencedName() }
                     .let(::FqName)
             }

idea/testData/resolve/references/ExternalCompanionObject.Data.kt

@@ -0,0 +1,5 @@
+package dependency
+
+class C {
+    companion object
+}

idea/testData/resolve/references/ExternalCompanionObject.kt

@@ -0,0 +1,5 @@
+fun usage() {
+    dependency.C.<caret>Companion
+}
+
+// REF: companion object of (dependency).C

idea/tests/org/jetbrains/kotlin/idea/resolve/ReferenceResolveTestGenerated.java

@@ -144,6 +144,11 @@ public class ReferenceResolveTestGenerated extends AbstractReferenceResolveTest
         runTest("idea/testData/resolve/references/EnumValues.kt");
     }
 
+    @TestMetadata("ExternalCompanionObject.kt")
+    public void testExternalCompanionObject() throws Exception {
+        runTest("idea/testData/resolve/references/ExternalCompanionObject.kt");
+    }
+
     @TestMetadata("FakeJavaLang1.kt")
     public void testFakeJavaLang1() throws Exception {
         runTest("idea/testData/resolve/references/FakeJavaLang1.kt");