[FIR] Remove dependent data flow variables after receiver reassignment

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

@@ -61,6 +61,8 @@ class RealVariable(
     variableIndexForDebug: Int,
     val stability: PropertyStability,
 ) : DataFlowVariable(variableIndexForDebug) {
+    val dependentVariables = mutableSetOf<RealVariable>()
+
     override fun equals(other: Any?): Boolean {
         return this === other
     }

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

@@ -22,6 +22,13 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
 
     abstract fun addImplication(flow: FLOW, implication: Implication)
 
+    fun removeAllAboutVariable(flow: FLOW, variable: RealVariable?) {
+        if (variable == null) return
+        removeTypeStatementsAboutVariable(flow, variable)
+        removeLogicStatementsAboutVariable(flow, variable)
+        removeAliasInformationAboutVariable(flow, variable)
+    }
+
     abstract fun removeTypeStatementsAboutVariable(flow: FLOW, variable: RealVariable)
     abstract fun removeLogicStatementsAboutVariable(flow: FLOW, variable: DataFlowVariable)
     abstract fun removeAliasInformationAboutVariable(flow: FLOW, variable: RealVariable)

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

@@ -135,11 +135,11 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         val variables = flows.flatMap { it.approvedTypeStatements.keys }.toSet()
         for (variable in variables) {
             val info = mergeOperation(flows.map { it.getApprovedTypeStatements(variable) }) ?: continue
-            removeTypeStatementsAboutVariable(commonFlow, variable)
+            commonFlow.approvedTypeStatements -= variable
+            commonFlow.approvedTypeStatementsDiff -= variable
             val thereWereReassignments = variable.hasDifferentReassignments(flows)
             if (thereWereReassignments) {
-                removeLogicStatementsAboutVariable(commonFlow, variable)
-                removeAliasInformationAboutVariable(commonFlow, variable)
+                removeAllAboutVariable(commonFlow, variable)
             }
             commonFlow.addApprovedStatements(info)
         }
@@ -222,18 +222,24 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
                     flow.logicStatements[backAlias]?.let { existing -> existing + newStatements } ?: newStatements.toPersistentList()
                 flow.logicStatements = flow.logicStatements.put(backAlias, replacedStatements)
             }
-            flow.approvedTypeStatements[alias]?.let { it ->
-                val newStatements = it.replaceVariable(alias, backAlias)
-                val replacedStatements =
-                    flow.approvedTypeStatements[backAlias]?.let { existing -> existing + newStatements } ?: newStatements
-                flow.approvedTypeStatements = flow.approvedTypeStatements.put(backAlias, replacedStatements.toPersistent())
-            }
-            flow.approvedTypeStatementsDiff[alias]?.let { it ->
-                val newStatements = it.replaceVariable(alias, backAlias)
-                val replacedStatements =
-                    flow.approvedTypeStatementsDiff[backAlias]?.let { existing -> existing + newStatements } ?: newStatements
-                flow.approvedTypeStatementsDiff = flow.approvedTypeStatementsDiff.put(backAlias, replacedStatements.toPersistent())
+
+            fun processApprovedStatements(isDiff: Boolean) {
+                val statements = if (isDiff) flow.approvedTypeStatementsDiff else flow.approvedTypeStatements
+                statements[alias]?.let { it ->
+                    val newStatements = it.replaceVariable(alias, backAlias)
+                    val replacedStatements =
+                        statements[backAlias]?.let { existing -> existing + newStatements } ?: newStatements
+                    val result = statements.put(backAlias, replacedStatements.toPersistent())
+                    if (isDiff) {
+                        flow.approvedTypeStatementsDiff = result
+                    } else {
+                        flow.approvedTypeStatements = result
+                    }
+                }
             }
+
+            processApprovedStatements(isDiff = false)
+            processApprovedStatements(isDiff = true)
         }
 
         val original = flow.directAliasMap[alias]?.variable
@@ -310,17 +316,31 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
     override fun removeTypeStatementsAboutVariable(flow: PersistentFlow, variable: RealVariable) {
         flow.approvedTypeStatements -= variable
         flow.approvedTypeStatementsDiff -= variable
+        variable.forEachTransitiveDependentVariable {
+            flow.approvedTypeStatements -= it
+            flow.approvedTypeStatementsDiff -= it
+        }
     }
 
     override fun removeLogicStatementsAboutVariable(flow: PersistentFlow, variable: DataFlowVariable) {
         flow.logicStatements -= variable
+        val realVariable = variable as? RealVariable
+        realVariable?.forEachTransitiveDependentVariable {
+            flow.logicStatements -= it
+        }
         var newLogicStatements = flow.logicStatements
         for ((key, implications) in flow.logicStatements) {
             val implicationsToDelete = mutableListOf<Implication>()
             implications.forEach { implication ->
-                if (implication.effect.variable == variable) {
+                val implicationVariable = implication.effect.variable
+                if (implicationVariable == variable) {
                     implicationsToDelete += implication
                 }
+                realVariable?.forEachTransitiveDependentVariable {
+                    if (implicationVariable == it) {
+                        implicationsToDelete += implication
+                    }
+                }
             }
             if (implicationsToDelete.isEmpty()) continue
             val newImplications = implications.removeAll(implicationsToDelete)
@@ -346,6 +366,13 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         }
     }
 
+    private fun RealVariable.forEachTransitiveDependentVariable(action: (RealVariable) -> Unit) {
+        dependentVariables.forEach {
+            it.forEachTransitiveDependentVariable(action)
+            action(it)
+        }
+    }
+
     override fun translateVariableFromConditionInStatements(
         flow: PersistentFlow,
         originalVariable: DataFlowVariable,

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

@@ -106,7 +106,9 @@ class VariableStorageImpl(private val session: FirSession) : VariableStorage() {
         }
 
         val receiverVariable = receiver?.let { getOrCreateVariable(flow, it) }
-        return RealVariable(identifier, isThisReference, receiverVariable, counter++, stability)
+        return RealVariable(identifier, isThisReference, receiverVariable, counter++, stability).also {
+            (receiverVariable as? RealVariable)?.dependentVariables?.add(it)
+        }
     }
 
     @JvmName("getOrCreateRealVariableOrNull")