KT-33211 Quickfix "add parameter" for method references should infer functional type instead of KFunction (#2664)

* "Add parameter to function" quick fix: add function type instead of KFunction for method references

#KT-33211 Fixed

* Support suspend function references

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

@@ -9,6 +9,7 @@ import com.intellij.psi.PsiElement
 import com.intellij.psi.tree.IElementType
 import org.jetbrains.kotlin.KtNodeTypes
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
+import org.jetbrains.kotlin.builtins.isKSuspendFunctionType
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.caches.resolve.resolveToCall
@@ -16,6 +17,7 @@ import org.jetbrains.kotlin.idea.core.getDeepestSuperDeclarations
 import org.jetbrains.kotlin.idea.core.replaced
 import org.jetbrains.kotlin.idea.core.setType
 import org.jetbrains.kotlin.idea.references.mainReference
+import org.jetbrains.kotlin.idea.references.resolveToDescriptors
 import org.jetbrains.kotlin.idea.search.usagesSearch.descriptor
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.load.java.descriptors.JavaMethodDescriptor
@@ -31,10 +33,11 @@ import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameUnsafe
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.resolve.source.getPsi
 import org.jetbrains.kotlin.types.KotlinType
+import org.jetbrains.kotlin.types.KotlinTypeFactory
 import org.jetbrains.kotlin.types.isFlexible
+import org.jetbrains.kotlin.types.typeUtil.builtIns
 import org.jetbrains.kotlin.types.typeUtil.isUnit
 import org.jetbrains.kotlin.util.OperatorChecks
-import org.jetbrains.kotlin.idea.references.resolveToDescriptors
 
 fun KtContainerNode.description(): String? {
     when (node.elementType) {
@@ -352,3 +355,11 @@ internal fun Sequence<PsiElement>.lastWithPersistedElementOrNull(elementShouldPe
 
     return if (checked) lastElement else null
 }
+
+fun KotlinType.reflectToRegularFunctionType(): KotlinType {
+    val isTypeAnnotatedWithExtensionFunctionType = annotations.findAnnotation(KotlinBuiltIns.FQ_NAMES.extensionFunctionType) != null
+    val parameterCount = if (isTypeAnnotatedWithExtensionFunctionType) arguments.size - 2 else arguments.size - 1
+    val classDescriptor =
+        if (isKSuspendFunctionType) builtIns.getSuspendFunction(parameterCount) else builtIns.getFunction(parameterCount)
+    return KotlinTypeFactory.simpleNotNullType(annotations, classDescriptor, arguments)
+}

idea/src/org/jetbrains/kotlin/idea/quickfix/AddFunctionParametersFix.kt

@@ -24,6 +24,8 @@ import com.intellij.psi.PsiNamedElement
 import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.psi.search.PsiSearchHelper
 import com.intellij.psi.search.searches.ReferencesSearch
+import org.jetbrains.kotlin.builtins.isKFunctionType
+import org.jetbrains.kotlin.builtins.isKSuspendFunctionType
 import org.jetbrains.kotlin.descriptors.ConstructorDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.diagnostics.Diagnostic
@@ -32,9 +34,10 @@ import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.idea.KotlinBundle
 import org.jetbrains.kotlin.idea.caches.resolve.analyze
 import org.jetbrains.kotlin.idea.core.CollectingNameValidator
-import org.jetbrains.kotlin.idea.util.getDataFlowAwareTypes
+import org.jetbrains.kotlin.idea.intentions.reflectToRegularFunctionType
 import org.jetbrains.kotlin.idea.refactoring.changeSignature.*
 import org.jetbrains.kotlin.idea.util.application.runReadAction
+import org.jetbrains.kotlin.idea.util.getDataFlowAwareTypes
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.getParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
@@ -191,7 +194,9 @@ class AddFunctionParametersFix(
     ): KotlinParameterInfo {
         val name = getNewArgumentName(argument, validator)
         val expression = argument.getArgumentExpression()
-        val type = expression?.let { getDataFlowAwareTypes(it).firstOrNull() } ?: functionDescriptor.builtIns.nullableAnyType
+        val type = (expression?.let { getDataFlowAwareTypes(it).firstOrNull() } ?: functionDescriptor.builtIns.nullableAnyType).let {
+            if (it.isKFunctionType || it.isKSuspendFunctionType) it.reflectToRegularFunctionType() else it
+        }
         return KotlinParameterInfo(functionDescriptor, -1, name, KotlinTypeInfo(false, null)).apply {
             currentTypeInfo = KotlinTypeInfo(false, type)
             originalTypeInfo.type?.let { typesToShorten.add(it) }

idea/src/org/jetbrains/kotlin/idea/quickfix/QuickFixFactoryForTypeMismatchError.kt

@@ -20,8 +20,7 @@ import com.intellij.codeInsight.intention.IntentionAction
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.psi.util.PsiTreeUtil
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
-import org.jetbrains.kotlin.builtins.functions.FunctionClassDescriptor
+import org.jetbrains.kotlin.builtins.isKFunctionType
-import org.jetbrains.kotlin.builtins.getFunctionalClassKind
 import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.diagnostics.Diagnostic
 import org.jetbrains.kotlin.diagnostics.Errors
@@ -30,6 +29,7 @@ import org.jetbrains.kotlin.idea.caches.resolve.analyzeWithContent
 import org.jetbrains.kotlin.idea.caches.resolve.findModuleDescriptor
 import org.jetbrains.kotlin.idea.caches.resolve.getResolutionFacade
 import org.jetbrains.kotlin.idea.core.quickfix.QuickFixUtil
+import org.jetbrains.kotlin.idea.intentions.reflectToRegularFunctionType
 import org.jetbrains.kotlin.idea.project.languageVersionSettings
 import org.jetbrains.kotlin.idea.util.approximateWithResolvableType
 import org.jetbrains.kotlin.idea.util.getResolutionScope
@@ -48,21 +48,17 @@ import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
 import org.jetbrains.kotlin.resolve.constants.IntegerValueTypeConstant
 import org.jetbrains.kotlin.resolve.descriptorUtil.fqNameSafe
 import org.jetbrains.kotlin.types.KotlinType
-import org.jetbrains.kotlin.types.KotlinTypeFactory
 import org.jetbrains.kotlin.types.TypeUtils
 import org.jetbrains.kotlin.types.TypeUtils.NO_EXPECTED_TYPE
-import org.jetbrains.kotlin.types.typeUtil.*
+import org.jetbrains.kotlin.types.typeUtil.isInterface
+import org.jetbrains.kotlin.types.typeUtil.isSignedOrUnsignedNumberType
+import org.jetbrains.kotlin.types.typeUtil.isSubtypeOf
+import org.jetbrains.kotlin.types.typeUtil.makeNullable
 import java.util.*
 
 //TODO: should use change signature to deal with cases of multiple overridden descriptors
 class QuickFixFactoryForTypeMismatchError : KotlinIntentionActionsFactory() {
 
-    private fun KotlinType.reflectToRegularFunctionType(): KotlinType {
-        val isTypeAnnotatedWithExtensionFunctionType = annotations.findAnnotation(KotlinBuiltIns.FQ_NAMES.extensionFunctionType) != null
-        val parameterCount = if (isTypeAnnotatedWithExtensionFunctionType) arguments.size - 2 else arguments.size - 1
-        return KotlinTypeFactory.simpleNotNullType(annotations, builtIns.getFunction(parameterCount), arguments)
-    }
-
     override fun doCreateActions(diagnostic: Diagnostic): List<IntentionAction> {
         val actions = LinkedList<IntentionAction>()
 
@@ -183,7 +179,7 @@ class QuickFixFactoryForTypeMismatchError : KotlinIntentionActionsFactory() {
         ) {
             val scope = callable.getResolutionScope(context, callable.getResolutionFacade())
             val typeToInsert = expressionType.approximateWithResolvableType(scope, false)
-            if (typeToInsert.constructor.declarationDescriptor?.getFunctionalClassKind() == FunctionClassDescriptor.Kind.KFunction) {
+            if (typeToInsert.isKFunctionType) {
                 actions.add(createFix(callable, typeToInsert.reflectToRegularFunctionType()))
             }
             actions.add(createFix(callable, typeToInsert))

idea/testData/quickfix/changeSignature/addFunctionParameterWithAnonymousFunction.kt

@@ -0,0 +1,6 @@
+// "Add parameter to function 'baz'" "true"
+fun baz() {}
+
+fun foo() {
+    baz(fun(i: Int): String { return i.toString() }<caret>)
+}

idea/testData/quickfix/changeSignature/addFunctionParameterWithAnonymousFunction.kt.after

@@ -0,0 +1,6 @@
+// "Add parameter to function 'baz'" "true"
+fun baz(function: (Int) -> String) {}
+
+fun foo() {
+    baz(fun(i: Int): String { return i.toString() }<caret>)
+}

idea/testData/quickfix/changeSignature/addFunctionParameterWithFunctionReference.kt

@@ -0,0 +1,8 @@
+// "Add parameter to function 'baz'" "true"
+fun bar(): Int = 42
+
+fun baz() {}
+
+fun foo() {
+    baz(::bar<caret>)
+}

idea/testData/quickfix/changeSignature/addFunctionParameterWithFunctionReference.kt.after

@@ -0,0 +1,8 @@
+// "Add parameter to function 'baz'" "true"
+fun bar(): Int = 42
+
+fun baz(kFunction0: () -> Int) {}
+
+fun foo() {
+    baz(::bar)
+}

idea/testData/quickfix/changeSignature/addFunctionParameterWithLambda.kt

@@ -0,0 +1,6 @@
+// "Add parameter to function 'baz'" "true"
+fun baz() {}
+
+fun foo() {
+    baz { i: Int -> i.toString() }<caret>
+}

idea/testData/quickfix/changeSignature/addFunctionParameterWithLambda.kt.after

@@ -0,0 +1,6 @@
+// "Add parameter to function 'baz'" "true"
+fun baz(function: (Int) -> String) {}
+
+fun foo() {
+    baz { i: Int -> i.toString() }<caret>
+}

idea/testData/quickfix/changeSignature/addFunctionParameterWithSuspendFunctionReference.kt

@@ -0,0 +1,8 @@
+// "Add parameter to function 'baz'" "true"
+suspend fun bar(): Int = 42
+
+suspend fun baz() {}
+
+suspend fun foo() {
+    baz(::bar<caret>)
+}

idea/testData/quickfix/changeSignature/addFunctionParameterWithSuspendFunctionReference.kt.after

@@ -0,0 +1,8 @@
+// "Add parameter to function 'baz'" "true"
+suspend fun bar(): Int = 42
+
+suspend fun baz(kSuspendFunction0: suspend () -> Int) {}
+
+suspend fun foo() {
+    baz(::bar)
+}

idea/tests/org/jetbrains/kotlin/idea/quickfix/QuickFixTestGenerated.java

@@ -2323,6 +2323,21 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             runTest("idea/testData/quickfix/changeSignature/addFunctionParameterLongNameRuntime.kt");
         }
 
+        @TestMetadata("addFunctionParameterWithAnonymousFunction.kt")
+        public void testAddFunctionParameterWithAnonymousFunction() throws Exception {
+            runTest("idea/testData/quickfix/changeSignature/addFunctionParameterWithAnonymousFunction.kt");
+        }
+
+        @TestMetadata("addFunctionParameterWithFunctionReference.kt")
+        public void testAddFunctionParameterWithFunctionReference() throws Exception {
+            runTest("idea/testData/quickfix/changeSignature/addFunctionParameterWithFunctionReference.kt");
+        }
+
+        @TestMetadata("addFunctionParameterWithLambda.kt")
+        public void testAddFunctionParameterWithLambda() throws Exception {
+            runTest("idea/testData/quickfix/changeSignature/addFunctionParameterWithLambda.kt");
+        }
+
         @TestMetadata("addFunctionParameterWithSmartCast.kt")
         public void testAddFunctionParameterWithSmartCast() throws Exception {
             runTest("idea/testData/quickfix/changeSignature/addFunctionParameterWithSmartCast.kt");
@@ -2333,6 +2348,11 @@ public class QuickFixTestGenerated extends AbstractQuickFixTest {
             runTest("idea/testData/quickfix/changeSignature/addFunctionParameterWithSmartcastedArgument.kt");
         }
 
+        @TestMetadata("addFunctionParameterWithSuspendFunctionReference.kt")
+        public void testAddFunctionParameterWithSuspendFunctionReference() throws Exception {
+            runTest("idea/testData/quickfix/changeSignature/addFunctionParameterWithSuspendFunctionReference.kt");
+        }
+
         @TestMetadata("addNothingReturnType.kt")
         public void testAddNothingReturnType() throws Exception {
             runTest("idea/testData/quickfix/changeSignature/addNothingReturnType.kt");