[FIR] Fix local RETURN_TYPE_MISMATCH with flexible Unit and multiple Unit returns

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

@@ -13433,6 +13433,12 @@ public class DiagnosticCompilerTestFE10TestdataTestGenerated extends AbstractDia
                     runTest("compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnWithExplicitUnit.kt");
                 }
 
+                @Test
+                @TestMetadata("LocalReturnWithFlexibleUnitType.kt")
+                public void testLocalReturnWithFlexibleUnitType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnWithFlexibleUnitType.kt");
+                }
+
                 @Test
                 @TestMetadata("LocalReturnsWithExplicitReturnType.kt")
                 public void testLocalReturnsWithExplicitReturnType() throws Exception {
@@ -13445,6 +13451,12 @@ public class DiagnosticCompilerTestFE10TestdataTestGenerated extends AbstractDia
                     runTest("compiler/testData/diagnostics/tests/functionLiterals/return/MixedReturnsFromLambda.kt");
                 }
 
+                @Test
+                @TestMetadata("MultipleLocalUnitReturns.kt")
+                public void testMultipleLocalUnitReturns() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/functionLiterals/return/MultipleLocalUnitReturns.kt");
+                }
+
                 @Test
                 @TestMetadata("NoCommonSystem.kt")
                 public void testNoCommonSystem() throws Exception {

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

@@ -13433,6 +13433,12 @@ public class LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated
                     runTest("compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnWithExplicitUnit.kt");
                 }
 
+                @Test
+                @TestMetadata("LocalReturnWithFlexibleUnitType.kt")
+                public void testLocalReturnWithFlexibleUnitType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnWithFlexibleUnitType.kt");
+                }
+
                 @Test
                 @TestMetadata("LocalReturnsWithExplicitReturnType.kt")
                 public void testLocalReturnsWithExplicitReturnType() throws Exception {
@@ -13445,6 +13451,12 @@ public class LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated
                     runTest("compiler/testData/diagnostics/tests/functionLiterals/return/MixedReturnsFromLambda.kt");
                 }
 
+                @Test
+                @TestMetadata("MultipleLocalUnitReturns.kt")
+                public void testMultipleLocalUnitReturns() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/functionLiterals/return/MultipleLocalUnitReturns.kt");
+                }
+
                 @Test
                 @TestMetadata("NoCommonSystem.kt")
                 public void testNoCommonSystem() throws Exception {

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

@@ -13433,6 +13433,12 @@ public class FirLightTreeOldFrontendDiagnosticsTestGenerated extends AbstractFir
                     runTest("compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnWithExplicitUnit.kt");
                 }
 
+                @Test
+                @TestMetadata("LocalReturnWithFlexibleUnitType.kt")
+                public void testLocalReturnWithFlexibleUnitType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnWithFlexibleUnitType.kt");
+                }
+
                 @Test
                 @TestMetadata("LocalReturnsWithExplicitReturnType.kt")
                 public void testLocalReturnsWithExplicitReturnType() throws Exception {
@@ -13445,6 +13451,12 @@ public class FirLightTreeOldFrontendDiagnosticsTestGenerated extends AbstractFir
                     runTest("compiler/testData/diagnostics/tests/functionLiterals/return/MixedReturnsFromLambda.kt");
                 }
 
+                @Test
+                @TestMetadata("MultipleLocalUnitReturns.kt")
+                public void testMultipleLocalUnitReturns() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/functionLiterals/return/MultipleLocalUnitReturns.kt");
+                }
+
                 @Test
                 @TestMetadata("NoCommonSystem.kt")
                 public void testNoCommonSystem() throws Exception {

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

@@ -13439,6 +13439,12 @@ public class FirPsiOldFrontendDiagnosticsTestGenerated extends AbstractFirPsiDia
                     runTest("compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnWithExplicitUnit.kt");
                 }
 
+                @Test
+                @TestMetadata("LocalReturnWithFlexibleUnitType.kt")
+                public void testLocalReturnWithFlexibleUnitType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnWithFlexibleUnitType.kt");
+                }
+
                 @Test
                 @TestMetadata("LocalReturnsWithExplicitReturnType.kt")
                 public void testLocalReturnsWithExplicitReturnType() throws Exception {
@@ -13451,6 +13457,12 @@ public class FirPsiOldFrontendDiagnosticsTestGenerated extends AbstractFirPsiDia
                     runTest("compiler/testData/diagnostics/tests/functionLiterals/return/MixedReturnsFromLambda.kt");
                 }
 
+                @Test
+                @TestMetadata("MultipleLocalUnitReturns.kt")
+                public void testMultipleLocalUnitReturns() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/functionLiterals/return/MultipleLocalUnitReturns.kt");
+                }
+
                 @Test
                 @TestMetadata("NoCommonSystem.kt")
                 public void testNoCommonSystem() throws Exception {

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

@@ -91,7 +91,7 @@ object FirFunctionReturnTypeMismatchChecker : FirReturnExpressionChecker() {
                 }
             }
         } else if (resultExpression.source?.kind is KtFakeSourceElementKind.ImplicitUnit &&
-            !functionReturnType.isUnit &&
+            !functionReturnType.lowerBoundIfFlexible().isUnit &&
             shouldCheckMismatchForAnonymousFunction(targetElement, expression)
         ) {
             // Disallow cases like
@@ -120,8 +120,11 @@ 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
+        // Check if any return expression other than the current is explicit and not Unit
-        return cfgNodes.any { if (it.fir === expression) false else it.fir is FirReturnExpression }
+        return cfgNodes.any {
+            val fir = it.fir
+            if (fir === expression) false else fir is FirReturnExpression && !fir.result.resolvedType.isUnit
+        }
     }
 }
 

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

@@ -0,0 +1,20 @@
+// FIR_IDENTICAL
+
+// FILE: JavaClass.java
+
+@FunctionalInterface
+public interface JavaClass<T> {
+    public T invoke();
+}
+
+// FILE: main.kt
+
+fun main() {
+    JavaClass {
+        if (true) {
+            return@JavaClass
+        } else {
+            return@JavaClass
+        }
+    }
+}

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

@@ -0,0 +1,11 @@
+// FIR_IDENTICAL
+
+fun <T> execute(block: () -> T): T = block()
+
+fun main() {
+    execute {
+        if (true) return@execute
+        if (true) return@execute Unit
+        execute { Any() }
+    }
+}

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

@@ -13439,6 +13439,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                     runTest("compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnWithExplicitUnit.kt");
                 }
 
+                @Test
+                @TestMetadata("LocalReturnWithFlexibleUnitType.kt")
+                public void testLocalReturnWithFlexibleUnitType() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/functionLiterals/return/LocalReturnWithFlexibleUnitType.kt");
+                }
+
                 @Test
                 @TestMetadata("LocalReturnsWithExplicitReturnType.kt")
                 public void testLocalReturnsWithExplicitReturnType() throws Exception {
@@ -13451,6 +13457,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                     runTest("compiler/testData/diagnostics/tests/functionLiterals/return/MixedReturnsFromLambda.kt");
                 }
 
+                @Test
+                @TestMetadata("MultipleLocalUnitReturns.kt")
+                public void testMultipleLocalUnitReturns() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/functionLiterals/return/MultipleLocalUnitReturns.kt");
+                }
+
                 @Test
                 @TestMetadata("NoCommonSystem.kt")
                 public void testNoCommonSystem() throws Exception {