Lambda --> reference: correct handling of companion references

Reference receivers are named more accurately now #KT-13341 Fixed

idea/idea-core/src/org/jetbrains/kotlin/idea/core/ShortenReferences.kt

@@ -56,7 +56,9 @@ import java.util.*
 class ShortenReferences(val options: (KtElement) -> Options = { Options.DEFAULT }) {
     data class Options(
             val removeThisLabels: Boolean = false,
-            val removeThis: Boolean = false
+            val removeThis: Boolean = false,
+            // TODO: remove this option and all related stuff (RETAIN_COMPANION etc.) after KT-13934 fixed
+            val removeExplicitCompanion: Boolean = true
     ) {
         companion object {
             val DEFAULT = Options()
@@ -68,6 +70,8 @@ class ShortenReferences(val options: (KtElement) -> Options = { Options.DEFAULT
         @JvmField
         val DEFAULT = ShortenReferences()
 
+        val RETAIN_COMPANION = ShortenReferences { Options(removeExplicitCompanion = false) }
+
         private fun DeclarationDescriptor.asString()
                 = DescriptorRenderer.FQ_NAMES_IN_TYPES.render(this)
 
@@ -152,6 +156,15 @@ class ShortenReferences(val options: (KtElement) -> Options = { Options.DEFAULT
 
         val failedToImportDescriptors = LinkedHashSet<DeclarationDescriptor>()
 
+        val companionElementFilter = { element: PsiElement ->
+            if (element is KtElement && !options(element).removeExplicitCompanion) {
+                FilterResult.SKIP
+            }
+            else {
+                elementFilter(element)
+            }
+        }
+
         while (true) {
             // Processors order is important here so that enclosing elements are not shortened before their children are, e.g.
             // test.foo(this@A) -> foo(this)
@@ -159,7 +172,7 @@ class ShortenReferences(val options: (KtElement) -> Options = { Options.DEFAULT
                     ShortenTypesProcessor(file, elementFilter, failedToImportDescriptors),
                     ShortenThisExpressionsProcessor(file, elementFilter, failedToImportDescriptors),
                     ShortenQualifiedExpressionsProcessor(file, elementFilter, failedToImportDescriptors),
-                    RemoveExplicitCompanionObjectReferenceProcessor(file, elementFilter, failedToImportDescriptors)
+                    RemoveExplicitCompanionObjectReferenceProcessor(file, companionElementFilter, failedToImportDescriptors)
             )
 
             // step 1: collect qualified elements to analyze (no resolve at this step)

idea/src/org/jetbrains/kotlin/idea/intentions/ConvertLambdaToReferenceIntention.kt

@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
 import org.jetbrains.kotlin.idea.core.ShortenReferences
 import org.jetbrains.kotlin.idea.core.isVisible
+import org.jetbrains.kotlin.idea.imports.importableFqName
 import org.jetbrains.kotlin.idea.inspections.IntentionBasedInspection
 import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
 import org.jetbrains.kotlin.idea.util.approximateFlexibleTypes
@@ -171,7 +172,7 @@ class ConvertLambdaToReferenceIntention : SelfTargetingOffsetIndependentIntentio
         if (lambdaArgument == null) {
             // Without lambda argument syntax, just replace lambda with reference
             val callableReferenceExpr = factory.createCallableReferenceExpression(referenceName) ?: return
-            (element.replace(callableReferenceExpr) as? KtElement)?.let { ShortenReferences.DEFAULT.process(it) }
+            (element.replace(callableReferenceExpr) as? KtElement)?.let { ShortenReferences.RETAIN_COMPANION.process(it) }
         }
         else {
             // Otherwise, replace the whole argument list for lambda argument-using call
@@ -204,11 +205,11 @@ class ConvertLambdaToReferenceIntention : SelfTargetingOffsetIndependentIntentio
             }
             val argumentList = outerCallExpression.valueArgumentList
             if (argumentList == null) {
-                (lambdaArgument.replace(newArgumentList) as? KtElement)?.let { ShortenReferences.DEFAULT.process(it) }
+                (lambdaArgument.replace(newArgumentList) as? KtElement)?.let { ShortenReferences.RETAIN_COMPANION.process(it) }
             }
             else {
                 (argumentList.replace(newArgumentList) as? KtValueArgumentList)?.let {
-                    ShortenReferences.DEFAULT.process(it.arguments.last())
+                    ShortenReferences.RETAIN_COMPANION.process(it.arguments.last())
                 }
                 lambdaArgument.delete()
             }
@@ -257,7 +258,8 @@ class ConvertLambdaToReferenceIntention : SelfTargetingOffsetIndependentIntentio
                                 }
                             }
                             else {
-                                "${receiverDescriptor.name}::$selectorReferenceName"
+                                val receiverName = receiverDescriptor.importableFqName ?: receiverDescriptor.name
+                                "$receiverName::$selectorReferenceName"
                             }
                         }
                         is KtThisExpression -> {

idea/testData/intentions/convertLambdaToReference/companion.kt

@@ -0,0 +1,11 @@
+// WITH_RUNTIME
+
+class Foo {
+    companion object {
+        fun create(x: String): Foo = Foo()
+    }
+}
+
+fun main(args: Array<String>) {
+    listOf("a").map {<caret> Foo.create(it) }
+}

idea/testData/intentions/convertLambdaToReference/companion.kt.after

@@ -0,0 +1,11 @@
+// WITH_RUNTIME
+
+class Foo {
+    companion object {
+        fun create(x: String): Foo = Foo()
+    }
+}
+
+fun main(args: Array<String>) {
+    listOf("a").map(Foo.Companion::create)
+}

idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java

@@ -4316,6 +4316,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("companion.kt")
+        public void testCompanion() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertLambdaToReference/companion.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("constructor.kt")
         public void testConstructor() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertLambdaToReference/constructor.kt");