From 0627dbcb781cb192d1ef1171151d9cd2524c085c Mon Sep 17 00:00:00 2001
From: Mark Punzalan <markpun@google.com>
Date: Thu, 29 Jul 2021 07:11:58 +0000
Subject: [PATCH] FIR IDE: Check stability of smartcast expressions in
 KtFirExpressionTypeProvider.getDefiniteNullability().

---
 .../components/KtFirExpressionTypeProvider.kt | 30 ++++++++++---------
 1 file changed, 16 insertions(+), 14 deletions(-)

diff --git a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirExpressionTypeProvider.kt b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirExpressionTypeProvider.kt
index 182ac631fe9..43b049756aa 100644
--- a/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirExpressionTypeProvider.kt
+++ b/idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirExpressionTypeProvider.kt
@@ -142,21 +142,23 @@ internal class KtFirExpressionTypeProvider(
         getDefiniteNullability(expression) == DefiniteNullability.DEFINITELY_NOT_NULL
 
     private fun getDefiniteNullability(expression: KtExpression): DefiniteNullability = withValidityAssertion {
-        // TODO: Check stability of smartcast (pending PR 4382)
-        return when (val fir = expression.getOrBuildFir(analysisSession.firResolveState)) {
-            is FirExpressionWithSmartcastToNull -> DefiniteNullability.DEFINITELY_NULL
-            is FirExpression -> {
-                // Note: This includes FirExpressionWithSmartcast
-                with(analysisSession.rootModuleSession.typeContext) {
-                    if (!fir.typeRef.coneType.isNullableType()) {
-                        DefiniteNullability.DEFINITELY_NOT_NULL
-                    } else {
-                        DefiniteNullability.UNKNOWN
-                    }
-                }
-            }
-            else -> DefiniteNullability.UNKNOWN
+        fun FirExpression.isNotNullable() = with(analysisSession.rootModuleSession.typeContext) {
+            !typeRef.coneType.isNullableType()
         }
+
+        when (val fir = expression.getOrBuildFir(analysisSession.firResolveState)) {
+            is FirExpressionWithSmartcastToNull -> if (fir.isStable) {
+                return DefiniteNullability.DEFINITELY_NULL
+            }
+            is FirExpressionWithSmartcast -> if (fir.isStable && fir.isNotNullable()) {
+                return DefiniteNullability.DEFINITELY_NOT_NULL
+            }
+            is FirExpression -> if (fir.isNotNullable()) {
+                return DefiniteNullability.DEFINITELY_NOT_NULL
+            }
+        }
+
+        return DefiniteNullability.UNKNOWN
     }
 }