FIR DFA: don't update receivers' types if flow is unchanged

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

@@ -5,10 +5,8 @@
 
 package org.jetbrains.kotlin.fir.resolve.dfa
 
-import org.jetbrains.kotlin.fir.types.ConeKotlinType
-
 abstract class Flow {
     abstract val approvedTypeStatements: TypeStatements
     abstract fun unwrapVariable(variable: RealVariable): RealVariable
-    abstract fun getType(variable: RealVariable): Set<ConeKotlinType>?
+    abstract fun getTypeStatement(variable: RealVariable): TypeStatement?
 }

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

@@ -15,7 +15,8 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
     abstract fun unionFlow(flows: Collection<FLOW>): FLOW
 
     // -------------------------------- Flow mutators --------------------------------
-    abstract fun addTypeStatement(flow: FLOW, statement: TypeStatement)
+    // Returns all known information about the variable, or null if unchanged by this statement:
+    abstract fun addTypeStatement(flow: FLOW, statement: TypeStatement): TypeStatement?
     abstract fun addImplication(flow: FLOW, implication: Implication)
     abstract fun addLocalVariableAlias(flow: FLOW, alias: RealVariable, underlyingVariable: RealVariable)
     abstract fun recordNewAssignment(flow: FLOW, variable: RealVariable, index: Int)
@@ -47,7 +48,7 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
         right.isEmpty() -> right
         else -> buildMap {
             for ((variable, leftStatement) in left) {
-                put(variable, or(listOf(leftStatement, right[variable] ?: continue)))
+                put(variable, or(listOf(leftStatement, right[variable] ?: continue))!!)
             }
         }
     }
@@ -57,14 +58,14 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
         right.isEmpty() -> left
         else -> left.toMutableMap().apply {
             for ((variable, rightStatement) in right) {
-                put(variable, { rightStatement }, { and(listOf(it, rightStatement)) })
+                put(variable, { rightStatement }, { and(listOf(it, rightStatement))!! })
             }
         }
     }
 
-    private inline fun Collection<TypeStatement>.singleOrNew(exactType: () -> MutableSet<ConeKotlinType>): TypeStatement =
+    private inline fun Collection<TypeStatement>.singleOrNew(exactType: () -> MutableSet<ConeKotlinType>): TypeStatement? =
         when (size) {
-            0 -> throw AssertionError("need at least one statement")
+            0 -> null
             1 -> first()
             else -> {
                 val variable = first().variable
@@ -84,9 +85,9 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
         }
     }
 
-    protected fun and(statements: Collection<TypeStatement>): TypeStatement =
+    protected fun and(statements: Collection<TypeStatement>): TypeStatement? =
         statements.singleOrNew { statements.flatMapTo(mutableSetOf()) { it.exactType } }
 
-    protected fun or(statements: Collection<TypeStatement>): TypeStatement =
+    protected fun or(statements: Collection<TypeStatement>): TypeStatement? =
         statements.singleOrNew { unifyTypes(statements.map { it.exactType })?.let { mutableSetOf(it) } ?: mutableSetOf() }
 }

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

@@ -64,11 +64,8 @@ class PersistentFlow : Flow {
     override fun unwrapVariable(variable: RealVariable): RealVariable =
         directAliasMap[variable] ?: variable
 
-    fun getTypeStatement(variable: RealVariable): TypeStatement =
-        approvedTypeStatements[unwrapVariable(variable)]?.copy(variable = variable) ?: MutableTypeStatement(variable)
-
-    override fun getType(variable: RealVariable): Set<ConeKotlinType>? =
-        approvedTypeStatements[unwrapVariable(variable)]?.exactType
+    override fun getTypeStatement(variable: RealVariable): TypeStatement? =
+        approvedTypeStatements[unwrapVariable(variable)]?.copy(variable = variable)
 }
 
 abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSystem<PersistentFlow>(context) {
@@ -108,7 +105,7 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         // Computing the statements for these aliases is redundant as the result is equal
         // to the statements for whatever they are aliasing.
         val variables = flows.flatMapTo(mutableSetOf()) { it.approvedTypeStatements.keys + it.directAliasMap.keys } - commonAliases.keys
-        val statements = variables.mapNotNull { variable ->
+        val statements = variables.mapNotNull computeStatement@{ variable ->
             val statement = if (allExecute) {
                 // All input flows execute in some order. If none of the flows reassign, i.e. the only key
                 // in `byAssignment` is `commonFlow`'s assignment index, then all statements are true.
@@ -118,12 +115,14 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
                 if (byAssignment.size > 1) {
                     byAssignment.remove(commonFlow.assignmentIndex[variable] ?: -1)
                 }
-                or(byAssignment.values.map { flowSubset -> and(flowSubset.map { it.getTypeStatement(variable) }) })
+                or(byAssignment.values.map { flowSubset ->
+                    and(flowSubset.mapNotNull { it.getTypeStatement(variable) }) ?: return@computeStatement null
+                })
             } else {
                 // One input flow executes - one set of statements is true, others might be false.
-                or(flows.map { it.getTypeStatement(variable) })
+                or(flows.map { it.getTypeStatement(variable) ?: return@computeStatement null })
             }
-            if (statement.isNotEmpty) variable to statement.toPersistent() else null
+            if (statement?.isEmpty == false) variable to statement.toPersistent() else null
         }
 
         // If a variable was reassigned in one branch, it was reassigned at the join point.
@@ -213,13 +212,15 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         }
     }
 
-    override fun addTypeStatement(flow: PersistentFlow, statement: TypeStatement) {
-        if (statement.exactType.isEmpty()) return
+    override fun addTypeStatement(flow: PersistentFlow, statement: TypeStatement): TypeStatement? {
+        if (statement.exactType.isEmpty()) return null
         val variable = statement.variable
         val oldExactType = flow.approvedTypeStatements[variable]?.exactType
         val newExactType = oldExactType?.addAll(statement.exactType) ?: statement.exactType.toPersistentSet()
-        if (newExactType === oldExactType) return
-        flow.approvedTypeStatements = flow.approvedTypeStatements.put(variable, PersistentTypeStatement(variable, newExactType))
+        if (newExactType === oldExactType) return null
+        val newStatement = PersistentTypeStatement(variable, newExactType)
+        flow.approvedTypeStatements = flow.approvedTypeStatements.put(variable, newStatement)
+        return newStatement
     }
 
     override fun addImplication(flow: PersistentFlow, implication: Implication) {
@@ -283,7 +284,7 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
                 flow.logicStatements = flow.logicStatements.put(variable, stillUnknown)
             }
         }
-        return approvedTypeStatements.asMap().mapValues { and(it.value) }
+        return approvedTypeStatements.asMap().mapValues { and(it.value)!! }
     }
 
     override fun recordNewAssignment(flow: PersistentFlow, variable: RealVariable, index: Int) {