Shorten References: Fix removal of explicit companion references in the position of callee expression

#KT-10102

idea/idea-analysis/src/org/jetbrains/kotlin/idea/util/ShortenReferences.kt

@@ -464,7 +464,25 @@ public class ShortenReferences(val options: (KtElement) -> Options = { Options.D
         }
 
         override fun shortenElement(element: KtDotQualifiedExpression): KtElement {
-            return element.replace(element.getReceiverExpression()) as KtElement
+            val receiver = element.receiverExpression
+            val selector = element.selectorExpression ?: return element
+
+            return when (receiver) {
+                is KtSimpleNameExpression -> {
+                    val identifier = receiver.getIdentifier() ?: return element
+                    (selector.getCalleeExpressionIfAny() as? KtSimpleNameExpression)?.getIdentifier()?.replace(identifier)
+                    element.replace(selector) as KtExpression
+                }
+
+                is KtQualifiedExpression -> {
+                    val identifier = (receiver.selectorExpression as? KtSimpleNameExpression)?.getIdentifier() ?: return element
+                    (selector.getCalleeExpressionIfAny() as? KtSimpleNameExpression)?.getIdentifier()?.replace(identifier)
+                    receiver.selectorExpression?.replace(selector)
+                    element.replace(receiver) as KtExpression
+                }
+
+                else -> element
+            }
         }
     }
 }

idea/testData/refactoring/move/kotlin/moveTopLevelDeclarations/misc/shortenCompanionObject2/after/b/boo.kt

@@ -0,0 +1,17 @@
+package b
+
+import java.util.function.IntPredicate
+
+interface Factory {
+    operator fun invoke(i: Int): IntPredicate
+
+    companion object {
+        inline operator fun invoke(crossinline f: (Int) -> IntPredicate) = object : Factory {
+            override fun invoke(i: Int) = f(i)
+        }
+    }
+}
+
+fun foo(): Factory = Factory { k ->
+    IntPredicate { n -> n % k == 0 }
+}

idea/testData/refactoring/move/kotlin/moveTopLevelDeclarations/misc/shortenCompanionObject2/after/b/foo.kt

@@ -0,0 +1,4 @@
+package b
+
+import java.util.function.IntPredicate
+

idea/testData/refactoring/move/kotlin/moveTopLevelDeclarations/misc/shortenCompanionObject2/before/b/boo.kt

@@ -0,0 +1,13 @@
+package b
+
+import java.util.function.IntPredicate
+
+interface Factory {
+    operator fun invoke(i: Int): IntPredicate
+
+    companion object {
+        inline operator fun invoke(crossinline f: (Int) -> IntPredicate) = object : Factory {
+            override fun invoke(i: Int) = f(i)
+        }
+    }
+}

idea/testData/refactoring/move/kotlin/moveTopLevelDeclarations/misc/shortenCompanionObject2/before/b/foo.kt

@@ -0,0 +1,7 @@
+package b
+
+import java.util.function.IntPredicate
+
+fun <caret>foo(): Factory = Factory { k ->
+    IntPredicate { n -> n % k == 0 }
+}

idea/testData/refactoring/move/kotlin/moveTopLevelDeclarations/misc/shortenCompanionObject2/shortenCompanionObject.test

@@ -0,0 +1,5 @@
+{
+  "mainFile": "b/foo.kt",
+  "type": "MOVE_KOTLIN_TOP_LEVEL_DECLARATIONS",
+  "targetFile": "b/boo.kt"
+}

idea/testData/shortenRefs/removeCompanionRefInCalleeExpression.kt

@@ -0,0 +1,11 @@
+interface Factory {
+    operator fun invoke(i: Int): IntPredicate
+
+    companion object {
+        inline operator fun invoke(crossinline f: (Int) -> IntPredicate) = object : Factory {
+            override fun invoke(i: Int) = f(i)
+        }
+    }
+}
+
+<selection>fun foo(): Factory = Factory.Companion { k -> IntPredicate { n -> n % k == 0 } }</selection>

