[kotlin] KTIJ-35945 Account for implicit companion object references in ImportMemberIntention

Fix is only for K2 Mode; K1 Mode would incorrectly import `Companion` instead of the class itself

^KTIJ-35945 Fixed

GitOrigin-RevId: 0249d0a5a9969dade611934bf1f47692e3e1717c

plugins/kotlin/code-insight/intentions-k2/src/org/jetbrains/kotlin/idea/k2/codeinsight/intentions/ImportMemberIntention.kt

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.analysis.api.symbols.*
 import org.jetbrains.kotlin.idea.base.analysis.api.utils.invokeShortening
 import org.jetbrains.kotlin.idea.base.resources.KotlinBundle
 import org.jetbrains.kotlin.idea.codeinsight.api.applicable.intentions.KotlinApplicableModCommandAction
+import org.jetbrains.kotlin.idea.codeinsight.utils.resolveCompanionObjectShortReferenceToContainingClassSymbol
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
 import org.jetbrains.kotlin.psi.KtElement
@@ -49,7 +50,14 @@ internal class ImportMemberIntention :
         (element is KtDotQualifiedExpression && !element.isInImportDirective()) || element is KtUserType
 
     override fun KaSession.prepareContext(element: KtElement): Context? {
-        val symbol = element.actualReference?.resolveToSymbol() ?: return null
+        val reference = element.actualReference ?: return null
+
+        val symbol =
+            // for implicit companion object references, we want to import the outer class
+            reference.resolveCompanionObjectShortReferenceToContainingClassSymbol()
+                ?: reference.resolveToSymbol()
+                ?: return null
+        
         val file = element.containingKtFile
         return computeContext(file, symbol)?.takeUnless {
             symbol.isTopLevel && symbol.containingFile?.isInSamePackage(file) == true

plugins/kotlin/code-insight/intentions-k2/tests/test/org/jetbrains/kotlin/idea/k2/intentions/tests/K2IntentionTestGenerated.java

@@ -1590,6 +1590,11 @@ public abstract class K2IntentionTestGenerated extends AbstractK2IntentionTest {
             runTest("../../../idea/tests/testData/intentions/importMember/ClassCallChain3.kt");
         }
 
+        @TestMetadata("CompanionObjectImplicitReference.kt")
+        public void testCompanionObjectImplicitReference() throws Exception {
+            runTest("../../../idea/tests/testData/intentions/importMember/CompanionObjectImplicitReference.kt");
+        }
+
         @TestMetadata("EnumMember.kt")
         public void testEnumMember() throws Exception {
             runTest("../../../idea/tests/testData/intentions/importMember/EnumMember.kt");

plugins/kotlin/idea/tests/test/org/jetbrains/kotlin/idea/intentions/K1IntentionTestGenerated.java

@@ -11324,6 +11324,11 @@ public abstract class K1IntentionTestGenerated extends AbstractK1IntentionTest {
             runTest("testData/intentions/importMember/ClassCallChain3.kt");
         }
 
+        @TestMetadata("CompanionObjectImplicitReference.kt")
+        public void testCompanionObjectImplicitReference() throws Exception {
+            runTest("testData/intentions/importMember/CompanionObjectImplicitReference.kt");
+        }
+
         @TestMetadata("EnumMember.kt")
         public void testEnumMember() throws Exception {
             runTest("testData/intentions/importMember/EnumMember.kt");

plugins/kotlin/idea/tests/testData/intentions/importMember/CompanionObjectImplicitReference.1.kt

@@ -0,0 +1,7 @@
+package dep
+
+class MyClass {
+    companion object {
+        fun fromCompanion() {}
+    }
+}

plugins/kotlin/idea/tests/testData/intentions/importMember/CompanionObjectImplicitReference.1.kt.after

@@ -0,0 +1,7 @@
+package dep
+
+class MyClass {
+    companion object {
+        fun fromCompanion() {}
+    }
+}

plugins/kotlin/idea/tests/testData/intentions/importMember/CompanionObjectImplicitReference.kt

@@ -0,0 +1,10 @@
+// PRIORITY: HIGH
+// INTENTION_TEXT: "Add import for 'dep.MyClass'"
+// WITH_STDLIB
+package test
+
+fun foo() {
+    dep.MyClass<caret>.fromCompanion()
+}
+
+// IGNORE_K1

plugins/kotlin/idea/tests/testData/intentions/importMember/CompanionObjectImplicitReference.kt.after

@@ -0,0 +1,12 @@
+// PRIORITY: HIGH
+// INTENTION_TEXT: "Add import for 'dep.MyClass'"
+// WITH_STDLIB
+package test
+
+import dep.MyClass
+
+fun foo() {
+    MyClass<caret>.fromCompanion()
+}
+
+// IGNORE_K1