FIR DFA: update assignment indices when joining flows

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

@@ -1221,7 +1221,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
             // TODO: add unstable smartcast for non-local var
             val variable = variableStorage.getRealVariableWithoutUnwrappingAlias(node.flow, property.symbol, assignment)
             if (variable != null) {
-                logicSystem.removeAllAboutVariable(node.flow, variable)
+                logicSystem.recordNewAssignment(node.flow, variable, context.newAssignmentIndex())
             }
         }
         processConditionalContract(assignment)
@@ -1243,7 +1243,6 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         )
         val isAssignment = assignment != null
         if (isAssignment) {
-            logicSystem.removeAllAboutVariable(flow, propertyVariable)
             logicSystem.recordNewAssignment(flow, propertyVariable, context.newAssignmentIndex())
         }
 

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

@@ -11,6 +11,7 @@ import org.jetbrains.kotlin.fir.types.ConeInferenceContext
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.utils.addIfNotNull
 import java.util.*
+import kotlin.math.max
 
 data class PersistentTypeStatement(
     override val variable: RealVariable,
@@ -138,13 +139,20 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         // the aliasing of `y` to `a.x` after `if (p) { y = a.x } else { y = a.x }`.
         val commonAliases = computeCommonAliases(flows)
 
+        // If a variable was reassigned in one branch, it was reassigned at the join point.
+        val reassignedVariables = mutableMapOf<RealVariable, Int>()
         for (flow in flows) {
             for ((variable, index) in flow.assignmentIndex) {
                 if (commonFlow.assignmentIndex[variable] != index) {
-                    removeAllAboutVariable(commonFlow, variable)
+                    // Ideally we should generate an entirely new index here, but it doesn't really
+                    // matter; the important part is that it's different from `commonFlow.previousFlow`.
+                    reassignedVariables[variable] = max(index, reassignedVariables[variable] ?: 0)
                 }
             }
         }
+        for ((variable, index) in reassignedVariables) {
+            recordNewAssignment(commonFlow, variable, index)
+        }
 
         val toReplace = statements.map { it.variable to it.toPersistent() }
         commonFlow.approvedTypeStatements += toReplace
@@ -158,7 +166,7 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         return commonFlow
     }
 
-    private fun computeCommonAliases(flows: Collection<PersistentFlow>): MutableMap<RealVariable, RealVariableAndType> =
+    private fun computeCommonAliases(flows: Collection<PersistentFlow>): Map<RealVariable, RealVariableAndType> =
         flows.first().directAliasMap.filterTo(mutableMapOf()) { (variable, alias) ->
             flows.all { it.directAliasMap[variable] == alias }
         }
@@ -471,6 +479,7 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
     }
 
     override fun recordNewAssignment(flow: PersistentFlow, variable: RealVariable, index: Int) {
+        removeAllAboutVariable(flow, variable)
         flow.assignmentIndex = flow.assignmentIndex.put(variable, index)
     }
 

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

@@ -38,3 +38,26 @@ fun test3(p: Boolean) {
     x.length // ok
     c.x<!UNSAFE_CALL!>.<!>length // bad
 }
+
+fun test4(p: Boolean, q: Boolean) {
+    var c = C("...")
+    val x = c.x
+    if (x == null) return
+    x.length // ok
+    c.x.length // ok
+    if (p) {
+        if (q) {
+            c = C(null)
+        } else {
+            c = C(null)
+        }
+    } else {
+        if (q) {
+            c = C(null)
+        } else {
+            c = C(null)
+        }
+    }
+    x.length // ok
+    c.x<!UNSAFE_CALL!>.<!>length // bad
+}

compiler/testData/diagnostics/tests/smartCasts/variables/reassignedDependency.kt

@@ -38,3 +38,26 @@ fun test3(p: Boolean) {
     <!DEBUG_INFO_SMARTCAST!>x<!>.length // ok
     c.x<!UNSAFE_CALL!>.<!>length // bad
 }
+
+fun test4(p: Boolean, q: Boolean) {
+    var c = C("...")
+    val x = c.x
+    if (x == null) return
+    <!DEBUG_INFO_SMARTCAST!>x<!>.length // ok
+    c.x<!UNSAFE_CALL!>.<!>length // ok
+    if (p) {
+        if (q) {
+            c = C(null)
+        } else {
+            c = C(null)
+        }
+    } else {
+        if (q) {
+            c = C(null)
+        } else {
+            c = C(null)
+        }
+    }
+    <!DEBUG_INFO_SMARTCAST!>x<!>.length // ok
+    c.x<!UNSAFE_CALL!>.<!>length // bad
+}