FIR DFA: avoid inverting type statements

By necessity, type statements underapproximate. Therefore, inverting
doesn't always produce an equal result. It's best to propagate negations
inwards to preserve precision as much as possible.

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/FirReturnsImpliesAnalyzer.kt

@@ -143,8 +143,8 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
         val variable = statement.variable
         if (!variable.isReal()) return newTypeStatements
         val extraStatement = when (statement.operation) {
-            Operation.NotEqNull -> simpleTypeStatement(variable, true, builtinTypes.anyType.type)
+            Operation.NotEqNull -> variable.nullabilityStatement(builtinTypes, isNull = false)
-            Operation.EqNull -> simpleTypeStatement(variable, false, builtinTypes.anyType.type)
+            Operation.EqNull -> variable.nullabilityStatement(builtinTypes, isNull = true)
             else -> return newTypeStatements
         }
         return andForTypeStatements(newTypeStatements, mapOf(variable to extraStatement))
@@ -157,39 +157,50 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
         flow: Flow,
         context: CheckerContext
     ): TypeStatements? {
-        fun buildTypeStatements(arg: ConeValueParameterReference, exactType: Boolean, type: ConeKotlinType): MutableTypeStatements? {
+        fun getOrCreateRealVariable(arg: ConeValueParameterReference): RealVariable? {
             val parameterSymbol = function.getParameterSymbol(arg.parameterIndex, context)
 
             @OptIn(SymbolInternals::class)
             val parameter = parameterSymbol.fir
-            val realVar = variableStorage.getOrCreateRealVariable(flow, parameterSymbol, parameter)
+            return variableStorage.getOrCreateRealVariable(flow, parameterSymbol, parameter)?.takeIf {
-                ?.takeIf {
+                it.stability == PropertyStability.STABLE_VALUE ||
-                    it.stability == PropertyStability.STABLE_VALUE ||
+                        // TODO: consider removing the part below
-                            // TODO: consider removing the part below
+                        it.stability == PropertyStability.LOCAL_VAR
-                            it.stability == PropertyStability.LOCAL_VAR
-                }
-            return realVar?.to(simpleTypeStatement(realVar, exactType, type))?.let { mutableMapOf(it) }
-        }
-        return when (this) {
-            is ConeBinaryLogicExpression -> {
-                val left = left.buildTypeStatements(function, logicSystem, variableStorage, flow, context)
-                val right = right.buildTypeStatements(function, logicSystem, variableStorage, flow, context)
-                if (left != null && right != null) {
-                    if (kind == LogicOperationKind.AND)
-                        logicSystem.andForTypeStatements(left, right)
-                    else
-                        logicSystem.orForTypeStatements(left, right)
-                } else (left ?: right)
             }
-            is ConeIsInstancePredicate -> buildTypeStatements(arg, !isNegated, type)
+        }
-            is ConeIsNullPredicate -> buildTypeStatements(arg, isNegated, context.session.builtinTypes.anyType.type)
-            is ConeLogicalNot -> arg.buildTypeStatements(function, logicSystem, variableStorage, flow, context)
-                ?.mapValuesTo(mutableMapOf()) { (_, value) -> value.invert() }
 
+        fun ConeBooleanExpression.toTypeStatements(inverted: Boolean): TypeStatements? = when (this) {
+            is ConeBinaryLogicExpression -> {
+                val left = left.toTypeStatements(inverted)
+                val right = right.toTypeStatements(inverted)
+                when {
+                    left == null -> right
+                    right == null -> left
+                    (kind == LogicOperationKind.AND) == !inverted -> logicSystem.andForTypeStatements(left, right)
+                    else -> logicSystem.orForTypeStatements(left, right)
+                }
+            }
+            is ConeIsInstancePredicate ->
+                getOrCreateRealVariable(arg)?.let {
+                    if (isNegated == inverted) it typeEq type else it typeNotEq type
+                }?.singleton()
+            is ConeIsNullPredicate ->
+                getOrCreateRealVariable(arg)?.nullabilityStatement(context.session.builtinTypes, isNull = isNegated == inverted)
+                    ?.singleton()
+            is ConeLogicalNot -> arg.toTypeStatements(!inverted)
             else -> null
         }
+
+        return toTypeStatements(inverted = false)
     }
 
+    private fun RealVariable.nullabilityStatement(builtinTypes: BuiltinTypes, isNull: Boolean) =
+        // TODO: opposite for builtinTypes.nullableNothingType.type
+        if (isNull) this typeNotEq builtinTypes.anyType.type else this typeEq builtinTypes.anyType.type
+
+    private fun TypeStatement.singleton(): TypeStatements =
+        mapOf(variable to this)
+
     private fun ConeKotlinType.isInapplicableWith(operation: Operation, session: FirSession): Boolean {
         return (operation == Operation.EqFalse || operation == Operation.EqTrue)
                 && !AbstractTypeChecker.isSubtypeOf(session.typeContext, session.builtinTypes.booleanType.type, this)
@@ -203,14 +214,6 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
         else -> true
     }
 
-    private fun simpleTypeStatement(realVar: RealVariable, exactType: Boolean, type: ConeKotlinType): MutableTypeStatement {
-        return MutableTypeStatement(
-            realVar,
-            if (exactType) linkedSetOf(type) else linkedSetOf(),
-            if (!exactType) linkedSetOf(type) else linkedSetOf()
-        )
-    }
-
     private fun CFGNode<*>.collectBranchExits(nodes: MutableList<CFGNode<*>> = mutableListOf()): List<CFGNode<*>> {
         if (this is BlockExitNode) {
             nodes += previousCfgNodes

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

@@ -541,8 +541,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
                             (expressionVariable eq isEq) implies (it.effect)
                         }
                         Operation.EqTrue, Operation.EqFalse -> {
-                            if (shouldInvert) (it.condition.invert()) implies (it.effect)
+                            if (shouldInvert) it.invertCondition() else it
-                            else it
                         }
                     }
                 }
@@ -937,6 +936,8 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
             is RealVariable -> variable.explicitReceiverVariable ?: return
             is SyntheticVariable -> variableStorage.getOrCreateVariable(flow, safeCall.receiver)
         }
