Convert lambda to reference: suggest on lambda argument for extension function parameter

#KT-22014 Fixed
#KT-22015 Fixed

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

@@ -6,10 +6,7 @@
 package org.jetbrains.kotlin.idea.intentions
 
 import com.intellij.openapi.editor.Editor
-import org.jetbrains.kotlin.builtins.getReturnTypeFromFunctionType
+import org.jetbrains.kotlin.builtins.*
-import org.jetbrains.kotlin.builtins.isExtensionFunctionType
-import org.jetbrains.kotlin.builtins.isFunctionType
-import org.jetbrains.kotlin.builtins.isSuspendFunctionType
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.idea.KotlinBundle
@@ -19,20 +16,26 @@ import org.jetbrains.kotlin.idea.core.ShortenReferences
 import org.jetbrains.kotlin.idea.imports.importableFqName
 import org.jetbrains.kotlin.idea.inspections.IntentionBasedInspection
 import org.jetbrains.kotlin.idea.project.languageVersionSettings
+import org.jetbrains.kotlin.idea.refactoring.fqName.fqName
 import org.jetbrains.kotlin.idea.util.IdeDescriptorRenderers
 import org.jetbrains.kotlin.idea.util.approximateFlexibleTypes
 import org.jetbrains.kotlin.idea.util.getResolutionScope
 import org.jetbrains.kotlin.psi.*
+import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.renderer.render
+import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.BindingContext.FUNCTION
 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.calls.components.hasDefaultValue
 import org.jetbrains.kotlin.resolve.calls.model.VarargValueArgument
 import org.jetbrains.kotlin.resolve.descriptorUtil.isCompanionObject
+import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExtensionReceiver
 import org.jetbrains.kotlin.resolve.scopes.utils.getImplicitReceiversHierarchy
 import org.jetbrains.kotlin.synthetic.SyntheticJavaPropertyDescriptor
+import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.types.isDynamic
 import org.jetbrains.kotlin.types.isError
 import org.jetbrains.kotlin.types.typeUtil.isTypeParameter
