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/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsTestGenerated.java

@@ -28217,6 +28217,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                     runTest("compiler/testData/diagnostics/tests/smartCasts/variables/accessorAndFunction.kt");
                 }
 
+                @Test
+                @TestMetadata("aliasing.kt")
+                public void testAliasing() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/smartCasts/variables/aliasing.kt");
+                }
+
                 @Test
                 public void testAllFilesPresentInVariables() throws Exception {
                     KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/smartCasts/variables"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);

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

@@ -28217,6 +28217,12 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                     runTest("compiler/testData/diagnostics/tests/smartCasts/variables/accessorAndFunction.kt");
                 }
 
+                @Test
+                @TestMetadata("aliasing.kt")
+                public void testAliasing() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/smartCasts/variables/aliasing.kt");
+                }
+
                 @Test
                 public void testAllFilesPresentInVariables() throws Exception {
                     KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/smartCasts/variables"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);

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
+
 }

compiler/fir/tree/src/org/jetbrains/kotlin/fir/types/FirTypeUtils.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.builtins.functions.FunctionClassKind
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
 import org.jetbrains.kotlin.fir.expressions.FirConstExpression
 import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.FirExpressionWithSmartcast
 import org.jetbrains.kotlin.fir.render
 import org.jetbrains.kotlin.fir.types.impl.FirImplicitBuiltinTypeRef
 import org.jetbrains.kotlin.name.ClassId
@@ -51,6 +52,14 @@ val FirExpression.isNullLiteral: Boolean
             this.value == null &&
             this.source != null
 
+@OptIn(ExperimentalContracts::class)
+fun FirExpression.isStableSmartcast(): Boolean {
+    contract {
+        returns(true) implies (this@isStableSmartcast is FirExpressionWithSmartcast)
+    }
+    return this is FirExpressionWithSmartcast && this.isStable
+}
+
 private val FirTypeRef.classLikeTypeOrNull: ConeClassLikeType?
     get() = when (this) {
         is FirImplicitBuiltinTypeRef -> type

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

@@ -0,0 +1,65 @@
+fun test() {
+    var a: Any? = null
+    var b = a
+    var c = b
+    // Now both `b` and `c` are aliases of `a`.
+
+    if (a is String) {
+        b.length // OK
+        c.length // OK
+    }
+    if (b is String) {
+        a.length // OK
+        c.length // OK
+    }
+    if (c is String) {
+        a.length // OK
+        b.length // OK
+    }
+
+    b = 3 // break `b` -> `a`
+    if (a is String) {
+        b.<!UNRESOLVED_REFERENCE!>length<!> // error
+        c.length // OK, since `c` is aliased to `a`
+    }
+    if (b is String) {
+        a.<!UNRESOLVED_REFERENCE!>length<!> // error
+        c.<!UNRESOLVED_REFERENCE!>length<!> // error
+    }
+    if (c is String) {
+        a.length // OK, since `c` is alised to `a`
+        b.<!UNRESOLVED_REFERENCE!>length<!> // error
+    }
+
+    a = 2 // break `c` -> `a`
+    if (a is String) {
+        b.<!UNRESOLVED_REFERENCE!>length<!> // error
+        c.<!UNRESOLVED_REFERENCE!>length<!> // error
+    }
+    if (b is String) {
+        a.<!UNRESOLVED_REFERENCE!>length<!> // error
+        c.<!UNRESOLVED_REFERENCE!>length<!> // error
+    }
+    if (c is String) {
+        a.<!UNRESOLVED_REFERENCE!>length<!> // error
+        b.<!UNRESOLVED_REFERENCE!>length<!> // error
+    }
+
+    c = b // create aliasing `c` -> `b`
+    c.unaryPlus() // OK due to aliasing
+    b = ""
+    c.unaryPlus() // OK since `c` should carry all typing information that was on `b` before `b = ""`.
+
+    c = ""
+    c.length // OK
+    c.<!UNRESOLVED_REFERENCE!>unaryPlus<!>() // error
+}
+
+fun test2() {
+    var a: Any? = null
+    a = ""
+    var b = a
+    a = 3
+    b.length // OK
+    b.<!UNRESOLVED_REFERENCE!>unaryPlus<!>() // error
+}

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

@@ -0,0 +1,65 @@
+fun test() {
+    var a: Any? = null
+    var b = a
+    var c = b
+    // Now both `b` and `c` are aliases of `a`.
+
+    if (a is String) {
+        b.<!UNRESOLVED_REFERENCE!>length<!> // OK
+        c.<!UNRESOLVED_REFERENCE!>length<!> // OK
+    }
+    if (b is String) {
+        a.<!UNRESOLVED_REFERENCE!>length<!> // OK
+        c.<!UNRESOLVED_REFERENCE!>length<!> // OK
+    }
+    if (c is String) {
+        a.<!UNRESOLVED_REFERENCE!>length<!> // OK
+        b.<!UNRESOLVED_REFERENCE!>length<!> // OK
+    }
+
+    b = 3 // break `b` -> `a`
+    if (a is String) {
+        b.<!UNRESOLVED_REFERENCE!>length<!> // error
+        c.<!UNRESOLVED_REFERENCE!>length<!> // OK, since `c` is aliased to `a`
+    }
+    if (<!USELESS_IS_CHECK!>b is String<!>) {
+        a.<!UNRESOLVED_REFERENCE!>length<!> // error
+        c.<!UNRESOLVED_REFERENCE!>length<!> // error
+    }
+    if (c is String) {
+        a.<!UNRESOLVED_REFERENCE!>length<!> // OK, since `c` is alised to `a`
+        b.<!UNRESOLVED_REFERENCE!>length<!> // error
+    }
+
+    a = 2 // break `c` -> `a`
+    if (<!USELESS_IS_CHECK!>a is String<!>) {
+        b.<!UNRESOLVED_REFERENCE!>length<!> // error
+        c.<!UNRESOLVED_REFERENCE!>length<!> // error
+    }
+    if (<!USELESS_IS_CHECK!>b is String<!>) {
+        a.<!UNRESOLVED_REFERENCE!>length<!> // error
+        c.<!UNRESOLVED_REFERENCE!>length<!> // error
+    }
+    if (c is String) {
+        a.<!UNRESOLVED_REFERENCE!>length<!> // error
+        b.<!UNRESOLVED_REFERENCE!>length<!> // error
+    }
+
+    c = b // create aliasing `c` -> `b`
+    <!DEBUG_INFO_SMARTCAST!>c<!>.unaryPlus() // OK due to aliasing
+    b = ""
+    <!DEBUG_INFO_SMARTCAST!>c<!>.unaryPlus() // OK since `c` should carry all typing information that was on `b` before `b = ""`.
+
+    c = ""
+    <!DEBUG_INFO_SMARTCAST!>c<!>.length // OK
+    c.<!UNRESOLVED_REFERENCE!>unaryPlus<!>() // error
+}
+
+fun test2() {
+    var a: Any? = <!VARIABLE_WITH_REDUNDANT_INITIALIZER!>null<!>
+    a = ""
+    var b = a
+    a = 3
+    <!DEBUG_INFO_SMARTCAST!>b<!>.length // OK
+    b.<!UNRESOLVED_REFERENCE!>unaryPlus<!>() // error
+}

compiler/testData/diagnostics/tests/smartCasts/variables/aliasing.txt

@@ -0,0 +1,4 @@
+package
+
+public fun test(): kotlin.Unit
+public fun test2(): kotlin.Unit

compiler/testData/diagnostics/tests/smartCasts/varnotnull/boundInitializerWrong.fir.kt

@@ -5,7 +5,7 @@ fun foo() {
     val y = x
     x = null
     if (y != null) {
-        x.hashCode()
+        x<!UNSAFE_CALL!>.<!>hashCode()
     }
 }
 

compiler/testData/diagnostics/tests/smartCasts/varnotnull/plusplusMinusminus.fir.kt

@@ -11,7 +11,7 @@ operator fun Long?.inc() = this?.let { it + 1 }
 fun bar(arg: Long?): Long {
     var i = arg
     if (i++ == 5L) {
-        return i<!UNSAFE_CALL!>--<!> + i
+        return i<!UNSAFE_CALL!>--<!> <!UNSAFE_OPERATOR_CALL!>+<!> i
     }
     if (i++ == 7L) {
         return i++ <!UNSAFE_OPERATOR_CALL!>+<!> i

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java

@@ -28307,6 +28307,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                     runTest("compiler/testData/diagnostics/tests/smartCasts/variables/accessorAndFunction.kt");
                 }
 
+                @Test
+                @TestMetadata("aliasing.kt")
+                public void testAliasing() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/smartCasts/variables/aliasing.kt");
+                }
+
                 @Test
                 public void testAllFilesPresentInVariables() throws Exception {
                     KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/smartCasts/variables"), Pattern.compile("^(.*)\\.kts?$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);

idea/idea-frontend-fir/idea-fir-low-level-api/tests/org/jetbrains/kotlin/idea/fir/low/level/api/diagnostic/compiler/based/DiagnosisCompilerTestFE10TestdataTestGenerated.java

@@ -28217,6 +28217,12 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                     runTest("compiler/testData/diagnostics/tests/smartCasts/variables/accessorAndFunction.kt");
                 }
 