+        // TODO? if the callee has non-null return type, then (variable eq null) implies (receiverVariable eq null)
+        //   if (x?.toString() == null) { /* x == null */ }
         flow.addImplication((variable notEq null) implies (receiverVariable notEq null))
     }
 

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

@@ -17,10 +17,7 @@ data class PersistentTypeStatement(
     override val variable: RealVariable,
     override val exactType: PersistentSet<ConeKotlinType>,
     override val exactNotType: PersistentSet<ConeKotlinType>
-) : TypeStatement() {
+) : TypeStatement()
-    override fun invert(): PersistentTypeStatement =
-        PersistentTypeStatement(variable, exactNotType, exactType)
-}
 
 typealias PersistentApprovedTypeStatements = PersistentMap<RealVariable, PersistentTypeStatement>
 typealias PersistentImplications = PersistentMap<DataFlowVariable, PersistentList<Implication>>
@@ -397,11 +394,6 @@ private fun TypeStatement.toPersistent(): PersistentTypeStatement = when (this)
     else -> PersistentTypeStatement(variable, exactType.toPersistentSet(), exactNotType.toPersistentSet())
 }
 
-fun TypeStatement.asMutableStatement(): MutableTypeStatement = when (this) {
-    is MutableTypeStatement -> this
-    else -> MutableTypeStatement(variable, exactType.toMutableSet(), exactNotType.toMutableSet())
-}
-
 @JvmName("replaceVariableInStatements")
 private fun PersistentApprovedTypeStatements.replaceVariable(from: RealVariable, to: RealVariable?): PersistentApprovedTypeStatements {
     val existing = this[from] ?: return this
@@ -440,10 +432,9 @@ private fun Implication.replaceVariable(from: RealVariable, to: RealVariable): I
     else -> this
 }
 
-private fun Statement<*>.replaceVariable(from: RealVariable, to: RealVariable): Statement<*> =
+private fun Statement.replaceVariable(from: RealVariable, to: RealVariable): Statement =
     if (variable != from) this else when (this) {
         is OperationStatement -> copy(variable = to)
         is PersistentTypeStatement -> copy(variable = to)
         is MutableTypeStatement -> MutableTypeStatement(to, exactType, exactNotType)
-        else -> throw IllegalArgumentException("unknown type of statement $this")
     }

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

@@ -17,14 +17,6 @@ class MutableTypeStatement(
     override val exactType: MutableSet<ConeKotlinType> = linkedSetOf(),
     override val exactNotType: MutableSet<ConeKotlinType> = linkedSetOf()
 ) : TypeStatement() {
-    override fun invert(): MutableTypeStatement {
-        return MutableTypeStatement(
-            variable,
-            LinkedHashSet(exactNotType),
-            LinkedHashSet(exactType)
-        )
-    }
-
     fun copy(): MutableTypeStatement = MutableTypeStatement(variable, LinkedHashSet(exactType), LinkedHashSet(exactNotType))
 }
 
@@ -33,7 +25,6 @@ fun Implication.invertCondition(): Implication = Implication(condition.invert(),
 // --------------------------------------- Aliases ---------------------------------------
 
 typealias TypeStatements = Map<RealVariable, TypeStatement>
-typealias MutableTypeStatements = MutableMap<RealVariable, MutableTypeStatement>
 
 // --------------------------------------- DSL ---------------------------------------
 
@@ -55,7 +46,7 @@ infix fun DataFlowVariable.notEq(constant: Boolean?): OperationStatement {
     return OperationStatement(this, condition)
 }
 
-infix fun OperationStatement.implies(effect: Statement<*>): Implication = Implication(this, effect)
+infix fun OperationStatement.implies(effect: Statement): Implication = Implication(this, effect)
 
 infix fun RealVariable.typeEq(type: ConeKotlinType): MutableTypeStatement =
     MutableTypeStatement(this) andTypeEq type

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

@@ -7,8 +7,7 @@ package org.jetbrains.kotlin.fir.resolve.dfa
 
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 
-sealed class Statement<T : Statement<T>> {
+sealed class Statement {
-    abstract fun invert(): T
     abstract val variable: DataFlowVariable
 }
 
@@ -19,8 +18,8 @@ sealed class Statement<T : Statement<T>> {
  * d == True
  * d == False
  */
-data class OperationStatement(override val variable: DataFlowVariable, val operation: Operation) : Statement<OperationStatement>() {
+data class OperationStatement(override val variable: DataFlowVariable, val operation: Operation) : Statement() {
-    override fun invert(): OperationStatement {
+    fun invert(): OperationStatement {
         return OperationStatement(variable, operation.invert())
     }
 
@@ -29,7 +28,7 @@ data class OperationStatement(override val variable: DataFlowVariable, val opera
     }
 }
 
-abstract class TypeStatement : Statement<TypeStatement>() {
+sealed class TypeStatement : Statement() {
     abstract override val variable: RealVariable
     abstract val exactType: Set<ConeKotlinType>
     abstract val exactNotType: Set<ConeKotlinType>
@@ -47,7 +46,7 @@ abstract class TypeStatement : Statement<TypeStatement>() {
 
 class Implication(
     val condition: OperationStatement,
-    val effect: Statement<*>
+    val effect: Statement
 ) {
     override fun toString(): String {
         return "$condition -> $effect"