[FIR] Fix clearing info about DF variable after reassignment

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

@@ -3141,6 +3141,11 @@ public class LazyBodyIsNotTouchedTilContractsPhaseTestGenerated extends Abstract
                 runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/safeCalls/boundSafeCallAndIsCheck.kt");
             }
 
+            @TestMetadata("complexIfWithOr.kt")
+            public void testComplexIfWithOr() throws Exception {
+                runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/safeCalls/complexIfWithOr.kt");
+            }
+
             @TestMetadata("safeCallAndEqualityToBool.kt")
             public void testSafeCallAndEqualityToBool() throws Exception {
                 runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/safeCalls/safeCallAndEqualityToBool.kt");

compiler/fir/analysis-tests/testData/resolve/smartcasts/safeCalls/complexIfWithOr.fir.txt

@@ -0,0 +1,19 @@
+FILE: complexIfWithOr.kt
+    public abstract interface State : R|kotlin/Any| {
+    }
+    public abstract interface Complex : R|kotlin/Any| {
+        public abstract val superClass: R|Complex?|
+            public get(): R|Complex?|
+
+    }
+    public abstract interface ExceptionState : R|State| {
+    }
+    public final fun test(qualifier: R|State?|): R|kotlin/Unit| {
+        when () {
+            ==(R|<local>/qualifier|, Null(null)) || (R|<local>/qualifier| is R|ExceptionState|) || ==((R|<local>/qualifier| as? R|Complex|)?.{ $subj$.R|/Complex.superClass| }, Null(null)) ->  {
+                ^test Unit
+            }
+        }
+
+        R|<local>/qualifier|.R|/Complex.superClass|
+    }

compiler/fir/analysis-tests/testData/resolve/smartcasts/safeCalls/complexIfWithOr.kt

@@ -0,0 +1,13 @@
+interface State
+interface Complex {
+    val superClass: Complex?
+}
+
+interface ExceptionState : State
+
+fun test(qualifier: State?) {
+    if (qualifier == null || qualifier is ExceptionState || (qualifier as? Complex)?.superClass == null) {
+        return
+    }
+    qualifier.superClass
+}

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

@@ -3520,6 +3520,12 @@ public class FirDiagnosticTestGenerated extends AbstractFirDiagnosticTest {
                     runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/safeCalls/boundSafeCallAndIsCheck.kt");
                 }
 
+                @Test
+                @TestMetadata("complexIfWithOr.kt")
+                public void testComplexIfWithOr() throws Exception {
+                    runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/safeCalls/complexIfWithOr.kt");
+                }
+
                 @Test
                 @TestMetadata("safeCallAndEqualityToBool.kt")
                 public void testSafeCallAndEqualityToBool() throws Exception {

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

@@ -3566,6 +3566,12 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
                     runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/safeCalls/boundSafeCallAndIsCheck.kt");
                 }
 
+                @Test
+                @TestMetadata("complexIfWithOr.kt")
+                public void testComplexIfWithOr() throws Exception {
+                    runTest("compiler/fir/analysis-tests/testData/resolve/smartcasts/safeCalls/complexIfWithOr.kt");
+                }
+
                 @Test
                 @TestMetadata("safeCallAndEqualityToBool.kt")
                 public void testSafeCallAndEqualityToBool() throws Exception {

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -14946,6 +14946,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
             runTest("compiler/testData/codegen/box/fir/FakeOverrideBuilder.kt");
         }
 
+        @Test
+        @TestMetadata("falsePositiveBoundSmartcast.kt")
+        public void testFalsePositiveBoundSmartcast() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/falsePositiveBoundSmartcast.kt");
+        }
+
         @Test
         @TestMetadata("Fir2IrClassifierStorage.kt")
         public void testFir2IrClassifierStorage() throws Exception {

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

@@ -49,6 +49,12 @@ class DataFlowAnalyzerContext<FLOW : Flow>(
     var variableStorage = variableStorage
         private set
 
+    private var assignmentCounter = 0
+
+    fun newAssignmentIndex(): Int {
+        return assignmentCounter++
+    }
+
     fun reset() {
         graphBuilder.reset()
         variablesForWhenConditions.clear()
@@ -84,7 +90,6 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
                 private val receiverStack: PersistentImplicitReceiverStack
                     get() = components.implicitReceiverStack as PersistentImplicitReceiverStack
 
-                private val symbolProvider = components.session.symbolProvider
                 private val visibilityChecker = components.session.visibilityChecker
 
                 override val logicSystem: PersistentLogicSystem =
@@ -673,7 +678,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         if (possiblyChangedVariables.isEmpty()) return
         val flow = node.flow
         for (variable in possiblyChangedVariables) {
-            logicSystem.removeAllAboutVariableIncludingAliasInformation(flow, variable)
+            flow.removeAllAboutVariable(variable)
         }
     }
 
@@ -946,6 +951,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         if (isAssignment) {
             logicSystem.removeLocalVariableAlias(flow, propertyVariable)
             flow.removeAllAboutVariable(propertyVariable)
+            logicSystem.recordNewAssignment(flow, propertyVariable, context.newAssignmentIndex())
         }
 
         variableStorage.getOrCreateRealVariable(flow, initializer.symbol, initializer)?.let { initializerVariable ->
@@ -1256,7 +1262,9 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
 
     private fun FLOW.removeAllAboutVariable(variable: RealVariable?) {
         if (variable == null) return
-        logicSystem.removeAllAboutVariable(this, variable)
+        logicSystem.removeTypeStatementsAboutVariable(this, variable)
+        logicSystem.removeLogicStatementsAboutVariable(this, variable)
+        logicSystem.removeAliasInformationAboutVariable(this, variable)
     }
 
     private fun FLOW.fork(): FLOW {

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

@@ -18,6 +18,7 @@ abstract class Flow {
 
     abstract val directAliasMap: Map<RealVariable, RealVariableAndType>
     abstract val backwardsAliasMap: Map<RealVariable, List<RealVariable>>
+    abstract val assignmentIndex: Map<RealVariable, Int>
 }
 
 fun Flow.unwrapVariable(variable: RealVariable): RealVariable {
@@ -36,8 +37,9 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
 
     abstract fun addImplication(flow: FLOW, implication: Implication)
 
-    abstract fun removeAllAboutVariable(flow: FLOW, variable: RealVariable)
-    abstract fun removeAllAboutVariableIncludingAliasInformation(flow: FLOW, variable: RealVariable)
+    abstract fun removeTypeStatementsAboutVariable(flow: FLOW, variable: RealVariable)
+    abstract fun removeLogicStatementsAboutVariable(flow: FLOW, variable: DataFlowVariable)
+    abstract fun removeAliasInformationAboutVariable(flow: FLOW, variable: RealVariable)
 
     abstract fun translateVariableFromConditionInStatements(
         flow: FLOW,
@@ -58,6 +60,8 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
     abstract fun addLocalVariableAlias(flow: FLOW, alias: RealVariable, underlyingVariable: RealVariableAndType)
     abstract fun removeLocalVariableAlias(flow: FLOW, alias: RealVariable)
 
+    abstract fun recordNewAssignment(flow: FLOW, variable: RealVariable, index: Int)
+
     protected abstract fun getImplicationsWithVariable(flow: FLOW, variable: DataFlowVariable): Collection<Implication>
 
     protected abstract fun ConeKotlinType.isAcceptableForSmartcast(): Boolean

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

@@ -11,7 +11,6 @@ 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.NoSuchElementException
 
 data class PersistentTypeStatement(
     override val variable: RealVariable,
@@ -55,6 +54,8 @@ class PersistentFlow : Flow {
     override var directAliasMap: PersistentMap<RealVariable, RealVariableAndType>
     override var backwardsAliasMap: PersistentMap<RealVariable, PersistentList<RealVariable>>
 
+    override var assignmentIndex: PersistentMap<RealVariable, Int>
+
     constructor(previousFlow: PersistentFlow) {
         this.previousFlow = previousFlow
         approvedTypeStatements = previousFlow.approvedTypeStatements
@@ -63,6 +64,7 @@ class PersistentFlow : Flow {
 
         directAliasMap = previousFlow.directAliasMap
         backwardsAliasMap = previousFlow.backwardsAliasMap
+        assignmentIndex = previousFlow.assignmentIndex
     }
 
     constructor() {
@@ -73,6 +75,7 @@ class PersistentFlow : Flow {
 
         directAliasMap = persistentMapOf()
         backwardsAliasMap = persistentMapOf()
+        assignmentIndex = persistentMapOf()
     }
 
     override fun getTypeStatement(variable: RealVariable): TypeStatement? {
@@ -133,7 +136,12 @@ 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, commonFlow) }) ?: continue
-            removeAllAboutVariable(commonFlow, variable)
+            removeTypeStatementsAboutVariable(commonFlow, variable)
+            val thereWereReassignments = variable.hasDifferentReassignments(flows)
+            if (thereWereReassignments) {
+                removeLogicStatementsAboutVariable(commonFlow, variable)
+                removeAliasInformationAboutVariable(commonFlow, variable)
+            }
             commonFlow.addApprovedStatements(info)
         }
 
@@ -141,6 +149,15 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         return commonFlow
     }
 
+    private fun RealVariable.hasDifferentReassignments(flows: Collection<PersistentFlow>): Boolean {
+        val firstIndex = flows.first().assignmentIndex[this] ?: -1
+        for (flow in flows) {
+            val index = flow.assignmentIndex[this] ?: -1
+            if (index != firstIndex) return true
+        }
+        return false
+    }
+
     private fun computeAliasesThatDontChange(
         flows: Collection<PersistentFlow>
     ): MutableMap<RealVariable, RealVariableAndType> {
@@ -247,15 +264,33 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         }
     }
 
-    override fun removeAllAboutVariable(flow: PersistentFlow, variable: RealVariable) {
-        flow.logicStatements -= variable
+    override fun removeTypeStatementsAboutVariable(flow: PersistentFlow, variable: RealVariable) {
         flow.approvedTypeStatements -= variable
         flow.approvedTypeStatementsDiff -= variable
-        // TODO: should we search variable in all logic statements?
     }
 
-    override fun removeAllAboutVariableIncludingAliasInformation(flow: PersistentFlow, variable: RealVariable) {
-        removeAllAboutVariable(flow, variable)
+    override fun removeLogicStatementsAboutVariable(flow: PersistentFlow, variable: DataFlowVariable) {
+        flow.logicStatements -= variable
+        var newLogicStatements = flow.logicStatements
+        for ((key, implications) in flow.logicStatements) {
+            val implicationsToDelete = mutableListOf<Implication>()
+            implications.forEach { implication ->
+                if (implication.effect.variable == variable) {
+                    implicationsToDelete += implication
+                }
+            }
+            if (implicationsToDelete.isEmpty()) continue
+            val newImplications = implications.removeAll(implicationsToDelete)
+            newLogicStatements = if (newImplications.isNotEmpty()) {
+                newLogicStatements.put(key, newImplications)
+            } else {
+                newLogicStatements.remove(key)
+            }
+        }
+        flow.logicStatements = newLogicStatements
+    }
+
+    override fun removeAliasInformationAboutVariable(flow: PersistentFlow, variable: RealVariable) {
         val existedAlias = flow.directAliasMap[variable]?.variable
         if (existedAlias != null) {
             flow.directAliasMap = flow.directAliasMap.remove(variable)
@@ -408,6 +443,10 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         return flow.logicStatements[variable] ?: emptyList()
     }
 
+    override fun recordNewAssignment(flow: PersistentFlow, variable: RealVariable, index: Int) {
+        flow.assignmentIndex = flow.assignmentIndex.put(variable, index)
+    }
+
     // --------------------------------------------------------------------\
 }
 

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

@@ -128,6 +128,7 @@ private infix fun FirElement.isEqualsTo(other: FirElement): Boolean {
 
 sealed class Statement<T : Statement<T>> {
     abstract fun invert(): T
+    abstract val variable: DataFlowVariable
 }
 
 /*
@@ -137,7 +138,7 @@ sealed class Statement<T : Statement<T>> {
  * d == True
  * d == False
  */
-data class OperationStatement(val variable: DataFlowVariable, val operation: Operation) : Statement<OperationStatement>() {
+data class OperationStatement(override val variable: DataFlowVariable, val operation: Operation) : Statement<OperationStatement>() {
     override fun invert(): OperationStatement {
         return OperationStatement(variable, operation.invert())
     }
@@ -148,7 +149,7 @@ data class OperationStatement(val variable: DataFlowVariable, val operation: Ope
 }
 
 abstract class TypeStatement : Statement<TypeStatement>() {
-    abstract val variable: RealVariable
+    abstract override val variable: RealVariable
     abstract val exactType: Set<ConeKotlinType>
     abstract val exactNotType: Set<ConeKotlinType>
 
@@ -209,6 +210,8 @@ fun Implication.invertCondition(): Implication = Implication(condition.invert(),
 typealias TypeStatements = Map<RealVariable, TypeStatement>
 typealias MutableTypeStatements = MutableMap<RealVariable, MutableTypeStatement>
 
+typealias MutableOperationStatements = MutableMap<RealVariable, MutableTypeStatement>
+
 fun MutableTypeStatements.addStatement(variable: RealVariable, statement: TypeStatement) {
     put(variable, statement.asMutableStatement()) { it.apply { this += statement } }
 }

compiler/testData/codegen/box/fir/falsePositiveBoundSmartcast.kt

@@ -0,0 +1,21 @@
+abstract class A {
+    abstract fun foo(): String
+}
+
+class B : A() {
+    override fun foo() = "OK"
+}
+
+class C : A() {
+    override fun foo() = "fail"
+}
+
+fun test(c: C, cond: Boolean): String {
+    var x: A = c
+    if (cond) {
+        x = B()
+    }
+    return x.foo()
+}
+
+fun box(): String = test(C(), true)

compiler/testData/codegen/boxInline/nonLocalReturns/tryFinally/variables/kt7792.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_FIR: JVM_IR
 // FILE: 1.kt
 
 package test

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

@@ -14946,6 +14946,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             runTest("compiler/testData/codegen/box/fir/FakeOverrideBuilder.kt");
         }
 
+        @Test
+        @TestMetadata("falsePositiveBoundSmartcast.kt")
+        public void testFalsePositiveBoundSmartcast() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/falsePositiveBoundSmartcast.kt");
+        }
+
         @Test
         @TestMetadata("Fir2IrClassifierStorage.kt")
         public void testFir2IrClassifierStorage() throws Exception {

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

@@ -14946,6 +14946,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             runTest("compiler/testData/codegen/box/fir/FakeOverrideBuilder.kt");
         }
 
+        @Test
+        @TestMetadata("falsePositiveBoundSmartcast.kt")
+        public void testFalsePositiveBoundSmartcast() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/falsePositiveBoundSmartcast.kt");
+        }
+
         @Test
         @TestMetadata("Fir2IrClassifierStorage.kt")
         public void testFir2IrClassifierStorage() throws Exception {

compiler/tests-gen/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -12338,6 +12338,11 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
             runTest("compiler/testData/codegen/box/fir/FakeOverrideBuilder.kt");
         }
 
+        @TestMetadata("falsePositiveBoundSmartcast.kt")
+        public void testFalsePositiveBoundSmartcast() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/falsePositiveBoundSmartcast.kt");
+        }
+
         @TestMetadata("Fir2IrClassifierStorage.kt")
         public void testFir2IrClassifierStorage() throws Exception {
             runTest("compiler/testData/codegen/box/fir/Fir2IrClassifierStorage.kt");

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/es6/semantics/IrJsCodegenBoxES6TestGenerated.java

@@ -10945,6 +10945,11 @@ public class IrJsCodegenBoxES6TestGenerated extends AbstractIrJsCodegenBoxES6Tes
         public void testClassCanNotBeCastedToVoid() throws Exception {
             runTest("compiler/testData/codegen/box/fir/classCanNotBeCastedToVoid.kt");
         }
+
+        @TestMetadata("falsePositiveBoundSmartcast.kt")
+        public void testFalsePositiveBoundSmartcast() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/falsePositiveBoundSmartcast.kt");
+        }
     }
 
     @TestMetadata("compiler/testData/codegen/box/fullJdk")

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/ir/semantics/IrJsCodegenBoxTestGenerated.java

@@ -10402,6 +10402,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
         public void testClassCanNotBeCastedToVoid() throws Exception {
             runTest("compiler/testData/codegen/box/fir/classCanNotBeCastedToVoid.kt");
         }
+
+        @TestMetadata("falsePositiveBoundSmartcast.kt")
+        public void testFalsePositiveBoundSmartcast() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/falsePositiveBoundSmartcast.kt");
+        }
     }
 
     @TestMetadata("compiler/testData/codegen/box/fullJdk")

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -10402,6 +10402,11 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
         public void testClassCanNotBeCastedToVoid() throws Exception {
             runTest("compiler/testData/codegen/box/fir/classCanNotBeCastedToVoid.kt");
         }
+
+        @TestMetadata("falsePositiveBoundSmartcast.kt")
+        public void testFalsePositiveBoundSmartcast() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/falsePositiveBoundSmartcast.kt");
+        }
     }
 
     @TestMetadata("compiler/testData/codegen/box/fullJdk")

js/js.tests/tests-gen/org/jetbrains/kotlin/js/test/wasm/semantics/IrCodegenBoxWasmTestGenerated.java

@@ -5233,6 +5233,11 @@ public class IrCodegenBoxWasmTestGenerated extends AbstractIrCodegenBoxWasmTest
         public void testClassCanNotBeCastedToVoid() throws Exception {
             runTest("compiler/testData/codegen/box/fir/classCanNotBeCastedToVoid.kt");
         }
+
+        @TestMetadata("falsePositiveBoundSmartcast.kt")
+        public void testFalsePositiveBoundSmartcast() throws Exception {
+            runTest("compiler/testData/codegen/box/fir/falsePositiveBoundSmartcast.kt");
+        }
     }
 
     @TestMetadata("compiler/testData/codegen/box/fullJdk")