[FIR] Fix missing RETURN_TYPE_MISMATCH for labeled return statements

in lambda functions, ^KT-59907 Fixed

compiler/fir/analysis-tests/testData/resolveWithStdlib/diagnostics/functionReturnTypeMismatchChecker/basic.kt

@@ -37,7 +37,7 @@ fun errorWithLambda(): String {
         return@foo
     } foo {
         bar()
-        return@foo <!ARGUMENT_TYPE_MISMATCH!>10<!>
+        return@foo <!ARGUMENT_TYPE_MISMATCH, RETURN_TYPE_MISMATCH!>10<!>
     }
 
     return ""

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirFunctionReturnTypeMismatchChecker.kt

@@ -9,20 +9,19 @@ import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.KtRealSourceElementKind
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.diagnostics.reportOn
+import org.jetbrains.kotlin.fir.analysis.cfa.util.previousCfgNodes
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.checkers.declaration.hasExplicitReturnType
 import org.jetbrains.kotlin.fir.analysis.checkers.isSubtypeForTypeMismatch
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.NULL_FOR_NONNULL_TYPE
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.RETURN_TYPE_MISMATCH
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors.SMARTCAST_IMPOSSIBLE
-import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
-import org.jetbrains.kotlin.fir.declarations.FirConstructor
-import org.jetbrains.kotlin.fir.declarations.FirErrorFunction
-import org.jetbrains.kotlin.fir.declarations.FirPropertyAccessor
+import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
 import org.jetbrains.kotlin.fir.expressions.FirSmartCastExpression
 import org.jetbrains.kotlin.fir.expressions.FirWhenExpression
 import org.jetbrains.kotlin.fir.expressions.isExhaustive
+import org.jetbrains.kotlin.fir.resolve.dfa.controlFlowGraph
 import org.jetbrains.kotlin.fir.types.*
 
 object FirFunctionReturnTypeMismatchChecker : FirReturnExpressionChecker() {
@@ -31,7 +30,7 @@ object FirFunctionReturnTypeMismatchChecker : FirReturnExpressionChecker() {
         if (expression.source?.kind == KtFakeSourceElementKind.DelegatedPropertyAccessor) return
 
         val targetElement = expression.target.labeledElement
-        if (targetElement is FirErrorFunction || targetElement is FirAnonymousFunction && targetElement.isLambda) {
+        if (targetElement is FirErrorFunction) {
             return
         }
 
@@ -91,7 +90,21 @@ object FirFunctionReturnTypeMismatchChecker : FirReturnExpressionChecker() {
                     )
                 }
             }
-        } else if (resultExpression.source?.kind is KtFakeSourceElementKind.ImplicitUnit && !functionReturnType.isUnit) {
+        } else if (resultExpression.source?.kind is KtFakeSourceElementKind.ImplicitUnit &&
+            !functionReturnType.isUnit &&
+            shouldCheckMismatchForAnonymousFunction(targetElement, expression)
+        ) {
+            // Disallow cases like
+            //     fun foo(): Any { return }
+            // Allow cases like
+            //     fun foo(): Unit { return }
+            //     fun foo() { return Unit }
+            // But ignore anonymous functions without explicit returns, the following code is valid (where type of parameter R is being inferenced):
+            //     run {
+            //        if (flag) return@run
+            //        "str"
+            //     }
+            // If type parameter is specified explicitly, checking is performed in the branch above and RETURN_TYPE_MISMATCH is reported
             reporter.reportOn(
                 resultExpression.source,
                 RETURN_TYPE_MISMATCH,
@@ -103,4 +116,12 @@ object FirFunctionReturnTypeMismatchChecker : FirReturnExpressionChecker() {
             )
         }
     }
+
+    private fun shouldCheckMismatchForAnonymousFunction(targetElement: FirFunction, expression: FirReturnExpression): Boolean {
+        if (targetElement !is FirAnonymousFunction || !targetElement.isLambda) return true
+        val cfgNodes = targetElement.controlFlowGraphReference?.controlFlowGraph?.exitNode?.previousCfgNodes ?: return true
+        // Check if any return expression other than the current is explicit
+        return cfgNodes.any { if (it.fir === expression) false else it.fir is FirReturnExpression }
+    }
 }
+

compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnNoCoercionToUnit.fir.kt

@@ -1,6 +1,6 @@
 fun test(a: Int) {
     run f@{
-      if (a > 0) return@f
+      if (a > 0) <!RETURN_TYPE_MISMATCH!>return@f<!>
       else return@f 1
     }
 }

compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnSecondUnit.fir.kt

@@ -1,6 +0,0 @@
-val flag = true
-
-val a = b@ {
-    if (flag) return@b 4
-    return@b
-}

compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnSecondUnit.kt

@@ -1,6 +1,7 @@
+// FIR_IDENTICAL
 val flag = true
 
 val a = b@ {
     if (flag) return@b <!RETURN_TYPE_MISMATCH!>4<!>
     return@b
-}
+}

compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnUnitAndDontCareType.fir.kt

@@ -14,7 +14,7 @@ val b = run {
 
 // Unit
 val c = run {
-    if (flag) return@run
+    if (flag) <!RETURN_TYPE_MISMATCH!>return@run<!>
 
     return@run 4
-}
+}

compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnsWithExplicitReturnType.fir.kt

@@ -1,6 +1,6 @@
 fun test(a: Int) {
     run<Int>f@{
-      if (a > 0) return@f <!ARGUMENT_TYPE_MISMATCH!>""<!>
+      if (a > 0) return@f <!ARGUMENT_TYPE_MISMATCH, RETURN_TYPE_MISMATCH!>""<!>
       return@f 1
     }
 

compiler/testData/diagnostics/tests/generics/projectionsScope/typeMismatchInLambda.fir.kt

@@ -20,7 +20,7 @@ fun test(a: A<out CharSequence>, z: Out<CharSequence>) {
     a.bar { <!ARGUMENT_TYPE_MISMATCH, TYPE_MISMATCH!>z.id()<!> }
 
     a.foo {
-        z.foobar(if (1 > 2) return@foo <!ARGUMENT_TYPE_MISMATCH!>""<!> else "")
+        z.foobar(if (1 > 2) return@foo <!ARGUMENT_TYPE_MISMATCH, RETURN_TYPE_MISMATCH!>""<!> else "")
         <!ARGUMENT_TYPE_MISMATCH!>""<!>
     }
 }

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

@@ -32,9 +32,9 @@ fun foo() {
     // 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 {
+            else -> return@myRun <!RETURN_TYPE_MISMATCH!>buildList {
                 add(4)
-            }
+            }<!>
         }
     }<!>
 }