K2: Fix undesirable NEW_INFERENCE_ERROR for case of return + buildList

While plainly repeating K1 behavior might be a dumb solution,
but inventing another one might be quite complicated
because we need to stop fixing variables into Nothing in many cases,
but probably not in all of them (see KT-58232).

^KT-58149 Fixed
^KT-58232 Related

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

@@ -16935,6 +16935,12 @@ public class DiagnosticCompilerTestFE10TestdataTestGenerated extends AbstractDia
                     runTest("compiler/testData/diagnostics/tests/inference/nothingType/reportImplicitNothingOnlyForOwnTypeParameters.kt");
                 }
 
+                @Test
+                @TestMetadata("returnAsLastStatementInLambda.kt")
+                public void testReturnAsLastStatementInLambda() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/nothingType/returnAsLastStatementInLambda.kt");
+                }
+
                 @Test
                 @TestMetadata("selectWithNull.kt")
                 public void testSelectWithNull() throws Exception {

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

@@ -16935,6 +16935,12 @@ public class LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated
                     runTest("compiler/testData/diagnostics/tests/inference/nothingType/reportImplicitNothingOnlyForOwnTypeParameters.kt");
                 }
 
+                @Test
+                @TestMetadata("returnAsLastStatementInLambda.kt")
+                public void testReturnAsLastStatementInLambda() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/nothingType/returnAsLastStatementInLambda.kt");
+                }
+
                 @Test
                 @TestMetadata("selectWithNull.kt")
                 public void testSelectWithNull() throws Exception {

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

@@ -16935,6 +16935,12 @@ public class FirLightTreeOldFrontendDiagnosticsTestGenerated extends AbstractFir
                     runTest("compiler/testData/diagnostics/tests/inference/nothingType/reportImplicitNothingOnlyForOwnTypeParameters.kt");
                 }
 
+                @Test
+                @TestMetadata("returnAsLastStatementInLambda.kt")
+                public void testReturnAsLastStatementInLambda() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/nothingType/returnAsLastStatementInLambda.kt");
+                }
+
                 @Test
                 @TestMetadata("selectWithNull.kt")
                 public void testSelectWithNull() throws Exception {

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

@@ -16941,6 +16941,12 @@ public class FirPsiOldFrontendDiagnosticsTestGenerated extends AbstractFirPsiDia
                     runTest("compiler/testData/diagnostics/tests/inference/nothingType/reportImplicitNothingOnlyForOwnTypeParameters.kt");
                 }
 
