FIR DFA: split Flow into PersistentFlow and MutableFlow

The distinction is similar to persistent data structures and their
builders. Only the MutableFlow can be passed to LogicSystem for
modification; when it's ready, it can be converted into PersistentFlow
and attached to a CFG node.

The result is that the API is cleaner, the implementation is a bit more
neat, and hopefully the use of PersistentHashMap.Builder improves
performance a little. Also the node-to-flow map can now be removed in
favor of just storing PersistentFlow inside a node, seeing as it's
explicitly immutable and all that.

analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/resolver/resolverUtils.kt

@@ -47,9 +47,9 @@ internal open class StubBodyResolveTransformerComponents(
     transformer,
     context,
 ) {
-    override val dataFlowAnalyzer: FirDataFlowAnalyzer<*>
+    override val dataFlowAnalyzer: FirDataFlowAnalyzer
-        get() = object : FirDataFlowAnalyzer<PersistentFlow>(this@StubBodyResolveTransformerComponents, context.dataFlowAnalyzerContext) {
+        get() = object : FirDataFlowAnalyzer(this@StubBodyResolveTransformerComponents, context.dataFlowAnalyzerContext) {
-            override val logicSystem: LogicSystem<PersistentFlow>
+            override val logicSystem: LogicSystem
                 get() = error("Should not be called")
 
             override val receiverStack: Iterable<ImplicitReceiverValue<*>>

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

@@ -67,7 +67,7 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
         node: CFGNode<*>,
         effectDeclaration: ConeConditionalEffectDeclaration,
         function: FirFunction,
-        logicSystem: LogicSystem<PersistentFlow>,
+        logicSystem: LogicSystem,
         dataFlowInfo: DataFlowInfo,
         context: CheckerContext
     ): Boolean {
@@ -87,7 +87,7 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
             }
         }
 
-        var flow = dataFlowInfo.flowOnNodes.getValue(node) as PersistentFlow
+        var flow = dataFlowInfo.flowOnNodes.getValue(node)
         val operation = effect.value.toOperation()
         if (operation != null) {
             if (resultExpression is FirConstExpression<*>) {
@@ -100,7 +100,7 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
                 if (resultVar != null) {
                     val impliedByReturnValue = logicSystem.approveOperationStatement(flow, OperationStatement(resultVar, operation))
                     if (impliedByReturnValue.isNotEmpty()) {
-                        flow = logicSystem.forkFlow(flow).also { logicSystem.addTypeStatements(it, impliedByReturnValue) }
+                        flow = flow.fork().also { logicSystem.addTypeStatements(it, impliedByReturnValue) }.freeze()
                     }
                 }
             }
@@ -126,8 +126,8 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
         }
 
         val conditionStatements = logicSystem.approveContractStatement(
-            flow, effectDeclaration.condition, argumentVariables, substitutor = null
+            effectDeclaration.condition, argumentVariables, substitutor = null
-        ) ?: return true
+        ) { logicSystem.approveOperationStatement(flow, it) } ?: return true
 
         return !conditionStatements.values.all { requirement ->
             val originalType = requirement.variable.identifier.symbol.correspondingParameterType ?: return@all true

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/BodyResolveComponents.kt

@@ -42,7 +42,7 @@ abstract class BodyResolveComponents : SessionHolder {
     abstract val callCompleter: FirCallCompleter
     abstract val doubleColonExpressionResolver: FirDoubleColonExpressionResolver
     abstract val syntheticCallGenerator: FirSyntheticCallGenerator
-    abstract val dataFlowAnalyzer: FirDataFlowAnalyzer<*>
+    abstract val dataFlowAnalyzer: FirDataFlowAnalyzer
     abstract val outerClassManager: FirOuterClassManager
     abstract val integerLiteralAndOperatorApproximationTransformer: IntegerLiteralAndOperatorApproximationTransformer
 }

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

@@ -36,10 +36,10 @@ import org.jetbrains.kotlin.types.ConstantValueKind
 import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 
-class DataFlowAnalyzerContext<FLOW : Flow>(
+class DataFlowAnalyzerContext(
     val graphBuilder: ControlFlowGraphBuilder,
     variableStorage: VariableStorageImpl,
-    flowOnNodes: MutableMap<CFGNode<*>, FLOW>,
+    flowOnNodes: MutableMap<CFGNode<*>, PersistentFlow>,
     val preliminaryLoopVisitor: PreliminaryLoopVisitor,
     val variablesClearedBeforeLoop: Stack<List<RealVariable>>,
 ) {
@@ -68,7 +68,7 @@ class DataFlowAnalyzerContext<FLOW : Flow>(
     }
 
     companion object {
-        fun <FLOW : Flow> empty(session: FirSession): DataFlowAnalyzerContext<FLOW> =
+        fun empty(session: FirSession): DataFlowAnalyzerContext =
             DataFlowAnalyzerContext(
                 ControlFlowGraphBuilder(), VariableStorageImpl(session),
                 mutableMapOf(), PreliminaryLoopVisitor(), stackOf()
@@ -77,16 +77,16 @@ class DataFlowAnalyzerContext<FLOW : Flow>(
 }
 
 @OptIn(DfaInternals::class)
-abstract class FirDataFlowAnalyzer<FLOW : Flow>(
+abstract class FirDataFlowAnalyzer(
     protected val components: FirAbstractBodyResolveTransformer.BodyResolveTransformerComponents,
-    private val context: DataFlowAnalyzerContext<FLOW>
+    private val context: DataFlowAnalyzerContext
 ) {
     companion object {
         fun createFirDataFlowAnalyzer(
             components: FirAbstractBodyResolveTransformer.BodyResolveTransformerComponents,
-            dataFlowAnalyzerContext: DataFlowAnalyzerContext<PersistentFlow>
+            dataFlowAnalyzerContext: DataFlowAnalyzerContext
-        ): FirDataFlowAnalyzer<*> =
+        ): FirDataFlowAnalyzer =
-            object : FirDataFlowAnalyzer<PersistentFlow>(components, dataFlowAnalyzerContext) {
+            object : FirDataFlowAnalyzer(components, dataFlowAnalyzerContext) {
                 override val receiverStack: PersistentImplicitReceiverStack
                     get() = components.implicitReceiverStack as PersistentImplicitReceiverStack
 
@@ -129,7 +129,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
             }
     }
 
-    protected abstract val logicSystem: LogicSystem<FLOW>
+    protected abstract val logicSystem: LogicSystem
     protected abstract val receiverStack: Iterable<ImplicitReceiverValue<*>>
     protected abstract fun receiverUpdated(symbol: FirBasedSymbol<*>, info: TypeStatement?)
 
@@ -230,16 +230,16 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         }
         val (postponedLambdaEnterNode, functionEnterNode) = graphBuilder.enterAnonymousFunction(anonymousFunction)
         postponedLambdaEnterNode?.mergeIncomingFlow()
-        val flowOnEntry = functionEnterNode.mergeIncomingFlow()
+        functionEnterNode.mergeIncomingFlow {
-        val invocationKind = anonymousFunction.invocationKind
+            if (anonymousFunction.invocationKind?.canBeRevisited() != false) {
-        if (invocationKind == null || invocationKind.canBeRevisited()) {
+                // TODO: if invocation can happen 0 times, there will be an edge from `functionEnterNode`
-            // TODO: if invocation can happen 0 times, there will be an edge from `functionEnterNode`
+                //  to `functionExitNode`, so erasing statements here causes all information to be lost
-            //  to `functionExitNode`, so erasing statements here causes all information to be lost
+                //  even though `statements from before && statements made inside the lambda` are correct.
-            //  even though `statements from before && statements made inside the lambda` are correct.
+                //     x = ""
-            //     x = ""
+                //     callUnknownNumberOfTimes { x = "" }
-            //     callUnknownNumberOfTimes { x = "" }
+                //     /* x is String no matter how many times the lambda is called, but that information got lost */
-            //     /* x is String no matter how many times the lambda is called, but that information got lost */
+                enterCapturingStatement(it, anonymousFunction)
-            enterCapturingStatement(flowOnEntry, anonymousFunction)
+            }
         }
     }
 
@@ -248,8 +248,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
             anonymousFunction
         )
         val (functionExitNode, postponedLambdaExitNode, graph) = graphBuilder.exitAnonymousFunction(anonymousFunction)
-        val invocationKind = anonymousFunction.invocationKind
+        if (anonymousFunction.invocationKind?.canBeRevisited() != false) {
-        if (invocationKind == null || invocationKind.canBeRevisited()) {
             exitCapturingStatement(anonymousFunction)
         }
         functionExitNode.mergeIncomingFlow()
@@ -379,10 +378,13 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     // ----------------------------------- Operator call -----------------------------------
 
     fun exitTypeOperatorCall(typeOperatorCall: FirTypeOperatorCall) {
-        val node = graphBuilder.exitTypeOperatorCall(typeOperatorCall)
+        graphBuilder.exitTypeOperatorCall(typeOperatorCall).mergeIncomingFlow {
-        val flow = node.mergeIncomingFlow()
+            if (typeOperatorCall.operation !in FirOperation.TYPES) return@mergeIncomingFlow
+            addTypeOperatorStatements(it, typeOperatorCall)
+        }
+    }
 
-        if (typeOperatorCall.operation !in FirOperation.TYPES) return
+    private fun addTypeOperatorStatements(flow: MutableFlow, typeOperatorCall: FirTypeOperatorCall) {
         val type = typeOperatorCall.conversionTypeRef.coneType
         val operandVariable = variableStorage.getOrCreateIfReal(flow, typeOperatorCall.argument) ?: return
         when (val operation = typeOperatorCall.operation) {
@@ -390,9 +392,9 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
                 val isType = operation == FirOperation.IS
                 when (type) {
                     // x is Nothing? <=> x == null
-                    nullableNothing -> processEqNull(node, typeOperatorCall.argument, isType)
+                    nullableNothing -> processEqNull(flow, typeOperatorCall, typeOperatorCall.argument, isType)
                     // x is Any <=> x != null
-                    any -> processEqNull(node, typeOperatorCall.argument, !isType)
+                    any -> processEqNull(flow, typeOperatorCall, typeOperatorCall.argument, !isType)
                     else -> {
                         val expressionVariable = variableStorage.createSynthetic(typeOperatorCall)
                         if (operandVariable.isReal()) {
@@ -440,7 +442,6 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
 
     fun exitEqualityOperatorCall(equalityOperatorCall: FirEqualityOperatorCall) {
         val node = graphBuilder.exitEqualityOperatorCall(equalityOperatorCall)
-        node.mergeIncomingFlow()
         val operation = equalityOperatorCall.operation
         val leftOperand = equalityOperatorCall.arguments[0]
         val rightOperand = equalityOperatorCall.arguments[1]
@@ -465,27 +466,31 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         val leftIsNull = leftIsNullConst || leftOperand.coneType.isNullableNothing && !rightIsNullConst
         val rightIsNull = rightIsNullConst || rightOperand.coneType.isNullableNothing && !leftIsNullConst
 
-        when {
+        node.mergeIncomingFlow { flow ->
-            leftConst != null && rightConst != null -> return
+            when {
-            leftIsNull -> processEqNull(node, rightOperand, operation.isEq())
+                leftConst != null && rightConst != null -> return@mergeIncomingFlow
-            rightIsNull -> processEqNull(node, leftOperand, operation.isEq())
+                leftIsNull -> processEqNull(flow, equalityOperatorCall, rightOperand, operation.isEq())
-            leftConst != null -> processEqWithConst(node, rightOperand, leftConst, operation)
+                rightIsNull -> processEqNull(flow, equalityOperatorCall, leftOperand, operation.isEq())
-            rightConst != null -> processEqWithConst(node, leftOperand, rightConst, operation)
+                leftConst != null -> processEqConst(flow, equalityOperatorCall, rightOperand, leftConst, operation.isEq())
-            else -> processEq(node, leftOperand, rightOperand, operation)
+                rightConst != null -> processEqConst(flow, equalityOperatorCall, leftOperand, rightConst, operation.isEq())
+                else -> processEq(flow, equalityOperatorCall, leftOperand, rightOperand, operation)
+            }
         }
     }
 
-    private fun processEqWithConst(
+    private fun processEqConst(
-        node: EqualityOperatorCallNode, operand: FirExpression, const: FirConstExpression<*>, operation: FirOperation
+        flow: MutableFlow,
+        expression: FirExpression,
+        operand: FirExpression,
+        const: FirConstExpression<*>,
+        isEq: Boolean
     ) {
-        val isEq = operation.isEq()
         if (const.kind == ConstantValueKind.Null) {
-            return processEqNull(node, operand, isEq)
+            return processEqNull(flow, expression, operand, isEq)
         }
 
-        val flow = node.flow
         val operandVariable = variableStorage.getOrCreateIfReal(flow, operand) ?: return
-        val expressionVariable = variableStorage.createSynthetic(node.fir)
+        val expressionVariable = variableStorage.createSynthetic(expression)
         if (const.kind == ConstantValueKind.Boolean && operand.coneType.isBooleanOrNullableBoolean) {
             val expected = (const.value as Boolean)
             flow.addImplication((expressionVariable eq isEq) implies (operandVariable eq expected))
@@ -498,16 +503,16 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         }
     }
 
-    private fun processEqNull(node: CFGNode<*>, operand: FirExpression, isEq: Boolean) {
+    private fun processEqNull(flow: MutableFlow, expression: FirExpression, operand: FirExpression, isEq: Boolean) {
-        val flow = node.flow
         val operandVariable = variableStorage.getOrCreateIfReal(flow, operand) ?: return
-        val expressionVariable = variableStorage.createSynthetic(node.fir)
+        val expressionVariable = variableStorage.createSynthetic(expression)
         flow.addImplication((expressionVariable eq isEq) implies (operandVariable eq null))
         flow.addImplication((expressionVariable eq !isEq) implies (operandVariable notEq null))
     }
 
     private fun processEq(
-        node: EqualityOperatorCallNode,
+        flow: MutableFlow,
+        expression: FirExpression,
         leftOperand: FirExpression,
         rightOperand: FirExpression,
         operation: FirOperation,
@@ -525,13 +530,12 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
             return
         }
 
-        val flow = node.flow
         // TODO: should be `getOrCreateIfRealAndUnchanged(flow from LHS, flow, leftOperand)`, otherwise the statement will
         //  be added even if the value has changed in the RHS. Currently the only previous node is the RHS.
         val leftOperandVariable = variableStorage.getOrCreateIfReal(flow, leftOperand)
         val rightOperandVariable = variableStorage.getOrCreateIfReal(flow, rightOperand)
         if (leftOperandVariable == null && rightOperandVariable == null) return
-        val expressionVariable = variableStorage.createSynthetic(node.fir)
+        val expressionVariable = variableStorage.createSynthetic(expression)
 
         if (leftIsNullable || rightIsNullable) {
             // `a == b:Any` => `a != null`; the inverse is not true - we don't know when `a` *is* `null`
@@ -605,11 +609,9 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     // ----------------------------------- Check not null call -----------------------------------
 
     fun exitCheckNotNullCall(checkNotNullCall: FirCheckNotNullCall, callCompleted: Boolean) {
-        // Add `Any` to the set of possible types; the intersection type `T? & Any` will be reduced to `T` after smartcast.
         val (node, unionNode) = graphBuilder.exitCheckNotNullCall(checkNotNullCall, callCompleted)
-        val flow = node.mergeIncomingFlow()
+        node.mergeIncomingFlow { flow ->
-        val argumentVariable = variableStorage.getOrCreateIfReal(flow, checkNotNullCall.argument)
+            val argumentVariable = variableStorage.getOrCreateIfReal(flow, checkNotNullCall.argument) ?: return@mergeIncomingFlow
-        if (argumentVariable != null) {
             flow.commitOperationStatement(argumentVariable notEq null)
         }
         unionNode?.unionFlowFromArguments()
@@ -625,28 +627,23 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         graphBuilder.enterWhenBranchCondition(whenBranch).mergeWhenBranchEntryFlow()
     }
 
-    private fun CFGNode<*>.mergeWhenBranchEntryFlow() {
+    private fun CFGNode<*>.mergeWhenBranchEntryFlow() = mergeIncomingFlow { flow ->
-        val previousConditionExitNode = previousNodes.singleOrNull()
+        val previousConditionExitNode = previousNodes.singleOrNull() as? WhenBranchConditionExitNode ?: return@mergeIncomingFlow
-        if (previousConditionExitNode is WhenBranchConditionExitNode) {
+        val previousCondition = previousConditionExitNode.fir.condition
-            val flow = mergeIncomingFlow()
+        if (!previousCondition.coneType.isBoolean) return@mergeIncomingFlow
-            val previousCondition = previousConditionExitNode.fir.condition
+        val previousConditionVariable = variableStorage.get(flow, previousCondition) ?: return@mergeIncomingFlow
-            if (previousCondition.coneType.isBoolean) {
+        flow.commitOperationStatement(previousConditionVariable eq false)
-                val previousConditionVariable = variableStorage.get(flow, previousCondition) ?: return
-                flow.commitOperationStatement(previousConditionVariable eq false)
-            }
-        } else { // first branch
-            mergeIncomingFlow()
-        }
     }
 
     fun exitWhenBranchCondition(whenBranch: FirWhenBranch) {
         val (conditionExitNode, resultEnterNode) = graphBuilder.exitWhenBranchCondition(whenBranch)
-        val conditionExitFlow = conditionExitNode.mergeIncomingFlow()
+        conditionExitNode.mergeIncomingFlow()
-        val resultEnterFlow = resultEnterNode.mergeIncomingFlow()
+        resultEnterNode.mergeIncomingFlow { flow ->
-        // If the condition is invalid, don't generate smart casts to Any or Boolean.
+            // If the condition is invalid, don't generate smart casts to Any or Boolean.
-        if (whenBranch.condition.coneType.isBoolean) {
+            if (whenBranch.condition.coneType.isBoolean) {
-            val conditionVariable = variableStorage.get(conditionExitFlow, whenBranch.condition) ?: return
+                val conditionVariable = variableStorage.get(flow, whenBranch.condition) ?: return@mergeIncomingFlow
-            resultEnterFlow.commitOperationStatement(conditionVariable eq true)
+                flow.commitOperationStatement(conditionVariable eq true)
+            }
         }
     }
 
@@ -670,56 +667,64 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     fun enterWhileLoop(loop: FirLoop) {
         val (loopEnterNode, loopConditionEnterNode) = graphBuilder.enterWhileLoop(loop)
         loopEnterNode.mergeIncomingFlow()
-        val loopConditionEnterFlow = loopConditionEnterNode.mergeIncomingFlow()
+        loopConditionEnterNode.mergeIncomingFlow { flow -> enterCapturingStatement(flow, loop) }
-        enterCapturingStatement(loopConditionEnterFlow, loop)
     }
 
     fun exitWhileLoopCondition(loop: FirLoop) {
         val (loopConditionExitNode, loopBlockEnterNode) = graphBuilder.exitWhileLoopCondition(loop)
-        val conditionExitFlow = loopConditionExitNode.mergeIncomingFlow()
+        loopConditionExitNode.mergeIncomingFlow()
-        val blockEnterFlow = loopBlockEnterNode.mergeIncomingFlow()
+        loopBlockEnterNode.mergeIncomingFlow { flow ->
-        if (loop.condition.coneType.isBoolean) {
+            if (loop.condition.coneType.isBoolean) {
-            val conditionVariable = variableStorage.get(conditionExitFlow, loop.condition) ?: return
+                val conditionVariable = variableStorage.get(flow, loop.condition) ?: return@mergeIncomingFlow
-            blockEnterFlow.commitOperationStatement(conditionVariable eq true)
+                flow.commitOperationStatement(conditionVariable eq true)
+            }
         }
     }
 
     fun exitWhileLoop(loop: FirLoop) {
         val (conditionEnterNode, blockExitNode, exitNode) = graphBuilder.exitWhileLoop(loop)
         blockExitNode.mergeIncomingFlow()
-        val possiblyChangedVariables = exitCapturingStatement(loop)
+        exitNode.mergeIncomingFlow {
+            processWhileLoopExit(it, exitNode, conditionEnterNode)
+            processLoopExit(it, exitNode, exitNode.firstPreviousNode as LoopConditionExitNode)
+        }
+    }
+
+    private fun processWhileLoopExit(flow: MutableFlow, node: LoopExitNode, conditionEnterNode: LoopConditionEnterNode) {
+        val possiblyChangedVariables = exitCapturingStatement(node.fir)
+        if (possiblyChangedVariables.isNullOrEmpty()) return
         // While analyzing the loop we might have added some backwards jumps to `conditionEnterNode` which weren't
         // there at the time its flow was computed - which is why we erased all information about `possiblyChangedVariables`
         // from it. Now that we have those edges, we can restore type information for the code after the loop.
-        if (!possiblyChangedVariables.isNullOrEmpty()) {
+        val conditionEnterFlow = conditionEnterNode.flow
-            val conditionEnterFlow = conditionEnterNode.flow
+        val loopEnterAndContinueFlows = conditionEnterNode.livePreviousFlows
-            val loopEnterAndContinueFlows = conditionEnterNode.livePreviousFlows
+        val conditionExitAndBreakFlows = node.livePreviousFlows
-            val conditionExitAndBreakFlows = exitNode.livePreviousFlows
+        possiblyChangedVariables.forEach { variable ->
-            possiblyChangedVariables.forEach { variable ->
+            // The statement about `variable` in `conditionEnterFlow` should be empty, so to obtain the new statement
-                // The statement about `variable` in `conditionEnterFlow` should be empty, so to obtain the new statement
+            // we can simply add the now-known input to whatever was inferred from nothing so long as the value is the same.
-                // we can simply add the now-known input to whatever was inferred from nothing so long as the value is the same.
+            val toAdd = logicSystem.or(loopEnterAndContinueFlows.map { it.getTypeStatement(variable) ?: return@forEach })
-                val statement = logicSystem.or(loopEnterAndContinueFlows.map { it.getTypeStatement(variable) ?: return@forEach })
+                ?.takeIf { it.isNotEmpty } ?: return@forEach
-                    ?: return@forEach
+            val newStatement = logicSystem.or(conditionExitAndBreakFlows.map {
-                for (beforeExitFlow in conditionExitAndBreakFlows) {
+                val atExit = it.getTypeStatement(variable)
-                    if (logicSystem.isSameValueIn(conditionEnterFlow, beforeExitFlow, variable)) {
+                if (logicSystem.isSameValueIn(conditionEnterFlow, it, variable)) {
-                        beforeExitFlow.addTypeStatement(statement)
+                    if (atExit != null) logicSystem.and(listOf(atExit, toAdd)) else toAdd
-                    }
+                } else {
-                }
+                    atExit
-            }
+                } ?: return@forEach
+            }) ?: return@forEach
+            flow.addTypeStatement(newStatement)
         }
-        exitNode.mergeLoopExitFlow(exitNode.firstPreviousNode as LoopConditionExitNode)
     }
 
-    private fun LoopExitNode.mergeLoopExitFlow(conditionExitNode: LoopConditionExitNode) {
+    private fun processLoopExit(flow: MutableFlow, node: LoopExitNode, conditionExitNode: LoopConditionExitNode) {
-        val flow = mergeIncomingFlow()
+        if (conditionExitNode.isDead || node.previousNodes.count { !it.isDead } > 1) return
-        if (conditionExitNode.isDead || previousNodes.count { !it.isDead } > 1) return
         if (conditionExitNode.fir.coneType.isBoolean) {
             val variable = variableStorage.get(flow, conditionExitNode.fir) ?: return
             flow.commitOperationStatement(variable eq false)
         }
     }
 
-    private fun enterCapturingStatement(flow: FLOW, statement: FirStatement) {
+    private fun enterCapturingStatement(flow: MutableFlow, statement: FirStatement) {
         val reassignedNames = context.preliminaryLoopVisitor.enterCapturingStatement(statement)
         if (reassignedNames.isEmpty()) return
         // TODO: only choose the innermost variable for each name
@@ -745,8 +750,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
 
     fun enterDoWhileLoop(loop: FirLoop) {
         val (loopEnterNode, loopBlockEnterNode) = graphBuilder.enterDoWhileLoop(loop)
-        val loopEnterFlow = loopEnterNode.mergeIncomingFlow()
+        loopEnterNode.mergeIncomingFlow { flow -> enterCapturingStatement(flow, loop) }
-        enterCapturingStatement(loopEnterFlow, loop)
         loopBlockEnterNode.mergeIncomingFlow()
     }
 
@@ -759,7 +763,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     fun exitDoWhileLoop(loop: FirLoop) {
         val (loopConditionExitNode, loopExitNode) = graphBuilder.exitDoWhileLoop(loop)
         loopConditionExitNode.mergeIncomingFlow()
-        loopExitNode.mergeLoopExitFlow(loopConditionExitNode)
+        loopExitNode.mergeIncomingFlow { processLoopExit(it, loopExitNode, loopConditionExitNode) }
         exitCapturingStatement(loop)
     }
 
@@ -803,8 +807,9 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     fun enterQualifiedAccessExpression() {}
 
     fun exitQualifiedAccessExpression(qualifiedAccessExpression: FirQualifiedAccessExpression) {
-        val flow = graphBuilder.exitQualifiedAccessExpression(qualifiedAccessExpression).mergeIncomingFlow()
+        graphBuilder.exitQualifiedAccessExpression(qualifiedAccessExpression).mergeIncomingFlow { flow ->
-        processConditionalContract(flow, qualifiedAccessExpression)
+            processConditionalContract(flow, qualifiedAccessExpression)
+        }
     }
 
     fun exitSmartCastExpression(smartCastExpression: FirSmartCastExpression) {
@@ -812,26 +817,28 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     }
 
     fun enterSafeCallAfterNullCheck(safeCall: FirSafeCallExpression) {
-        val flow = graphBuilder.enterSafeCall(safeCall).mergeIncomingFlow()
+        graphBuilder.enterSafeCall(safeCall).mergeIncomingFlow { flow ->
-        val receiverVariable = variableStorage.getOrCreateIfReal(flow, safeCall.receiver) ?: return
+            val receiverVariable = variableStorage.getOrCreateIfReal(flow, safeCall.receiver) ?: return@mergeIncomingFlow
-        flow.commitOperationStatement(receiverVariable notEq null)
+            flow.commitOperationStatement(receiverVariable notEq null)
+        }
     }
 
     fun exitSafeCall(safeCall: FirSafeCallExpression) {
         val (node, mergePostponedLambdaExitsNode) = graphBuilder.exitSafeCall()
-        val flow = node.mergeIncomingFlow()
+        node.mergeIncomingFlow { flow ->
+            // If there is only 1 previous node, then this is LHS of `a?.b ?: c`; then the null-case
+            // edge from `a` goes directly to `c` and this node's flow already assumes `b` executed.
+            if (node.previousNodes.size < 2) return@mergeIncomingFlow
+            // Otherwise if the result is non-null, then `b` executed, which implies `a` is not null
+            // and every statement from `b` holds.
+            val expressionVariable = variableStorage.getOrCreate(flow, safeCall)
+            // TODO? if the callee has non-null return type, then safe-call == null => receiver == null
+            //   if (x?.toString() == null) { /* x == null */ }
+            // TODO? all new implications in previous node's flow are valid here if receiver != null
+            //  (that requires a second level of implications: receiver != null => condition => effect).
+            flow.addAllConditionally(expressionVariable notEq null, node.lastPreviousNode.flow)
+        }
         mergePostponedLambdaExitsNode?.mergeIncomingFlow()
-        // If there is only 1 previous node, then this is LHS of `a?.b ?: c`; then the null-case
-        // edge from `a` goes directly to `c` and this node's flow already assumes `b` executed.
-        if (node.previousNodes.size < 2) return
-        // Otherwise if the result is non-null, then `b` executed, which implies `a` is not null
-        // and every statement from `b` holds.
-        val expressionVariable = variableStorage.getOrCreate(flow, safeCall)
-        // TODO? if the callee has non-null return type, then safe-call == null => receiver == null
-        //   if (x?.toString() == null) { /* x == null */ }
-        // TODO? all new implications in previous node's flow are valid here if receiver != null
-        //  (that requires a second level of implications: receiver != null => condition => effect).
-        flow.addAllConditionally(expressionVariable notEq null, node.lastPreviousNode.flow)
     }
 
     fun exitResolvedQualifierNode(resolvedQualifier: FirResolvedQualifier) {
@@ -868,8 +875,9 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         }
         val (functionCallNode, unionNode) = graphBuilder.exitFunctionCall(functionCall, callCompleted)
         unionNode?.unionFlowFromArguments()
-        val flow = functionCallNode.mergeIncomingFlow()
+        functionCallNode.mergeIncomingFlow {
-        processConditionalContract(flow, functionCall)
+            processConditionalContract(it, functionCall)
+        }
     }
 
     fun exitDelegatedConstructorCall(call: FirDelegatedConstructorCall, callCompleted: Boolean) {
@@ -884,10 +892,6 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         callNode.mergeIncomingFlow()
     }
 
-    private fun UnionFunctionCallArgumentsNode.unionFlowFromArguments() {
-        flow = logicSystem.unionFlow(previousNodes.map { it.flow })
-    }
-
     private fun FirQualifiedAccess.orderedArguments(callee: FirFunction): Array<out FirExpression?>? {
         val receiver = extensionReceiver.takeIf { it != FirNoReceiverExpression }
             ?: dispatchReceiver.takeIf { it != FirNoReceiverExpression }
@@ -905,7 +909,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         }
     }
 
-    private fun processConditionalContract(flow: FLOW, qualifiedAccess: FirQualifiedAccess) {
+    private fun processConditionalContract(flow: MutableFlow, qualifiedAccess: FirQualifiedAccess) {
         val callee = when (qualifiedAccess) {
             is FirFunctionCall -> qualifiedAccess.toResolvedCallableSymbol()?.fir as? FirSimpleFunction
             is FirQualifiedAccessExpression -> (qualifiedAccess.calleeReference.resolvedSymbol?.fir as? FirProperty)?.getter
@@ -944,9 +948,9 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         for (conditionalEffect in conditionalEffects) {
             val effect = conditionalEffect.effect as? ConeReturnsEffectDeclaration ?: continue
             val operation = effect.value.toOperation()
-            val statements = logicSystem.approveContractStatement(
+            val statements = logicSystem.approveContractStatement(conditionalEffect.condition, argumentVariables, substitutor) {
-                flow, conditionalEffect.condition, argumentVariables, substitutor, removeApprovedOrImpossible = operation == null
+                logicSystem.approveOperationStatement(flow, it, removeApprovedOrImpossible = operation == null)
-            ) ?: continue // TODO: do what if the result is known to be false?
+            } ?: continue // TODO: do what if the result is known to be false?
             if (operation == null) {
                 flow.addAllStatements(statements)
             } else {
@@ -962,28 +966,30 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     }
 
     fun exitLocalVariableDeclaration(variable: FirProperty, hadExplicitType: Boolean) {
-        val flow = graphBuilder.exitVariableDeclaration(variable).mergeIncomingFlow()
+        graphBuilder.exitVariableDeclaration(variable).mergeIncomingFlow { flow ->
-        val initializer = variable.initializer ?: return
+            val initializer = variable.initializer ?: return@mergeIncomingFlow
-        exitVariableInitialization(flow, initializer, variable, assignment = null, hadExplicitType)
+            exitVariableInitialization(flow, initializer, variable, assignment = null, hadExplicitType)
+        }
     }
 
     fun exitVariableAssignment(assignment: FirVariableAssignment) {
-        val flow = graphBuilder.exitVariableAssignment(assignment).mergeIncomingFlow()
+        graphBuilder.exitVariableAssignment(assignment).mergeIncomingFlow { flow ->
-        val property = assignment.lValue.resolvedSymbol?.fir as? FirProperty ?: return
+            val property = assignment.lValue.resolvedSymbol?.fir as? FirProperty ?: return@mergeIncomingFlow
-        if (property.isLocal || property.isVal) {
+            if (property.isLocal || property.isVal) {
-            exitVariableInitialization(flow, assignment.rValue, property, assignment, hasExplicitType = false)
+                exitVariableInitialization(flow, assignment.rValue, property, assignment, hasExplicitType = false)
-        } else {
+            } else {
-            // TODO: add unstable smartcast for non-local var
+                // TODO: add unstable smartcast for non-local var
-            val variable = variableStorage.getRealVariableWithoutUnwrappingAlias(flow, assignment)
+                val variable = variableStorage.getRealVariableWithoutUnwrappingAlias(flow, assignment)
-            if (variable != null) {
+                if (variable != null) {
-                logicSystem.recordNewAssignment(flow, variable, context.newAssignmentIndex())
+                    logicSystem.recordNewAssignment(flow, variable, context.newAssignmentIndex())
+                }
             }
+            processConditionalContract(flow, assignment)
         }
-        processConditionalContract(flow, assignment)
     }
 
     private fun exitVariableInitialization(
-        flow: FLOW,
+        flow: MutableFlow,
         initializer: FirExpression,
         property: FirProperty,
         assignment: FirVariableAssignment?,
@@ -1044,23 +1050,26 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
 
     fun exitLeftBinaryLogicExpressionArgument(binaryLogicExpression: FirBinaryLogicExpression) {
         val (leftExitNode, rightEnterNode) = graphBuilder.exitLeftBinaryLogicExpressionArgument(binaryLogicExpression)
-        val leftExitFlow = leftExitNode.mergeIncomingFlow()
+        leftExitNode.mergeIncomingFlow()
-        val rightEnterFlow = rightEnterNode.mergeIncomingFlow()
+        rightEnterNode.mergeIncomingFlow { flow ->
-        val leftOperandVariable = variableStorage.get(leftExitFlow, leftExitNode.firstPreviousNode.fir) ?: return
+            val leftOperandVariable = variableStorage.get(flow, binaryLogicExpression.leftOperand) ?: return@mergeIncomingFlow
-        val isAnd = binaryLogicExpression.kind == LogicOperationKind.AND
+            val isAnd = binaryLogicExpression.kind == LogicOperationKind.AND
-        rightEnterFlow.commitOperationStatement(leftOperandVariable eq isAnd)
+            flow.commitOperationStatement(leftOperandVariable eq isAnd)
+        }
     }
 
-    fun exitBinaryLogicExpression(binaryLogicExpression: FirBinaryLogicExpression) {
+    fun exitBinaryLogicExpression() {
-        val node = graphBuilder.exitBinaryLogicExpression()
+        graphBuilder.exitBinaryLogicExpression().mergeBinaryLogicOperatorFlow()
-        val isAnd = binaryLogicExpression.kind == LogicOperationKind.AND
+    }
-        val flowFromLeft = node.leftOperandNode.flow
-        val flowFromRight = node.rightOperandNode.flow
-        val flow = node.mergeIncomingFlow()
 
-        val leftVariable = variableStorage.get(flowFromLeft, binaryLogicExpression.leftOperand)
+    private fun AbstractBinaryExitNode<FirBinaryLogicExpression>.mergeBinaryLogicOperatorFlow() = mergeIncomingFlow { flow ->
-        val leftIsBoolean = leftVariable != null && binaryLogicExpression.leftOperand.coneType.isBoolean
+        val isAnd = fir.kind == LogicOperationKind.AND
-        if (!node.leftOperandNode.isDead && node.rightOperandNode.isDead) {
+        val flowFromLeft = leftOperandNode.flow
+        val flowFromRight = rightOperandNode.flow
+
+        val leftVariable = variableStorage.get(flowFromLeft, fir.leftOperand)
+        val leftIsBoolean = leftVariable != null && fir.leftOperand.coneType.isBoolean
+        if (!leftOperandNode.isDead && rightOperandNode.isDead) {
             // If the right operand does not terminate, then we know that the value of the entire expression
             // has to be saturating (true for or, false for and), and it has to be produced by the left operand.
             if (leftIsBoolean) {
@@ -1068,9 +1077,9 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
                 flow.commitOperationStatement(leftVariable!! eq !isAnd)
             }
         } else {
-            val rightVariable = variableStorage.get(flowFromRight, binaryLogicExpression.rightOperand)
+            val rightVariable = variableStorage.get(flowFromRight, fir.rightOperand)
-            val rightIsBoolean = rightVariable != null && binaryLogicExpression.rightOperand.coneType.isBoolean
+            val rightIsBoolean = rightVariable != null && fir.rightOperand.coneType.isBoolean
-            val operatorVariable = variableStorage.createSynthetic(binaryLogicExpression)
+            val operatorVariable = variableStorage.createSynthetic(fir)
             // If `left && right` is true, then both are evaluated to true. If `left || right` is false, then both are false.
             // Approved type statements for RHS already contain everything implied by the corresponding value of LHS.
             val bothEvaluated = operatorVariable eq isAnd
@@ -1098,7 +1107,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         }
     }
 
-    private fun exitBooleanNot(flow: FLOW, expression: FirFunctionCall) {
+    private fun exitBooleanNot(flow: MutableFlow, expression: FirFunctionCall) {
         val argumentVariable = variableStorage.get(flow, expression.dispatchReceiver) ?: return
         val expressionVariable = variableStorage.createSynthetic(expression)
         // Alternatively: (expression == true => argument == false) && (expression == false => argument == true)
@@ -1153,26 +1162,28 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
 
     fun exitElvisLhs(elvisExpression: FirElvisExpression) {
         val (lhsExitNode, lhsIsNotNullNode, rhsEnterNode) = graphBuilder.exitElvisLhs(elvisExpression)
-        val flow = lhsExitNode.mergeIncomingFlow()
+        val lhsVariable = variableStorage.getOrCreateIfReal(lhsExitNode.mergeIncomingFlow(), elvisExpression.lhs)
-        val lhsIsNotNullFlow = lhsIsNotNullNode.mergeIncomingFlow()
+        lhsIsNotNullNode.mergeIncomingFlow { flow ->
-        val rhsEnterFlow = rhsEnterNode.mergeIncomingFlow()
+            lhsVariable?.let { flow.commitOperationStatement(it notEq null) }
-        val lhsVariable = variableStorage.getOrCreateIfReal(flow, elvisExpression.lhs) ?: return
+        }
-        lhsIsNotNullFlow.commitOperationStatement(lhsVariable notEq null)
+        rhsEnterNode.mergeIncomingFlow { flow ->
-        rhsEnterFlow.commitOperationStatement(lhsVariable eq null)
+            lhsVariable?.let { flow.commitOperationStatement(it eq null) }
+        }
     }
 
     fun exitElvis(elvisExpression: FirElvisExpression, isLhsNotNull: Boolean) {
         val (node, mergePostponedLambdaExitsNode) = graphBuilder.exitElvis(isLhsNotNull)
-        val flow = node.mergeIncomingFlow()
+        node.mergeIncomingFlow { flow ->
+            // If LHS is never null, then the edge from RHS is dead and this node's flow already contains
+            // all statements from LHS unconditionally.
+            if (isLhsNotNull) return@mergeIncomingFlow
+            // For any predicate P(x), if P(v) != P(u ?: v) then u != null. In general this requires two levels of
+            // implications, but for constant v the logic system can handle some basic cases of P(x).
+            val rhs = (elvisExpression.rhs as? FirConstExpression<*>)?.value as? Boolean ?: return@mergeIncomingFlow
+            val elvisVariable = variableStorage.createSynthetic(elvisExpression)
+            flow.addAllConditionally(elvisVariable eq !rhs, node.firstPreviousNode.flow)
+        }
         mergePostponedLambdaExitsNode?.mergeIncomingFlow()
-        // If LHS is never null, then the edge from RHS is dead and this node's flow already contains
-        // all statements from LHS unconditionally.
-        if (isLhsNotNull) return
-        // For any predicate P(x), if P(v) != P(u ?: v) then u != null. In general this requires two levels of
-        // implications, but for constant v the logic system can handle some basic cases of P(x).
-        val rhs = (elvisExpression.rhs as? FirConstExpression<*>)?.value as? Boolean ?: return
-        val elvisVariable = variableStorage.createSynthetic(elvisExpression)
-        flow.addAllConditionally(elvisVariable eq !rhs, node.firstPreviousNode.flow)
     }
 
     // Callable reference
@@ -1197,15 +1208,12 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         }
     }
 
-    private var CFGNode<*>.flow: FLOW
+    private val CFGNode<*>.flow: PersistentFlow
         get() = context.flowOnNodes.getValue(this.origin)
-        set(value) {
-            context.flowOnNodes[this.origin] = value
-        }
 
     private val CFGNode<*>.origin: CFGNode<*> get() = if (this is StubNode) firstPreviousNode else this
 
-    private val CFGNode<*>.livePreviousFlows: List<FLOW>
+    private val CFGNode<*>.livePreviousFlows: List<PersistentFlow>
         get() = previousNodes.mapNotNull { it.takeIf { this.isDead || !it.isDead }?.flow }
 
     // Smart cast information is taken from `graphBuilder.lastNode`, but the problem with receivers specifically
@@ -1213,16 +1221,16 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
     // we explicitly patch up the stack by calling `receiverUpdated` in a way that maintains consistency with
     // `getTypeUsingSmartcastInfo`; i.e. at any point between calls to this class' methods the types in the implicit
     // receiver stack also correspond to the data flow information attached to `graphBuilder.lastNode`.
-    private var currentReceiverState: FLOW? = null
+    private var currentReceiverState: Flow? = null
 
     // Generally when calling some method on `graphBuilder`, one of the nodes it returns is the new `lastNode`.
     // In that case `mergeIncomingFlow` will automatically ensure consistency once called on that node.
-    private fun CFGNode<*>.mergeIncomingFlow(): FLOW {
+    private fun CFGNode<*>.buildIncomingFlow(): MutableFlow {
         val previousFlows = previousNodes.mapNotNull {
             val incomingEdgeKind = incomingEdges.getValue(it).kind
             it.takeIf { incomingEdgeKind.usedInDfa || (isDead && incomingEdgeKind.usedInDeadDfa) }?.flow
         }
-        val result = logicSystem.joinFlow(previousFlows).also { flow = it }
+        val result = logicSystem.joinFlow(previousFlows, union = false)
         if (graphBuilder.lastNodeOrNull == this) {
             // Here it is, the new `lastNode`. If the previous state is the only predecessor, then there is actually
             // nothing to update; `addTypeStatement` has already ensured we have the correct information.
@@ -1234,6 +1242,25 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         return result
     }
 
+    private fun UnionFunctionCallArgumentsNode.unionFlowFromArguments() {
+        context.flowOnNodes[this] = logicSystem.joinFlow(previousNodes.map { it.flow }, union = true).freeze()
+    }
+
+    private fun CFGNode<*>.setFlow(builder: MutableFlow): PersistentFlow {
+        val flow = builder.freeze()
+        context.flowOnNodes[this] = flow
+        if (currentReceiverState === builder) {
+            currentReceiverState = flow
+        }
+        return flow
+    }
+
+    private inline fun CFGNode<*>.mergeIncomingFlow(crossinline builder: (MutableFlow) -> Unit): PersistentFlow =
+        setFlow(buildIncomingFlow().also(builder))
+
+    private fun CFGNode<*>.mergeIncomingFlow() =
+        setFlow(buildIncomingFlow())
+
     // In rare cases (like after exiting functions) after adding more nodes `graphBuilder` will revert the current
     // state to a previously created node, so none of the nodes it returned are `lastNode` and `mergeIncomingFlow`
     // will not ensure consistency. In that case an explicit call to `resetReceivers` is needed to roll back the stack
@@ -1244,7 +1271,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         currentReceiverState = currentFlow
     }
 
-    private fun updateAllReceivers(from: FLOW?, to: FLOW?) {
+    private fun updateAllReceivers(from: Flow?, to: Flow?) {
         receiverStack.forEach {
             variableStorage.getLocalVariable(it.boundSymbol)?.let { variable ->
                 val newStatement = to?.getTypeStatement(variable)
@@ -1255,40 +1282,32 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         }
     }
 
-    private fun FLOW.addImplication(statement: Implication) {
+    private fun MutableFlow.addImplication(statement: Implication) {
         logicSystem.addImplication(this, statement)
     }
 
-    private fun FLOW.addTypeStatement(info: TypeStatement) {
+    private fun MutableFlow.addTypeStatement(info: TypeStatement) {
         val newStatement = logicSystem.addTypeStatement(this, info) ?: return
         if (newStatement.variable.isThisReference && this === currentReceiverState) {
             receiverUpdated(newStatement.variable.identifier.symbol, newStatement)
         }
     }
 
-    private fun FLOW.addAllStatements(statements: TypeStatements) {
+    private fun MutableFlow.addAllStatements(statements: TypeStatements) =
-        val newStatements = logicSystem.addTypeStatements(this, statements)
+        statements.values.forEach { addTypeStatement(it) }
-        if (this === currentReceiverState) {
-            for (newStatement in newStatements) {
-                if (newStatement.variable.isThisReference) {
-                    receiverUpdated(newStatement.variable.identifier.symbol, newStatement)
-                }
-            }
-        }
-    }
 
-    private fun FLOW.addAllConditionally(condition: OperationStatement, statements: TypeStatements) =
+    private fun MutableFlow.addAllConditionally(condition: OperationStatement, statements: TypeStatements) =
         statements.values.forEach { addImplication(condition implies it) }
 
-    private fun FLOW.addAllConditionally(condition: OperationStatement, from: FLOW) =
+    private fun MutableFlow.addAllConditionally(condition: OperationStatement, from: Flow) =
         from.knownVariables.forEach {
             // Only add the statement if this variable is not aliasing another in `this` (but it could be aliasing in `from`).
             if (unwrapVariable(it) == it) addImplication(condition implies (from.getTypeStatement(it) ?: return@forEach))
         }
 
-    private fun FLOW.commitOperationStatement(statement: OperationStatement) =
+    private fun MutableFlow.commitOperationStatement(statement: OperationStatement) =
         addAllStatements(logicSystem.approveOperationStatement(this, statement, removeApprovedOrImpossible = true))
 
-    private fun VariableStorageImpl.getOrCreateIfRealAndUnchanged(originalFlow: FLOW, currentFlow: FLOW, fir: FirElement) =
+    private fun VariableStorageImpl.getOrCreateIfRealAndUnchanged(originalFlow: PersistentFlow, currentFlow: MutableFlow, fir: FirElement) =
         getOrCreateIfReal(originalFlow, fir)?.takeIf { !it.isReal() || logicSystem.isSameValueIn(originalFlow, currentFlow, it) }
 }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt

@@ -65,7 +65,7 @@ class FirCallCompletionResultsWriterTransformer(
     private val finalSubstitutor: ConeSubstitutor,
     private val typeCalculator: ReturnTypeCalculator,
     private val typeApproximator: ConeTypeApproximator,
-    private val dataFlowAnalyzer: FirDataFlowAnalyzer<*>,
+    private val dataFlowAnalyzer: FirDataFlowAnalyzer,
     private val integerOperatorApproximator: IntegerLiteralAndOperatorApproximationTransformer,
     private val context: BodyResolveContext,
     private val mode: Mode = Mode.Normal

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt

@@ -21,7 +21,6 @@ import org.jetbrains.kotlin.fir.resolve.calls.ImplicitReceiverValue
 import org.jetbrains.kotlin.fir.resolve.calls.InaccessibleImplicitReceiverValue
 import org.jetbrains.kotlin.fir.resolve.calls.ResolutionContext
 import org.jetbrains.kotlin.fir.resolve.dfa.DataFlowAnalyzerContext
-import org.jetbrains.kotlin.fir.resolve.dfa.PersistentFlow
 import org.jetbrains.kotlin.fir.resolve.inference.FirBuilderInferenceSession
 import org.jetbrains.kotlin.fir.resolve.inference.FirCallCompleter
 import org.jetbrains.kotlin.fir.resolve.inference.FirDelegatedPropertyInferenceSession
@@ -42,7 +41,7 @@ import org.jetbrains.kotlin.name.SpecialNames.UNDERSCORE_FOR_UNUSED_VAR
 
 class BodyResolveContext(
     val returnTypeCalculator: ReturnTypeCalculator,
-    val dataFlowAnalyzerContext: DataFlowAnalyzerContext<PersistentFlow>,
+    val dataFlowAnalyzerContext: DataFlowAnalyzerContext,
     val targetedLocalClasses: Set<FirClassLikeDeclaration> = emptySet(),
     val outerLocalClassForNested: MutableMap<FirClassLikeSymbol<*>, FirClassLikeSymbol<*>> = mutableMapOf()
 ) {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirAbstractBodyResolveTransformer.kt

@@ -66,7 +66,7 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
     protected inline val typeResolverTransformer: FirSpecificTypeResolverTransformer get() = components.typeResolverTransformer
     protected inline val callResolver: FirCallResolver get() = components.callResolver
     protected inline val callCompleter: FirCallCompleter get() = components.callCompleter
-    protected inline val dataFlowAnalyzer: FirDataFlowAnalyzer<*> get() = components.dataFlowAnalyzer
+    protected inline val dataFlowAnalyzer: FirDataFlowAnalyzer get() = components.dataFlowAnalyzer
     protected inline val scopeSession: ScopeSession get() = components.scopeSession
     protected inline val file: FirFile get() = components.file
 
@@ -103,7 +103,7 @@ abstract class FirAbstractBodyResolveTransformer(phase: FirResolvePhase) : FirAb
             session
         )
         override val callCompleter: FirCallCompleter = FirCallCompleter(transformer, this)
-        override val dataFlowAnalyzer: FirDataFlowAnalyzer<*> =
+        override val dataFlowAnalyzer: FirDataFlowAnalyzer =
             FirDataFlowAnalyzer.createFirDataFlowAnalyzer(this, context.dataFlowAnalyzerContext)
         override val syntheticCallGenerator: FirSyntheticCallGenerator = FirSyntheticCallGenerator(this)
         override val doubleColonExpressionResolver: FirDoubleColonExpressionResolver = FirDoubleColonExpressionResolver(session)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -829,7 +829,7 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
             .transformLeftOperand(this, ResolutionMode.WithExpectedType(booleanType))
             .also(dataFlowAnalyzer::exitLeftBinaryLogicExpressionArgument)
             .transformRightOperand(this, ResolutionMode.WithExpectedType(booleanType))
-            .also(dataFlowAnalyzer::exitBinaryLogicExpression)
+            .also { dataFlowAnalyzer.exitBinaryLogicExpression() }
             .transformOtherChildren(transformer, ResolutionMode.WithExpectedType(booleanType))
             .also { it.resultType = booleanType }
     }

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

@@ -25,7 +25,7 @@ class FirControlFlowGraphReferenceImpl(
     }
 }
 
-class DataFlowInfo(val variableStorage: VariableStorage, val flowOnNodes: Map<CFGNode<*>, Flow>)
+class DataFlowInfo(val variableStorage: VariableStorage, val flowOnNodes: Map<CFGNode<*>, PersistentFlow>)
 
 val FirControlFlowGraphReference.controlFlowGraph: ControlFlowGraph?
     get() = (this as? FirControlFlowGraphReferenceImpl)?.controlFlowGraph

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

@@ -5,8 +5,103 @@
 
 package org.jetbrains.kotlin.fir.resolve.dfa
 
+import kotlinx.collections.immutable.PersistentList
+import kotlinx.collections.immutable.PersistentMap
+import kotlinx.collections.immutable.PersistentSet
+import kotlinx.collections.immutable.persistentHashMapOf
+
 abstract class Flow {
     abstract val knownVariables: Set<RealVariable>
     abstract fun unwrapVariable(variable: RealVariable): RealVariable
     abstract fun getTypeStatement(variable: RealVariable): TypeStatement?
 }
+
+class PersistentFlow internal constructor(
+    private val previousFlow: PersistentFlow?,
+    private val approvedTypeStatements: PersistentMap<RealVariable, PersistentTypeStatement>,
+    internal val logicStatements: PersistentMap<DataFlowVariable, PersistentList<Implication>>,
+    // RealVariable describes a storage in memory; a pair of RealVariable with its assignment
+    // index at a particular execution point forms an SSA value corresponding to the result of
+    // an initializer.
+    internal val assignmentIndex: PersistentMap<RealVariable, Int>,
+    // RealVariables thus form equivalence sets by values they reference. One is chosen
+    // as a representative of that set, while the rest are mapped to that representative
+    // in `directAliasMap`. `backwardsAliasMap` maps each representative to the rest of the set.
+    internal val directAliasMap: PersistentMap<RealVariable, RealVariable>,
+    private val backwardsAliasMap: PersistentMap<RealVariable, PersistentSet<RealVariable>>,
+) : Flow() {
+    private val level: Int = if (previousFlow != null) previousFlow.level + 1 else 0
+
+    override val knownVariables: Set<RealVariable>
+        get() = approvedTypeStatements.keys + directAliasMap.keys
+
+    override fun unwrapVariable(variable: RealVariable): RealVariable =
+        directAliasMap[variable] ?: variable
+
+    override fun getTypeStatement(variable: RealVariable): TypeStatement? =
+        approvedTypeStatements[unwrapVariable(variable)]?.copy(variable = variable)
+
+    fun lowestCommonAncestor(other: PersistentFlow): PersistentFlow? {
+        var left = this
+        var right = other
+        while (left.level > right.level) {
+            left = left.previousFlow ?: return null
+        }
+        while (right.level > left.level) {
+            right = right.previousFlow ?: return null
+        }
+        while (left != right) {
+            left = left.previousFlow ?: return null
+            right = right.previousFlow ?: return null
+        }
+        return left
+    }
+
+    fun fork(): MutableFlow = MutableFlow(
+        this,
+        approvedTypeStatements.builder(),
+        logicStatements.builder(),
+        assignmentIndex.builder(),
+        directAliasMap.builder(),
+        backwardsAliasMap.builder(),
+    )
+}
+
+class MutableFlow internal constructor(
+    private val previousFlow: PersistentFlow?,
+    internal val approvedTypeStatements: PersistentMap.Builder<RealVariable, PersistentTypeStatement>,
+    internal val logicStatements: PersistentMap.Builder<DataFlowVariable, PersistentList<Implication>>,
+    internal val assignmentIndex: PersistentMap.Builder<RealVariable, Int>,
+    internal val directAliasMap: PersistentMap.Builder<RealVariable, RealVariable>,
+    internal val backwardsAliasMap: PersistentMap.Builder<RealVariable, PersistentSet<RealVariable>>,
+) : Flow() {
+    constructor() : this(
+        null,
+        emptyPersistentHashMapBuilder(),
+        emptyPersistentHashMapBuilder(),
+        emptyPersistentHashMapBuilder(),
+        emptyPersistentHashMapBuilder(),
+        emptyPersistentHashMapBuilder(),
+    )
+
+    override val knownVariables: Set<RealVariable>
+        get() = approvedTypeStatements.keys + directAliasMap.keys
+
+    override fun unwrapVariable(variable: RealVariable): RealVariable =
+        directAliasMap[variable] ?: variable
+
+    override fun getTypeStatement(variable: RealVariable): TypeStatement? =
+        approvedTypeStatements[unwrapVariable(variable)]?.copy(variable = variable)
+
+    fun freeze(): PersistentFlow = PersistentFlow(
+        previousFlow,
+        approvedTypeStatements.build(),
+        logicStatements.build(),
+        assignmentIndex.build(),
+        directAliasMap.build(),
+        backwardsAliasMap.build(),
+    )
+}
+
+private fun <K, V> emptyPersistentHashMapBuilder(): PersistentMap.Builder<K, V> =
+    persistentHashMapOf<K, V>().builder()

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

@@ -7,35 +7,32 @@ package org.jetbrains.kotlin.fir.resolve.dfa
 
 import org.jetbrains.kotlin.fir.types.*
 
-abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceContext) {
+abstract class LogicSystem(protected val context: ConeInferenceContext) {
-    // --------------------------- Flow graph constructors ---------------------------
+    abstract fun joinFlow(flows: Collection<PersistentFlow>, union: Boolean): MutableFlow
-    abstract fun createEmptyFlow(): FLOW
-    abstract fun forkFlow(flow: FLOW): FLOW
-    abstract fun joinFlow(flows: Collection<FLOW>): FLOW
-    abstract fun unionFlow(flows: Collection<FLOW>): FLOW
 
-    // -------------------------------- Flow mutators --------------------------------
     // Returns all known information about the variable, or null if unchanged by this statement:
-    abstract fun addTypeStatement(flow: FLOW, statement: TypeStatement): TypeStatement?
+    abstract fun addTypeStatement(flow: MutableFlow, statement: TypeStatement): TypeStatement?
-    abstract fun addTypeStatements(flow: FLOW, statements: TypeStatements): List<TypeStatement>
+    abstract fun addImplication(flow: MutableFlow, implication: Implication)
-    abstract fun addImplication(flow: FLOW, implication: Implication)
+    abstract fun addLocalVariableAlias(flow: MutableFlow, alias: RealVariable, underlyingVariable: RealVariable)
-    abstract fun addLocalVariableAlias(flow: FLOW, alias: RealVariable, underlyingVariable: RealVariable)
+    abstract fun recordNewAssignment(flow: MutableFlow, variable: RealVariable, index: Int)
-    abstract fun recordNewAssignment(flow: FLOW, variable: RealVariable, index: Int)
+    abstract fun isSameValueIn(a: PersistentFlow, b: MutableFlow, variable: RealVariable): Boolean
-    abstract fun isSameValueIn(a: FLOW, b: FLOW, variable: RealVariable): Boolean
+    abstract fun isSameValueIn(a: PersistentFlow, b: PersistentFlow, variable: RealVariable): Boolean
+
+    fun addTypeStatements(flow: MutableFlow, statements: TypeStatements): List<TypeStatement> =
+        statements.values.mapNotNull { addTypeStatement(flow, it) }
 
     abstract fun translateVariableFromConditionInStatements(
-        flow: FLOW,
+        flow: MutableFlow,
         originalVariable: DataFlowVariable,
         newVariable: DataFlowVariable,
-        shouldRemoveOriginalStatements: Boolean = originalVariable.isSynthetic(),
         transform: (Implication) -> Implication? = { it },
     )
 
-    // This does *not* commit the results to the flow (but it does mutate the flow if removeApprovedOrImpossible=true)
+    abstract fun approveOperationStatement(flow: PersistentFlow, statement: OperationStatement): TypeStatements
     abstract fun approveOperationStatement(
-        flow: FLOW,
+        flow: MutableFlow,
-        approvedStatement: OperationStatement,
+        statement: OperationStatement,
-        removeApprovedOrImpossible: Boolean = false
+        removeApprovedOrImpossible: Boolean,
     ): TypeStatements
 
     protected abstract fun ConeKotlinType.isAcceptableForSmartcast(): Boolean
@@ -57,7 +54,7 @@ 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))!! })
+                this[variable] = this[variable]?.let { and(listOf(it, rightStatement))!! } ?: rightStatement
             }
         }
     }

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

@@ -7,98 +7,32 @@ package org.jetbrains.kotlin.fir.resolve.dfa
 
 import kotlinx.collections.immutable.*
 import org.jetbrains.kotlin.fir.types.ConeInferenceContext
-import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.types.AbstractTypeChecker
 import java.util.*
 import kotlin.math.max
 
-data class PersistentTypeStatement(
+abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSystem(context) {
-    override val variable: RealVariable,
-    override val exactType: PersistentSet<ConeKotlinType>,
-) : TypeStatement()
-
-typealias PersistentApprovedTypeStatements = PersistentMap<RealVariable, PersistentTypeStatement>
-typealias PersistentImplications = PersistentMap<DataFlowVariable, PersistentList<Implication>>
-
-class PersistentFlow : Flow {
-    val previousFlow: PersistentFlow?
-    var approvedTypeStatements: PersistentApprovedTypeStatements
-    var logicStatements: PersistentImplications
-    val level: Int
-
-    /*
-     * val x = a
-     * val y = a
-     *
-     * directAliasMap: { x -> a, y -> a}
-     * backwardsAliasMap: { a -> [x, y] }
-     */
-    var directAliasMap: PersistentMap<RealVariable, RealVariable>
-    var backwardsAliasMap: PersistentMap<RealVariable, PersistentSet<RealVariable>>
-
-    var assignmentIndex: PersistentMap<RealVariable, Int>
-
-    constructor(previousFlow: PersistentFlow) {
-        this.previousFlow = previousFlow
-        approvedTypeStatements = previousFlow.approvedTypeStatements
-        logicStatements = previousFlow.logicStatements
-        level = previousFlow.level + 1
-
-        directAliasMap = previousFlow.directAliasMap
-        backwardsAliasMap = previousFlow.backwardsAliasMap
-        assignmentIndex = previousFlow.assignmentIndex
-    }
-
-    constructor() {
-        previousFlow = null
-        approvedTypeStatements = persistentHashMapOf()
-        logicStatements = persistentHashMapOf()
-        level = 1
-
-        directAliasMap = persistentMapOf()
-        backwardsAliasMap = persistentMapOf()
-        assignmentIndex = persistentMapOf()
-    }
-
-    override val knownVariables: Set<RealVariable>
-        get() = approvedTypeStatements.keys + directAliasMap.keys
-
-    override fun unwrapVariable(variable: RealVariable): RealVariable =
-        directAliasMap[variable] ?: variable
-
-    override fun getTypeStatement(variable: RealVariable): TypeStatement? =
-        approvedTypeStatements[unwrapVariable(variable)]?.copy(variable = variable)
-}
-
-abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSystem<PersistentFlow>(context) {
     abstract val variableStorage: VariableStorageImpl
 
-    override fun createEmptyFlow(): PersistentFlow =
+    override fun joinFlow(flows: Collection<PersistentFlow>, union: Boolean): MutableFlow {
-        PersistentFlow()
-
-    override fun forkFlow(flow: PersistentFlow): PersistentFlow =
-        PersistentFlow(flow)
-
-    override fun joinFlow(flows: Collection<PersistentFlow>): PersistentFlow =
-        foldFlow(flows, allExecute = false)
-
-    override fun unionFlow(flows: Collection<PersistentFlow>): PersistentFlow =
-        foldFlow(flows, allExecute = true)
-
-    private fun foldFlow(flows: Collection<PersistentFlow>, allExecute: Boolean): PersistentFlow {
         when (flows.size) {
-            0 -> return createEmptyFlow()
+            0 -> return MutableFlow()
-            1 -> return forkFlow(flows.first())
+            1 -> return flows.first().fork()
         }
 
-        val commonFlow = flows.reduce(::lowestCommonFlow)
+        val commonFlow = flows.reduce { a, b ->
-        val result = forkFlow(commonFlow)
+            // If you're debugging this assertion error, most likely cause is that a node is not
+            // marked as dead when all its input edges are dead. In that case it will have an empty flow,
+            // and joining that with a non-empty flow from another branch will fail.
+            a.lowestCommonAncestor(b) ?: error("no common ancestor in $a, $b")
+        }
+        val result = commonFlow.fork()
 
         // If a variable was reassigned in one branch, it was reassigned at the join point.
         val reassignedVariables = mutableMapOf<RealVariable, Int>()
         for (flow in flows) {
             for ((variable, index) in flow.assignmentIndex) {
-                if (commonFlow.assignmentIndex[variable] != index) {
+                if (result.assignmentIndex[variable] != index) {
                     // Ideally we should generate an entirely new index here, but it doesn't really
                     // matter; the important part is that it's different from `commonFlow.previousFlow`.
                     reassignedVariables[variable] = max(index, reassignedVariables[variable] ?: 0)
@@ -109,22 +43,21 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
             recordNewAssignment(result, variable, index)
         }
 
-        // TODO: if `allExecute`, then all aliases from all flows are valid so long as other flows don't have a contradicting one
+        // TODO: if `union`, then all aliases from all flows are valid so long as other flows don't have a contradicting one
         for ((from, to) in flows.first().directAliasMap) {
             // If `from -> to` is still in `result` (was not removed by the above code), then it is also in all `flows`,
             // as the only way to break aliasing is by reassignment.
-            if (result.directAliasMap[from] != to && flows.all { it.directAliasMap[from] == to }) {
+            if (result.directAliasMap[from] != to && flows.all { it.unwrapVariable(from) == to }) {
                 // if (p) { y = x } else { y = x } <-- after `if`, `y -> x` is in all `flows`, but not in `result`
                 // (which was forked from the flow before the `if`)
                 addLocalVariableAlias(result, from, to)
             }
         }
 
-        val approvedTypeStatements = result.approvedTypeStatements.builder()
         flows.flatMapTo(mutableSetOf()) { it.knownVariables }.forEach computeStatement@{ variable ->
             val statement = if (variable in result.directAliasMap) {
                 return@computeStatement // statements about alias == statements about aliased variable
-            } else if (!allExecute) {
+            } else if (!union) {
                 // if (condition) { /* x: S1 */ } else { /* x: S2 */ } // -> x: S1 | S2
                 or(flows.mapTo(mutableSetOf()) { it.getTypeStatement(variable) ?: return@computeStatement })
             } else if (variable !in reassignedVariables) {
@@ -145,19 +78,18 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
                 or(byAssignment.values.mapTo(mutableSetOf()) { and(it.filterNotNull()) ?: return@computeStatement })
             }
             if (statement?.isNotEmpty == true) {
-                approvedTypeStatements[variable] = statement.toPersistent()
+                result.approvedTypeStatements[variable] = statement.toPersistent()
             }
         }
-        result.approvedTypeStatements = approvedTypeStatements.build()
         // TODO: compute common implications?
         return result
     }
 
-    override fun addLocalVariableAlias(flow: PersistentFlow, alias: RealVariable, underlyingVariable: RealVariable) {
+    override fun addLocalVariableAlias(flow: MutableFlow, alias: RealVariable, underlyingVariable: RealVariable) {
         flow.addAliases(persistentSetOf(alias), flow.unwrapVariable(underlyingVariable))
     }
 
-    private fun PersistentFlow.replaceVariable(variable: RealVariable, replacement: RealVariable?) {
+    private fun MutableFlow.replaceVariable(variable: RealVariable, replacement: RealVariable?) {
         val original = directAliasMap[variable]
         if (original != null) {
             // All statements should've been made about whatever variable this is an alias to. There is nothing to replace.
@@ -170,9 +102,9 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
             // was broken in another flow. However, in *this* flow dependent variables should have no information attached to them.
 
             val siblings = backwardsAliasMap.getValue(original).remove(variable)
-            directAliasMap -= variable
+            directAliasMap.remove(variable)
-            backwardsAliasMap = if (siblings.isNotEmpty()) {
+            if (siblings.isNotEmpty()) {
-                backwardsAliasMap.put(original, siblings)
+                backwardsAliasMap[original] = siblings
             } else {
                 backwardsAliasMap.remove(original)
             }
@@ -188,8 +120,8 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
                     variableStorage.copyRealVariableWithRemapping(dependent, variable, it)
                 })
             }
-            logicStatements = logicStatements.replaceVariable(variable, replacementOrNext)
+            logicStatements.replaceVariable(variable, replacementOrNext)
-            approvedTypeStatements = approvedTypeStatements.replaceVariable(variable, replacementOrNext)
+            approvedTypeStatements.replaceVariable(variable, replacementOrNext)
             if (aliases != null) {
                 backwardsAliasMap -= variable
                 if (replacementOrNext != null) {
@@ -200,79 +132,72 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         }
     }
 
-    private fun PersistentFlow.addAliases(aliases: PersistentSet<RealVariable>, target: RealVariable) {
+    private fun MutableFlow.addAliases(aliases: PersistentSet<RealVariable>, target: RealVariable) {
         val withoutSelf = aliases - target
         if (withoutSelf.isNotEmpty()) {
-            directAliasMap = withoutSelf.associateWithTo(directAliasMap.builder()) { target }.build()
+            withoutSelf.associateWithTo(directAliasMap) { target }
-            backwardsAliasMap = backwardsAliasMap.put(target, { withoutSelf }, { it + withoutSelf })
+            backwardsAliasMap[target] = backwardsAliasMap[target]?.addAll(withoutSelf) ?: withoutSelf
         }
     }
 
-    private fun PersistentFlow.completeStatement(statement: TypeStatement): PersistentTypeStatement? {
+    override fun addTypeStatement(flow: MutableFlow, statement: TypeStatement): TypeStatement? {
         if (statement.exactType.isEmpty()) return null
         val variable = statement.variable
-        val oldExactType = approvedTypeStatements[variable]?.exactType
+        val oldExactType = flow.approvedTypeStatements[variable]?.exactType
         val newExactType = oldExactType?.addAll(statement.exactType) ?: statement.exactType.toPersistentSet()
         if (newExactType === oldExactType) return null
-        return PersistentTypeStatement(variable, newExactType)
+        return PersistentTypeStatement(variable, newExactType).also { flow.approvedTypeStatements[variable] = it }
     }
 
-    override fun addTypeStatement(flow: PersistentFlow, statement: TypeStatement): TypeStatement? =
+    override fun addImplication(flow: MutableFlow, implication: Implication) {
-        flow.completeStatement(statement)?.also {
-            flow.approvedTypeStatements = flow.approvedTypeStatements.put(it.variable, it)
-        }
-
-    override fun addTypeStatements(flow: PersistentFlow, statements: TypeStatements): List<TypeStatement> {
-        val builder = flow.approvedTypeStatements.builder()
-        val result = statements.values.mapNotNull { statement ->
-            flow.completeStatement(statement)?.also { builder[it.variable] = it }
-        }
-        flow.approvedTypeStatements = builder.build()
-        return result
-    }
-
-    override fun addImplication(flow: PersistentFlow, implication: Implication) {
         val effect = implication.effect
         if (effect == implication.condition) return
         if (effect is TypeStatement && (effect.isEmpty ||
                     flow.approvedTypeStatements[effect.variable]?.exactType?.containsAll(effect.exactType) == true)
         ) return
         val variable = implication.condition.variable
-        flow.logicStatements = flow.logicStatements.put(variable, { persistentListOf(implication) }, { it + implication })
+        flow.logicStatements[variable] = flow.logicStatements[variable]?.add(implication) ?: persistentListOf(implication)
     }
 
     override fun translateVariableFromConditionInStatements(
-        flow: PersistentFlow,
+        flow: MutableFlow,
         originalVariable: DataFlowVariable,
         newVariable: DataFlowVariable,
-        shouldRemoveOriginalStatements: Boolean,
         transform: (Implication) -> Implication?
     ) {
-        val statements = flow.logicStatements[originalVariable]?.takeIf { it.isNotEmpty() } ?: return
+        val statements = if (originalVariable.isSynthetic())
+            flow.logicStatements.remove(originalVariable)
+        else
+            flow.logicStatements[originalVariable]
+        if (statements.isNullOrEmpty()) return
         val existing = flow.logicStatements[newVariable] ?: persistentListOf()
-        val result = statements.mapNotNullTo(existing.builder()) {
+        flow.logicStatements[newVariable] = statements.mapNotNullTo(existing.builder()) {
             // TODO: rethink this API - technically it permits constructing invalid flows
             //  (transform can replace the variable in the condition with anything)
             transform(OperationStatement(newVariable, it.condition.operation) implies it.effect)
         }.build()
-        if (shouldRemoveOriginalStatements) {
-            flow.logicStatements -= originalVariable
-        }
-        flow.logicStatements = flow.logicStatements.put(newVariable, result)
     }
 
     private val nullableNothingType = context.session.builtinTypes.nullableNothingType.type
     private val anyType = context.session.builtinTypes.anyType.type
 
+    override fun approveOperationStatement(flow: PersistentFlow, statement: OperationStatement): TypeStatements =
+        approveOperationStatement(flow.logicStatements, statement, removeApprovedOrImpossible = false)
+
     override fun approveOperationStatement(
-        flow: PersistentFlow,
+        flow: MutableFlow,
-        approvedStatement: OperationStatement,
+        statement: OperationStatement,
         removeApprovedOrImpossible: Boolean,
+    ): TypeStatements = approveOperationStatement(flow.logicStatements, statement, removeApprovedOrImpossible)
+
+    private fun approveOperationStatement(
+        logicStatements: Map<DataFlowVariable, PersistentList<Implication>>,
+        approvedStatement: OperationStatement,
+        removeApprovedOrImpossible: Boolean
     ): TypeStatements {
         val result = mutableMapOf<RealVariable, MutableTypeStatement>()
         val queue = LinkedList<OperationStatement>().apply { this += approvedStatement }
         val approved = mutableSetOf<OperationStatement>()
-        val logicStatements = flow.logicStatements.builder()
         while (queue.isNotEmpty()) {
             val next = queue.removeFirst()
             // Defense from cycles in facts
@@ -297,73 +222,64 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
                 }
                 removeApprovedOrImpossible && knownValue != null
             }
-            if (stillUnknown.isEmpty()) {
+            if (stillUnknown != statements && logicStatements is MutableMap) {
-                logicStatements.remove(variable)
+                if (stillUnknown.isEmpty()) {
-            } else if (stillUnknown != statements) {
+                    logicStatements.remove(variable)
-                logicStatements[variable] = stillUnknown
+                } else {
+                    logicStatements[variable] = stillUnknown
+                }
             }
         }
-        flow.logicStatements = logicStatements.build()
         return result
     }
 
-    override fun recordNewAssignment(flow: PersistentFlow, variable: RealVariable, index: Int) {
+    override fun recordNewAssignment(flow: MutableFlow, variable: RealVariable, index: Int) {
         flow.replaceVariable(variable, null)
-        flow.assignmentIndex = flow.assignmentIndex.put(variable, index)
+        flow.assignmentIndex[variable] = index
     }
 
+    override fun isSameValueIn(a: PersistentFlow, b: MutableFlow, variable: RealVariable): Boolean =
+        a.assignmentIndex[variable] == b.assignmentIndex[variable]
+
     override fun isSameValueIn(a: PersistentFlow, b: PersistentFlow, variable: RealVariable): Boolean =
         a.assignmentIndex[variable] == b.assignmentIndex[variable]
 }
 
-private fun lowestCommonFlow(left: PersistentFlow, right: PersistentFlow): PersistentFlow {
-    val level = minOf(left.level, right.level)
-
-    @Suppress("NAME_SHADOWING")
-    var left = left
-    while (left.level > level) {
-        left = left.previousFlow!!
-    }
-    @Suppress("NAME_SHADOWING")
-    var right = right
-    while (right.level > level) {
-        right = right.previousFlow!!
-    }
-    while (left != right) {
-        left = left.previousFlow!!
-        right = right.previousFlow!!
-    }
-    return left
-}
-
 private fun TypeStatement.toPersistent(): PersistentTypeStatement = when (this) {
     is PersistentTypeStatement -> this
     else -> PersistentTypeStatement(variable, exactType.toPersistentSet())
 }
 
 @JvmName("replaceVariableInStatements")
-private fun PersistentApprovedTypeStatements.replaceVariable(from: RealVariable, to: RealVariable?): PersistentApprovedTypeStatements {
+private fun MutableMap<RealVariable, PersistentTypeStatement>.replaceVariable(from: RealVariable, to: RealVariable?) {
-    val existing = this[from] ?: return this
+    val existing = remove(from) ?: return
-    return if (to != null) remove(from).put(to, existing.copy(variable = to)) else remove(from)
+    if (to != null) {
+        put(to, existing.copy(variable = to))
+    }
 }
 
 @JvmName("replaceVariableInImplications")
-private fun PersistentImplications.replaceVariable(from: RealVariable, to: RealVariable?): PersistentImplications =
+private fun MutableMap<DataFlowVariable, PersistentList<Implication>>.replaceVariable(from: RealVariable, to: RealVariable?) {
-    mutate { result ->
+    val existing = remove(from)
-        for ((variable, implications) in this) {
+    val toReplace = entries.mapNotNull { (variable, implications) ->
-            val newImplications = if (to != null) {
+        val newImplications = if (to != null) {
-                implications.replaceAll { it.replaceVariable(from, to) }
+            implications.replaceAll { it.replaceVariable(from, to) }
-            } else {
+        } else {
-                implications.removeAll { it.effect.variable == from }
+            implications.removeAll { it.effect.variable == from }
-            }
-            if (newImplications.isNotEmpty()) {
-                result[implications.first().condition.variable] = newImplications
-            } else {
-                result.remove(variable)
-            }
         }
-        result.remove(from)
+        if (newImplications != implications) variable to newImplications else null
     }
+    for ((variable, implications) in toReplace) {
+        if (implications.isEmpty()) {
+            remove(variable)
+        } else {
+            put(variable, implications)
+        }
+    }
+    if (existing != null && to != null) {
+        put(to, existing.replaceAll { it.replaceVariable(from, to) })
+    }
+}
 
 private inline fun <T> PersistentList<T>.replaceAll(block: (T) -> T): PersistentList<T> =
     mutate { result ->

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

@@ -20,18 +20,14 @@ fun ConeConstantReference.toOperation(): Operation? = when (this) {
 }
 
 // Returns `null` if the statement is always false.
-fun <F : Flow> LogicSystem<F>.approveContractStatement(
+fun LogicSystem.approveContractStatement(
-    flow: F,
     statement: ConeBooleanExpression,
     arguments: Array<out DataFlowVariable?>, // 0 = receiver (null if doesn't exist)
     substitutor: ConeSubstitutor?,
-    removeApprovedOrImpossible: Boolean = false,
+    approveOperationStatement: (OperationStatement) -> TypeStatements
 ): TypeStatements? {
-    fun OperationStatement.approve() =
-        approveOperationStatement(flow, this, removeApprovedOrImpossible)
-
     fun DataFlowVariable.processEqNull(isEq: Boolean): TypeStatements =
-        OperationStatement(this, if (isEq) Operation.EqNull else Operation.NotEqNull).approve()
+        approveOperationStatement(OperationStatement(this, if (isEq) Operation.EqNull else Operation.NotEqNull))
 
     fun ConeBooleanExpression.visit(inverted: Boolean): TypeStatements? = when (this) {
         is ConeBooleanConstantReference ->
@@ -61,9 +57,7 @@ fun <F : Flow> LogicSystem<F>.approveContractStatement(
         is ConeIsNullPredicate ->
             arguments.getOrNull(arg.parameterIndex + 1)?.processEqNull(inverted == isNegated) ?: mapOf()
         is ConeBooleanValueParameterReference ->
-            arguments.getOrNull(parameterIndex + 1)?.let {
+            arguments.getOrNull(parameterIndex + 1)?.let { approveOperationStatement(it eq !inverted) } ?: mapOf()
-                OperationStatement(it, if (inverted) Operation.EqFalse else Operation.EqTrue).approve()
-            } ?: mapOf()
         is ConeBinaryLogicExpression -> {
             val a = left.visit(inverted)
             val b = right.visit(inverted)

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.resolve.dfa
 
+import kotlinx.collections.immutable.PersistentSet
 import org.jetbrains.kotlin.fir.types.ConeErrorType
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import kotlin.contracts.ExperimentalContracts
@@ -12,6 +13,11 @@ import kotlin.contracts.contract
 
 // --------------------------------------- Facts ---------------------------------------
 
+data class PersistentTypeStatement(
+    override val variable: RealVariable,
+    override val exactType: PersistentSet<ConeKotlinType>,
+) : TypeStatement()
+
 class MutableTypeStatement(
     override val variable: RealVariable,
     override val exactType: MutableSet<ConeKotlinType> = linkedSetOf(),

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

@@ -5,7 +5,6 @@
 
 package org.jetbrains.kotlin.fir.resolve.dfa
 
-import kotlinx.collections.immutable.PersistentMap
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.expressions.*
@@ -16,41 +15,6 @@ import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirSyntheticPropertySymbol
 import org.jetbrains.kotlin.fir.types.*
-import kotlin.contracts.ExperimentalContracts
-import kotlin.contracts.InvocationKind
-import kotlin.contracts.contract
-
-@OptIn(ExperimentalContracts::class)
-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, valueProducer())
-    } else {
-        put(key, remappingFunction(existing))
-    }
-}
-
-@OptIn(ExperimentalContracts::class)
-internal inline fun <K, V> PersistentMap<K, V>.put(
-    key: K,
-    valueProducer: () -> V,
-    remappingFunction: (existing: V) -> V
-): PersistentMap<K, V> {
-    contract {
-        callsInPlace(remappingFunction, InvocationKind.AT_MOST_ONCE)
-        callsInPlace(valueProducer, InvocationKind.AT_MOST_ONCE)
-    }
-    val existing = this[key]
-    return if (existing == null) {
-        put(key, valueProducer())
-    } else {
-        put(key, remappingFunction(existing))
-    }
-}
 
 fun TypeStatement?.smartCastedType(context: ConeTypeContext, originalType: ConeKotlinType): ConeKotlinType =
     if (this != null && exactType.isNotEmpty()) {