[FIR] Extract function type kind for lambda in argument position

^KT-56381 Fixed

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArguments.kt

@@ -63,8 +63,11 @@ fun Candidate.preprocessLambdaArgument(
 
     if (expectedType != null) {
         val parameters = resolvedArgument.parameters
+        val functionTypeKind = context.session.functionTypeService.extractSingleSpecialKindForFunction(anonymousFunction.symbol)
+            ?: resolvedArgument.expectedFunctionTypeKind
+            ?: FunctionTypeKind.Function
         val lambdaType = createFunctionType(
-            resolvedArgument.expectedFunctionTypeKind ?: FunctionTypeKind.Function,
+            functionTypeKind,
             if (resolvedArgument.coerceFirstParameterToExtensionReceiver) parameters.drop(1) else parameters,
             resolvedArgument.receiver,
             resolvedArgument.returnType,

plugins/fir-plugin-prototype/testData/diagnostics/functionalTypes/inference.fir.txt

@@ -0,0 +1,38 @@
+FILE: inference.kt
+    public final fun <T : R|suspend () -> kotlin/Unit|> suspendId(f: R|T|): R|T| {
+        ^suspendId R|<local>/f|
+    }
+    public final fun <T : R|@R|org/jetbrains/kotlin/fir/plugin/MyComposable|()  some/MyComposableFunction0<kotlin/Unit>|> composableId(f: R|T|): R|T| {
+        ^composableId R|<local>/f|
+    }
+    public final fun test_1(): R|kotlin/Unit| {
+        R|/id|<R|some/MyComposableFunction0<kotlin/Unit>|>(@R|org/jetbrains/kotlin/fir/plugin/MyComposable|() id@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
+            ^@id Unit
+        }
+        )
+        lval f: R|() -> kotlin/Unit| = R|/id<CS errors: /id>#|<R|() -> kotlin/Unit|>(@R|org/jetbrains/kotlin/fir/plugin/MyComposable|() id@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
+            ^@id Unit
+        }
+        )
+        R|/composableId|<R|@R|org/jetbrains/kotlin/fir/plugin/MyComposable|()  some/MyComposableFunction0<kotlin/Unit>|>(composableId@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
+            ^@composableId Unit
+        }
+        )
+        R|/composableId|<R|@R|org/jetbrains/kotlin/fir/plugin/MyComposable|()  some/MyComposableFunction0<kotlin/Unit>|>(@R|org/jetbrains/kotlin/fir/plugin/MyComposable|() composableId@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
+            ^@composableId Unit
+        }
+        )
+        R|/suspendId<CS errors: /suspendId>#|<R|suspend () -> kotlin/Unit|>(@R|org/jetbrains/kotlin/fir/plugin/MyComposable|() suspendId@fun <anonymous>(): R|kotlin/Unit| <inline=NoInline>  {
+            ^@suspendId Unit
+        }
+        )
+    }
+    public final fun test_2(): R|kotlin/Unit| {
+        R|/select|<R|some/MyComposableFunction0<kotlin/String>|>(vararg(@R|org/jetbrains/kotlin/fir/plugin/MyComposable|() select@fun <anonymous>(): R|kotlin/String| <inline=Unknown>  {
+            ^ String(a)
+        }
+        , select@fun <anonymous>(): R|kotlin/String| <inline=Unknown>  {
+            ^ String(b)
+        }
+        ))
+    }

plugins/fir-plugin-prototype/testData/diagnostics/functionalTypes/inference.kt

@@ -0,0 +1,17 @@
+// INFERENCE_HELPERS
+import org.jetbrains.kotlin.fir.plugin.MyComposable
+
+fun <T : suspend () -> Unit> suspendId(f: T): T = f
+fun <T : @MyComposable () -> Unit> composableId(f: T): T = f
+
+fun test_1() {
+    <!DEBUG_INFO_EXPRESSION_TYPE("some.MyComposableFunction0<kotlin.Unit>")!>id(@MyComposable {})<!>
+    val f: () -> Unit = id(@MyComposable <!ARGUMENT_TYPE_MISMATCH!>{}<!>) // should be an error
+    <!DEBUG_INFO_EXPRESSION_TYPE("some.MyComposableFunction0<kotlin.Unit>")!>composableId({})<!>
+    <!DEBUG_INFO_EXPRESSION_TYPE("some.MyComposableFunction0<kotlin.Unit>")!>composableId(@MyComposable {})<!>
+    <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.coroutines.SuspendFunction0<kotlin.Unit>")!>suspendId(@MyComposable <!ARGUMENT_TYPE_MISMATCH!>{}<!>)<!>
+}
+
+fun test_2() {
+    <!DEBUG_INFO_EXPRESSION_TYPE("some.MyComposableFunction0<kotlin.String>")!>select(@MyComposable { "a" }, { "b" })<!>
+}

plugins/fir-plugin-prototype/tests-gen/org/jetbrains/kotlin/fir/plugin/runners/FirPluginDiagnosticTestGenerated.java

@@ -61,6 +61,12 @@ public class FirPluginDiagnosticTestGenerated extends AbstractFirPluginDiagnosti
             runTest("plugins/fir-plugin-prototype/testData/diagnostics/functionalTypes/ambigousKinds.kt");
         }
 
+        @Test
+        @TestMetadata("inference.kt")
+        public void testInference() throws Exception {
+            runTest("plugins/fir-plugin-prototype/testData/diagnostics/functionalTypes/inference.kt");
+        }
+
         @Test
         @TestMetadata("simple.kt")
         public void testSimple() throws Exception {