FIR DFA: don't assume != true/false => == false/true

This also fixes some returnsNotNull contracts because the old code added
an implication that `== true` => `!= null` then promptly removed any
statement that this could've affected if the argument was a synthetic
variable.

^KT-26612 tag fixed-in-k2

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

@@ -497,46 +497,24 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         }
     }
 
-    // const != null
     private fun processEqWithConst(
         node: EqualityOperatorCallNode, operand: FirExpression, const: FirConstExpression<*>, operation: FirOperation
     ) {
         val isEq = operation.isEq()
-        val expressionVariable = variableStorage.createSyntheticVariable(node.fir)
+        if (const.kind == ConstantValueKind.Null) {
+            return processEqNull(node, operand, isEq)
+        }
+
         val flow = node.flow
+        val expressionVariable = variableStorage.createSyntheticVariable(node.fir)
         val operandVariable = variableStorage.getOrCreateVariable(node.previousFlow, operand)
-        // expression == const -> expression != null
+        // expression == non-null const -> expression != null
         flow.addImplication((expressionVariable eq isEq) implies (operandVariable notEq null))
-
+        if (const.kind == ConstantValueKind.Boolean) {
-        // propagating facts for (... == true) and (... == false)
+            val expected = (const.value as Boolean)
-        when (const.kind) {
+            flow.addImplication((expressionVariable eq isEq) implies (operandVariable eq expected))
-            ConstantValueKind.Boolean -> {
+            if (operand.coneType.isBoolean) {
-                val constValue = const.value as Boolean
+                flow.addImplication((expressionVariable eq !isEq) implies (operandVariable eq !expected))
-                val shouldInvert = isEq xor constValue
-
-                logicSystem.translateVariableFromConditionInStatements(flow, operandVariable, expressionVariable) {
-                    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.invertCondition() else it
-                    }
-                }
-            }
-            ConstantValueKind.Null -> {
-                logicSystem.translateVariableFromConditionInStatements(flow, operandVariable, expressionVariable) {
-                    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.
             }
         }
     }
@@ -546,7 +524,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         val expressionVariable = variableStorage.createSyntheticVariable(node.fir)
         val operandVariable = variableStorage.getOrCreateVariable(node.previousFlow, operand)
         flow.addImplication((expressionVariable eq isEq) implies (operandVariable eq null))
-        flow.addImplication((expressionVariable notEq isEq) implies (operandVariable notEq null))
+        flow.addImplication((expressionVariable eq !isEq) implies (operandVariable notEq null))
     }
 
     private fun processEq(
@@ -961,24 +939,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
                     lastNode.flow.commitOperationStatement(argumentVariable eq true)
                 }
 
-                is ConeBooleanConstantReference -> {
+                else -> {
-                    logicSystem.translateVariableFromConditionInStatements(
-                        lastNode.flow,
-                        argumentVariable,
-                        functionCallVariable,
-                        shouldRemoveOriginalStatements = true,
-                        filter = { it.condition.operation == Operation.EqTrue },
-                        transform = {
-                            when (value) {
-                                ConeBooleanConstantReference.TRUE -> it
-                                ConeBooleanConstantReference.FALSE -> it.invertCondition()
-                                else -> throw IllegalStateException()
-                            }
-                        }
-                    )
-                }
-
-                ConeConstantReference.NOT_NULL, ConeConstantReference.NULL -> {
                     logicSystem.translateVariableFromConditionInStatements(
                         lastNode.flow,
                         argumentVariable,
@@ -988,8 +949,6 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
                         transform = { OperationStatement(it.condition.variable, value.toOperation()) implies it.effect }
                     )
                 }
-
-                else -> throw IllegalArgumentException("Unsupported constant reference: $value")
             }
         }
         graphBuilder.exitContract(qualifiedAccess).mergeIncomingFlow(updateReceivers = true)
