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

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

@@ -14060,6 +14060,12 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                     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/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

compiler/testData/diagnostics/tests/inference/coercionToUnit/lambdaWithNullableUnitInElvis.kt

@@ -0,0 +1,8 @@
+// FIR_IDENTICAL
+// SKIP_TXT
+
+fun nullableF(): (() -> Unit)?= null
+
+fun String.unit() {}
+
+fun foo(x: String?): () -> Unit = nullableF() ?: { x?.unit() }

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

@@ -14066,6 +14066,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                     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 {