FIR IDE: Check function call before trying to drop the receiver

idea/idea-fir/tests/org/jetbrains/kotlin/shortenRefs/FirShortenRefsTestGenerated.java

@@ -46,6 +46,26 @@ public class FirShortenRefsTestGenerated extends AbstractFirShortenRefsTest {
             runTest("idea/testData/shortenRefsFir/calls/classInSameFile.kt");
         }
 
+        @TestMetadata("explicitlyImportedFunctionFromLocalObject.kt")
+        public void testExplicitlyImportedFunctionFromLocalObject() throws Exception {
+            runTest("idea/testData/shortenRefsFir/calls/explicitlyImportedFunctionFromLocalObject.kt");
+        }
+
+        @TestMetadata("extenstionFunctionOnCompanionObjectReceiverNotShortened.kt")
+        public void testExtenstionFunctionOnCompanionObjectReceiverNotShortened() throws Exception {
+            runTest("idea/testData/shortenRefsFir/calls/extenstionFunctionOnCompanionObjectReceiverNotShortened.kt");
+        }
+
+        @TestMetadata("extenstionFunctionOnObjectReceiverNotShortened.kt")
+        public void testExtenstionFunctionOnObjectReceiverNotShortened() throws Exception {
+            runTest("idea/testData/shortenRefsFir/calls/extenstionFunctionOnObjectReceiverNotShortened.kt");
+        }
+
+        @TestMetadata("extenstionFunctionReceiverNotShortened.kt")
+        public void testExtenstionFunctionReceiverNotShortened() throws Exception {
+            runTest("idea/testData/shortenRefsFir/calls/extenstionFunctionReceiverNotShortened.kt");
+        }
+
         @TestMetadata("functionInSameFile.kt")
         public void testFunctionInSameFile() throws Exception {
             runTest("idea/testData/shortenRefsFir/calls/functionInSameFile.kt");

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirReferenceShortener.kt

@@ -11,11 +11,11 @@ import com.intellij.psi.SmartPsiElementPointer
 import com.intellij.util.containers.addIfNotNull
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.declarations.FirClass
-import org.jetbrains.kotlin.fir.declarations.FirFile
-import org.jetbrains.kotlin.fir.declarations.FirResolvedImport
+import org.jetbrains.kotlin.fir.analysis.checkers.toRegularClass
+import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirFunctionCall
 import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
+import org.jetbrains.kotlin.fir.expressions.impl.FirNoReceiverExpression
 import org.jetbrains.kotlin.fir.psi
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
@@ -247,6 +247,8 @@ internal class KtFirReferenceShortener(
         override fun visitFunctionCall(functionCall: FirFunctionCall) {
             super.visitFunctionCall(functionCall)
 
+            if (!canBePossibleToDropReceiver(functionCall)) return
+
             val callExpression = functionCall.psi as? KtCallExpression ?: return
             val qualifiedCallExpression = callExpression.getDotQualifiedExpressionForSelector() ?: return
 
@@ -261,6 +263,16 @@ internal class KtFirReferenceShortener(
             }
         }
 
+        private fun canBePossibleToDropReceiver(functionCall: FirFunctionCall): Boolean {
+            // we can remove receiver only if it is a qualifier
+            val explicitReceiver = functionCall.explicitReceiver as? FirResolvedQualifier ?: return false
+
+            // if there is no extension receiver necessary, then it can be removed
+            if (functionCall.extensionReceiver is FirNoReceiverExpression) return true
+
+            val receiverType = explicitReceiver.typeRef.toRegularClass(firResolveState.rootModuleSession) ?: return true
+            return receiverType.classKind != ClassKind.OBJECT
+        }
 
         override fun visitResolvedQualifier(resolvedQualifier: FirResolvedQualifier) {
             super.visitResolvedQualifier(resolvedQualifier)

idea/testData/shortenRefsFir/calls/explicitlyImportedFunctionFromLocalObject.kt

@@ -0,0 +1,12 @@
+// FIR_COMPARISON
+package test
+
+import test.Obj.funFromObj
+
+object Obj {
+    fun funFromObj() {}
+}
+
+fun usage() {
+    <selection>Obj.funFromObj()</selection>
+}

idea/testData/shortenRefsFir/calls/explicitlyImportedFunctionFromLocalObject.kt.after

@@ -0,0 +1,12 @@
+// FIR_COMPARISON
+package test
+
+import test.Obj.funFromObj
+
+object Obj {
+    fun funFromObj() {}
+}
+
+fun usage() {
+    funFromObj()
+}

idea/testData/shortenRefsFir/calls/extenstionFunctionOnCompanionObjectReceiverNotShortened.kt

@@ -0,0 +1,12 @@
+// FIR_COMPARISON
+package test
+
+class T {
+    companion object
+}
+
+fun T.Companion.ext() {}
+
+fun usage() {
+    <selection>T.ext()</selection>
+}

idea/testData/shortenRefsFir/calls/extenstionFunctionOnCompanionObjectReceiverNotShortened.kt.after

@@ -0,0 +1,12 @@
+// FIR_COMPARISON
+package test
+
+class T {
+    companion object
+}
+
+fun T.Companion.ext() {}
+
+fun usage() {
+    T.ext()
+}

idea/testData/shortenRefsFir/calls/extenstionFunctionOnObjectReceiverNotShortened.kt

@@ -0,0 +1,10 @@
+// FIR_COMPARISON
+package test
+
+object T
+
+fun T.ext() {}
+
+fun usage() {
+    <selection>T.ext()</selection>
+}

idea/testData/shortenRefsFir/calls/extenstionFunctionOnObjectReceiverNotShortened.kt.after

@@ -0,0 +1,10 @@
+// FIR_COMPARISON
+package test
+
+object T
+
+fun T.ext() {}
+
+fun usage() {
+    T.ext()
+}

idea/testData/shortenRefsFir/calls/extenstionFunctionReceiverNotShortened.kt

@@ -0,0 +1,10 @@
+// FIR_COMPARISON
+package test
+
+class T
+
+fun T.ext() {}
+
+fun usage(t: T) {
+    <selection>t.ext()</selection>
+}

idea/testData/shortenRefsFir/calls/extenstionFunctionReceiverNotShortened.kt.after

@@ -0,0 +1,10 @@
+// FIR_COMPARISON
+package test
+
+class T
+
+fun T.ext() {}
+
+fun usage(t: T) {
+    t.ext()
+}