FIR DFA: bring into compliance with De Morgan's laws

I don't even know if this affects any real code, but it did uncover
some deeper issues, like the fact that `!is` did not add any statements
when `true`.

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

@@ -139,21 +139,17 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
         statement: OperationStatement,
         builtinTypes: BuiltinTypes
     ): MutableTypeStatements {
-        val newTypeStatements: MutableTypeStatements = mutableMapOf()
+        val newTypeStatements = flow.approvedTypeStatements.asMutableStatements()
-
         approveStatementsTo(newTypeStatements, flow, statement, flow.logicStatements.flatMap { it.value })
-        newTypeStatements.mergeTypeStatements(flow.approvedTypeStatements)
 
         val variable = statement.variable
-        if (variable.isReal()) {
+        if (!variable.isReal()) return newTypeStatements
-            if (statement.operation == Operation.NotEqNull) {
+        val extraStatement = when (statement.operation) {
-                newTypeStatements.addStatement(variable, simpleTypeStatement(variable, true, builtinTypes.anyType.type))
+            Operation.NotEqNull -> simpleTypeStatement(variable, true, builtinTypes.anyType.type)
-            } else if (statement.operation == Operation.EqNull) {
+            Operation.EqNull -> simpleTypeStatement(variable, false, builtinTypes.anyType.type)
-                newTypeStatements.addStatement(variable, simpleTypeStatement(variable, false, builtinTypes.anyType.type))
+            else -> return newTypeStatements
-            }
         }
-
+        return andForTypeStatements(newTypeStatements, mapOf(variable to extraStatement))
-        return newTypeStatements
     }
 
     private fun ConeBooleanExpression.buildTypeStatements(
@@ -181,9 +177,10 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
                 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) {
+                    if (kind == LogicOperationKind.AND)
-                        left.apply { mergeTypeStatements(right) }
+                        logicSystem.andForTypeStatements(left, right)
-                    } else logicSystem.orForTypeStatements(left, right)
+                    else
+                        logicSystem.orForTypeStatements(left, right)
                 } else (left ?: right)
             }
             is ConeIsInstancePredicate -> buildTypeStatements(arg, !isNegated, type)

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

@@ -632,6 +632,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
 
         // left == right && right not null -> left != null
         // [processEqNull] adds both implications: operator call could be true or false. We definitely need the matched case only.
+        // TODO: this is incomprehensible - the comments below say what the equivalent expression is
         fun shouldAddImplicationForStatement(operationStatement: OperationStatement): Boolean {
             if (!checkAddImplicationForStatement) return true
             // Only if operation statement is == True, i.e., left == right
@@ -639,17 +640,22 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
             return !isEq && operationStatementOp == Operation.EqTrue || isEq && operationStatementOp == Operation.EqFalse
         }
 