@@ -58,8 +61,7 @@ open class ConvertLambdaToReferenceIntention(textGetter: () -> String) : SelfTar
     private fun isConvertibleCallInLambda(
         callableExpression: KtExpression,
         explicitReceiver: KtExpression? = null,
-        lambdaExpression: KtLambdaExpression,
+        lambdaExpression: KtLambdaExpression
-        lambdaMustReturnUnit: Boolean
     ): Boolean {
         val context = callableExpression.analyze()
         val calleeReferenceExpression = when (callableExpression) {
@@ -68,6 +70,15 @@ open class ConvertLambdaToReferenceIntention(textGetter: () -> String) : SelfTar
             else -> return false
         }
         val calleeDescriptor = context[REFERENCE_TARGET, calleeReferenceExpression] as? CallableMemberDescriptor ?: return false
+
+        val lambdaParameterType = lambdaExpression.lambdaParameterType(context)
+        if (lambdaParameterType?.isSuspendFunctionType == true) return false
+        if (lambdaParameterType?.isExtensionFunctionType == true) {
+            if (explicitReceiver != null && explicitReceiver !is KtThisExpression) return false
+            val receiver = calleeDescriptor.dispatchReceiverParameter ?: calleeDescriptor.extensionReceiverParameter
+            if (lambdaParameterType.getReceiverTypeFromFunctionType() != receiver?.type) return false
+        }
+
         // No references with type parameters
         if (calleeDescriptor.typeParameters.isNotEmpty()) return false
         // No references to Java synthetic properties
@@ -87,6 +98,8 @@ open class ConvertLambdaToReferenceIntention(textGetter: () -> String) : SelfTar
 
         val callableArgumentsCount = (callableExpression as? KtCallExpression)?.valueArguments?.size ?: 0
         if (calleeDescriptor.valueParameters.size != callableArgumentsCount) return false
+        val lambdaMustReturnUnit =
+            if (lambdaParameterType?.isFunctionType == true) lambdaParameterType.getReturnTypeFromFunctionType().isUnit() else false
         if (lambdaMustReturnUnit) {
             calleeDescriptor.returnType.let {
                 // If Unit required, no references to non-Unit callables
@@ -135,26 +148,11 @@ open class ConvertLambdaToReferenceIntention(textGetter: () -> String) : SelfTar
 
     override fun isApplicableTo(element: KtLambdaExpression): Boolean {
         val singleStatement = element.singleStatementOrNull() ?: return false
-        val lambdaParent = element.parent
-        var lambdaMustReturnUnit = false
-        if (lambdaParent is KtLambdaArgument) {
-            val outerCalleeDescriptor = lambdaParent.outerCalleeDescriptor() ?: return false
-            val lambdaParameterType = outerCalleeDescriptor.valueParameters.lastOrNull()?.type
-            if (lambdaParameterType?.isSuspendFunctionType == true) return false
-            if (lambdaParameterType?.isFunctionType == true) {
-                // For lambda parameter with receiver, conversion is not allowed
-                if (lambdaParameterType.isExtensionFunctionType) return false
-                // Special Unit case (non-Unit returning lambda is accepted here, but non-Unit returning reference is not)
-                lambdaMustReturnUnit = lambdaParameterType.getReturnTypeFromFunctionType().isUnit()
-            }
-        }
-
         return when (singleStatement) {
             is KtCallExpression -> {
                 isConvertibleCallInLambda(
                     callableExpression = singleStatement,
-                    lambdaExpression = element,
+                    lambdaExpression = element
-                    lambdaMustReturnUnit = lambdaMustReturnUnit
                 )
             }
             is KtNameReferenceExpression -> false // Global property reference is not possible (?!)
@@ -163,8 +161,7 @@ open class ConvertLambdaToReferenceIntention(textGetter: () -> String) : SelfTar
                 isConvertibleCallInLambda(
                     callableExpression = selector,
                     explicitReceiver = singleStatement.receiverExpression,
-                    lambdaExpression = element,
+                    lambdaExpression = element
-                    lambdaMustReturnUnit = lambdaMustReturnUnit
                 )
             }
             else -> false
@@ -221,7 +218,16 @@ open class ConvertLambdaToReferenceIntention(textGetter: () -> String) : SelfTar
     }
 
     companion object {
+        private fun KtLambdaExpression.lambdaParameterType(context: BindingContext? = null): KotlinType? {
+            val argument = parent as? KtValueArgument ?: return null
+            val callExpression = argument.getStrictParentOfType<KtCallExpression>() ?: return null
+            return callExpression
+                .getResolvedCall(context ?: analyze(BodyResolveMode.PARTIAL))
+                ?.getParameterForArgument(argument)?.type
+        }
+
         private fun buildReferenceText(lambdaExpression: KtLambdaExpression, shortTypes: Boolean): String? {
+            val lambdaParameterType = lambdaExpression.lambdaParameterType()
             return when (val singleStatement = lambdaExpression.singleStatementOrNull()) {
                 is KtCallExpression -> {
                     val calleeReferenceExpression = singleStatement.calleeExpression as? KtNameReferenceExpression ?: return null
@@ -229,12 +235,13 @@ open class ConvertLambdaToReferenceIntention(textGetter: () -> String) : SelfTar
                     val receiver = resolvedCall.dispatchReceiver ?: resolvedCall.extensionReceiver
                     val descriptor by lazy { receiver?.type?.constructor?.declarationDescriptor }
                     val receiverText = when {
+                        lambdaParameterType?.isExtensionFunctionType == true ->
+                            lambdaParameterType.getReceiverTypeFromFunctionType()?.fqName?.asString()
                         receiver == null || descriptor?.isCompanionObject() == true -> ""
                         receiver is ExtensionReceiver ||
                                 lambdaExpression.getResolutionScope().getImplicitReceiversHierarchy().size == 1 -> "this"
-                        else -> descriptor?.name?.let { "this@$it" } ?: return null
+                        else -> descriptor?.name?.let { "this@$it" }
-                    }
+                    } ?: return null
-
                     "$receiverText::${singleStatement.getCallReferencedName()}"
                 }
                 is KtDotQualifiedExpression -> {
@@ -264,7 +271,12 @@ open class ConvertLambdaToReferenceIntention(textGetter: () -> String) : SelfTar
                             }
                         }
                         else -> {
-                            "${receiver.text}::$selectorReferenceName"
+                            val receiverText = if (lambdaParameterType?.isExtensionFunctionType == true) {
+                                lambdaParameterType.getReceiverTypeFromFunctionType()?.fqName?.asString()
+                            } else {
+                                receiver.text
+                            } ?: return null
+                            "$receiverText::$selectorReferenceName"
                         }
                     }
                 }

idea/testData/intentions/convertLambdaToReference/argumentWithReceiver2.kt

@@ -0,0 +1,16 @@
+package p
+
+class A {
+    fun f1() {
+    }
+}
+
+fun myA(func: A.() -> Unit) {
+    A().func()
+}
+
+fun main() {
+    myA {<caret>
+        f1()
+    }
+}

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

@@ -0,0 +1,14 @@
+package p
+
+class A {
+    fun f1() {
+    }
+}
+
+fun myA(func: A.() -> Unit) {
+    A().func()
+}
+
+fun main() {
+    myA(A::f1)
+}

idea/testData/intentions/convertLambdaToReference/argumentWithReceiver3.kt

@@ -0,0 +1,16 @@
+package p
+
+class A {
+    fun f1() {
+    }
+}
+
+fun myA(func: A.() -> Unit) {
+    A().func()
+}
+
+fun main() {
+    myA {<caret>
+        this.f1()
+    }
+}

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

@@ -0,0 +1,14 @@
+package p
+
+class A {
+    fun f1() {
+    }
+}
+
+fun myA(func: A.() -> Unit) {
+    A().func()
+}
+
+fun main() {
+    myA(A::f1)
+}

idea/testData/intentions/convertLambdaToReference/argumentWithReceiver4.kt

@@ -0,0 +1,14 @@
+package p
+
+class A {
+    fun f1() {
+    }
+}
+
+fun myA(func: A.() -> Unit) {
+    A().func()
+}
+
+fun main() {
+    myA({<caret> f1() })
+}

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

@@ -0,0 +1,14 @@
+package p
+
+class A {
+    fun f1() {
+    }
+}
+
+fun myA(func: A.() -> Unit) {
+    A().func()
+}
+
+fun main() {
+    myA(A::f1)
+}

idea/testData/intentions/convertLambdaToReference/argumentWithReceiver5.kt

@@ -0,0 +1,14 @@
+package p
+
+class A {
+    fun f1() {
+    }
+}
+
+fun myA(func: A.() -> Unit) {
+    A().func()
+}
+
+fun main() {
+    myA({<caret> this.f1() })
+}

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

@@ -0,0 +1,14 @@
+package p
+
+class A {
+    fun f1() {
+    }
+}
+
+fun myA(func: A.() -> Unit) {
+    A().func()
+}
+
+fun main() {
+    myA(A::f1)
+}

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

@@ -4875,6 +4875,26 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             runTest("idea/testData/intentions/convertLambdaToReference/argumentWithReceiver.kt");
         }
 
+        @TestMetadata("argumentWithReceiver2.kt")
+        public void testArgumentWithReceiver2() throws Exception {
+            runTest("idea/testData/intentions/convertLambdaToReference/argumentWithReceiver2.kt");
+        }
+
+        @TestMetadata("argumentWithReceiver3.kt")
+        public void testArgumentWithReceiver3() throws Exception {
+            runTest("idea/testData/intentions/convertLambdaToReference/argumentWithReceiver3.kt");
+        }
+
+        @TestMetadata("argumentWithReceiver4.kt")
+        public void testArgumentWithReceiver4() throws Exception {
+            runTest("idea/testData/intentions/convertLambdaToReference/argumentWithReceiver4.kt");
+        }
+
+        @TestMetadata("argumentWithReceiver5.kt")
+        public void testArgumentWithReceiver5() throws Exception {
+            runTest("idea/testData/intentions/convertLambdaToReference/argumentWithReceiver5.kt");
+        }
+
         @TestMetadata("backtickEscaped.kt")
         public void testBacktickEscaped() throws Exception {
             runTest("idea/testData/intentions/convertLambdaToReference/backtickEscaped.kt");