Convert lambda to reference: add type parameter to outer call expression if needed

#KT-37744 Fixed

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

@@ -20,6 +20,7 @@ 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.incremental.components.NoLookupLocation
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.renderer.render
@@ -30,9 +31,12 @@ import org.jetbrains.kotlin.resolve.DescriptorUtils
 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.components.isVararg
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 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.lazy.descriptors.LazyPackageDescriptor
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExtensionReceiver
 import org.jetbrains.kotlin.resolve.scopes.utils.getImplicitReceiversHierarchy
 import org.jetbrains.kotlin.synthetic.SyntheticJavaPropertyDescriptor
@@ -54,11 +58,6 @@ open class ConvertLambdaToReferenceIntention(textGetter: () -> String) : SelfTar
 
     open fun buildReferenceText(element: KtLambdaExpression) = buildReferenceText(lambdaExpression = element, shortTypes = false)
 
-    private fun KtLambdaArgument.outerCalleeDescriptor(): FunctionDescriptor? {
-        val outerCallExpression = parent as? KtCallExpression ?: return null
-        return outerCallExpression.resolveToCall()?.resultingDescriptor as? FunctionDescriptor
-    }
-
     private fun isConvertibleCallInLambda(
         callableExpression: KtExpression,
         explicitReceiver: KtExpression? = null,
@@ -173,6 +172,15 @@ open class ConvertLambdaToReferenceIntention(textGetter: () -> String) : SelfTar
     override fun applyTo(element: KtLambdaExpression, editor: Editor?) {
         val referenceName = buildReferenceText(element) ?: return
         val factory = KtPsiFactory(element)
+        val parent = element.parent
+
+        val outerCallExpression = parent.getStrictParentOfType<KtCallExpression>()
+        val outerCallContext = outerCallExpression?.analyze(BodyResolveMode.PARTIAL)
+        val resolvedOuterCall = outerCallContext?.let { outerCallExpression.getResolvedCall(it) }
+        if (parent is KtValueArgument && resolvedOuterCall != null) {
+            outerCallExpression.addTypeArgumentsIfNeeded(element, parent, resolvedOuterCall, outerCallContext)
+        }
+
         val lambdaArgument = element.parentValueArgument() as? KtLambdaArgument
         if (lambdaArgument == null) {
             // Without lambda argument syntax, just replace lambda with reference
@@ -180,8 +188,7 @@ open class ConvertLambdaToReferenceIntention(textGetter: () -> String) : SelfTar
             (element.replace(callableReferenceExpr) as? KtElement)?.let { ShortenReferences.RETAIN_COMPANION.process(it) }
         } else {
             // Otherwise, replace the whole argument list for lambda argument-using call
-            val outerCallExpression = lambdaArgument.parent as? KtCallExpression ?: return
+            val outerCalleeDescriptor = resolvedOuterCall?.resultingDescriptor ?: return
-            val outerCalleeDescriptor = lambdaArgument.outerCalleeDescriptor() ?: return
             // Parameters with default value
             val valueParameters = outerCalleeDescriptor.valueParameters
             val arguments = outerCallExpression.valueArguments.filter { it !is KtLambdaArgument }
@@ -219,6 +226,39 @@ open class ConvertLambdaToReferenceIntention(textGetter: () -> String) : SelfTar
         }
     }
 
+    private fun KtCallExpression.addTypeArgumentsIfNeeded(
+        lambda: KtLambdaExpression,
+        valueArgument: KtValueArgument,
+        resolvedCall: ResolvedCall<out CallableDescriptor>,
+        context: BindingContext,
+    ) {
+        val parameter = resolvedCall.getParameterForArgument(valueArgument) ?: return
+        val parameterType = if (parameter.isVararg) {
+            parameter.original.type.arguments.firstOrNull()?.type
+        } else {
+            parameter.original.type
+        } ?: return
+        if (parameterType.arguments.none { it.type.isTypeParameter() }) return
+
+        val calledFunctionInLambda = lambda.singleStatementOrNull()
+            ?.getResolvedCall(context)?.resultingDescriptor as? FunctionDescriptor ?: return
+        val overloadedFunctions = calledFunctionInLambda.overloadedFunctions()
+        if (overloadedFunctions.count { it.valueParameters.size == calledFunctionInLambda.valueParameters.size } < 2) return
+
+        if (InsertExplicitTypeArgumentsIntention.isApplicableTo(this, context)) {
+            InsertExplicitTypeArgumentsIntention.applyTo(this)
+        }
+    }
+
+    private fun FunctionDescriptor.overloadedFunctions(): Collection<SimpleFunctionDescriptor> {
+        val memberScope = when (val containingDeclaration = this.containingDeclaration) {
+            is ClassDescriptor -> containingDeclaration.unsubstitutedMemberScope
+            is LazyPackageDescriptor -> containingDeclaration.getMemberScope()
+            else -> null
+        }
+        return memberScope?.getContributedFunctions(name, NoLookupLocation.FROM_IDE).orEmpty()
+    }
+
     companion object {
         private fun KtLambdaExpression.lambdaParameterType(context: BindingContext? = null): KotlinType? {
             val argument = parentValueArgument() ?: return null

idea/testData/intentions/convertLambdaToReference/typeInferenceFromLambdaParameters.kt

@@ -0,0 +1,8 @@
+fun overloadFun(p: Int, q: Long) {}
+fun overloadFun(p: String, q: Long) {}
+
+fun <T, U> foo(fn: (T, U) -> Unit) {}
+
+fun test() {
+    foo {<caret> x: String, y: Long -> overloadFun(x, y) }
+}

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

@@ -0,0 +1,8 @@
+fun overloadFun(p: Int, q: Long) {}
+fun overloadFun(p: String, q: Long) {}
+
+fun <T, U> foo(fn: (T, U) -> Unit) {}
+
+fun test() {
+    foo<String, Long>(::overloadFun)
+}

idea/testData/intentions/convertLambdaToReference/typeInferenceFromLambdaParameters2.kt

@@ -0,0 +1,8 @@
+fun overloadFun(p: Int, q: Long) {}
+fun overloadFun(p: String, q: Long) {}
+
+fun <T, U> foo(fn: (T, U) -> Unit) {}
+
+fun test() {
+    foo({<caret> x: String, y: Long -> overloadFun(x, y) })
+}

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

@@ -0,0 +1,8 @@
+fun overloadFun(p: Int, q: Long) {}
+fun overloadFun(p: String, q: Long) {}
+
+fun <T, U> foo(fn: (T, U) -> Unit) {}
+
+fun test() {
+    foo<String, Long>(::overloadFun)
+}

idea/testData/intentions/convertLambdaToReference/typeInferenceFromLambdaParameters3.kt

@@ -0,0 +1,12 @@
+fun overloadFun(p: Int, q: Long) {}
+fun overloadFun(p: String, q: Long) {}
+
+fun interface Action<T, U> {
+    fun bar(t: T, u: U)
+}
+
+fun <T, U> foo(a: Action<T, U>) {}
+
+fun test() {
+    foo {<caret> x: String, y: Long -> overloadFun(x, y) }
+}

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

@@ -0,0 +1,12 @@
+fun overloadFun(p: Int, q: Long) {}
+fun overloadFun(p: String, q: Long) {}
+
+fun interface Action<T, U> {
+    fun bar(t: T, u: U)
+}
+
+fun <T, U> foo(a: Action<T, U>) {}
+
+fun test() {
+    foo<String, Long>(::overloadFun)
+}

idea/testData/intentions/convertLambdaToReference/typeInferenceFromLambdaParameters4.kt

@@ -0,0 +1,7 @@
+fun overloadFun(p: Int, q: Long) {}
+
+fun <T, U> foo(fn: (T, U) -> Unit) {}
+
+fun test() {
+    foo {<caret> x: Int, y: Long -> overloadFun(x, y) }
+}

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

@@ -0,0 +1,7 @@
+fun overloadFun(p: Int, q: Long) {}
+
+fun <T, U> foo(fn: (T, U) -> Unit) {}
+
+fun test() {
+    foo(::overloadFun)
+}

idea/testData/intentions/convertLambdaToReference/typeInferenceFromLambdaParameters5.kt

@@ -0,0 +1,8 @@
+fun overloadFun(p: Int, q: Long) {}
+fun overloadFun(p: Int) {}
+
+fun <T, U> foo(fn: (T, U) -> Unit) {}
+
+fun test() {
+    foo {<caret> x: Int, y: Long -> overloadFun(x, y) }
+}

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

@@ -0,0 +1,8 @@
+fun overloadFun(p: Int, q: Long) {}
+fun overloadFun(p: Int) {}
+
+fun <T, U> foo(fn: (T, U) -> Unit) {}
+
+fun test() {
+    foo(::overloadFun)
+}

idea/testData/intentions/convertLambdaToReference/typeInferenceFromLambdaParameters6.kt

@@ -0,0 +1,8 @@
+fun overloadFun(p: Int) {}
+fun overloadFun(p: String) {}
+
+fun <T> ambiguityFun(vararg fn: (T) -> Unit) {}
+
+fun overloadContext() {
+    ambiguityFun({<caret> x: String -> overloadFun(x) }, ::overloadFun)
+}

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

@@ -0,0 +1,8 @@
+fun overloadFun(p: Int) {}
+fun overloadFun(p: String) {}
+
+fun <T> ambiguityFun(vararg fn: (T) -> Unit) {}
+
+fun overloadContext() {
+    ambiguityFun<String>(::overloadFun, ::overloadFun)
+}

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

@@ -5230,6 +5230,36 @@ public class IntentionTestGenerated extends AbstractIntentionTest {
             runTest("idea/testData/intentions/convertLambdaToReference/typeFromJavaFlexibleRecursive.kt");
         }
 
+        @TestMetadata("typeInferenceFromLambdaParameters.kt")
+        public void testTypeInferenceFromLambdaParameters() throws Exception {
+            runTest("idea/testData/intentions/convertLambdaToReference/typeInferenceFromLambdaParameters.kt");
+        }
+
+        @TestMetadata("typeInferenceFromLambdaParameters2.kt")
+        public void testTypeInferenceFromLambdaParameters2() throws Exception {
+            runTest("idea/testData/intentions/convertLambdaToReference/typeInferenceFromLambdaParameters2.kt");
+        }
+
+        @TestMetadata("typeInferenceFromLambdaParameters3.kt")
+        public void testTypeInferenceFromLambdaParameters3() throws Exception {
+            runTest("idea/testData/intentions/convertLambdaToReference/typeInferenceFromLambdaParameters3.kt");
+        }
+
+        @TestMetadata("typeInferenceFromLambdaParameters4.kt")
+        public void testTypeInferenceFromLambdaParameters4() throws Exception {
+            runTest("idea/testData/intentions/convertLambdaToReference/typeInferenceFromLambdaParameters4.kt");
+        }
+
+        @TestMetadata("typeInferenceFromLambdaParameters5.kt")
+        public void testTypeInferenceFromLambdaParameters5() throws Exception {
+            runTest("idea/testData/intentions/convertLambdaToReference/typeInferenceFromLambdaParameters5.kt");
+        }
+
+        @TestMetadata("typeInferenceFromLambdaParameters6.kt")
+        public void testTypeInferenceFromLambdaParameters6() throws Exception {
+            runTest("idea/testData/intentions/convertLambdaToReference/typeInferenceFromLambdaParameters6.kt");
+        }
+
         @TestMetadata("unit.kt")
         public void testUnit() throws Exception {
             runTest("idea/testData/intentions/convertLambdaToReference/unit.kt");