FIR DFA: use PersistentSet.Builder for exactType

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

@@ -641,10 +641,10 @@ abstract class FirDataFlowAnalyzer(
             val newStatement = logicSystem.or(conditionExitAndBreakFlows.map {
                 val atExit = it.getTypeStatement(variable)
                 if (logicSystem.isSameValueIn(conditionEnterFlow, it, variable)) {
-                    if (atExit != null) logicSystem.and(listOf(atExit, toAdd)) else toAdd
+                    logicSystem.and(atExit, toAdd)
                 } else {
-                    atExit
+                    atExit ?: return@forEach
-                } ?: return@forEach
+                }
             }) ?: return@forEach
             flow.addTypeStatement(newStatement)
         }

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

@@ -278,7 +278,7 @@ abstract class LogicSystem(private val context: ConeInferenceContext) {
         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)) ?: continue)
             }
         }
     }
@@ -288,50 +288,62 @@ abstract class LogicSystem(private val context: ConeInferenceContext) {
         right.isEmpty() -> left
         else -> left.toMutableMap().apply {
             for ((variable, rightStatement) in right) {
-                this[variable] = this[variable]?.let { and(listOf(it, rightStatement))!! } ?: rightStatement
+                this[variable] = and(this[variable], rightStatement)
             }
         }
     }
 
-    private inline fun Collection<TypeStatement>.singleOrNew(exactType: () -> MutableSet<ConeKotlinType>): TypeStatement? =
-        when (size) {
-            0 -> null
-            1 -> first()
-            else -> {
-                val variable = first().variable
-                assert(all { it.variable == variable }) { "folding statements for different variables" }
-                MutableTypeStatement(variable, exactType())
-            }
-        }
-
-    private fun unifyTypes(types: Collection<Set<ConeKotlinType>>): ConeKotlinType? {
-        if (types.any { it.isEmpty() }) return null
-        val intersected = types.map { ConeTypeIntersector.intersectTypes(context, it.toList()) }
-        val unified = context.commonSuperTypeOrNull(intersected) ?: return null
-        return when {
-            unified.isNullableAny -> null
-            unified.isAcceptableForSmartcast() -> unified
-            unified.canBeNull -> null
-            else -> context.anyType()
-        }
-    }
-
     private operator fun MutableTypeStatement.plusAssign(other: TypeStatement) {
         exactType += other.exactType
     }
 
-    fun and(statements: Collection<TypeStatement>): TypeStatement? =
+    fun and(a: TypeStatement?, b: TypeStatement): TypeStatement =
-        statements.singleOrNew { statements.flatMapTo(mutableSetOf()) { it.exactType } }
+        a?.toMutable()?.apply { this += b } ?: b
 
-    fun or(statements: Collection<TypeStatement>): TypeStatement? =
+    fun and(statements: Collection<TypeStatement>): TypeStatement? {
-        statements.singleOrNew { unifyTypes(statements.map { it.exactType })?.let { mutableSetOf(it) } ?: mutableSetOf() }
+        when (statements.size) {
+            0 -> return null
+            1 -> return statements.first()
+        }
+        val iterator = statements.iterator()
+        val result = iterator.next().toMutable()
+        while (iterator.hasNext()) {
+            result += iterator.next()
+        }
+        return result
+    }
+
+    fun or(statements: Collection<TypeStatement>): TypeStatement? {
+        when (statements.size) {
+            0 -> return null
+            1 -> return statements.first()
+        }
+        val variable = statements.first().variable
+        assert(statements.all { it.variable == variable }) { "folding statements for different variables" }
+        if (statements.any { it.isEmpty }) return null
+        val intersected = statements.map { ConeTypeIntersector.intersectTypes(context, it.exactType.toList()) }
+        val unified = context.commonSuperTypeOrNull(intersected) ?: return null
+        val result = when {
+            unified.isNullableAny -> return null
+            unified.isAcceptableForSmartcast() -> unified
+            unified.canBeNull -> return null
+            else -> context.anyType()
+        }
+        return PersistentTypeStatement(variable, persistentSetOf(result))
+    }
 }
 
 private fun TypeStatement.toPersistent(): PersistentTypeStatement = when (this) {
     is PersistentTypeStatement -> this
+    // If this statement was obtained via `toMutable`, `toPersistentSet` will call `build`.
     else -> PersistentTypeStatement(variable, exactType.toPersistentSet())
 }
 
+private fun TypeStatement.toMutable(): MutableTypeStatement = when (this) {
+    is PersistentTypeStatement -> MutableTypeStatement(variable, exactType.builder())
+    else -> MutableTypeStatement(variable, LinkedHashSet(exactType))
+}
+
 @JvmName("replaceVariableInStatements")
 private fun MutableMap<RealVariable, PersistentTypeStatement>.replaceVariable(from: RealVariable, to: RealVariable?) {
     val existing = remove(from) ?: return

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

@@ -21,9 +21,7 @@ data class PersistentTypeStatement(
 class MutableTypeStatement(
     override val variable: RealVariable,
     override val exactType: MutableSet<ConeKotlinType> = linkedSetOf(),
-) : TypeStatement() {
+) : TypeStatement()
-    fun copy(): MutableTypeStatement = MutableTypeStatement(variable, LinkedHashSet(exactType))
-}
 
 // --------------------------------------- Aliases ---------------------------------------