K2: don't report RETURN_TYPE_MISMATCH on last statement returns

#KT-53198 Fixed

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

@@ -37734,6 +37734,12 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/testsWithStdLib/functionLiterals"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
             }
 
+            @Test
+            @TestMetadata("nullableUnitLiteral.kt")
+            public void testNullableUnitLiteral() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/functionLiterals/nullableUnitLiteral.kt");
+            }
+
             @Test
             @TestMetadata("pseudocodeMemoryOverhead.kt")
             public void testPseudocodeMemoryOverhead() throws Exception {

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

@@ -37734,6 +37734,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/testsWithStdLib/functionLiterals"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
             }
 
+            @Test
+            @TestMetadata("nullableUnitLiteral.kt")
+            public void testNullableUnitLiteral() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/functionLiterals/nullableUnitLiteral.kt");
+            }
+
             @Test
             @TestMetadata("pseudocodeMemoryOverhead.kt")
             public void testPseudocodeMemoryOverhead() throws Exception {

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

@@ -37734,6 +37734,12 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/testsWithStdLib/functionLiterals"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
             }
 
+            @Test
+            @TestMetadata("nullableUnitLiteral.kt")
+            public void testNullableUnitLiteral() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/functionLiterals/nullableUnitLiteral.kt");
+            }
+
             @Test
             @TestMetadata("pseudocodeMemoryOverhead.kt")
             public void testPseudocodeMemoryOverhead() throws Exception {

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.analysis.checkers.expression
 
+import org.jetbrains.kotlin.KtFakeSourceElementKind
 import org.jetbrains.kotlin.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.diagnostics.reportOn
 import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
@@ -36,6 +37,13 @@ object FirFunctionReturnTypeMismatchChecker : FirReturnExpressionChecker() {
             context.session.builtinTypes.unitType.coneType
         else
             targetElement.returnTypeRef.coneType
+        if (targetElement is FirAnonymousFunction &&
+            expression.source?.kind is KtFakeSourceElementKind.ImplicitReturn.FromLastStatement &&
+            functionReturnType.isUnit
+        ) {
+            return
+        }
+
         val typeContext = context.session.typeContext
         val returnExpressionType = resultExpression.typeRef.coneType
 

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

@@ -24,7 +24,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 object FirReturnSyntaxAndLabelChecker : FirReturnExpressionChecker() {
     override fun check(expression: FirReturnExpression, context: CheckerContext, reporter: DiagnosticReporter) {
         val source = expression.source
-        if (source?.kind == KtFakeSourceElementKind.ImplicitReturn) return
+        if (source?.kind is KtFakeSourceElementKind.ImplicitReturn) return
 
         val labeledElement = expression.target.labeledElement
         val targetSymbol = labeledElement.symbol

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/extended/UnreachableCodeChecker.kt

@@ -40,7 +40,8 @@ object UnreachableCodeChecker : FirControlFlowChecker() {
     }
 
     private val sourceKindsToSkip = setOf(
-        KtFakeSourceElementKind.ImplicitReturn,
+        KtFakeSourceElementKind.ImplicitReturn.FromExpressionBody,
+        KtFakeSourceElementKind.ImplicitReturn.FromLastStatement,
         KtFakeSourceElementKind.DesugaredForLoop
     )
 

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -324,7 +324,7 @@ class Fir2IrVisitor(
         return returnExpression.convertWithOffsets { startOffset, endOffset ->
             val result = returnExpression.result
             // For implicit returns, use the expression endOffset to generate the expected line number for debugging.
-            val returnStartOffset = if (returnExpression.source?.kind == KtFakeSourceElementKind.ImplicitReturn) endOffset else startOffset
+            val returnStartOffset = if (returnExpression.source?.kind is KtFakeSourceElementKind.ImplicitReturn) endOffset else startOffset
             IrReturnImpl(
                 returnStartOffset, endOffset, irBuiltIns.nothingType,
                 when (irTarget) {

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt

@@ -195,7 +195,7 @@ class ExpressionsConverter(
                         if (statements.isEmpty()) {
                             statements.add(
                                 buildReturnExpression {
-                                    source = expressionSource.fakeElement(KtFakeSourceElementKind.ImplicitReturn)
+                                    source = expressionSource.fakeElement(KtFakeSourceElementKind.ImplicitReturn.FromExpressionBody)
                                     this.target = target
                                     result = buildUnitExpression {
                                         source = expressionSource.fakeElement(KtFakeSourceElementKind.ImplicitUnit)

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -1435,7 +1435,7 @@ open class RawFirBuilder(
                         if (statements.isEmpty()) {
                             statements.add(
                                 buildReturnExpression {
-                                    source = expressionSource.fakeElement(KtFakeSourceElementKind.ImplicitReturn)
+                                    source = expressionSource.fakeElement(KtFakeSourceElementKind.ImplicitReturn.FromExpressionBody)
                                     this.target = target
                                     result = buildUnitExpression {
                                         source = expressionSource.fakeElement(KtFakeSourceElementKind.ImplicitUnit)

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/BaseFirBuilder.kt

@@ -192,7 +192,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
 
             source =
                 if (fromKtReturnExpression) baseSource?.realElement()
-                else baseSource?.fakeElement(KtFakeSourceElementKind.ImplicitReturn)
+                else baseSource?.fakeElement(KtFakeSourceElementKind.ImplicitReturn.FromExpressionBody)
             result = this@toReturn
             if (labelName == null) {
                 target = context.firFunctionTargets.lastOrNull { !it.isLambda } ?: FirFunctionTarget(labelName, isLambda = false).apply {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -944,7 +944,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                     override fun <E : FirElement> transformElement(element: E, data: FirExpression): E {
                         if (element == lastStatement) {
                             val returnExpression = buildReturnExpression {
-                                source = element.source?.fakeElement(KtFakeSourceElementKind.ImplicitReturn)
+                                source = element.source?.fakeElement(KtFakeSourceElementKind.ImplicitReturn.FromLastStatement)
                                 result = lastStatement
                                 target = FirFunctionTarget(null, isLambda = this@addReturn.isLambda).also {
                                     it.bind(this@addReturn)

compiler/frontend.common/src/org/jetbrains/kotlin/KtSourceElement.kt

@@ -59,7 +59,11 @@ sealed class KtFakeSourceElementKind : KtSourceElementKind() {
 
     // for lambdas & functions with expression bodies the return statement is added
     // with a fake sources which refers to the return target
-    object ImplicitReturn : KtFakeSourceElementKind()
+    sealed class ImplicitReturn : KtFakeSourceElementKind() {
+        object FromExpressionBody : ImplicitReturn()
+
+        object FromLastStatement : ImplicitReturn()
+    }
 
     // return expression in procedures -> return Unit
     // with a fake sources which refers to the return statement

compiler/testData/diagnostics/testsWithStdLib/functionLiterals/nullableUnitLiteral.fir.kt

@@ -0,0 +1,10 @@
+// FULL_JDK
+
+fun test() {
+    val closeable: java.io.Closeable? = null
+    val closeF = fun() { closeable?.close() }
+    val closeFB = fun(): Unit = <!RETURN_TYPE_MISMATCH!>closeable?.close()<!>
+    val closeFR = fun() { return <!RETURN_TYPE_MISMATCH!>closeable?.close()<!> }
+    val closeL = { closeable?.close() }
+    val closeLR = label@ { return@label closeable?.close() }
+}

compiler/testData/diagnostics/testsWithStdLib/functionLiterals/nullableUnitLiteral.kt

@@ -0,0 +1,10 @@
+// FULL_JDK
+
+fun test() {
+    val closeable: java.io.Closeable? = null
+    val closeF = fun() { closeable?.close() }
+    val closeFB = fun(): Unit = <!TYPE_MISMATCH!>closeable?.close()<!>
+    val closeFR = fun() { return <!TYPE_MISMATCH!>closeable?.close()<!> }
+    val closeL = { closeable?.close() }
+    val closeLR = label@ { return@label closeable?.close() }
+}

compiler/testData/diagnostics/testsWithStdLib/functionLiterals/nullableUnitLiteral.txt

@@ -0,0 +1,3 @@
+package
+
+public fun test(): kotlin.Unit

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

@@ -37824,6 +37824,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/testsWithStdLib/functionLiterals"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
             }
 
+            @Test
+            @TestMetadata("nullableUnitLiteral.kt")
+            public void testNullableUnitLiteral() throws Exception {
+                runTest("compiler/testData/diagnostics/testsWithStdLib/functionLiterals/nullableUnitLiteral.kt");
+            }
+
             @Test
             @TestMetadata("pseudocodeMemoryOverhead.kt")
             public void testPseudocodeMemoryOverhead() throws Exception {

compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirDiagnosticTest.kt

@@ -201,7 +201,7 @@ abstract class AbstractKtDiagnosticsTest : AbstractFirBaseDiagnosticsTest() {
         val sourceElement = element.source ?: return null
         val sourceKind = sourceElement.kind
         if (sourceKind !in allowedKindsForDebugInfo) {
-            if (sourceKind != KtFakeSourceElementKind.ImplicitReturn || sourceElement.elementType != KtNodeTypes.RETURN) {
+            if (sourceKind !is KtFakeSourceElementKind.ImplicitReturn || sourceElement.elementType != KtNodeTypes.RETURN) {
                 return null
             }
         }