+                @Test
+                @TestMetadata("returnAsLastStatementInLambda.kt")
+                public void testReturnAsLastStatementInLambda() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/nothingType/returnAsLastStatementInLambda.kt");
+                }
+
                 @Test
                 @TestMetadata("selectWithNull.kt")
                 public void testSelectWithNull() throws Exception {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt

@@ -5,8 +5,8 @@
 
 package org.jetbrains.kotlin.fir.resolve.dfa.cfg
 
+import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.contracts.description.*
-import org.jetbrains.kotlin.contracts.description.LogicOperationKind
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.declarations.utils.hasExplicitBackingField
@@ -89,8 +89,27 @@ class ControlFlowGraphBuilder {
             is BlockExitNode -> when {
                 // lambda@{ x } -> x
                 // lambda@{ class C } -> Unit-returning stub
-                function.isLambda -> fir.statements.lastOrNull() as? FirExpression
-                    ?: buildUnitExpression { source = fir.statements.lastOrNull()?.source ?: fir.source }
+                function.isLambda -> {
+                    val lastStatement = fir.statements.lastOrNull()
+
+                    when {
+                        // Skip last return statement because otherwise it add Nothing constraint on the lambda return type.
+                        // That might lead to preliminary variable fixation to Nothing that is kind of undesirable in most cases.
+                        // Note, that the expression that is going to be returned would be used as another element of `returnValues`.
+                        //
+                        // While that remains a bit questionable why adding such trivial Nothing constraint makes variable being fixed there
+                        // but currently it doesn't look like we've got easy answers to those questions, so we just repeat K1 behavior
+                        // (see `val lastExpressionArgument` at KotlinResolutionCallbacksImpl.analyzeAndGetLambdaReturnArguments)
+                        // Probably, that might be removed once KT-58232 is fixed
+                        lastStatement is FirReturnExpression &&
+                                lastStatement.target.labeledElement.symbol == function.symbol &&
+                                lastStatement.source?.kind != KtFakeSourceElementKind.ImplicitReturn.FromLastStatement ->
+                            null
+                        else ->
+                            lastStatement as? FirExpression
+                                ?: buildUnitExpression { source = fir.statements.lastOrNull()?.source ?: fir.source }
+                    }
+                }
                 // fun() { terminatingExpression } -> nothing (checker will emit an error if return type is not Unit)
                 // fun() { throw } or fun() { returnsNothing() } -> Nothing-returning stub
                 else -> FirStub.takeIf { _ -> previousNodes.all { it is StubNode } }

compiler/testData/diagnostics/tests/inference/nothingType/returnAsLastStatementInLambda.fir.kt

@@ -0,0 +1,40 @@
+// WITH_STDLIB
+// ISSUE: KT-58149
+
+fun <U> myRun(calc: () -> U): U {
+    return calc()
+}
+
+fun foo() {
+    // OK, works because we explicitly avoid adding last "return" as lambda result
+    val dates1 = myRun {
+        return@myRun buildList {
+            add(1)
+        }
+    }
+
+    <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.collections.List<kotlin.Int>")!>dates1<!>
+
+    // OK, because the type of when is inferred to List<*>
+    val dates2 = myRun {
+        when {
+            true -> return@myRun buildList {
+                add(2)
+            }
+            else -> return@myRun buildList {
+                add(3)
+            }
+        }
+    }
+
+    <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.collections.List<*>")!>dates2<!>
+
+    // Doesn't work both in K1 and K2, but probably should (KT-58232 for tracking)
+    val dates3 = <!NEW_INFERENCE_ERROR!>myRun {
+        when {
+            else -> return@myRun buildList {
+                add(4)
+            }
+        }
+    }<!>
+}

compiler/testData/diagnostics/tests/inference/nothingType/returnAsLastStatementInLambda.kt

@@ -0,0 +1,40 @@
+// WITH_STDLIB
+// ISSUE: KT-58149
+
+fun <U> myRun(calc: () -> U): U {
+    return calc()
+}
+
+fun foo() {
+    // OK, works because we explicitly avoid adding last "return" as lambda result
+    val dates1 = myRun {
+        return@myRun buildList {
+            add(1)
+        }
+    }
+
+    <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.collections.List<kotlin.Int>")!>dates1<!>
+
+    // OK, because the type of when is inferred to List<*>
+    val dates2 = myRun {
+        when {
+            true -> return@myRun buildList {
+                add(2)
+            }
+            else -> return@myRun buildList {
+                add(3)
+            }
+        }
+    }
+
+    <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.collections.List<*>")!>dates2<!>
+
+    // Doesn't work both in K1 and K2, but probably should (KT-58232 for tracking)
+    <!UNREACHABLE_CODE!>val dates3 =<!> <!TYPE_MISMATCH!>myRun<!> {
+        when {
+            else -> return@myRun <!TYPE_MISMATCH, TYPE_MISMATCH!>buildList {
+                add(4)
+            }<!>
+        }
+    }
+}

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

@@ -16941,6 +16941,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                     runTest("compiler/testData/diagnostics/tests/inference/nothingType/reportImplicitNothingOnlyForOwnTypeParameters.kt");
                 }
 
+                @Test
+                @TestMetadata("returnAsLastStatementInLambda.kt")
+                public void testReturnAsLastStatementInLambda() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/nothingType/returnAsLastStatementInLambda.kt");
+                }
+
                 @Test
                 @TestMetadata("selectWithNull.kt")
                 public void testSelectWithNull() throws Exception {