[FIR] Remove dependent data flow variables after receiver reassignment

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosisCompilerFirTestdataTestGenerated.java

@@ -3474,6 +3474,12 @@ public class DiagnosisCompilerFirTestdataTestGenerated extends AbstractDiagnosis
         @TestMetadata("compiler/fir/analysis-tests/testData/resolve/smartcasts")
         @TestDataPath("$PROJECT_ROOT")
         public class Smartcasts {
+            @Test
+            @TestMetadata("accessToMemberAfterReceiverReassignment.kt")
+            public void testAccessToMemberAfterReceiverReassignment() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/accessToMemberAfterReceiverReassignment.kt");
+            }
+
             @Test
             public void testAllFilesPresentInSmartcasts() throws Exception {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/analysis-tests/testData/resolve/smartcasts"), Pattern.compile("^([^.]+)\\.kt$"), null, true);

compiler/fir/analysis-tests/legacy-fir-tests/tests-gen/org/jetbrains/kotlin/fir/LazyBodyIsNotTouchedTilContractsPhaseTestGenerated.java

@@ -3075,6 +3075,11 @@ public class LazyBodyIsNotTouchedTilContractsPhaseTestGenerated extends Abstract
             KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
         }
 
+        @TestMetadata("accessToMemberAfterReceiverReassignment.kt")
+        public void testAccessToMemberAfterReceiverReassignment() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/accessToMemberAfterReceiverReassignment.kt");
+        }
+
         public void testAllFilesPresentInSmartcasts() throws Exception {
             KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/analysis-tests/testData/resolve/smartcasts"), Pattern.compile("^([^.]+)\\.kt$"), null, true);
         }

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

@@ -0,0 +1,64 @@
+FILE: accessToMemberAfterReceiverReassignment.kt
+    public final class My : R|kotlin/Any| {
+        public constructor(x: R|My?|, z: R|My?| = Null(null)): R|My| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final val x: R|My?| = R|<local>/x|
+            public get(): R|My?|
+
+        public final val z: R|My?| = R|<local>/z|
+            public get(): R|My?|
+
+    }
+    public final fun baseTest(): R|kotlin/Unit| {
+        lvar y: R|My?| = R|/My.My|(R|/My.My|(Null(null)))
+        when () {
+            !=(R|<local>/y|?.{ $subj$.R|/My.x| }, Null(null)) ->  {
+                R|<local>/y|.R|/My.x|.R|/My.x|
+                R|<local>/y| = R|/My.My|(Null(null))
+                R|<local>/y|.R|/My.x|.<Inapplicable(UNSAFE_CALL): /My.x>#
+            }
+        }
+
+    }
+    public final fun deepChainTest(): R|kotlin/Unit| {
+        lvar y: R|My?| = R|/My.My|(R|/My.My|(Null(null)))
+        when () {
+            !=(R|<local>/y|?.{ $subj$.R|/My.x| }?.{ $subj$.R|/My.x| }, Null(null)) ->  {
+                R|<local>/y|.R|/My.x|.R|/My.x|.R|/My.x|
+                R|<local>/y| = R|/My.My|(Null(null))
+                R|<local>/y|.R|/My.x|.<Inapplicable(UNSAFE_CALL): /My.x>#.<Inapplicable(UNSAFE_CALL): /My.x>#
+            }
+        }
+
+    }
+    public final fun backwardAliasTest(z: R|My|): R|kotlin/Unit| {
+        lvar y: R|My| = R|<local>/z|
+        when () {
+            !=(R|<local>/y|.R|/My.x|, Null(null)) ->  {
+                R|<local>/y|.R|/My.x|.R|/My.x|
+                R|<local>/y| = R|/My.My|(Null(null))
+                R|<local>/y|.R|/My.x|.<Inapplicable(UNSAFE_CALL): /My.x>#
+            }
+        }
+
+    }
+    public final fun severalMembersTest(): R|kotlin/Unit| {
+        lvar y: R|My| = R|/My.My|(R|/My.My|(Null(null)), R|/My.My|(Null(null)))
+        when () {
+            !=(R|<local>/y|.R|/My.x|, Null(null)) ->  {
+                when () {
+                    !=(R|<local>/y|.R|/My.z|, Null(null)) ->  {
+                        R|<local>/y|.R|/My.x|.R|/My.x|
+                        R|<local>/y|.R|/My.z|.R|/My.z|
+                        R|<local>/y| = R|/My.My|(Null(null))
+                        R|<local>/y|.R|/My.x|.<Inapplicable(UNSAFE_CALL): /My.x>#
+                        R|<local>/y|.R|/My.z|.<Inapplicable(UNSAFE_CALL): /My.z>#
+                    }
+                }
+
+            }
+        }
+
+    }

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

