FIR DFA: fix handling of equality operation

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

@@ -483,26 +483,50 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         }
 
         // propagating facts for (... == true) and (... == false)
-        if (const.kind == ConstantValueKind.Boolean) {
-            val constValue = const.value as Boolean
-            val shouldInvert = isEq xor constValue
+        when (const.kind) {
+            ConstantValueKind.Boolean -> {
+                val constValue = const.value as Boolean
+                val shouldInvert = isEq xor constValue
 
-            logicSystem.translateVariableFromConditionInStatements(
-                flow,
-                operandVariable,
-                expressionVariable,
-                shouldRemoveOriginalStatements = operandVariable.isSynthetic()
-            ) {
-                when (it.condition.operation) {
-                    Operation.EqNull, Operation.NotEqNull -> {
-                        (expressionVariable eq isEq) implies (it.effect)
-                    }
-                    Operation.EqTrue, Operation.EqFalse -> {
-                        if (shouldInvert) (it.condition.invert()) implies (it.effect)
-                        else it
+                logicSystem.translateVariableFromConditionInStatements(
+                    flow,
+                    operandVariable,
+                    expressionVariable,
+                    shouldRemoveOriginalStatements = operandVariable.isSynthetic()
+                ) {
+                    when (it.condition.operation) {
+                        // Whatever the result is after comparing operandVariable with `true` or `false` cannot let you imply effects that apply
+                        // when the operandVariable is null. Hence we return null here.
+                        Operation.EqNull -> null
+                        Operation.NotEqNull -> {
+                            (expressionVariable eq isEq) implies (it.effect)
+                        }
+                        Operation.EqTrue, Operation.EqFalse -> {
+                            if (shouldInvert) (it.condition.invert()) implies (it.effect)
+                            else it
+                        }
                     }
                 }
             }
+            ConstantValueKind.Null -> {
+                logicSystem.translateVariableFromConditionInStatements(
+                    flow,
+                    operandVariable,
+                    expressionVariable,
+                    shouldRemoveOriginalStatements = operandVariable.isSynthetic()
+                ) {
+                    when (it.condition.operation) {
+                        Operation.EqNull -> (expressionVariable eq isEq) implies (it.effect)
+                        Operation.NotEqNull -> (expressionVariable eq !isEq) implies (it.effect)
+                        // Whatever the result is after comparing operandVariable with `null` cannot let you imply effects that apply when the
+                        // operandVariable is true or false.
+                        Operation.EqTrue, Operation.EqFalse -> null
+                    }
+                }
+            }
+            else -> {
+                // Inconclusive if the user code compares with other constants.
+            }
         }
     }
 

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

@@ -47,7 +47,7 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
         newVariable: DataFlowVariable,
         shouldRemoveOriginalStatements: Boolean,
         filter: (Implication) -> Boolean = { true },
-        transform: (Implication) -> Implication = { it },
+        transform: (Implication) -> Implication? = { it },
     )
 
     abstract fun approveStatementsInsideFlow(

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

@@ -297,11 +297,11 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         newVariable: DataFlowVariable,
         shouldRemoveOriginalStatements: Boolean,
         filter: (Implication) -> Boolean,
-        transform: (Implication) -> Implication
+        transform: (Implication) -> Implication?
     ) {
         with(flow) {
             val statements = logicStatements[originalVariable]?.takeIf { it.isNotEmpty() } ?: return
-            val newStatements = statements.filter(filter).map {
+            val newStatements = statements.filter(filter).mapNotNull {
                 val newStatement = OperationStatement(newVariable, it.condition.operation) implies it.effect
                 transform(newStatement)
             }.toPersistentList()

compiler/tests-spec/testData/diagnostics/notLinked/contracts/analysis/smartcasts/neg/8.fir.kt

@@ -116,7 +116,7 @@ fun case_5(value_1: Number?, value_2: String?) {
 fun case_6(value_1: Number, value_2: String?, value_3: Any?) {
     when (value_3.case_6(value_1, value_2)) {
         true -> {
-            println(value_3.equals(""))
+            println(value_3<!UNSAFE_CALL!>.<!>equals(""))
             println(value_2<!UNSAFE_CALL!>.<!>length)
         }
         false -> {