K2: Fix false-positive TYPE_MISMATCH for suspend lambdas

Why did the problem existed?
At first, lambdas were analyzed with suspend function expected type,
because it's the WHEN-case and we propagate expected type info to
the branches.

Then, after the lambdas was introduced to the containing inference
system, we're creating ResolvedLambdaAtom using the information
from analyzed lambda's shape, but didn't use known lambda resulting
type (from which we might infer FunctionTypeKind).

So, the fix is just using that already obtained information.

^KT-57446 Fixed

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosticCompilerTestFE10TestdataTestGenerated.java

@@ -31927,6 +31927,12 @@ public class DiagnosticCompilerTestFE10TestdataTestGenerated extends AbstractDia
             public void testSuspendConversionWithReferenceAdaptation() throws Exception {
                 runTest("compiler/testData/diagnostics/tests/suspendConversion/suspendConversionWithReferenceAdaptation.kt");
             }
+
+            @Test
+            @TestMetadata("suspendFunctionExpectedTypeAndWhen.kt")
+            public void testSuspendFunctionExpectedTypeAndWhen() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/suspendConversion/suspendFunctionExpectedTypeAndWhen.kt");
+            }
         }
 
         @Nested

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated.java

@@ -31927,6 +31927,12 @@ public class LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated
             public void testSuspendConversionWithReferenceAdaptation() throws Exception {
                 runTest("compiler/testData/diagnostics/tests/suspendConversion/suspendConversionWithReferenceAdaptation.kt");
             }
+
+            @Test
+            @TestMetadata("suspendFunctionExpectedTypeAndWhen.kt")
+            public void testSuspendFunctionExpectedTypeAndWhen() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/suspendConversion/suspendFunctionExpectedTypeAndWhen.kt");
+            }
         }
 
         @Nested

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirLightTreeOldFrontendDiagnosticsTestGenerated.java

@@ -31927,6 +31927,12 @@ public class FirLightTreeOldFrontendDiagnosticsTestGenerated extends AbstractFir
             public void testSuspendConversionWithReferenceAdaptation() throws Exception {
                 runTest("compiler/testData/diagnostics/tests/suspendConversion/suspendConversionWithReferenceAdaptation.kt");
             }
+
+            @Test
+            @TestMetadata("suspendFunctionExpectedTypeAndWhen.kt")
+            public void testSuspendFunctionExpectedTypeAndWhen() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/suspendConversion/suspendFunctionExpectedTypeAndWhen.kt");
+            }
         }
 
         @Nested

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirPsiOldFrontendDiagnosticsTestGenerated.java

@@ -32023,6 +32023,12 @@ public class FirPsiOldFrontendDiagnosticsTestGenerated extends AbstractFirPsiDia
             public void testSuspendConversionWithReferenceAdaptation() throws Exception {
                 runTest("compiler/testData/diagnostics/tests/suspendConversion/suspendConversionWithReferenceAdaptation.kt");
             }
+
+            @Test
+            @TestMetadata("suspendFunctionExpectedTypeAndWhen.kt")
+            public void testSuspendFunctionExpectedTypeAndWhen() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/suspendConversion/suspendFunctionExpectedTypeAndWhen.kt");
+            }
         }
 
         @Nested

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

@@ -134,7 +134,7 @@ private fun extractLambdaInfo(
     return ResolvedLambdaAtom(
         argument,
         expectedType,
-        expectedFunctionTypeKind = null,
+        expectedFunctionTypeKind = argument.typeRef.coneTypeSafe<ConeKotlinType>()?.lowerBoundIfFlexible()?.functionTypeKind(session),
         receiverType,
         contextReceivers,
         parameters,

compiler/testData/diagnostics/tests/suspendConversion/suspendFunctionExpectedTypeAndWhen.kt

@@ -0,0 +1,14 @@
+// FIR_IDENTICAL
+// ISSUE: KT-57446
+
+fun test1() : suspend (Int) -> Unit = when {
+    true -> { _ -> }
+    else -> { _ -> }
+}
+
+fun test2() : suspend (Int) -> Unit = when {
+    true -> { x -> foo(x) }
+    else -> { y -> foo(y) }
+}
+
+suspend fun foo(x: Int) {}

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java

@@ -32771,6 +32771,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
             public void testSuspendConversionWithReferenceAdaptation() throws Exception {
                 runTest("compiler/testData/diagnostics/tests/suspendConversion/suspendConversionWithReferenceAdaptation.kt");
             }
+
+            @Test
+            @TestMetadata("suspendFunctionExpectedTypeAndWhen.kt")
+            public void testSuspendFunctionExpectedTypeAndWhen() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/suspendConversion/suspendFunctionExpectedTypeAndWhen.kt");
+            }
         }
 
         @Nested