[FIR] Infer Unit type for generic lambda function if implicit return exists

Make behavior more consistent with K1

^KT-63563 Fixed

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

@@ -14357,6 +14357,12 @@ public class DiagnosticCompilerTestFE10TestdataTestGenerated extends AbstractDia
                     runTest("compiler/testData/diagnostics/tests/functionLiterals/return/IfWithoutElseWithExplicitType.kt");
                 }
 
+                @Test
+                @TestMetadata("LambdaReturnTypeIsUnitIfImplicitReturnExists.kt")
+                public void testLambdaReturnTypeIsUnitIfImplicitReturnExists() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/functionLiterals/return/LambdaReturnTypeIsUnitIfImplicitReturnExists.kt");
+                }
+
                 @Test
                 @TestMetadata("LambdaWithParameter.kt")
                 public void testLambdaWithParameter() throws Exception {

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

@@ -14357,6 +14357,12 @@ public class LLFirPreresolvedReversedDiagnosticCompilerFE10TestDataTestGenerated
                     runTest("compiler/testData/diagnostics/tests/functionLiterals/return/IfWithoutElseWithExplicitType.kt");
                 }
 
+                @Test
+                @TestMetadata("LambdaReturnTypeIsUnitIfImplicitReturnExists.kt")
+                public void testLambdaReturnTypeIsUnitIfImplicitReturnExists() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/functionLiterals/return/LambdaReturnTypeIsUnitIfImplicitReturnExists.kt");
+                }
+
                 @Test
                 @TestMetadata("LambdaWithParameter.kt")
                 public void testLambdaWithParameter() throws Exception {

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

@@ -14351,6 +14351,12 @@ public class FirLightTreeOldFrontendDiagnosticsTestGenerated extends AbstractFir
                     runTest("compiler/testData/diagnostics/tests/functionLiterals/return/IfWithoutElseWithExplicitType.kt");
                 }
 
+                @Test
+                @TestMetadata("LambdaReturnTypeIsUnitIfImplicitReturnExists.kt")
+                public void testLambdaReturnTypeIsUnitIfImplicitReturnExists() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/functionLiterals/return/LambdaReturnTypeIsUnitIfImplicitReturnExists.kt");
+                }
+
                 @Test
                 @TestMetadata("LambdaWithParameter.kt")
                 public void testLambdaWithParameter() throws Exception {

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

@@ -14357,6 +14357,12 @@ public class FirPsiOldFrontendDiagnosticsTestGenerated extends AbstractFirPsiDia
                     runTest("compiler/testData/diagnostics/tests/functionLiterals/return/IfWithoutElseWithExplicitType.kt");
                 }
 
+                @Test
+                @TestMetadata("LambdaReturnTypeIsUnitIfImplicitReturnExists.kt")
+                public void testLambdaReturnTypeIsUnitIfImplicitReturnExists() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/functionLiterals/return/LambdaReturnTypeIsUnitIfImplicitReturnExists.kt");
+                }
+
                 @Test
                 @TestMetadata("LambdaWithParameter.kt")
                 public void testLambdaWithParameter() throws Exception {

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

@@ -9,7 +9,6 @@ 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
@@ -21,7 +20,6 @@ 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() {
@@ -91,19 +89,13 @@ object FirFunctionReturnTypeMismatchChecker : FirReturnExpressionChecker() {
                 }
             }
         } else if (resultExpression.source?.kind is KtFakeSourceElementKind.ImplicitUnit &&
-            !functionReturnType.lowerBoundIfFlexible().isUnit &&
+            !functionReturnType.lowerBoundIfFlexible().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,
@@ -116,15 +108,5 @@ 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 and not Unit
-        return cfgNodes.any {
-            val fir = it.fir
-            if (fir === expression) false else fir is FirReturnExpression && !fir.result.resolvedType.isUnit
-        }
-    }
 }
 

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

@@ -15,7 +15,6 @@ import org.jetbrains.kotlin.fir.diagnostics.ConeSimpleDiagnostic
 import org.jetbrains.kotlin.fir.diagnostics.DiagnosticKind
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.builder.buildSamConversionExpression
-import org.jetbrains.kotlin.fir.expressions.impl.FirPropertyAccessExpressionImpl
 import org.jetbrains.kotlin.fir.references.FirNamedReference
 import org.jetbrains.kotlin.fir.references.builder.buildResolvedCallableReference
 import org.jetbrains.kotlin.fir.references.builder.buildResolvedNamedReference
@@ -43,7 +42,6 @@ import org.jetbrains.kotlin.fir.visitors.FirDefaultTransformer
 import org.jetbrains.kotlin.fir.visitors.FirTransformer
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.resolve.calls.inference.model.InferredEmptyIntersection
-import org.jetbrains.kotlin.resolve.calls.tower.CandidateApplicability
 import org.jetbrains.kotlin.resolve.calls.tower.isSuccess
 import org.jetbrains.kotlin.types.TypeApproximatorConfiguration
 import org.jetbrains.kotlin.types.Variance
@@ -707,6 +705,8 @@ class FirCallCompletionResultsWriterTransformer(
 
         val initialReturnType = anonymousFunction.returnTypeRef.coneTypeSafe<ConeKotlinType>()
         val expectedReturnType = initialReturnType?.let { finallySubstituteOrSelf(it) }
+            ?: runIf(returnExpressions.any { it.source?.kind is KtFakeSourceElementKind.ImplicitUnit.Return })
+            { session.builtinTypes.unitType.coneType }
             ?: expectedType?.returnType(session) as? ConeClassLikeType
             ?: (data as? ExpectedArgumentType.ArgumentsMap)?.lambdasReturnTypes?.get(anonymousFunction)
 

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

@@ -0,0 +1,22 @@
+// FIR_IDENTICAL
+// ISSUE: KT-63563
+
+fun foo(x: () -> Any) = x()
+fun foo2(x: () -> Unit) = x()
+fun <T> foo3(x: () -> T): T = x()
+
+fun main2() {
+    foo {
+        if ("0".hashCode() == 42) <!RETURN_TYPE_MISMATCH!>return@foo<!>
+        ""
+    }
+    foo2 {
+        if ("1".hashCode() == 42) return@foo2
+        ""
+    }
+    foo3 { // Infer T to Unit -> no error
+        if ("2".hashCode() == 42) return@foo3
+        ""
+    }
+}
+

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

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

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

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 fun test(a: Int) {
     run f@{
       if (a > 0) return@f

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

@@ -1,20 +0,0 @@
-// NI_EXPECTED_FILE
-val flag = true
-
-// type of a was checked by txt
-val a = run { // () -> Unit
-    return@run
-}
-
-// Unit
-val b = run {
-    if (flag) return@run
-    5
-}
-
-// Unit
-val c = run {
-    if (flag) <!RETURN_TYPE_MISMATCH!>return@run<!>
-
-    return@run 4
-}

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

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // NI_EXPECTED_FILE
 val flag = true
 

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

@@ -14357,6 +14357,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                     runTest("compiler/testData/diagnostics/tests/functionLiterals/return/IfWithoutElseWithExplicitType.kt");
                 }
 
+                @Test
+                @TestMetadata("LambdaReturnTypeIsUnitIfImplicitReturnExists.kt")
+                public void testLambdaReturnTypeIsUnitIfImplicitReturnExists() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/functionLiterals/return/LambdaReturnTypeIsUnitIfImplicitReturnExists.kt");
+                }
+
                 @Test
                 @TestMetadata("LambdaWithParameter.kt")
                 public void testLambdaWithParameter() throws Exception {