@@ -1177,14 +1136,16 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
 
     private fun exitBooleanNot(functionCall: FirFunctionCall, node: FunctionCallNode) {
         val previousFlow = node.previousFlow
-        val booleanExpressionVariable = variableStorage.getOrCreateVariable(previousFlow, node.firstPreviousNode.fir)
+        val argumentVariable = variableStorage.getOrCreateVariable(previousFlow, node.firstPreviousNode.fir)
-        val variable = variableStorage.getOrCreateVariable(previousFlow, functionCall)
+        val expressionVariable = variableStorage.getOrCreateVariable(previousFlow, functionCall)
-        logicSystem.translateVariableFromConditionInStatements(
+        logicSystem.translateVariableFromConditionInStatements(node.flow, argumentVariable, expressionVariable) {
-            node.flow,
+            when (it.condition.operation) {
-            booleanExpressionVariable,
+                Operation.EqTrue -> expressionVariable eq false implies it.effect
-            variable,
+                Operation.EqFalse -> expressionVariable eq true implies it.effect
-            transform = { it.invertCondition() }
+                // `argumentVariable eq/notEq null` shouldn't exist since `argumentVariable` is presumably `Boolean`
-        )
+                else -> null
+            }
+        }
     }
 
     // ----------------------------------- Annotations -----------------------------------

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

@@ -19,31 +19,22 @@ class MutableTypeStatement(
     fun copy(): MutableTypeStatement = MutableTypeStatement(variable, LinkedHashSet(exactType))
 }
 
-fun Implication.invertCondition(): Implication = Implication(condition.invert(), effect)
-
 // --------------------------------------- Aliases ---------------------------------------
 
 typealias TypeStatements = Map<RealVariable, TypeStatement>
 
 // --------------------------------------- DSL ---------------------------------------
 
-infix fun DataFlowVariable.eq(constant: Boolean?): OperationStatement {
+infix fun DataFlowVariable.eq(constant: Boolean): OperationStatement =
-    val condition = when (constant) {
+    OperationStatement(this, if (constant) Operation.EqTrue else Operation.EqFalse)
-        true -> Operation.EqTrue
-        false -> Operation.EqFalse
-        null -> Operation.EqNull
-    }
-    return OperationStatement(this, condition)
-}
 
-infix fun DataFlowVariable.notEq(constant: Boolean?): OperationStatement {
+@Suppress("UNUSED_PARAMETER")
-    val condition = when (constant) {
+infix fun DataFlowVariable.eq(constant: Nothing?): OperationStatement =
-        true -> Operation.EqFalse
+    OperationStatement(this, Operation.EqNull)
-        false -> Operation.EqTrue
+
-        null -> Operation.NotEqNull
+@Suppress("UNUSED_PARAMETER")
-    }
+infix fun DataFlowVariable.notEq(constant: Nothing?): OperationStatement =
-    return OperationStatement(this, condition)
+    OperationStatement(this, Operation.NotEqNull)
-}
 
 infix fun OperationStatement.implies(effect: Statement): Implication = Implication(this, effect)
 

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

