[FIR] Infer type of elvis from lhs if rhs is Nothing

^KT-50875 Fixed

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

@@ -3645,6 +3645,12 @@ public class DiagnosisCompilerFirTestdataTestGenerated extends AbstractDiagnosis
                 runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/kt50788.kt");
             }
 
+            @Test
+            @TestMetadata("kt50875.kt")
+            public void testKt50875() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/kt50875.kt");
+            }
+
             @Test
             @TestMetadata("literalInWhenSubject.kt")
             public void testLiteralInWhenSubject() throws Exception {

compiler/fir/analysis-tests/legacy-fir-tests/tests-gen/org/jetbrains/kotlin/fir/LazyBodyIsNotTouchedTilContractsPhaseTestGenerated.java

@@ -3222,6 +3222,11 @@ public class LazyBodyIsNotTouchedTilContractsPhaseTestGenerated extends Abstract
             runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/kt50788.kt");
         }
 
+        @TestMetadata("kt50875.kt")
+        public void testKt50875() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/kt50875.kt");
+        }
+
         @TestMetadata("literalInWhenSubject.kt")
         public void testLiteralInWhenSubject() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/literalInWhenSubject.kt");

compiler/fir/analysis-tests/testData/resolve/smartcasts/kt50875.fir.txt

@@ -0,0 +1,17 @@
+FILE: kt50875.kt
+    public abstract interface A : R|kotlin/Any| {
+        public abstract fun foo(): R|kotlin/Unit|
+
+    }
+    @R|kotlin/Suppress|(names = vararg(String(INVISIBLE_MEMBER), String(INVISIBLE_REFERENCE))) public final fun <K> checkNotNull(x: R|K?|, y: R|K|): R|@Exact K| {
+        ^checkNotNull R|<local>/x| ?: R|<local>/y|
+    }
+    public final fun test_1(any: R|kotlin/Any|, a: R|A|, nullableA: R|A?|): R|kotlin/Unit| {
+        lvar x: R|kotlin/Any| = R|<local>/any|
+        R|<local>/x| = R|/checkNotNull|<R|kotlin/Any|>(R|<local>/nullableA|, R|<local>/a|)
+        R|<local>/x|.<Unresolved name: foo>#()
+        R|<local>/x| = R|<local>/nullableA| ?: R|<local>/a|
+        R|<local>/x|.<Unresolved name: foo>#()
+        R|<local>/x| = R|<local>/nullableA| ?: ^test_1 Unit
+        R|<local>/x|.R|/A.foo|()
+    }

compiler/fir/analysis-tests/testData/resolve/smartcasts/kt50875.kt

@@ -0,0 +1,23 @@
+// ISSUE: KT-50875
+
+interface A {
+    fun foo()
+}
+
+@Suppress("INVISIBLE_MEMBER", "INVISIBLE_REFERENCE")
+fun <K> checkNotNull(x: K?, y: K): @kotlin.internal.Exact K {
+    return x ?: y
+}
+
+fun test_1(any: Any, a: A, nullableA: A?) {
+    var x = any // x has type Any
+
+    x = checkNotNull(nullableA, a) // inferred to Any
+    x.<!UNRESOLVED_REFERENCE!>foo<!>() // no smartcast
+
+    x = nullableA ?: a // inferred to Any
+    x.<!UNRESOLVED_REFERENCE!>foo<!>() // no smartcast
+
+    x = nullableA ?: return // inferred to Any, but has A type
+    x.foo() // smartcast
+}

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

@@ -3645,6 +3645,12 @@ public class FirDiagnosticTestGenerated extends AbstractFirDiagnosticTest {
                 runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/kt50788.kt");
             }
 
+            @Test
+            @TestMetadata("kt50875.kt")
+            public void testKt50875() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/kt50875.kt");
+            }
+
             @Test
             @TestMetadata("literalInWhenSubject.kt")
             public void testLiteralInWhenSubject() throws Exception {

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

@@ -3645,6 +3645,12 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
                 runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/kt50788.kt");
             }
 
+            @Test
+            @TestMetadata("kt50875.kt")
+            public void testKt50875() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/kt50875.kt");
+            }
+
             @Test
             @TestMetadata("literalInWhenSubject.kt")
             public void testLiteralInWhenSubject() throws Exception {

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

@@ -243,6 +243,14 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
             }
         }
 
+        if (result.rhs.typeRef.coneTypeSafe<ConeKotlinType>()?.isNothing == true) {
+            val lhsType = result.lhs.typeRef.coneTypeSafe<ConeKotlinType>()
+            if (lhsType != null) {
+                val newReturnType = lhsType.makeConeTypeDefinitelyNotNullOrNotNull(session.typeContext)
+                result.replaceTypeRef(result.typeRef.resolvedTypeFromPrototype(newReturnType))
+            }
+        }
+
         session.typeContext.run {
             if (result.typeRef.coneTypeSafe<ConeKotlinType>()?.isNullableType() == true
                 && result.rhs.typeRef.coneTypeSafe<ConeKotlinType>()?.isNullableType() == false
@@ -255,8 +263,6 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
             }
         }
 
-
-
         dataFlowAnalyzer.exitElvis(elvisExpression)
         return result
     }

compiler/testData/diagnostics/tests/when/kt9972.fir.kt

@@ -20,9 +20,9 @@ fun test1(): Int {
 }
 
 fun test2(): Int {
-    val x: String = when {
+    val x: String = <!INITIALIZER_TYPE_MISMATCH!>when {
                         true -> Any()
                         else -> null
-                    } ?: return 0
+                    } ?: return 0<!>
     return x.hashCode()
 }