-        logicSystem.approveOperationStatement(flow, predicate).forEach { effect ->
+        // !checkAddImplicationForStatement || !isEq
-            if (shouldAddImplicationForStatement(expressionVariable eq true)) {
+        if (shouldAddImplicationForStatement(expressionVariable eq true)) {
+            logicSystem.approveOperationStatement(flow, predicate).forEach { effect ->
                 flow.addImplication((expressionVariable eq true) implies effect)
             }
-            if (shouldAddImplicationForStatement(expressionVariable eq false)) {
-                flow.addImplication((expressionVariable eq false) implies effect.invert())
-            }
         }
 
-        val expressionVariableIsEq = shouldAddImplicationForStatement(expressionVariable eq isEq)
+        // !checkAddImplicationForStatement || isEq
-        val expressionVariableIsNotEq = shouldAddImplicationForStatement(expressionVariable notEq isEq)
+        if (shouldAddImplicationForStatement(expressionVariable eq false)) {
+            logicSystem.approveOperationStatement(flow, predicate.invert()).forEach { effect ->
+                flow.addImplication((expressionVariable eq false) implies effect)
+            }
+        }
+
+        val expressionVariableIsEq = shouldAddImplicationForStatement(expressionVariable eq isEq) // !checkAddImplicationForStatement
+        val expressionVariableIsNotEq = shouldAddImplicationForStatement(expressionVariable notEq isEq) // true
 
         if (expressionVariableIsEq) {
             flow.addImplication((expressionVariable eq isEq) implies (operandVariable eq null))
@@ -665,10 +671,11 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
             if (expressionVariableIsNotEq) {
                 flow.addImplication((expressionVariable notEq isEq) implies (operandVariable typeEq any))
             }
-
+            // true
             if (shouldAddImplicationForStatement(expressionVariable eq !isEq)) {
                 flow.addImplication((expressionVariable eq !isEq) implies (operandVariable typeNotEq nullableNothing))
             }
+            // !checkAddImplicationForStatement
             if (shouldAddImplicationForStatement(expressionVariable notEq !isEq)) {
                 flow.addImplication((expressionVariable notEq !isEq) implies (operandVariable typeEq nullableNothing))
             }
@@ -1374,11 +1381,8 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
             val approvedIfTrue: MutableTypeStatements = mutableMapOf()
             logicSystem.approveStatementsTo(approvedIfTrue, flowFromRight, leftVariable eq bothEvaluated, conditionalFromLeft)
             logicSystem.approveStatementsTo(approvedIfTrue, flowFromRight, rightVariable eq bothEvaluated, conditionalFromRight)
-            approvedFromRight.forEach { (variable, info) ->
+            logicSystem.andForTypeStatements(approvedIfTrue, approvedFromRight).values.forEach {
-                approvedIfTrue.addStatement(variable, info)
+                flow.addImplication((operatorVariable eq bothEvaluated) implies it)
-            }
-            approvedIfTrue.values.forEach { info ->
-                flow.addImplication((operatorVariable eq bothEvaluated) implies info)
             }
 
             // left && right == False
@@ -1387,9 +1391,8 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
             val leftIsFalse = logicSystem.approveOperationStatement(flowFromLeft, leftVariable eq onlyLeftEvaluated, conditionalFromLeft)
             val rightIsFalse =
                 logicSystem.approveOperationStatement(flowFromRight, rightVariable eq onlyLeftEvaluated, conditionalFromRight)
-            approvedIfFalse.mergeTypeStatements(logicSystem.orForTypeStatements(leftIsFalse, rightIsFalse))
+            logicSystem.andForTypeStatements(approvedIfFalse, logicSystem.orForTypeStatements(leftIsFalse, rightIsFalse)).values.forEach {
-            approvedIfFalse.values.forEach { info ->
+                flow.addImplication((operatorVariable eq onlyLeftEvaluated) implies it)
-                flow.addImplication((operatorVariable eq onlyLeftEvaluated) implies info)
             }
         }
 

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

@@ -5,10 +5,7 @@
 
 package org.jetbrains.kotlin.fir.resolve.dfa
 
-import org.jetbrains.kotlin.fir.types.ConeInferenceContext
+import org.jetbrains.kotlin.fir.types.*
-import org.jetbrains.kotlin.fir.types.ConeKotlinType
-import org.jetbrains.kotlin.fir.types.canBeNull
-import org.jetbrains.kotlin.fir.types.commonSuperTypeOrNull
 
 abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceContext) {
     // ------------------------------- Flow operations -------------------------------
@@ -85,11 +82,8 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
         return approveOperationStatement(flow, approvedStatement, statements).values
     }
 
-    fun orForTypeStatements(
+    fun orForTypeStatements(left: TypeStatements, right: TypeStatements): MutableTypeStatements {
-        left: TypeStatements,
+        if (left.isEmpty() || right.isEmpty()) return mutableMapOf()
-        right: TypeStatements,
-    ): MutableTypeStatements {
-        if (left.isNullOrEmpty() || right.isNullOrEmpty()) return mutableMapOf()
         val map = mutableMapOf<RealVariable, MutableTypeStatement>()
         for (variable in left.keys.intersect(right.keys)) {
             val leftStatement = left.getValue(variable)
@@ -99,63 +93,55 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
         return map
     }
 
-    // ------------------------------- Util functions -------------------------------
+    fun andForTypeStatements(left: TypeStatements, right: TypeStatements): MutableTypeStatements {
-
+        if (left.isEmpty() && right.isEmpty()) return mutableMapOf()
-    // TODO
+        val map = left.asMutableStatements()
-    protected fun <E> Collection<Collection<E>>.intersectSets(): Set<E> {
+        for ((variable, rightStatement) in right) {
-        if (isEmpty()) return emptySet()
+            map[variable] = and(listOfNotNull(map[variable], rightStatement))
-        val iterator = iterator()
-        val result = LinkedHashSet<E>(iterator.next())
-        while (iterator.hasNext()) {
-            result.retainAll(iterator.next())
         }
-        return result
+        return map
     }
 
-    private inline fun manipulateTypeStatements(
+    // ------------------------------- Util functions -------------------------------
-        statements: Collection<TypeStatement>,
+
-        op: (Collection<Set<ConeKotlinType>>) -> MutableSet<ConeKotlinType>
+    private fun foldStatements(statements: Collection<TypeStatement>, all: Boolean): MutableTypeStatement {
-    ): MutableTypeStatement {
         require(statements.isNotEmpty())
         statements.singleOrNull()?.let { return it.asMutableStatement() }
         val variable = statements.first().variable
         assert(statements.all { it.variable == variable })
-        val exactType = op.invoke(statements.map { it.exactType })
+        // TypeStatement(variable, exactType, exactNotType) =
-        val exactNotType = op.invoke(statements.map { it.exactNotType })
+        //   variable is intersect(exactType) && variable !is intersect(exactNotType)
-        return MutableTypeStatement(variable, exactType, exactNotType)
+        // So `and` of two type statements computes `and` of exactType and `or` of `exactNotType`,
+        // while `or` is the opposite.
+        return if (all) {
+            val exactType = statements.flatMapTo(mutableSetOf()) { it.exactType }
+            // variable !is a && variable !is b =/=> variable !is commonSuperType(a, b)
+            // So in this case we can only take the union if either type is a subtype of the other.
+            val exactNotType = unifyTypes(statements.map { it.exactNotType }, onlyInputTypes = true)
+            MutableTypeStatement(variable, exactType, exactNotType?.let { mutableSetOf(it) } ?: mutableSetOf())
+        } else {
+            val exactType = unifyTypes(statements.map { it.exactType }, onlyInputTypes = false)
+            val exactNotType = statements.flatMapTo(mutableSetOf()) { it.exactNotType }
+            MutableTypeStatement(variable, exactType?.let { mutableSetOf(it) } ?: mutableSetOf(), exactNotType)
+        }
     }
 
-    protected fun or(statements: Collection<TypeStatement>): MutableTypeStatement =
+    private fun unifyTypes(types: Collection<Set<ConeKotlinType>>, onlyInputTypes: Boolean): ConeKotlinType? {
-        manipulateTypeStatements(statements, ::orForTypes)
+        if (types.any { it.isEmpty() }) return null
-
+        val intersected = types.map { ConeTypeIntersector.intersectTypes(context, it.toList()) }
-    private fun orForTypes(types: Collection<Set<ConeKotlinType>>): MutableSet<ConeKotlinType> {
+        val unified = context.commonSuperTypeOrNull(intersected) ?: return null
-        if (types.any { it.isEmpty() }) return mutableSetOf()
+        return when {
-        val intersectedTypes = types.map {
+            unified.isAcceptableForSmartcast() -> unified
-            if (it.size > 1) {
+            unified.canBeNull -> null
-                context.intersectTypes(it.toList())
+            else -> context.anyType()
-            } else {
+        }.takeIf { !onlyInputTypes || it in intersected }
-                assert(it.size == 1) { "We've already checked each set of types is not empty." }
-                it.single()
-            }
-        }
-        val result = mutableSetOf<ConeKotlinType>()
-        context.commonSuperTypeOrNull(intersectedTypes)?.let {
-            if (it.isAcceptableForSmartcast()) {
-                result.add(it)
-            } else if (!it.canBeNull) {
-                result.add(context.anyType())
-            }
-            Unit
-        }
-        return result
     }
 
     protected fun and(statements: Collection<TypeStatement>): MutableTypeStatement =
-        manipulateTypeStatements(statements, ::andForTypes)
+        foldStatements(statements, all = true)
 
-    private fun andForTypes(types: Collection<Set<ConeKotlinType>>): MutableSet<ConeKotlinType> {
+    protected fun or(statements: Collection<TypeStatement>): MutableTypeStatement =
-        return types.flatMapTo(mutableSetOf()) { it }
+        foldStatements(statements, all = false)
-    }
 }
 
 fun <FLOW : Flow> LogicSystem<FLOW>.approveOperationStatement(

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

@@ -18,20 +18,8 @@ data class PersistentTypeStatement(
     override val exactType: PersistentSet<ConeKotlinType>,
     override val exactNotType: PersistentSet<ConeKotlinType>
 ) : TypeStatement() {
-    override operator fun plus(other: TypeStatement): PersistentTypeStatement {
+    override fun invert(): PersistentTypeStatement =
-        return PersistentTypeStatement(
+        PersistentTypeStatement(variable, exactNotType, exactType)
-            variable,
-            exactType + other.exactType,
-            exactNotType + other.exactNotType
-        )
-    }
-
-    override val isEmpty: Boolean
-        get() = exactType.isEmpty() && exactNotType.isEmpty()
-
-    override fun invert(): PersistentTypeStatement {
-        return PersistentTypeStatement(variable, exactNotType, exactType)
-    }
 }
 
 typealias PersistentApprovedTypeStatements = PersistentMap<RealVariable, PersistentTypeStatement>
@@ -242,6 +230,9 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         }
     }
 
+    private fun PersistentApprovedTypeStatements.addTypeStatement(info: TypeStatement): PersistentApprovedTypeStatements =
+        put(info.variable, { info.toPersistent() }, { and(listOf(it, info)).toPersistent() })
+
     override fun addImplication(flow: PersistentFlow, implication: Implication) {
         if ((implication.effect as? TypeStatement)?.isEmpty == true) return
         if (implication.condition == implication.effect) return
@@ -288,14 +279,10 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
 
         val updatedReceivers = mutableSetOf<RealVariable>()
         approvedFacts.asMap().forEach { (variable, infos) ->
-            var resultInfo = PersistentTypeStatement(variable, persistentSetOf(), persistentSetOf())
-            for (info in infos) {
-                resultInfo += info
-            }
             if (variable.isThisReference) {
                 updatedReceivers += variable
             }
-            addTypeStatement(resultFlow, resultInfo)
+            addTypeStatement(resultFlow, and(infos))
         }
 
         updatedReceivers.forEach {
@@ -361,13 +348,7 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         val approveOperationStatements =
             approveOperationStatementsInternal(flow, approvedStatement, statements, shouldRemoveSynthetics = false)
         approveOperationStatements.asMap().forEach { (variable, infos) ->
-            for (info in infos) {
+            destination.put(variable, { and(infos) }, { and(listOf(it) + infos) })
-                val mutableInfo = info.asMutableStatement()
-                destination.put(variable, mutableInfo) {
-                    it += mutableInfo
-                    it
-                }
-            }
         }
     }
 
@@ -416,9 +397,6 @@ private fun lowestCommonFlow(left: PersistentFlow, right: PersistentFlow): Persi
     return left
 }
 
-private fun PersistentApprovedTypeStatements.addTypeStatement(info: TypeStatement): PersistentApprovedTypeStatements =
-    put(info.variable, { info.toPersistent() }, { it + info })
-
 private fun TypeStatement.toPersistent(): PersistentTypeStatement = when (this) {
     is PersistentTypeStatement -> this
     else -> PersistentTypeStatement(variable, exactType.toPersistentSet(), exactNotType.toPersistentSet())

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

@@ -12,22 +12,11 @@ import kotlin.contracts.contract
 
 // --------------------------------------- Facts ---------------------------------------
 
-operator fun TypeStatement.plus(other: TypeStatement?): TypeStatement = other?.let { this + other } ?: this
-
 class MutableTypeStatement(
     override val variable: RealVariable,
     override val exactType: MutableSet<ConeKotlinType> = linkedSetOf(),
     override val exactNotType: MutableSet<ConeKotlinType> = linkedSetOf()
 ) : TypeStatement() {
-    override fun plus(other: TypeStatement): MutableTypeStatement = MutableTypeStatement(
-        variable,
-        LinkedHashSet(exactType).apply { addAll(other.exactType) },
-        LinkedHashSet(exactNotType).apply { addAll(other.exactNotType) }
-    )
-
-    override val isEmpty: Boolean
-        get() = exactType.isEmpty() && exactType.isEmpty()
-
     override fun invert(): MutableTypeStatement {
         return MutableTypeStatement(
             variable,
@@ -36,11 +25,6 @@ class MutableTypeStatement(
         )
     }
 
-    operator fun plusAssign(info: TypeStatement) {
-        exactType += info.exactType
-        exactNotType += info.exactNotType
-    }
-
     fun copy(): MutableTypeStatement = MutableTypeStatement(variable, LinkedHashSet(exactType), LinkedHashSet(exactNotType))
 }
 
@@ -51,17 +35,8 @@ fun Implication.invertCondition(): Implication = Implication(condition.invert(),
 typealias TypeStatements = Map<RealVariable, TypeStatement>
 typealias MutableTypeStatements = MutableMap<RealVariable, MutableTypeStatement>
 
-typealias MutableOperationStatements = MutableMap<RealVariable, MutableTypeStatement>
+fun TypeStatements.asMutableStatements(): MutableTypeStatements =
-
+    mapValuesTo(mutableMapOf()) { it.value.asMutableStatement() }
-fun MutableTypeStatements.addStatement(variable: RealVariable, statement: TypeStatement) {
-    put(variable, statement.asMutableStatement()) { it.apply { this += statement } }
-}
-
-fun MutableTypeStatements.mergeTypeStatements(other: TypeStatements) {
-    other.forEach { (variable, info) ->
-        addStatement(variable, info)
-    }
-}
 
 // --------------------------------------- DSL ---------------------------------------
 

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

@@ -34,9 +34,11 @@ abstract class TypeStatement : Statement<TypeStatement>() {
     abstract val exactType: Set<ConeKotlinType>
     abstract val exactNotType: Set<ConeKotlinType>
 
-    abstract operator fun plus(other: TypeStatement): TypeStatement
+    val isEmpty: Boolean
-    abstract val isEmpty: Boolean
+        get() = exactType.isEmpty() && exactNotType.isEmpty()
-    val isNotEmpty: Boolean get() = !isEmpty
+
+    val isNotEmpty: Boolean
+        get() = !isEmpty
 
     override fun toString(): String {
         return "$variable: $exactType, $exactNotType"

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

@@ -26,13 +26,14 @@ import kotlin.contracts.InvocationKind
 import kotlin.contracts.contract
 
 @OptIn(ExperimentalContracts::class)
-internal inline fun <K, V> MutableMap<K, V>.put(key: K, value: V, remappingFunction: (existing: V) -> V) {
+internal inline fun <K, V> MutableMap<K, V>.put(key: K, valueProducer: () -> V, remappingFunction: (existing: V) -> V) {
     contract {
         callsInPlace(remappingFunction, InvocationKind.AT_MOST_ONCE)
+        callsInPlace(valueProducer, InvocationKind.AT_MOST_ONCE)
     }
     val existing = this[key]
     if (existing == null) {
-        put(key, value)
+        put(key, valueProducer())
     } else {
         put(key, remappingFunction(existing))
     }