@@ -19,10 +19,6 @@ sealed class Statement {
  * d == False
  */
 data class OperationStatement(override val variable: DataFlowVariable, val operation: Operation) : Statement() {
-    fun invert(): OperationStatement {
-        return OperationStatement(variable, operation.invert())
-    }
-
     override fun toString(): String {
         return "$variable $operation"
     }
@@ -55,13 +51,6 @@ class Implication(
 enum class Operation {
     EqTrue, EqFalse, EqNull, NotEqNull;
 
-    fun invert(): Operation = when (this) {
-        EqTrue -> EqFalse
-        EqFalse -> EqTrue
-        EqNull -> NotEqNull
-        NotEqNull -> EqNull
-    }
-
     fun valueIfKnown(given: Operation): Boolean? = when (this) {
         EqTrue, EqFalse -> if (given == NotEqNull) null else given == this
         EqNull -> given == EqNull

compiler/testData/diagnostics/testsWithStdLib/contracts/smartcasts/operatorsTests/equalsWithNullableBoolean.fir.kt

@@ -28,7 +28,7 @@ fun equalsFalse(x: Any?) {
         x.<!UNRESOLVED_REFERENCE!>length<!>
     }
     else {
-        x.length
+        x.<!UNRESOLVED_REFERENCE!>length<!>
     }
 }
 
@@ -52,7 +52,7 @@ fun notEqualsTrue(x: Any?) {
 
 fun notEqualsFalse(x: Any?) {
     if (safeIsString(x) != false) {
-        x.length
+        x.<!UNRESOLVED_REFERENCE!>length<!>
     }
     else {
         x.<!UNRESOLVED_REFERENCE!>length<!>
@@ -66,4 +66,4 @@ fun notEqualsNull(x: Any?) {
     else {
         x.<!UNRESOLVED_REFERENCE!>length<!>
     }
-}
+}

compiler/testData/diagnostics/testsWithStdLib/contracts/smartcasts/when/withSubjectNullableBoolean.fir.kt

@@ -14,7 +14,7 @@ fun safeIsString(x: Any?): Boolean? {
 fun elseWithNullableResult(x: Any?) {
     when (safeIsString(x)) {
         false -> x.<!UNRESOLVED_REFERENCE!>length<!>
-        else -> x.length
+        else -> x.<!UNRESOLVED_REFERENCE!>length<!>
     }
 
     when (safeIsString(x)) {
@@ -45,12 +45,12 @@ fun exhaustiveWithNullableResult(x: Any?) {
     when (safeIsString(x)) {
         false -> x.<!UNRESOLVED_REFERENCE!>length<!>
         true -> x.length
-        null -> x.length
+        null -> x.<!UNRESOLVED_REFERENCE!>length<!>
     }
 
     when (safeIsString(x)) {
         false -> x.<!UNRESOLVED_REFERENCE!>length<!>
-        null -> x.length
+        null -> x.<!UNRESOLVED_REFERENCE!>length<!>
         true -> x.length
     }
-}
+}

compiler/tests-spec/testData/diagnostics/notLinked/contracts/analysis/common/neg/1.fir.kt

@@ -66,7 +66,7 @@ fun case_2(value_1: Int?, value_2: Int?, value_3: Any?) {
             println(value_2)
         }
         null -> {
-            println(value_3?.xor(true))
+            <!OVERLOAD_RESOLUTION_AMBIGUITY!>println<!>(value_3?.<!UNRESOLVED_REFERENCE!>xor<!>(true))
             println(<!UNINITIALIZED_VARIABLE!>value_4<!>)
             println(value_1)
             println(value_2)

compiler/tests-spec/testData/diagnostics/notLinked/contracts/analysis/common/pos/1.fir.kt

@@ -61,7 +61,7 @@ fun case_2(value_1: Int?, value_2: Int?, value_3: Any?) {
             println(value_2)
         }
         null -> {
-            <!OVERLOAD_RESOLUTION_AMBIGUITY!>println<!>(value_3?.<!UNRESOLVED_REFERENCE!>xor<!>(true))
+            println(value_3?.xor(true))
             println(value_4)
             println(value_1)
             println(value_2)

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

@@ -94,7 +94,7 @@ fun case_4(value_1: Number, value_2: (() -> Unit)?) {
     } else if (contracts.case_4(value_1, value_2) == false) {
         println(value_2)
     } else if (contracts.case_4(value_1, value_2) == null) {
-        value_2()
+        <!UNSAFE_IMPLICIT_INVOKE_CALL!>value_2<!>()
     }
 }
 
@@ -107,7 +107,7 @@ fun case_5(value_1: Number?, value_2: String?) {
         }
         false -> {
             println(value_2<!UNSAFE_CALL!>.<!>length)
-            <!OVERLOAD_RESOLUTION_AMBIGUITY!>println<!>(value_1.<!UNRESOLVED_REFERENCE!>inv<!>())
+            println(value_1.inv())
         }
     }
 }
@@ -124,7 +124,7 @@ fun case_6(value_1: Number, value_2: String?, value_3: Any?) {
             println(value_2<!UNSAFE_CALL!>.<!>length)
         }
         null -> {
-            println(value_1.inv())
+            <!OVERLOAD_RESOLUTION_AMBIGUITY!>println<!>(value_1.<!UNRESOLVED_REFERENCE!>inv<!>())
         }
     }
 }

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

@@ -110,7 +110,7 @@ fun case_5(value_1: Number?, value_2: String?) {
             println(value_1.toByte())
         }
         false -> {
-            println(value_2<!UNSAFE_CALL!>.<!>length)
+            println(value_2.length)
             println(value_1.inv())
         }
     }
@@ -129,10 +129,10 @@ fun case_6(value_1: Number, value_2: String?, value_3: Any?) {
         }
         false -> {
             println(value_3.length)
-            println(value_2<!UNSAFE_CALL!>.<!>length)
+            println(value_2.length)
         }
         null -> {
-            <!OVERLOAD_RESOLUTION_AMBIGUITY!>println<!>(value_1.<!UNRESOLVED_REFERENCE!>inv<!>())
+            println(value_1.inv())
         }
     }
 }