FIR DFA: slightly simplify {and,or}ForTypeStatements

And use them less in boolean expression analysis. This uncovers a fun
new case that doesn't work:

    interface I { val x: String }
    class C(override val x: String) : I
    class D(override val x: String) : I

    fun foo(x: Any?) {
      if (x is D || (x as C).x != "") {
        println(x.x.length) // actually OK (x is D || x is C => x is I)
      }
    }

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

@@ -138,10 +138,8 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
         flow: PersistentFlow,
         statement: OperationStatement,
         builtinTypes: BuiltinTypes
-    ): MutableTypeStatements {
+    ): TypeStatements {
-        val newTypeStatements = flow.approvedTypeStatements.asMutableStatements()
+        val newTypeStatements = andForTypeStatements(flow.approvedTypeStatements, approveOperationStatement(flow, statement, null))
-        approveStatementsTo(newTypeStatements, flow, statement, flow.logicStatements.flatMap { it.value })
-
         val variable = statement.variable
         if (!variable.isReal()) return newTypeStatements
         val extraStatement = when (statement.operation) {
@@ -158,7 +156,7 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
         variableStorage: VariableStorageImpl,
         flow: Flow,
         context: CheckerContext
-    ): MutableTypeStatements? {
+    ): TypeStatements? {
         fun buildTypeStatements(arg: ConeValueParameterReference, exactType: Boolean, type: ConeKotlinType): MutableTypeStatements? {
             val parameterSymbol = function.getParameterSymbol(arg.parameterIndex, context)
 

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

@@ -1232,51 +1232,41 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
          * TODO: Here we should handle case when one of arguments is dead (e.g. in cases `false && expr` or `true || expr`)
          *  But since conditions with const are rare it can be delayed
          */
-
         val leftVariable = variableStorage.getOrCreateVariable(flow, binaryLogicExpression.leftOperand)
         val rightVariable = variableStorage.getOrCreateVariable(flow, binaryLogicExpression.rightOperand)
         val operatorVariable = variableStorage.getOrCreateVariable(flow, binaryLogicExpression)
 
         if (!node.leftOperandNode.isDead && node.rightOperandNode.isDead) {
-            /*
+            // If the right operand does not terminate, then we know that the value of the entire expression
-             * If there was a jump from right argument then we know that we well exit from
+            // has to be `onlyLeftEvaluated`, and it has to be produced by the left operand.
-             *   boolean operator only if right operand was not executed
-             *
-             *   a && return => a == false
-             *   a || return => a == true
-             */
             logicSystem.approveStatementsInsideFlow(
                 flow,
-                leftVariable eq !isAnd,
+                leftVariable eq onlyLeftEvaluated,
                 shouldForkFlow = false,
                 shouldRemoveSynthetics = true
             )
         } else {
-            val (conditionalFromLeft, conditionalFromRight, approvedFromRight) = logicSystem.collectInfoForBooleanOperator(
+            val (conditionalFromLeft, conditionalFromRight, approvedFromRight) =
-                flowFromLeft,
+                logicSystem.collectInfoForBooleanOperator(flowFromLeft, leftVariable, flowFromRight, rightVariable)
-                leftVariable,
-                flowFromRight,
-                rightVariable
-            )
 
-            // left && right == True
+            // If `left && right` is true, then both are true (and evaluated).
-            // left || right == False
+            // If `left || right` is false, then both are false.
-            val approvedIfTrue: MutableTypeStatements = mutableMapOf()
+            arrayOf(
-            logicSystem.approveStatementsTo(approvedIfTrue, flowFromRight, leftVariable eq bothEvaluated, conditionalFromLeft)
+                approvedFromRight,
-            logicSystem.approveStatementsTo(approvedIfTrue, flowFromRight, rightVariable eq bothEvaluated, conditionalFromRight)
+                logicSystem.approveOperationStatement(flowFromRight, leftVariable eq bothEvaluated, conditionalFromLeft),
-            logicSystem.andForTypeStatements(approvedIfTrue, approvedFromRight).values.forEach {
+                logicSystem.approveOperationStatement(flowFromRight, rightVariable eq bothEvaluated, conditionalFromRight),
-                flow.addImplication((operatorVariable eq bothEvaluated) implies it)
+            ).forEach { statements ->
+                statements.values.forEach { flow.addImplication((operatorVariable eq bothEvaluated) implies it) }
             }
 
-            // left && right == False
+            // If `left && right` is false, then either `left` is false, or both were evaluated and `right` is false.
-            // left || right == True
+            // If `left || right` is true, then either `left` is true, or both were evaluated and `right` is true.
-            val approvedIfFalse: MutableTypeStatements = mutableMapOf()
+            logicSystem.orForTypeStatements(
-            val leftIsFalse = logicSystem.approveOperationStatement(flowFromLeft, leftVariable eq onlyLeftEvaluated, conditionalFromLeft)
+                logicSystem.approveOperationStatement(flowFromLeft, leftVariable eq onlyLeftEvaluated, conditionalFromLeft),
-            val rightIsFalse =
+                // TODO: and(approvedFromRight, ...)? FE1.0 doesn't seem to handle that correctly either.
-                logicSystem.approveOperationStatement(flowFromRight, rightVariable eq onlyLeftEvaluated, conditionalFromRight)
+                //   if (x is A || whatever(x as B)) { /* x is (A | B) */ }
-            logicSystem.andForTypeStatements(approvedIfFalse, logicSystem.orForTypeStatements(leftIsFalse, rightIsFalse)).values.forEach {
+                logicSystem.approveOperationStatement(flowFromRight, rightVariable eq onlyLeftEvaluated, conditionalFromRight),
-                flow.addImplication((operatorVariable eq onlyLeftEvaluated) implies it)
+            ).values.forEach { flow.addImplication((operatorVariable eq onlyLeftEvaluated) implies it) }
-            }
         }
 
         logicSystem.updateAllReceivers(flow)

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

@@ -65,38 +65,37 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
         rightVariable: DataFlowVariable,
     ): InfoForBooleanOperator
 
-    abstract fun approveStatementsTo(
+    abstract fun approveOperationStatement(
-        destination: MutableTypeStatements,
         flow: FLOW,
         approvedStatement: OperationStatement,
-        statements: Collection<Implication>,
+        statementsForVariable: Collection<Implication>?
-    )
+    ): TypeStatements
 
-    fun orForTypeStatements(left: TypeStatements, right: TypeStatements): MutableTypeStatements {
+    fun orForTypeStatements(left: TypeStatements, right: TypeStatements): TypeStatements = when {
-        if (left.isEmpty() || right.isEmpty()) return mutableMapOf()
+        left.isEmpty() -> left
-        val map = mutableMapOf<RealVariable, MutableTypeStatement>()
+        right.isEmpty() -> right
-        for (variable in left.keys.intersect(right.keys)) {
+        else -> buildMap {
-            val leftStatement = left.getValue(variable)
+            for ((variable, leftStatement) in left) {
-            val rightStatement = right.getValue(variable)
+                put(variable, or(listOf(leftStatement, right[variable] ?: continue)))
-            map[variable] = or(listOf(leftStatement, rightStatement))
+            }
         }
-        return map
     }
 
-    fun andForTypeStatements(left: TypeStatements, right: TypeStatements): MutableTypeStatements {
+    fun andForTypeStatements(left: TypeStatements, right: TypeStatements): TypeStatements = when {
-        if (left.isEmpty() && right.isEmpty()) return mutableMapOf()
+        left.isEmpty() -> right
-        val map = left.asMutableStatements()
+        right.isEmpty() -> left
-        for ((variable, rightStatement) in right) {
+        else -> left.toMutableMap().apply {
-            map[variable] = and(listOfNotNull(map[variable], rightStatement))
+            for ((variable, rightStatement) in right) {
+                put(variable, { rightStatement }, { and(listOf(it, rightStatement)) })
+            }
         }
-        return map
     }
 
     // ------------------------------- Util functions -------------------------------
 
-    private fun foldStatements(statements: Collection<TypeStatement>, all: Boolean): MutableTypeStatement {
+    private fun foldStatements(statements: Collection<TypeStatement>, all: Boolean): TypeStatement {
         require(statements.isNotEmpty())
-        statements.singleOrNull()?.let { return it.asMutableStatement() }
+        statements.singleOrNull()?.let { return it }
         val variable = statements.first().variable
         assert(statements.all { it.variable == variable })
         // TypeStatement(variable, exactType, exactNotType) =
@@ -127,23 +126,13 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
         }.takeIf { !onlyInputTypes || it in intersected }
     }
 
-    protected fun and(statements: Collection<TypeStatement>): MutableTypeStatement =
+    protected fun and(statements: Collection<TypeStatement>): TypeStatement =
         foldStatements(statements, all = true)
 
-    protected fun or(statements: Collection<TypeStatement>): MutableTypeStatement =
+    protected fun or(statements: Collection<TypeStatement>): TypeStatement =
         foldStatements(statements, all = false)
 }
 
-fun <FLOW : Flow> LogicSystem<FLOW>.approveOperationStatement(
-    flow: FLOW,
-    approvedStatement: OperationStatement,
-    statements: Collection<Implication>,
-): MutableTypeStatements {
-    return mutableMapOf<RealVariable, MutableTypeStatement>().apply {
-        approveStatementsTo(this, flow, approvedStatement, statements)
-    }
-}
-
 /*
  *  used for:
  *   1. val b = x is String

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

@@ -339,18 +339,13 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         }
     }
 
-    override fun approveStatementsTo(
+    override fun approveOperationStatement(
-        destination: MutableTypeStatements,
         flow: PersistentFlow,
         approvedStatement: OperationStatement,
-        statements: Collection<Implication>
+        statementsForVariable: Collection<Implication>?
-    ) {
+    ): TypeStatements =
-        val approveOperationStatements =
+        approveOperationStatementsInternal(flow, approvedStatement, statementsForVariable, shouldRemoveSynthetics = false)
-            approveOperationStatementsInternal(flow, approvedStatement, statements, shouldRemoveSynthetics = false)
+            .asMap().mapValues { and(it.value) }
-        approveOperationStatements.asMap().forEach { (variable, infos) ->
-            destination.put(variable, { and(infos) }, { and(listOf(it) + infos) })
-        }
-    }
 
     override fun collectInfoForBooleanOperator(
         leftFlow: PersistentFlow,

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

@@ -35,9 +35,6 @@ fun Implication.invertCondition(): Implication = Implication(condition.invert(),
 typealias TypeStatements = Map<RealVariable, TypeStatement>
 typealias MutableTypeStatements = MutableMap<RealVariable, MutableTypeStatement>
 
-fun TypeStatements.asMutableStatements(): MutableTypeStatements =
-    mapValuesTo(mutableMapOf()) { it.value.asMutableStatement() }
-
 // --------------------------------------- DSL ---------------------------------------
 
 infix fun DataFlowVariable.eq(constant: Boolean?): OperationStatement {