[FIR] Infer labeled type for lambda in return in catch ^KT-51009 Fixed

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

@@ -14809,6 +14809,12 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                     runTest("compiler/testData/diagnostics/tests/inference/completion/flexibleType.kt");
                 }
 
+                @Test
+                @TestMetadata("inferLabeledElementTypeForLambdaInReturnInCatch.kt")
+                public void testInferLabeledElementTypeForLambdaInReturnInCatch() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/completion/inferLabeledElementTypeForLambdaInReturnInCatch.kt");
+                }
+
                 @Test
                 @TestMetadata("intersectionType.kt")
                 public void testIntersectionType() throws Exception {

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

@@ -14809,6 +14809,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                     runTest("compiler/testData/diagnostics/tests/inference/completion/flexibleType.kt");
                 }
 
+                @Test
+                @TestMetadata("inferLabeledElementTypeForLambdaInReturnInCatch.kt")
+                public void testInferLabeledElementTypeForLambdaInReturnInCatch() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/completion/inferLabeledElementTypeForLambdaInReturnInCatch.kt");
+                }
+
                 @Test
                 @TestMetadata("intersectionType.kt")
                 public void testIntersectionType() throws Exception {

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

@@ -14809,6 +14809,12 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                     runTest("compiler/testData/diagnostics/tests/inference/completion/flexibleType.kt");
                 }
 
+                @Test
+                @TestMetadata("inferLabeledElementTypeForLambdaInReturnInCatch.kt")
+                public void testInferLabeledElementTypeForLambdaInReturnInCatch() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/completion/inferLabeledElementTypeForLambdaInReturnInCatch.kt");
+                }
+
                 @Test
                 @TestMetadata("intersectionType.kt")
                 public void testIntersectionType() throws Exception {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt

@@ -702,7 +702,8 @@ class FirCallCompletionResultsWriterTransformer(
             return returnExpression
         }
 
-        return super.transformReturnExpression(returnExpression, data)
+        val newData = labeledElement.returnTypeRef.coneTypeSafe<ConeKotlinType>()?.toExpectedType()
+        return super.transformReturnExpression(returnExpression, newData)
     }
 
     override fun transformBlock(block: FirBlock, data: ExpectedArgumentType?): FirStatement {

compiler/testData/diagnostics/tests/inference/completion/inferLabeledElementTypeForLambdaInReturnInCatch.kt

@@ -0,0 +1,20 @@
+// FIR_IDENTICAL
+// ISSUE: KT-51009
+
+fun test(b: Boolean, f: () -> String?): () -> String {
+    val foo = try {
+        f
+    } catch (e: Exception) {
+        return { "1" } // Infer return type of test () -> String instead of type of `f`
+    }
+    return { "2" }
+}
+
+fun test2(b: Boolean, f: () -> String?) = run { // implicit return type
+    val foo = try {
+        f
+    } catch (e: Exception) {
+        return { "1" }
+    }
+    { "2" }
+}

compiler/testData/diagnostics/tests/inference/completion/inferLabeledElementTypeForLambdaInReturnInCatch.txt

@@ -0,0 +1,5 @@
+package
+
+public fun test(/*0*/ b: kotlin.Boolean, /*1*/ f: () -> kotlin.String?): () -> kotlin.String
+public fun test2(/*0*/ b: kotlin.Boolean, /*1*/ f: () -> kotlin.String?): () -> kotlin.String
+

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

@@ -14815,6 +14815,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                     runTest("compiler/testData/diagnostics/tests/inference/completion/flexibleType.kt");
                 }
 
+                @Test
+                @TestMetadata("inferLabeledElementTypeForLambdaInReturnInCatch.kt")
+                public void testInferLabeledElementTypeForLambdaInReturnInCatch() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/completion/inferLabeledElementTypeForLambdaInReturnInCatch.kt");
+                }
+
                 @Test
                 @TestMetadata("intersectionType.kt")
                 public void testIntersectionType() throws Exception {