From 791a7b1c310f9e78eb89e67a3b8a801b118a9b5f Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Thu, 2 Feb 2023 15:04:03 +0200
Subject: [PATCH] [FIR] Extract function type kind for lambda in argument
 position

^KT-56381 Fixed
---
 .../resolve/inference/PostponedArguments.kt   |  5 ++-
 .../functionalTypes/inference.fir.txt         | 38 +++++++++++++++++++
 .../diagnostics/functionalTypes/inference.kt  | 17 +++++++++
 .../FirPluginDiagnosticTestGenerated.java     |  6 +++
 4 files changed, 65 insertions(+), 1 deletion(-)
 create mode 100644 plugins/fir-plugin-prototype/testData/diagnostics/functionalTypes/inference.fir.txt
 create mode 100644 plugins/fir-plugin-prototype/testData/diagnostics/functionalTypes/inference.kt

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArguments.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArguments.kt
index 213644db40d..ad5b44f1fe5 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/inference/PostponedArguments.kt
+++ b/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,
diff --git a/plugins/fir-plugin-prototype/testData/diagnostics/functionalTypes/inference.fir.txt b/plugins/fir-plugin-prototype/testData/diagnostics/functionalTypes/inference.fir.txt
new file mode 100644
index 00000000000..8b08f8f32cf
--- /dev/null
+++ b/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)
+        }
+        ))
+    }
diff --git a/plugins/fir-plugin-prototype/testData/diagnostics/functionalTypes/inference.kt b/plugins/fir-plugin-prototype/testData/diagnostics/functionalTypes/inference.kt
new file mode 100644
index 00000000000..63621ef7678
--- /dev/null
+++ 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" })<!>
+}
diff --git a/plugins/fir-plugin-prototype/tests-gen/org/jetbrains/kotlin/fir/plugin/runners/FirPluginDiagnosticTestGenerated.java b/plugins/fir-plugin-prototype/tests-gen/org/jetbrains/kotlin/fir/plugin/runners/FirPluginDiagnosticTestGenerated.java
index 96ed142b53b..79495e9da36 100644
--- a/plugins/fir-plugin-prototype/tests-gen/org/jetbrains/kotlin/fir/plugin/runners/FirPluginDiagnosticTestGenerated.java
+++ 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 {