Convert reference to lambda: handle special case of extension functional type #KT-14985 Fixed

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

@@ -17,6 +17,7 @@
 package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.openapi.editor.Editor
+import org.jetbrains.kotlin.builtins.isExtensionFunctionType
 import org.jetbrains.kotlin.descriptors.CallableMemberDescriptor
 import org.jetbrains.kotlin.descriptors.PropertyDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
@@ -28,6 +29,8 @@ import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
 import org.jetbrains.kotlin.resolve.BindingContext.DOUBLE_COLON_LHS
 import org.jetbrains.kotlin.resolve.BindingContext.REFERENCE_TARGET
+import org.jetbrains.kotlin.resolve.calls.callUtil.getParameterForArgument
+import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.types.expressions.DoubleColonLHS
 
@@ -50,23 +53,30 @@ class ConvertReferenceToLambdaIntention : SelfTargetingOffsetIndependentIntentio
         val receiverNameAndType = receiverType?.let { KotlinNameSuggester.suggestNamesByType(it, validator = {
             name -> name !in parameterNamesAndTypes.map { it.first }
         }, defaultName = "receiver").first() to it }
-        val acceptsReceiverAsParameter = receiverNameAndType != null &&
+
+        val valueArgumentParent = element.parent as? KtValueArgument
+        val callGrandParent = valueArgumentParent?.parent?.parent as? KtCallExpression
+        val resolvedCall = callGrandParent?.getResolvedCall(context)
+        val matchingParameterType = resolvedCall?.getParameterForArgument(valueArgumentParent)?.type
+        val matchingParameterIsExtension = matchingParameterType?.isExtensionFunctionType ?: false
+
+        val acceptsReceiverAsParameter = receiverNameAndType != null && !matchingParameterIsExtension &&
                                          (targetDescriptor.dispatchReceiverParameter != null ||
                                           targetDescriptor.extensionReceiverParameter != null)
 
-        val referenceParent = element.parent
-        val insideCall = referenceParent is KtValueArgument
-
         val factory = KtPsiFactory(element)
         val targetName = reference.text
         val lambdaParameterNamesAndTypes =
                 if (acceptsReceiverAsParameter) listOf(receiverNameAndType!!) + parameterNamesAndTypes
                 else parameterNamesAndTypes
 
-        val receiverPrefix =
-                if (acceptsReceiverAsParameter) receiverNameAndType!!.first + "."
-                else receiverExpression?.let { it.text + "." } ?: ""
-        val lambdaExpression = if (insideCall && lambdaParameterNamesAndTypes.size == 1) {
+        val receiverPrefix = when {
+            acceptsReceiverAsParameter -> receiverNameAndType!!.first + "."
+            matchingParameterIsExtension -> ""
+            else -> receiverExpression?.let { it.text + "." } ?: ""
+        }
+
+        val lambdaExpression = if (valueArgumentParent != null && lambdaParameterNamesAndTypes.size == 1) {
             factory.createLambdaExpression(
                     parameters = "",
                     body = when {
@@ -81,7 +91,7 @@ class ConvertReferenceToLambdaIntention : SelfTargetingOffsetIndependentIntentio
         else {
             factory.createLambdaExpression(
                     parameters = lambdaParameterNamesAndTypes.joinToString(separator = ", ") {
-                        if (insideCall) it.first
+                        if (valueArgumentParent != null) it.first
                         else it.first + ": " + SOURCE_RENDERER.renderType(it.second)
                     },
                     body = if (targetDescriptor is PropertyDescriptor) {
@@ -99,11 +109,10 @@ class ConvertReferenceToLambdaIntention : SelfTargetingOffsetIndependentIntentio
         val lambdaResult = element.replace(lambdaExpression) as KtLambdaExpression
         ShortenReferences.DEFAULT.process(lambdaResult)
 
-        if (insideCall) {
-            val call = referenceParent?.parent?.parent as? KtCallExpression ?: return
+        if (valueArgumentParent != null && callGrandParent != null) {
             val moveOutOfParenthesis = MoveLambdaOutsideParenthesesIntention()
-            if (moveOutOfParenthesis.isApplicableTo(call, referenceParent.startOffset)) {
-                moveOutOfParenthesis.applyTo(call, editor)
+            if (moveOutOfParenthesis.isApplicableTo(callGrandParent, valueArgumentParent.startOffset)) {
+                moveOutOfParenthesis.applyTo(callGrandParent, editor)
             }
         }
     }

idea/testData/intentions/convertReferenceToLambda/apply.kt

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+
+class My {
+    fun foo() {}
+}
+
+val x = My().apply(<caret>My::foo)

idea/testData/intentions/convertReferenceToLambda/apply.kt.after

@@ -0,0 +1,7 @@
+// WITH_RUNTIME
+
+class My {
+    fun foo() {}
+}
+
+val x = My().apply { foo() }

idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType.kt

@@ -0,0 +1,8 @@
+class My {
+    fun foo(a: Int, b: String) = "$b$a"
+
+}
+
+fun baz(f: My.(Int, String) -> String) = My().f(42, "")
+
+val x = baz(<caret>My::foo)

idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType.kt.after

@@ -0,0 +1,8 @@
+class My {
+    fun foo(a: Int, b: String) = "$b$a"
+
+}
+
+fun baz(f: My.(Int, String) -> String) = My().f(42, "")
+
+val x = baz { a, b -> foo(a, b) }

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

@@ -4961,6 +4961,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             KotlinTestUtils.assertAllTestsPresentByMetadata(this.getClass(), new File("idea/testData/intentions/convertReferenceToLambda"), Pattern.compile("^([\\w\\-_]+)\\.kt$"), TargetBackend.ANY, true);
         }
 
+        @TestMetadata("apply.kt")
+        public void testApply() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertReferenceToLambda/apply.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("boundReference.kt")
         public void testBoundReference() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertReferenceToLambda/boundReference.kt");
@@ -4985,6 +4991,12 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             doTest(fileName);
         }
 
+        @TestMetadata("extensionFunctionalType.kt")
+        public void testExtensionFunctionalType() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("extensionProperty.kt")
         public void testExtensionProperty() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("idea/testData/intentions/convertReferenceToLambda/extensionProperty.kt");