FIR DFA: when removing a variable, move type statements about dependents

val c = C("...")
  val d = c
  if (c.x == null) return
  c.x.length // c.x has type String
  d.x.length // d.x -> c.x through d -> c
  c = C(null) // remove alias d -> c
  d.x.length // info from c.x moved to d.x

^KT-54824 Fixed

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/FirReturnsImpliesAnalyzer.kt

@@ -48,6 +48,9 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
         if (effects.isNullOrEmpty()) return
 
         val logicSystem = object : PersistentLogicSystem(context.session.typeContext) {
+            override val variableStorage: VariableStorageImpl
+                get() = dataFlowInfo.variableStorage as VariableStorageImpl
+
             override fun processUpdatedReceiverVariable(flow: PersistentFlow, variable: RealVariable) =
                 throw IllegalStateException("Receiver variable update is not possible for this logic system")
 

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

@@ -99,6 +99,9 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
 
                 override val logicSystem: PersistentLogicSystem =
                     object : PersistentLogicSystem(components.session.typeContext) {
+                        override val variableStorage: VariableStorageImpl
+                            get() = dataFlowAnalyzerContext.variableStorage
+
                         override fun processUpdatedReceiverVariable(flow: PersistentFlow, variable: RealVariable) {
                             val symbol = variable.identifier.symbol
 

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

@@ -74,6 +74,12 @@ class RealVariable(
     override fun hashCode(): Int {
         return _hashCode
     }
+
+    init {
+        if (explicitReceiverVariable is RealVariable) {
+            explicitReceiverVariable.dependentVariables.add(this)
+        }
+    }
 }
 
 class SyntheticVariable(val fir: FirElement, variableIndexForDebug: Int) : DataFlowVariable(variableIndexForDebug) {

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

@@ -94,6 +94,8 @@ class PersistentFlow : Flow {
 }
 
 abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSystem<PersistentFlow>(context) {
+    abstract val variableStorage: VariableStorageImpl
+
     override fun createEmptyFlow(): PersistentFlow {
         return PersistentFlow()
     }
@@ -174,45 +176,50 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
     }
 
     override fun removeAllAboutVariable(flow: PersistentFlow, variable: RealVariable) {
-        val original = flow.directAliasMap[variable]
+        flow.replaceVariable(variable, null)
+    }
+
+    private fun PersistentFlow.replaceVariable(variable: RealVariable, replacement: RealVariable?) {
+        val original = directAliasMap[variable]
         if (original != null) {
-            // All statements should've been made about whatever variable this is an alias to. There is nothing to remove.
+            // All statements should've been made about whatever variable this is an alias to. There is nothing to replace.
             if (AbstractTypeChecker.RUN_SLOW_ASSERTIONS) {
-                assert(variable !in flow.backwardsAliasMap)
+                assert(variable !in backwardsAliasMap)
-                assert(variable !in flow.logicStatements)
+                assert(variable !in logicStatements)
-                assert(variable !in flow.approvedTypeStatements)
+                assert(variable !in approvedTypeStatements)
-                assert(variable !in flow.approvedTypeStatementsDiff)
+                assert(variable !in approvedTypeStatementsDiff)
             }
             // `variable.dependentVariables` is not separated by flow, so it may be non-empty if aliasing of this variable
             // was broken in another flow. However, in *this* flow dependent variables should have no information attached to them.
 
-            val siblings = flow.backwardsAliasMap.getValue(original).remove(variable)
+            val siblings = backwardsAliasMap.getValue(original).remove(variable)
-            flow.directAliasMap -= variable
+            directAliasMap -= variable
-            flow.backwardsAliasMap = if (siblings.isNotEmpty()) {
+            backwardsAliasMap = if (siblings.isNotEmpty()) {
-                flow.backwardsAliasMap.put(original, siblings)
+                backwardsAliasMap.put(original, siblings)
             } else {
-                flow.backwardsAliasMap.remove(original)
+                backwardsAliasMap.remove(original)
+            }
+            if (replacement != null) {
+                addLocalVariableAlias(this, replacement, original)
             }
         } else {
-            val aliases = flow.backwardsAliasMap[variable]
+            val aliases = backwardsAliasMap[variable]
-            val replacement = aliases?.first()
+            // If asked to remove the variable but there are aliases, replace with a new representative for the alias group instead.
+            val replacementOrNext = replacement ?: aliases?.first()
             for (dependent in variable.dependentVariables) {
-                // TODO: replace the identifier in dependent variables instead (need to somehow interface with VariableStorage for that)
+                replaceVariable(dependent, replacementOrNext?.let {
-                //   var x = something
+                    variableStorage.copyRealVariableWithRemapping(dependent, variable, it)
-                //   val y = x
+                })
-                //   if (x.stableProperty !is A) return
-                //   y.stableProperty // A
-                //   x = somethingElse
-                //   y.stableProperty // still A
-                removeAllAboutVariable(flow, dependent)
             }
-            flow.logicStatements = flow.logicStatements.replaceVariable(variable, replacement)
+            logicStatements = logicStatements.replaceVariable(variable, replacementOrNext)
-            flow.approvedTypeStatements = flow.approvedTypeStatements.replaceVariable(variable, replacement)
+            approvedTypeStatements = approvedTypeStatements.replaceVariable(variable, replacementOrNext)
-            flow.approvedTypeStatementsDiff = flow.approvedTypeStatementsDiff.replaceVariable(variable, replacement)
+            approvedTypeStatementsDiff = approvedTypeStatementsDiff.replaceVariable(variable, replacementOrNext)
-            if (replacement != null) {
+            if (aliases != null) {
-                flow.directAliasMap -= replacement
+                backwardsAliasMap -= variable
-                flow.backwardsAliasMap -= variable
+                if (replacementOrNext != null) {
-                flow.addAliases(aliases, replacement)
+                    directAliasMap -= replacementOrNext
+                    addAliases(aliases, replacementOrNext)
+                }
             }
         }
     }

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

@@ -108,8 +108,23 @@ class VariableStorageImpl(private val session: FirSession) : VariableStorage() {
         }
 
         val receiverVariable = receiver?.let { getOrCreateVariable(flow, it) }
-        return RealVariable(identifier, isThisReference, receiverVariable, counter++, stability).also {
+        return RealVariable(identifier, isThisReference, receiverVariable, counter++, stability)
-            (receiverVariable as? RealVariable)?.dependentVariables?.add(it)
+    }
+
+    fun copyRealVariableWithRemapping(variable: RealVariable, from: RealVariable, to: RealVariable): RealVariable {
+        val newIdentifier = with(variable.identifier) {
+            copy(
+                dispatchReceiver = if (dispatchReceiver == from) to else dispatchReceiver,
+                extensionReceiver = if (extensionReceiver == from) to else extensionReceiver,
+            )
+        }
+        return _realVariables.getOrPut(newIdentifier) {
+            with(variable) {
+                RealVariable(
+                    newIdentifier, isThisReference, if (explicitReceiverVariable == from) to else explicitReceiverVariable,
+                    counter++, stability
+                )
+            }
         }
     }
 

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

@@ -73,5 +73,5 @@ fun test5() {
     c = C(null)
     x.length // ok
     c.x<!UNSAFE_CALL!>.<!>length // bad
-    d.x<!UNSAFE_CALL!>.<!>length // ok
+    d.x.length // ok
 }