+                @Test
+                @TestMetadata("aliasing.kt")
+                public void testAliasing() throws Exception {
+                    runTest("compiler/testData/diagnostics/tests/smartCasts/variables/aliasing.kt");
+                }
+
                 @Test
                 public void testAllFilesPresentInVariables() throws Exception {
                     KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/smartCasts/variables"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirSmartcastProvider.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.fir.expressions.FirExpressionWithSmartcast
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccessExpression
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.coneTypeSafe
+import org.jetbrains.kotlin.fir.types.isStableSmartcast
 import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirSafe
 import org.jetbrains.kotlin.idea.frontend.api.ImplicitReceiverSmartCast
 import org.jetbrains.kotlin.idea.frontend.api.ImplicitReceiverSmartcastKind
@@ -41,13 +42,13 @@ internal class KtFirSmartcastProvider(
             (extensionReceiver !is FirExpressionWithSmartcast || !extensionReceiver.isStable)
         ) return emptyList()
         buildList {
-            (dispatchReceiver as? FirExpressionWithSmartcast)?.takeIf { it.isStable }?.let { smartCasted ->
+            dispatchReceiver.takeIf { it.isStableSmartcast() }?.let { smartCasted ->
                 ImplicitReceiverSmartCast(
                     smartCasted.typeRef.coneTypeSafe<ConeKotlinType>()?.asKtType() ?: return@let null,
                     ImplicitReceiverSmartcastKind.DISPATCH
                 )
             }?.let(::add)
-            (extensionReceiver as? FirExpressionWithSmartcast)?.takeIf { it.isStable }?.let { smartCasted ->
+            extensionReceiver.takeIf { it.isStableSmartcast() }?.let { smartCasted ->
                 ImplicitReceiverSmartCast(
                     smartCasted.typeRef.coneTypeSafe<ConeKotlinType>()?.asKtType() ?: return@let null,
                     ImplicitReceiverSmartcastKind.EXTENSION