[FIR] Support smartcast after reference equality check

#KT-39000 Fixed

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

@@ -0,0 +1,12 @@
+interface A
+class B: A
+
+val X = B()
+
+fun foo(b: B) {}
+
+fun main(a: A) {
+    if (a === X) {
+        foo(a)
+    }
+}

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

@@ -0,0 +1,21 @@
+FILE: kt39000.kt
+    public abstract interface A : R|kotlin/Any| {
+    }
+    public final class B : R|A| {
+        public constructor(): R|B| {
+            super<R|kotlin/Any|>()
+        }
+
+    }
+    public final val X: R|B| = R|/B.B|()
+        public get(): R|B|
+    public final fun foo(b: R|B|): R|kotlin/Unit| {
+    }
+    public final fun main(a: R|A|): R|kotlin/Unit| {
+        when () {
+            ===(R|<local>/a|, R|/X|) ->  {
+                R|/foo|(R|<local>/a|)
+            }
+        }
+
+    }

compiler/fir/analysis-tests/tests/org/jetbrains/kotlin/fir/FirDiagnosticsTestGenerated.java

@@ -1989,6 +1989,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
             runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/kt37327.kt");
         }
 
+        @TestMetadata("kt39000.kt")
+        public void testKt39000() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/kt39000.kt");
+        }
+
         @TestMetadata("multipleCasts.kt")
         public void testMultipleCasts() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/multipleCasts.kt");

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

@@ -1989,6 +1989,11 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
             runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/kt37327.kt");
         }
 
+        @TestMetadata("kt39000.kt")
+        public void testKt39000() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/kt39000.kt");
+        }
+
         @TestMetadata("multipleCasts.kt")
         public void testMultipleCasts() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/multipleCasts.kt");

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt

@@ -411,6 +411,10 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
             leftIsNullable -> processEqNull(node, leftOperand, operation.invert())
             rightIsNullable -> processEqNull(node, rightOperand, operation.invert())
         }
+
+        if (operation == FirOperation.IDENTITY || operation == FirOperation.NOT_IDENTITY) {
+            processIdentity(node, leftOperand, rightOperand, operation)
+        }
     }
 
     private fun processEqNull(node: OperatorCallNode, operand: FirExpression, operation: FirOperation) {
@@ -444,6 +448,28 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         node.flow = flow
     }
 
+    private fun processIdentity(node: OperatorCallNode, leftOperand: FirExpression, rightOperand: FirExpression, operation: FirOperation) {
+        val flow = node.flow
+        val expressionVariable = variableStorage.getOrCreateVariable(node.previousFlow, node.fir)
+        val leftOperandVariable = variableStorage.getOrCreateVariable(node.previousFlow, leftOperand)
+        val rightOperandVariable = variableStorage.getOrCreateVariable(node.previousFlow, rightOperand)
+        val leftOperandType = leftOperand.coneType
+        val rightOperandType = rightOperand.coneType
+        val isEq = operation.isEq()
+
+        if (leftOperandVariable.isReal()) {
+            flow.addImplication((expressionVariable eq isEq) implies (leftOperandVariable typeEq rightOperandType))
+            flow.addImplication((expressionVariable notEq isEq) implies (leftOperandVariable typeNotEq rightOperandType))
+        }
+
+        if (rightOperandVariable.isReal()) {
+            flow.addImplication((expressionVariable eq isEq) implies (rightOperandVariable typeEq leftOperandType))
+            flow.addImplication((expressionVariable notEq isEq) implies (rightOperandVariable typeNotEq leftOperandType))
+        }
+
+        node.flow = flow
+    }
+
     // ----------------------------------- Jump -----------------------------------
 
     fun exitJump(jump: FirJump<*>) {

compiler/testData/diagnostics/tests/smartCasts/fakeSmartCastOnEquality.fir.kt

@@ -15,11 +15,11 @@ fun check(x1: Derived1, x: Base) {
     }
     if (x1 === x) {
         // OK
-        x.<!UNRESOLVED_REFERENCE!>foo<!>()
+        x.foo()
     }
     if (x1 !== x) {} else {
         // OK
-        x.<!UNRESOLVED_REFERENCE!>foo<!>()
+        x.foo()
     }
 }
 

compiler/tests-spec/testData/diagnostics/notLinked/dfa/pos/39.fir.kt

@@ -21,8 +21,8 @@ fun case_2(x: Number) {
     val y: Int? = null
 
     if (x === y) {
-        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Number")!>x<!>
+        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int & kotlin.Number")!>x<!>
-        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Number")!>x<!>.inv()
+        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int & kotlin.Number")!>x<!>.inv()
     }
 }
 
@@ -31,8 +31,8 @@ fun case_3(x: Number) {
     var y: Int? = null
 
     if (x === y) {
-        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Number")!>x<!>
+        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int & kotlin.Number")!>x<!>
-        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Number")!>x<!>.inv()
+        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int & kotlin.Number")!>x<!>.inv()
     }
 }
 

compiler/tests-spec/testData/diagnostics/notLinked/dfa/pos/48.fir.kt

@@ -25,8 +25,8 @@ fun case_2(x: Any?, y: Any?) {
     if (x as Int === y) {
         <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int & kotlin.Any?")!>x<!>
         <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int & kotlin.Any?")!>x<!>.<!INAPPLICABLE_CANDIDATE!>inv<!>(10)
-        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Any & kotlin.Any?")!>y<!>
+        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int & kotlin.Any?")!>y<!>
-        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Any & kotlin.Any?")!>y<!>.<!INAPPLICABLE_CANDIDATE!>inv<!>(10)
+        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int & kotlin.Any?")!>y<!>.<!INAPPLICABLE_CANDIDATE!>inv<!>(10)
     }
 }
 
@@ -53,7 +53,7 @@ fun case_4(x: Any?, y: Any?) {
     if (y === x as Int) {
         <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int & kotlin.Any?")!>x<!>
         <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int & kotlin.Any?")!>x<!>.<!INAPPLICABLE_CANDIDATE!>inv<!>(10)
-        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Any & kotlin.Any?")!>y<!>
+        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int & kotlin.Any?")!>y<!>
-        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Any & kotlin.Any?")!>y<!>.<!INAPPLICABLE_CANDIDATE!>inv<!>(10)
+        <!DEBUG_INFO_EXPRESSION_TYPE("kotlin.Int & kotlin.Any?")!>y<!>.<!INAPPLICABLE_CANDIDATE!>inv<!>(10)
     }
 }