[FIR] RHS assignment during equality comparison should invalidate DFA

If the right-hand side of an equality comparison contains an assignment
to the variable used in the left-hand side, then implications about the
equality comparison within data-flow analysis cannot be applied, as the
value of the variable will be different after the comparison.

^KT-55096 Fixed

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

@@ -398,8 +398,12 @@ abstract class FirDataFlowAnalyzer(
         graphBuilder.exitComparisonExpression(comparisonExpression).mergeIncomingFlow()
     }
 
+    fun exitEqualityOperatorLhs() {
+        graphBuilder.exitEqualityOperatorLhs()
+    }
+
     fun exitEqualityOperatorCall(equalityOperatorCall: FirEqualityOperatorCall) {
-        val node = graphBuilder.exitEqualityOperatorCall(equalityOperatorCall)
+        val (lhsExitNode, node) = graphBuilder.exitEqualityOperatorCall(equalityOperatorCall)
         val operation = equalityOperatorCall.operation
         val leftOperand = equalityOperatorCall.arguments[0]
         val rightOperand = equalityOperatorCall.arguments[1]
@@ -431,7 +435,7 @@ abstract class FirDataFlowAnalyzer(
                 rightIsNull -> processEqNull(flow, equalityOperatorCall, leftOperand, operation.isEq())
                 leftConst != null -> processEqConst(flow, equalityOperatorCall, rightOperand, leftConst, operation.isEq())
                 rightConst != null -> processEqConst(flow, equalityOperatorCall, leftOperand, rightConst, operation.isEq())
-                else -> processEq(flow, equalityOperatorCall, leftOperand, rightOperand, operation)
+                else -> processEq(flow, lhsExitNode.flow, equalityOperatorCall, leftOperand, rightOperand, operation)
             }
         }
     }
@@ -476,6 +480,7 @@ abstract class FirDataFlowAnalyzer(
 
     private fun processEq(
         flow: MutableFlow,
+        lhsExitFlow: PersistentFlow,
         expression: FirExpression,
         leftOperand: FirExpression,
         rightOperand: FirExpression,
@@ -494,10 +499,11 @@ abstract class FirDataFlowAnalyzer(
             return
         }
 
-        // Ideally it should be `getOrCreateIfRealAndUnchanged(flow from LHS, flow, leftOperand)`, otherwise the statement will
-        //  be added even if the value has changed in the RHS. Currently, the only previous node is the RHS.
-        // But seems like everything works and with current implementation
-        val leftOperandVariable = variableStorage.getOrCreateIfReal(flow, leftOperand)
+        // Only consider the LHS variable if it has not been reassigned in the RHS.
+        val leftOperandVariable = variableStorage.getOrCreateIfReal(flow, leftOperand).takeIf {
+            val variable = variableStorage.getRealVariableWithoutUnwrappingAlias(lhsExitFlow, leftOperand)
+            variable == null || logicSystem.isSameValueIn(lhsExitFlow, flow, variable)
+        }
         val rightOperandVariable = variableStorage.getOrCreateIfReal(flow, rightOperand)
         if (leftOperandVariable == null && rightOperandVariable == null) return
         val expressionVariable = variableStorage.createSynthetic(expression)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt

@@ -80,6 +80,7 @@ class ControlFlowGraphBuilder {
     private val exitSafeCallNodes: Stack<ExitSafeCallNode> = stackOf()
     private val exitElvisExpressionNodes: Stack<ElvisExitNode> = stackOf()
     private val elvisRhsEnterNodes: Stack<ElvisRhsEnterNode> = stackOf()
+    private val equalityOperatorCallLhsExitNodes: Stack<CFGNode<*>> = stackOf()
 
     private val notCompletedFunctionCalls: Stack<MutableList<FunctionCallNode>> = stackOf()
 
@@ -735,8 +736,19 @@ class ControlFlowGraphBuilder {
         return createComparisonExpressionNode(comparisonExpression).also { addNewSimpleNode(it) }
     }
 
-    fun exitEqualityOperatorCall(equalityOperatorCall: FirEqualityOperatorCall): EqualityOperatorCallNode {
-        return createEqualityOperatorCallNode(equalityOperatorCall).also { addNewSimpleNode(it) }
+    fun exitEqualityOperatorLhs() {
+        equalityOperatorCallLhsExitNodes.push(lastNode)
+    }
+
+    /**
+     * Returns a pair of nodes, where the first is the last node of the LHS of the equality operator
+     * call, and the second is the exit node of the equality operator call. This allows DFA to
+     * determine if an assignment took place within the RHS of the equality operator call.
+     */
+    fun exitEqualityOperatorCall(equalityOperatorCall: FirEqualityOperatorCall): Pair<CFGNode<*>, EqualityOperatorCallNode> {
+        val lhsExitNode = equalityOperatorCallLhsExitNodes.pop()
+        val node = createEqualityOperatorCallNode(equalityOperatorCall).also { addNewSimpleNode(it) }
+        return lhsExitNode to node
     }
 
     // ----------------------------------- Jump -----------------------------------

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

@@ -803,6 +803,7 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
         // In cases like materialize1() == materialize2() we add expected type just for the right argument.
         // One of the reasons is just consistency with K1 and with the desugared form `a.equals(b)`. See KT-47409 for clarifications.
         val leftArgumentTransformed: FirExpression = arguments[0].transform(transformer, ResolutionMode.ContextIndependent)
+        dataFlowAnalyzer.exitEqualityOperatorLhs()
         val rightArgumentTransformed: FirExpression = arguments[1].transform(transformer, withExpectedType(builtinTypes.nullableAnyType))
 
         equalityOperatorCall

compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/LogicSystem.kt

@@ -109,6 +109,10 @@ abstract class LogicSystem(private val context: ConeInferenceContext) {
     fun isSameValueIn(a: PersistentFlow, b: PersistentFlow, variable: RealVariable): Boolean =
         a.assignmentIndex[variable] == b.assignmentIndex[variable]
 
+    fun isSameValueIn(a: PersistentFlow, b: MutableFlow, variable: RealVariable): Boolean {
+        return a.assignmentIndex[variable] == b.assignmentIndex[variable]
+    }
+
     private fun MutableFlow.mergeAssignments(flows: Collection<PersistentFlow>) {
         // If a variable was reassigned in one branch, it was reassigned at the join point.
         val reassignedVariables = mutableMapOf<RealVariable, Int>()

compiler/testData/diagnostics/tests/smartCasts/variables/reassignedInRhs.fir.kt

@@ -27,9 +27,9 @@ fun unoverriddenEquals(a: Any?) {
     var b: Any?
     b = c
     if (b == c.also { b = a }) {
-        a.x // bad (b#0 is C, b#1 = a)
-        b.x // bad (b#0 is C, this is b#1)
-        if (<!USELESS_IS_CHECK!>b is C<!>) { // b#1
+        a.<!UNRESOLVED_REFERENCE!>x<!> // bad (b#0 is C, b#1 = a)
+        b.<!UNRESOLVED_REFERENCE!>x<!> // bad (b#0 is C, this is b#1)
+        if (b is C) { // b#1
             a.x // ok (b#1 = a)
             b.x // ok
         }