FIR: Fix inference of nullable unit lambda at RHS of elvis

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

@@ -14060,6 +14060,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                     runTest("compiler/testData/diagnostics/tests/inference/coercionToUnit/kt30242.kt");
                 }
 
+                @Test
+                @TestMetadata("lambdaWithNullableUnitInElvis.kt")
+                public void testLambdaWithNullableUnitInElvis() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/coercionToUnit/lambdaWithNullableUnitInElvis.kt");
+                }
+
                 @Test
                 @TestMetadata("noCoercion.kt")
                 public void testNoCoercion() throws Exception {

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

@@ -14060,6 +14060,12 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                     runTest("compiler/testData/diagnostics/tests/inference/coercionToUnit/kt30242.kt");
                 }
 
+                @Test
+                @TestMetadata("lambdaWithNullableUnitInElvis.kt")
+                public void testLambdaWithNullableUnitInElvis() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/inference/coercionToUnit/lambdaWithNullableUnitInElvis.kt");
+                }
+
                 @Test
                 @TestMetadata("noCoercion.kt")
                 public void testNoCoercion() throws Exception {

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

@@ -787,8 +787,9 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         val implicitReturns = returnExpressionsExceptLast.filter {
             (it as? FirExpression)?.typeRef is FirImplicitUnitTypeRef
         }
+
         val returnType =
-            if (implicitReturns.isNotEmpty()) {
+            if (implicitReturns.isNotEmpty() || lambda.returnType?.isUnit == true) {
                 // i.e., early return, e.g., l@{ ... return@l ... }
                 // Note that the last statement will be coerced to Unit if needed.
                 session.builtinTypes.unitType.type