From 3ec7866ead779d1771e81e30ed6f8113c22b8d2a Mon Sep 17 00:00:00 2001
From: "Denis.Zharkov" <denis.zharkov@jetbrains.com>
Date: Mon, 15 Nov 2021 16:16:19 +0300
Subject: [PATCH] FIR: Fix inference for case of non-nullable RHS of elvis

Some existing tests start failing after previous commits adding @Exact
attribute to `?:`

They have a form:
var x: String? = nullable()

if (x == null) {
   x = nullable() ?: "" // considering @Exact the whole elvis is inferred to nullable from expect type
}

x.length // should be smart cast
---
 .../inference/ifElvisReturn.fir.txt                |  2 +-
 .../FirControlFlowStatementsResolveTransformer.kt  | 14 ++++++++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/compiler/fir/analysis-tests/testData/resolveWithStdlib/inference/ifElvisReturn.fir.txt b/compiler/fir/analysis-tests/testData/resolveWithStdlib/inference/ifElvisReturn.fir.txt
index 6f331f220ed..26c712dade1 100644
--- a/compiler/fir/analysis-tests/testData/resolveWithStdlib/inference/ifElvisReturn.fir.txt
+++ b/compiler/fir/analysis-tests/testData/resolveWithStdlib/inference/ifElvisReturn.fir.txt
@@ -14,7 +14,7 @@ FILE: ifElvisReturn.kt
                 R|<local>/b|
             }
             else ->  {
-                R|/materialize|<R|kotlin/Nothing|>() ?: ^foo Unit
+                R|/materialize|<R|B|>() ?: ^foo Unit
             }
         }
 
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirControlFlowStatementsResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirControlFlowStatementsResolveTransformer.kt
index 9c831996154..2a1256d7cfb 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirControlFlowStatementsResolveTransformer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirControlFlowStatementsResolveTransformer.kt
@@ -243,6 +243,20 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
             }
         }
 
+        session.typeContext.run {
+            if (result.typeRef.coneTypeSafe<ConeKotlinType>()?.isNullableType() == true
+                && result.rhs.typeRef.coneTypeSafe<ConeKotlinType>()?.isNullableType() == false
+            ) {
+                // Sometimes return type for special call for elvis operator might be nullable,
+                // but result is not nullable if the right type is not nullable
+                result.replaceTypeRef(
+                    result.typeRef.withReplacedConeType(result.typeRef.coneType.makeConeTypeDefinitelyNotNullOrNotNull(session.typeContext))
+                )
+            }
+        }
+
+
+
         dataFlowAnalyzer.exitElvis(elvisExpression)
         return result
     }