FIR DFA: fix logic clear aliasing

The logic should clear back aliases as well. To ensure all back aliases
don't lose any information, statements on the original variable are
copied over to its aliases.

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/Arguments.kt

@@ -396,7 +396,7 @@ private fun checkApplicabilityForArgumentType(
 
     if (!csBuilder.addSubtypeConstraintIfCompatible(argumentType, expectedType, position)) {
         val smartcastExpression = argument as? FirExpressionWithSmartcast
-        if (smartcastExpression != null && smartcastExpression.smartcastStability != SmartcastStability.STABLE_VALUE) {
+        if (smartcastExpression != null && !smartcastExpression.isStable) {
             val unstableType = smartcastExpression.smartcastType.coneType
             if (csBuilder.addSubtypeConstraintIfCompatible(unstableType, expectedType, position)) {
                 sink.reportDiagnostic(UnstableSmartCast(smartcastExpression, expectedType))

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

@@ -202,13 +202,68 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         )
     }
 
+    /**
+     * For example, consider `var b = a`. In this case, `b` is an alias of `a`. In other words, `a` is the original variable of `b`.
+     *
+     * For back alias, consider the following.
+     * ```
+     * var b = a
+     * var c = b
+     * ```
+     * Here, if the current `alias` references `b`, `c` is a back alias of `b`. So if one calls this method with `b`, we must also
+     * remove aliasing between `b` and `c`. But before removing aliasing, we need to copy any statements that apply to `b` to `c`.
+     */
     override fun removeLocalVariableAlias(flow: PersistentFlow, alias: RealVariable) {
-        val original = flow.directAliasMap[alias]?.variable ?: return
-        flow.directAliasMap = flow.directAliasMap.remove(alias)
-        val variables = flow.backwardsAliasMap.getValue(original)
-        flow.backwardsAliasMap = flow.backwardsAliasMap.put(original, variables - alias)
+        val backAliases = flow.backwardsAliasMap[alias] ?: emptyList()
+        for (backAlias in backAliases) {
+            flow.logicStatements[alias]?.let { it ->
+                val newStatements = it.map { it.replaceVariable(alias, backAlias) }
+                val replacedStatements =
+                    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())
+            }
+        }
+
+        val original = flow.directAliasMap[alias]?.variable
+        if (original != null) {
+            flow.directAliasMap = flow.directAliasMap.remove(alias)
+            val variables = flow.backwardsAliasMap.getValue(original)
+            flow.backwardsAliasMap = flow.backwardsAliasMap.put(original, variables - alias)
+        }
+        flow.backwardsAliasMap = flow.backwardsAliasMap.remove(alias)
+        for (backAlias in backAliases) {
+            flow.directAliasMap = flow.directAliasMap.remove(backAlias)
+        }
     }
 
+    private fun Implication.replaceVariable(from: RealVariable, to: RealVariable): Implication {
+        return Implication(condition.replaceVariable(from, to), effect.replaceVariable(from, to))
+    }
+
+    private fun <T : Statement<T>> Statement<T>.replaceVariable(from: RealVariable, to: RealVariable): T {
+        val statement = when (this) {
+            is OperationStatement -> if (variable == from) copy(variable = to) else this
+            is PersistentTypeStatement -> if (variable == from) copy(variable = to) else this
+            is MutableTypeStatement -> if (variable == from) MutableTypeStatement(to, exactType, exactNotType) else this
+            else -> throw IllegalArgumentException("unknown type of statement $this")
+        }
+        @Suppress("UNCHECKED_CAST")
+        return statement as T
+    }
+
+
     @OptIn(DfaInternals::class)
     private fun PersistentFlow.getApprovedTypeStatements(variable: RealVariable): MutableTypeStatement {
         var flow = this

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

@@ -882,4 +882,5 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
 
     private val FirFunction.bodyResolved: Boolean
         get() = body !is FirLazyBlock && body?.typeRef is FirResolvedTypeRef
+
 }