idea/testData/shortenRefs/removeCompanionRefInCalleeExpression.kt.after

@@ -0,0 +1,11 @@
+interface Factory {
+    operator fun invoke(i: Int): IntPredicate
+
+    companion object {
+        inline operator fun invoke(crossinline f: (Int) -> IntPredicate) = object : Factory {
+            override fun invoke(i: Int) = f(i)
+        }
+    }
+}
+
+fun foo(): Factory = Factory { k -> IntPredicate { n -> n % k == 0 } }

idea/testData/shortenRefs/removeCompanionRefWithQualifiedReceiverInCalleeExpression.kt

@@ -0,0 +1,13 @@
+class T {
+    interface Factory {
+        operator fun invoke(i: Int): IntPredicate
+
+        companion object {
+            inline operator fun invoke(crossinline f: (Int) -> IntPredicate) = object : Factory {
+                override fun invoke(i: Int) = f(i)
+            }
+        }
+    }
+}
+
+<selection>fun foo(): T.Factory = T.Factory.Companion { k -> IntPredicate { n -> n % k == 0 } }</selection>

idea/testData/shortenRefs/removeCompanionRefWithQualifiedReceiverInCalleeExpression.kt.after

@@ -0,0 +1,13 @@
+class T {
+    interface Factory {
+        operator fun invoke(i: Int): IntPredicate
+
+        companion object {
+            inline operator fun invoke(crossinline f: (Int) -> IntPredicate) = object : Factory {
+                override fun invoke(i: Int) = f(i)
+            }
+        }
+    }
+}
+
+fun foo(): T.Factory = T.Factory { k -> IntPredicate { n -> n % k == 0 } }

idea/tests/org/jetbrains/kotlin/idea/refactoring/move/MoveTestGenerated.java

@@ -329,6 +329,12 @@ public class MoveTestGenerated extends AbstractMoveTest {
         doTest(fileName);
     }
 
+    @TestMetadata("kotlin/moveTopLevelDeclarations/misc/shortenCompanionObject2/shortenCompanionObject.test")
+    public void testKotlin_moveTopLevelDeclarations_misc_shortenCompanionObject2_ShortenCompanionObject() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/move/kotlin/moveTopLevelDeclarations/misc/shortenCompanionObject2/shortenCompanionObject.test");
+        doTest(fileName);
+    }
+
     @TestMetadata("kotlin/moveTopLevelDeclarations/misc/singletonsAndStatics/singletonsAndStatics.test")
     public void testKotlin_moveTopLevelDeclarations_misc_singletonsAndStatics_SingletonsAndStatics() throws Exception {
         String fileName = KotlinTestUtils.navigationMetadata("idea/testData/refactoring/move/kotlin/moveTopLevelDeclarations/misc/singletonsAndStatics/singletonsAndStatics.test");

idea/tests/org/jetbrains/kotlin/shortenRefs/ShortenRefsTestGenerated.java

@@ -83,6 +83,18 @@ public class ShortenRefsTestGenerated extends AbstractShortenRefsTest {
         doTest(fileName);
     }
 
+    @TestMetadata("removeCompanionRefInCalleeExpression.kt")
+    public void testRemoveCompanionRefInCalleeExpression() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/shortenRefs/removeCompanionRefInCalleeExpression.kt");
+        doTest(fileName);
+    }
+
+    @TestMetadata("removeCompanionRefWithQualifiedReceiverInCalleeExpression.kt")
+    public void testRemoveCompanionRefWithQualifiedReceiverInCalleeExpression() throws Exception {
+        String fileName = KotlinTestUtils.navigationMetadata("idea/testData/shortenRefs/removeCompanionRefWithQualifiedReceiverInCalleeExpression.kt");
+        doTest(fileName);
+    }
+
     @TestMetadata("idea/testData/shortenRefs/constructor")
     @TestDataPath("$PROJECT_ROOT")
     @RunWith(JUnit3RunnerWithInners.class)