From b1e8238ea2acc2d7c8a9cb771711216e1b48c3dc Mon Sep 17 00:00:00 2001
From: Toshiaki Kameyama <kmymtsak@gmail.com>
Date: Thu, 19 Dec 2019 19:53:51 +0900
Subject: [PATCH] "Convert reference to lambda" intention: handle extension
 function reference with extension function call

#KT-35558 Fixed
---
 .../ConvertReferenceToLambdaIntention.kt         | 16 ++++++++++++----
 .../extensionFunctionalType2.kt                  |  9 +++++++++
 .../extensionFunctionalType2.kt.after            |  9 +++++++++
 .../extensionFunctionalType3.kt                  |  9 +++++++++
 .../extensionFunctionalType3.kt.after            |  9 +++++++++
 .../extensionFunctionalType4.kt                  |  9 +++++++++
 .../extensionFunctionalType4.kt.after            |  9 +++++++++
 .../idea/intentions/IntentionTestGenerated.java  | 15 +++++++++++++++
 8 files changed, 81 insertions(+), 4 deletions(-)
 create mode 100644 idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType2.kt
 create mode 100644 idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType2.kt.after
 create mode 100644 idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType3.kt
 create mode 100644 idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType3.kt.after
 create mode 100644 idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType4.kt
 create mode 100644 idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType4.kt.after

diff --git a/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertReferenceToLambdaIntention.kt b/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertReferenceToLambdaIntention.kt
index e4ec9319c59..609ca8e292a 100644
--- a/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertReferenceToLambdaIntention.kt
+++ b/idea/src/org/jetbrains/kotlin/idea/intentions/ConvertReferenceToLambdaIntention.kt
@@ -24,6 +24,7 @@ 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.descriptorUtil.fqNameSafe
+import org.jetbrains.kotlin.resolve.descriptorUtil.isExtension
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.types.expressions.DoubleColonLHS
 
@@ -42,6 +43,7 @@ class ConvertReferenceToLambdaIntention : SelfTargetingOffsetIndependentIntentio
         val resolvedCall = callGrandParent?.getResolvedCall(context)
         val matchingParameterType = resolvedCall?.getParameterForArgument(valueArgumentParent)?.type
         val matchingParameterIsExtension = matchingParameterType?.isExtensionFunctionType ?: false
+        val firstArgumentIsThis = matchingParameterIsExtension && resolvedCall?.resultingDescriptor?.isExtension == true
 
         val receiverExpression = element.receiverExpression
         val receiverType = receiverExpression?.let {
@@ -87,23 +89,29 @@ class ConvertReferenceToLambdaIntention : SelfTargetingOffsetIndependentIntentio
                         if (targetDescriptor is PropertyDescriptor) "it.$targetName"
                         else "it.$targetName()"
                     else ->
-                        "$receiverPrefix$targetName(it)"
+                        "$receiverPrefix$targetName(${if (firstArgumentIsThis) "this" else "it"})"
                 }
             )
         } else {
+            val (lambdaParams, args) = if (firstArgumentIsThis) {
+                val thisArgument = if (parameterNamesAndTypes.isNotEmpty()) listOf("this") else emptyList()
+                lambdaParameterNamesAndTypes.drop(1) to (thisArgument + parameterNamesAndTypes.drop(1).map { it.first })
+            } else {
+                lambdaParameterNamesAndTypes to parameterNamesAndTypes.map { it.first }
+            }
             factory.createLambdaExpression(
-                parameters = lambdaParameterNamesAndTypes.joinToString(separator = ", ") {
+                parameters = lambdaParams.joinToString(separator = ", ") {
                     if (valueArgumentParent != null) it.first
                     else it.first + ": " + SOURCE_RENDERER.renderType(it.second)
                 },
                 body = if (targetDescriptor is PropertyDescriptor) {
                     "$receiverPrefix$targetName"
                 } else {
-                    parameterNamesAndTypes.joinToString(
+                    args.joinToString(
                         prefix = "$receiverPrefix$targetName(",
                         separator = ", ",
                         postfix = ")"
-                    ) { it.first }
+                    )
                 }
             )
         }
diff --git a/idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType2.kt b/idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType2.kt
new file mode 100644
index 00000000000..644ab2e7b82
--- /dev/null
+++ b/idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType2.kt
@@ -0,0 +1,9 @@
+fun func(a: Int) {
+}
+
+fun Int.foo(block: (Int) -> Unit) {
+}
+
+fun main() {
+    42.foo(::func<caret>)
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType2.kt.after b/idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType2.kt.after
new file mode 100644
index 00000000000..d8e393bad68
--- /dev/null
+++ b/idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType2.kt.after
@@ -0,0 +1,9 @@
+fun func(a: Int) {
+}
+
+fun Int.foo(block: (Int) -> Unit) {
+}
+
+fun main() {
+    42.foo { func(it) }
+}
\ No newline at end of file
diff --git a/idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType3.kt b/idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType3.kt
new file mode 100644
index 00000000000..23362468acc
--- /dev/null
+++ b/idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType3.kt
@@ -0,0 +1,9 @@
+fun func(a: Int) {
+}
+
+fun Int.foo(block: Int.() -> Unit) {
+}
+
+fun main() {
+    42.foo(::func<caret>)
+}
diff --git a/idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType3.kt.after b/idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType3.kt.after
new file mode 100644
index 00000000000..934a48956b8
--- /dev/null
+++ b/idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType3.kt.after
@@ -0,0 +1,9 @@
+fun func(a: Int) {
+}
+
+fun Int.foo(block: Int.() -> Unit) {
+}
+
+fun main() {
+    42.foo { func(this) }
+}
diff --git a/idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType4.kt b/idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType4.kt
new file mode 100644
index 00000000000..2c8e4c92927
--- /dev/null
+++ b/idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType4.kt
@@ -0,0 +1,9 @@
+fun func2(a: Int, b: Int) {
+}
+
+fun Int.foo(block: Int.(Int) -> Unit) {
+}
+
+fun main() {
+    42.foo(::func2<caret>)
+}
diff --git a/idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType4.kt.after b/idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType4.kt.after
new file mode 100644
index 00000000000..2645eb5213e
--- /dev/null
+++ b/idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType4.kt.after
@@ -0,0 +1,9 @@
+fun func2(a: Int, b: Int) {
+}
+
+fun Int.foo(block: Int.(Int) -> Unit) {
+}
+
+fun main() {
+    42.foo { b -> func2(this, b) }
+}
diff --git a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
index 4f8fb14ada1..8df4d86e983 100644
--- a/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
+++ b/idea/tests/org/jetbrains/kotlin/idea/intentions/IntentionTestGenerated.java
@@ -6455,6 +6455,21 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             runTest("idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType.kt");
         }
 
+        @TestMetadata("extensionFunctionalType2.kt")
+        public void testExtensionFunctionalType2() throws Exception {
+            runTest("idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType2.kt");
+        }
+
+        @TestMetadata("extensionFunctionalType3.kt")
+        public void testExtensionFunctionalType3() throws Exception {
+            runTest("idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType3.kt");
+        }
+
+        @TestMetadata("extensionFunctionalType4.kt")
+        public void testExtensionFunctionalType4() throws Exception {
+            runTest("idea/testData/intentions/convertReferenceToLambda/extensionFunctionalType4.kt");
+        }
+
         @TestMetadata("extensionProperty.kt")
         public void testExtensionProperty() throws Exception {
             runTest("idea/testData/intentions/convertReferenceToLambda/extensionProperty.kt");