From 69f74b2a5122e566b16f0a7b42c288cc72b32cbb Mon Sep 17 00:00:00 2001
From: Valentin Kipyatkov <valentin@jetbrains.com>
Date: Tue, 19 Sep 2017 16:45:57 +0300
Subject: [PATCH] Fixed ShortenReferences for extension on companion object

---
 .../kotlin/idea/core/ShortenReferences.kt     | 26 ++++++++++++++-----
 .../extensionForCompanionObject.dependency.kt |  8 ++++++
 .../extensionForCompanionObject.kt            |  5 ++++
 .../extensionForCompanionObject.kt.after      |  6 +++++
 .../shortenRefs/ShortenRefsTestGenerated.java |  6 +++++
 5 files changed, 44 insertions(+), 7 deletions(-)
 create mode 100644 idea/testData/shortenRefs/extensionForCompanionObject.dependency.kt
 create mode 100644 idea/testData/shortenRefs/extensionForCompanionObject.kt
 create mode 100644 idea/testData/shortenRefs/extensionForCompanionObject.kt.after

diff --git a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/ShortenReferences.kt b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/ShortenReferences.kt
index eb437ca0f19..b5aa4ee257e 100644
--- a/idea/idea-core/src/org/jetbrains/kotlin/idea/core/ShortenReferences.kt
+++ b/idea/idea-core/src/org/jetbrains/kotlin/idea/core/ShortenReferences.kt
@@ -538,13 +538,25 @@ class ShortenReferences(val options: (KtElement) -> Options = { Options.DEFAULT
 
         private fun canBePossibleToDropReceiver(element: KtDotQualifiedExpression, bindingContext: BindingContext): Boolean {
             val receiver = element.receiverExpression
-            if (receiver is KtThisExpression) return true
-            val qualifier = bindingContext[BindingContext.QUALIFIER, receiver] ?: return false
-            val classDescriptor = qualifier.descriptor as? ClassDescriptor ?: return true
-            if (classDescriptor.kind != ClassKind.OBJECT) return true
-            // for object receiver we should additionally check that it's dispatch receiver (that is the member is inside the object)
-            val resolvedCall = element.getResolvedCall(bindingContext) ?: return false
-            return resolvedCall.explicitReceiverKind == ExplicitReceiverKind.DISPATCH_RECEIVER
+            val nameRef = when (receiver) {
+                is KtThisExpression -> return true
+                is KtNameReferenceExpression -> receiver
+                is KtDotQualifiedExpression -> receiver.selectorExpression as? KtNameReferenceExpression ?: return false
+                else -> return false
+            }
+            val targetDescriptor = bindingContext[BindingContext.REFERENCE_TARGET, nameRef]
+            when (targetDescriptor) {
+                is ClassDescriptor -> {
+                    if (targetDescriptor.kind != ClassKind.OBJECT) return true
+                    // for object receiver we should additionally check that it's dispatch receiver (that is the member is inside the object)
+                    val resolvedCall = element.getResolvedCall(bindingContext) ?: return false
+                    return resolvedCall.explicitReceiverKind == ExplicitReceiverKind.DISPATCH_RECEIVER
+                }
+
+                is PackageViewDescriptor -> return true
+
+                else -> return false
+            }
         }
 
         private fun copyShortenAndAnalyze(
diff --git a/idea/testData/shortenRefs/extensionForCompanionObject.dependency.kt b/idea/testData/shortenRefs/extensionForCompanionObject.dependency.kt
new file mode 100644
index 00000000000..3761bbe2ad9
--- /dev/null
+++ b/idea/testData/shortenRefs/extensionForCompanionObject.dependency.kt
@@ -0,0 +1,8 @@
+package dependency
+
+class AClass {
+    companion object {
+    }
+}
+
+fun AClass.Companion.extFun(){}
diff --git a/idea/testData/shortenRefs/extensionForCompanionObject.kt b/idea/testData/shortenRefs/extensionForCompanionObject.kt
new file mode 100644
index 00000000000..d70a6731a41
--- /dev/null
+++ b/idea/testData/shortenRefs/extensionForCompanionObject.kt
@@ -0,0 +1,5 @@
+import dependency.extFun
+
+fun foo() {
+    <selection>dependency.AClass.extFun()</selection>
+}
\ No newline at end of file
diff --git a/idea/testData/shortenRefs/extensionForCompanionObject.kt.after b/idea/testData/shortenRefs/extensionForCompanionObject.kt.after
new file mode 100644
index 00000000000..a251c9ad3e8
--- /dev/null
+++ b/idea/testData/shortenRefs/extensionForCompanionObject.kt.after
@@ -0,0 +1,6 @@
+import dependency.AClass
+import dependency.extFun
+
+fun foo() {
+    AClass.extFun()
+}
\ No newline at end of file
diff --git a/idea/tests/org/jetbrains/kotlin/shortenRefs/ShortenRefsTestGenerated.java b/idea/tests/org/jetbrains/kotlin/shortenRefs/ShortenRefsTestGenerated.java
index 4644dea5d4b..5a2f7b80797 100644
--- a/idea/tests/org/jetbrains/kotlin/shortenRefs/ShortenRefsTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/shortenRefs/ShortenRefsTestGenerated.java
@@ -59,6 +59,12 @@ public class ShortenRefsTestGenerated extends AbstractShortenRefsTest {
         runTest("idea/testData/shortenRefs/descriptorsChangeAfterImportInsertion.kt");
     }
 
+    @TestMetadata("extensionForCompanionObject.kt")
+    public void testExtensionForCompanionObject() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/shortenRefs/extensionForCompanionObject.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("extensionForObject.kt")
     public void testExtensionForObject() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/shortenRefs/extensionForObject.kt");