@@ -0,0 +1,41 @@
+class My(val x: My?, val z: My? = null)
+
+fun baseTest() {
+    var y: My? = My(My(null))
+    if (y?.x != null) {
+        y.x.x
+        y = My(null)
+        y.x<!UNSAFE_CALL!>.<!>x
+    }
+}
+
+fun deepChainTest() {
+    var y: My? = My(My(null))
+    if (y?.x?.x != null) {
+        y.x.x.x
+        y = My(null)
+        y.x<!UNSAFE_CALL!>.<!>x<!UNSAFE_CALL!>.<!>x
+    }
+}
+
+fun backwardAliasTest(z: My) {
+    var y = z
+    if (y.x != null) {
+        y.x.x
+        y = My(null)
+        y.x<!UNSAFE_CALL!>.<!>x
+    }
+}
+
+fun severalMembersTest() {
+    var y = My(My(null), My(null))
+    if (y.x != null) {
+        if (y.z != null) {
+            y.x.x
+            y.z.z
+            y = My(null)
+            y.x<!UNSAFE_CALL!>.<!>x
+            y.z<!UNSAFE_CALL!>.<!>z
+        }
+    }
+}

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirDiagnosticTestGenerated.java

@@ -3474,6 +3474,12 @@ public class FirDiagnosticTestGenerated extends AbstractFirDiagnosticTest {
         @TestMetadata("compiler/fir/analysis-tests/testData/resolve/smartcasts")
         @TestDataPath("$PROJECT_ROOT")
         public class Smartcasts {
+            @Test
+            @TestMetadata("accessToMemberAfterReceiverReassignment.kt")
+            public void testAccessToMemberAfterReceiverReassignment() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/accessToMemberAfterReceiverReassignment.kt");
+            }
+
             @Test
             public void testAllFilesPresentInSmartcasts() throws Exception {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/analysis-tests/testData/resolve/smartcasts"), Pattern.compile("^([^.]+)\\.kt$"), null, true);

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

@@ -3474,6 +3474,12 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
         @TestMetadata("compiler/fir/analysis-tests/testData/resolve/smartcasts")
         @TestDataPath("$PROJECT_ROOT")
         public class Smartcasts {
+            @Test
+            @TestMetadata("accessToMemberAfterReceiverReassignment.kt")
+            public void testAccessToMemberAfterReceiverReassignment() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/accessToMemberAfterReceiverReassignment.kt");
+            }
+
             @Test
             public void testAllFilesPresentInSmartcasts() throws Exception {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/fir/analysis-tests/testData/resolve/smartcasts"), Pattern.compile("^([^.]+)\\.kt$"), null, true);

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

@@ -863,7 +863,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         if (possiblyChangedVariables.isEmpty()) return
         val flow = node.flow
         for (variable in possiblyChangedVariables) {
-            flow.removeAllAboutVariable(variable)
+            logicSystem.removeAllAboutVariable(flow, variable)
         }
     }
 
@@ -1170,7 +1170,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         val isAssignment = assignment != null
         if (isAssignment) {
             logicSystem.removeLocalVariableAlias(flow, propertyVariable)
-            flow.removeAllAboutVariable(propertyVariable)
+            logicSystem.removeAllAboutVariable(flow, propertyVariable)
             logicSystem.recordNewAssignment(flow, propertyVariable, context.newAssignmentIndex())
         }
 
@@ -1508,13 +1508,6 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         logicSystem.addTypeStatement(this, info)
     }
 
-    private fun FLOW.removeAllAboutVariable(variable: RealVariable?) {
-        if (variable == null) return
-        logicSystem.removeTypeStatementsAboutVariable(this, variable)
-        logicSystem.removeLogicStatementsAboutVariable(this, variable)
-        logicSystem.removeAliasInformationAboutVariable(this, variable)
-    }
-
     private fun FLOW.fork(): FLOW {
         return logicSystem.forkFlow(this)
     }

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)
+                removeAllAboutVariable(commonFlow, variable)
-                removeAliasInformationAboutVariable(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)
+            fun processApprovedStatements(isDiff: Boolean) {
-                val replacedStatements =
+                val statements = if (isDiff) flow.approvedTypeStatementsDiff else flow.approvedTypeStatements
-                    flow.approvedTypeStatements[backAlias]?.let { existing -> existing + newStatements } ?: newStatements
+                statements[alias]?.let { it ->
-                flow.approvedTypeStatements = flow.approvedTypeStatements.put(backAlias, replacedStatements.toPersistent())
+                    val newStatements = it.replaceVariable(alias, backAlias)
-            }
+                    val replacedStatements =
-            flow.approvedTypeStatementsDiff[alias]?.let { it ->
+                        statements[backAlias]?.let { existing -> existing + newStatements } ?: newStatements
-                val newStatements = it.replaceVariable(alias, backAlias)
+                    val result = statements.put(backAlias, replacedStatements.toPersistent())
-                val replacedStatements =
+                    if (isDiff) {
-                    flow.approvedTypeStatementsDiff[backAlias]?.let { existing -> existing + newStatements } ?: newStatements
+                        flow.approvedTypeStatementsDiff = result
-                flow.approvedTypeStatementsDiff = flow.approvedTypeStatementsDiff.put(backAlias, replacedStatements.toPersistent())
+                    } 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")