[FIR] Replace Flow with DelegatingFlow

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

@@ -111,6 +111,13 @@ class DataFlowVariableStorage {
         }
     }
 
+    fun removeVariableIfSynthetic(variable: DataFlowVariable): FirElement? {
+        if (variable is SyntheticDataFlowVariable) {
+            return removeVariable(variable)
+        }
+        return null
+    }
+
     operator fun get(variable: DataFlowVariable): FirElement? {
         return variable.fir
     }

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

@@ -0,0 +1,364 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.fir.resolve.dfa
+
+import com.google.common.collect.ArrayListMultimap
+import com.google.common.collect.Multimap
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+
+interface DelegatingFlow {
+    companion object {
+        fun createDefaultFlow(): DelegatingFlow = DelegatingFlowImpl(null)
+    }
+
+    fun getApprovedInfo(variable: RealDataFlowVariable): FirDataFlowInfo?
+    fun getConditionalInfos(variable: DataFlowVariable): Collection<ConditionalFirDataFlowInfo>
+
+    fun getVariablesInApprovedInfos(): Collection<RealDataFlowVariable>
+
+    fun createNewFlow(): DelegatingFlow
+}
+
+private class DelegatingFlowImpl(
+    val previousFlow: DelegatingFlowImpl?,
+    val approvedInfos: ApprovedInfos = mutableMapOf(),
+    val conditionalInfos: ConditionalInfos = ArrayListMultimap.create()
+) : DelegatingFlow {
+    val level: Int = previousFlow?.level?.plus(1) ?: 0
+
+    override fun getApprovedInfo(variable: RealDataFlowVariable): FirDataFlowInfo? {
+        val info = MutableFirDataFlowInfo(mutableSetOf(), mutableSetOf())
+        collect { flow ->
+            flow.approvedInfos[variable]?.let { info += it }
+            true
+        }
+        return info.takeIf { it.isNotEmpty }
+    }
+
+    override fun getConditionalInfos(variable: DataFlowVariable): Collection<ConditionalFirDataFlowInfo> {
+        val result = mutableSetOf<ConditionalFirDataFlowInfo>()
+        val isReal = !variable.isSynthetic()
+        collect {
+            val infos = it.conditionalInfos[variable]
+            result += infos
+            /*
+             * Here we use invariant that if we found some conditional info about synthetic variable
+             *   on some level then there is definitely no another info on other levels of flow
+             * It's correct because we can add condtional info about synthetic variable only
+             *   when we first met expression corresponding to that variable
+             */
+            isReal || infos.isEmpty()
+        }
+        return result
+    }
+
+    override fun getVariablesInApprovedInfos(): Collection<RealDataFlowVariable> {
+        val result = mutableSetOf<RealDataFlowVariable>()
+        collect {
+            result += approvedInfos.keys
+            true
+        }
+        return result
+    }
+
+    override fun createNewFlow(): DelegatingFlowImpl {
+        return DelegatingFlowImpl(this)
+    }
+
+    private inline fun collect(block: (DelegatingFlowImpl) -> Boolean) {
+        var flow: DelegatingFlowImpl? = this
+        while (flow != null) {
+            val shouldContinue = block(flow)
+            if (!shouldContinue) break
+            flow = flow.previousFlow
+        }
+    }
+}
+
+fun DelegatingFlow.approvedInfosFromTopFlow(): Map<RealDataFlowVariable, FirDataFlowInfo> {
+    require(this is DelegatingFlowImpl)
+    return approvedInfos
+}
+
+fun DelegatingFlow.conditionalInfosFromTopFlow(): Multimap<DataFlowVariable, ConditionalFirDataFlowInfo> {
+    require(this is DelegatingFlowImpl)
+    return ImmutableMultimap(conditionalInfos)
+}
+
+fun ApprovedInfos.mergeInfo(other: Map<RealDataFlowVariable, FirDataFlowInfo>) {
+    other.forEach { (variable, info) ->
+        merge(variable, info) { existingInfo, newInfo ->
+            (existingInfo as MutableFirDataFlowInfo) += newInfo
+            existingInfo
+        }
+    }
+}
+
+private class ImmutableMultimap<K, V>(private val original: Multimap<K, V>) : Multimap<K, V> by original {
+    override fun put(p0: K?, p1: V?): Boolean {
+        throw IllegalStateException()
+    }
+
+    override fun remove(p0: Any?, p1: Any?): Boolean {
+        throw IllegalStateException()
+    }
+
+    override fun putAll(p0: Multimap<out K, out V>): Boolean {
+        throw IllegalStateException()
+    }
+
+    override fun putAll(p0: K?, p1: MutableIterable<V>): Boolean {
+        throw IllegalStateException()
+    }
+
+    override fun removeAll(p0: Any?): MutableCollection<V> {
+        throw IllegalStateException()
+    }
+
+    override fun replaceValues(p0: K?, p1: MutableIterable<V>): MutableCollection<V> {
+        throw IllegalStateException()
+    }
+}
+
+abstract class DelegatingLogicSystem(private val context: DataFlowInferenceContext) {
+
+    // ------------------------------- Flow operations -------------------------------
+
+    fun mergeFlows(flows: Collection<DelegatingFlow>): DelegatingFlow {
+        if (flows.isEmpty()) return DelegatingFlow.createDefaultFlow()
+        flows.singleOrNull()?.let { return it }
+
+        @Suppress("UNCHECKED_CAST", "NAME_SHADOWING")
+        val flows = flows as Collection<DelegatingFlowImpl>
+        val commonFlow = flows.reduce(this::lowestCommonFlow)
+
+        val approvedInfosFromAllFlows = mutableListOf<ApprovedInfos>()
+        flows.forEach {
+            approvedInfosFromAllFlows += it.collectInfos(commonFlow)
+        }
+
+        // approved info
+        val variablesFromApprovedInfos = approvedInfosFromAllFlows.map { it.keys }.intersectSets()
+        for (variable in variablesFromApprovedInfos) {
+            val infos = approvedInfosFromAllFlows.mapNotNull { it[variable] }
+            if (infos.size != flows.size) continue
+            val intersectedInfo = or(infos)
+            if (intersectedInfo.isEmpty) continue
+            val existingInfo = commonFlow.approvedInfos[variable]
+            if (existingInfo == null) {
+                commonFlow.approvedInfos[variable] = intersectedInfo
+            } else {
+                (existingInfo as MutableFirDataFlowInfo) += intersectedInfo
+            }
+        }
+
+        updateAllReceivers(commonFlow)
+
+        return commonFlow
+    }
+
+    /*
+     *  used for:
+     *   1. val b = x is String
+     *   2. b = x is String
+     *   3. !b | b.not()   for Booleans
+     */
+    fun changeVariableForConditionFlow(
+        flow: DelegatingFlow,
+        sourceVariable: DataFlowVariable,
+        newVariable: DataFlowVariable,
+        transform: ((ConditionalFirDataFlowInfo) -> ConditionalFirDataFlowInfo)? = null
+    ) {
+        require(flow is DelegatingFlowImpl)
+        var infos = flow.getConditionalInfos(sourceVariable)
+        if (transform != null) {
+            infos = infos.map(transform)
+        }
+        flow.conditionalInfos.putAll(newVariable, infos)
+        if (sourceVariable.isSynthetic()) {
+            flow.conditionalInfos.removeAll(sourceVariable)
+        }
+    }
+
+    fun approveFactsInsideFlow(
+        variable: DataFlowVariable,
+        condition: Condition,
+        flow: DelegatingFlow,
+        shouldCreateNewFlow: Boolean,
+        shouldRemoveSynthetics: Boolean
+    ): DelegatingFlow {
+        require(flow is DelegatingFlowImpl)
+
+        val notApprovedFacts: Collection<ConditionalFirDataFlowInfo> = if (shouldRemoveSynthetics && variable.isSynthetic()) {
+            flow.conditionalInfos.removeAll(variable)
+        } else {
+            flow.getConditionalInfos(variable)
+        }
+
+        val resultFlow = if (shouldCreateNewFlow) flow.createNewFlow() else flow
+        if (notApprovedFacts.isEmpty()) {
+            return resultFlow
+        }
+
+        val newFacts = ArrayListMultimap.create<RealDataFlowVariable, FirDataFlowInfo>()
+        notApprovedFacts.forEach {
+            if (it.condition == condition) {
+                newFacts.put(it.variable, it.info)
+            }
+        }
+
+        val updatedReceivers = mutableSetOf<RealDataFlowVariable>()
+
+        newFacts.asMap().forEach { (variable, infos) ->
+            @Suppress("NAME_SHADOWING")
+            val info = MutableFirDataFlowInfo()
+            infos.forEach {
+                info += it
+            }
+            if (variable.isThisReference) {
+                updatedReceivers += variable
+            }
+            addApprovedInfo(resultFlow, variable, info)
+        }
+        updatedReceivers.forEach {
+            processUpdatedReceiverVariable(resultFlow, it)
+        }
+        return resultFlow
+    }
+
+    fun addApprovedInfo(flow: DelegatingFlow, variable: RealDataFlowVariable, info: FirDataFlowInfo) {
+        assert(info is MutableFirDataFlowInfo)
+        require(flow is DelegatingFlowImpl)
+        flow.approvedInfos.addInfo(variable, info)
+        if (variable.isThisReference) {
+            processUpdatedReceiverVariable(flow, variable)
+        }
+    }
+
+    fun addConditionalInfo(flow: DelegatingFlow, variable: DataFlowVariable, info: ConditionalFirDataFlowInfo) {
+        require(flow is DelegatingFlowImpl)
+        flow.conditionalInfos.put(variable, info)
+    }
+
+    // ------------------------------- DataFlowInfo operations -------------------------------
+
+    fun orForVerifiedFacts(
+        left: Map<RealDataFlowVariable, FirDataFlowInfo>,
+        right: Map<RealDataFlowVariable, FirDataFlowInfo>
+    ): ApprovedInfos {
+        if (left.isNullOrEmpty() || right.isNullOrEmpty()) return mutableMapOf()
+        val map = mutableMapOf<RealDataFlowVariable, FirDataFlowInfo>()
+        for (variable in left.keys.intersect(right.keys)) {
+            val leftInfo = left.getValue(variable)
+            val rightInfo = right.getValue(variable)
+            map[variable] = or(listOf(leftInfo, rightInfo))
+        }
+        return map
+    }
+
+    fun approveFactTo(
+        destination: ApprovedInfos,
+        condition: Condition,
+        notApprovedFacts: Collection<ConditionalFirDataFlowInfo>
+    ) {
+        if (notApprovedFacts.isEmpty()) {
+            return
+        }
+        notApprovedFacts.forEach {
+            if (it.condition == condition) {
+                destination.addInfo(it.variable, it.info)
+            }
+        }
+    }
+
+    fun approveFactTo(
+        destination: ApprovedInfos,
+        variable: DataFlowVariable,
+        condition: Condition,
+        flow: DelegatingFlow
+    ) {
+        require(flow is DelegatingFlowImpl)
+        val notApprovedFacts: Collection<ConditionalFirDataFlowInfo> = flow.getConditionalInfos(variable)
+        approveFactTo(destination, condition, notApprovedFacts)
+    }
+
+    fun approveFact(condition: Condition, notApprovedFacts: Collection<ConditionalFirDataFlowInfo>): ApprovedInfos {
+        return mutableMapOf<RealDataFlowVariable, FirDataFlowInfo>().apply { approveFactTo(this, condition, notApprovedFacts)}
+    }
+
+    fun approveFact(
+        variable: DataFlowVariable,
+        condition: Condition,
+        flow: DelegatingFlow
+    ): ApprovedInfos {
+        return mutableMapOf<RealDataFlowVariable, FirDataFlowInfo>().apply { approveFactTo(this, variable, condition, flow)}
+    }
+
+    // ------------------------------- Util functions -------------------------------
+
+    abstract fun processUpdatedReceiverVariable(flow: DelegatingFlow, variable: RealDataFlowVariable)
+
+    abstract fun updateAllReceivers(flow: DelegatingFlow)
+
+    // ------------------------------- Util functions -------------------------------
+
+    private fun <E> Collection<Set<E>>.intersectSets(): Set<E> = takeIf { isNotEmpty() }?.reduce { x, y -> x.intersect(y) } ?: emptySet()
+
+    private fun or(infos: Collection<FirDataFlowInfo>): FirDataFlowInfo {
+        infos.singleOrNull()?.let { return it }
+        val exactType = orTypes(infos.map { it.exactType })
+        val exactNotType = orTypes(infos.map { it.exactNotType })
+        return FirDataFlowInfo(exactType, exactNotType)
+    }
+
+    private fun orTypes(types: Collection<Set<ConeKotlinType>>): Set<ConeKotlinType> {
+        if (types.any { it.isEmpty() }) return emptySet()
+        val allTypes = types.flatMapTo(mutableSetOf()) { it }
+        val commonTypes = allTypes.toMutableSet()
+        types.forEach { commonTypes.retainAll(it) }
+        val differentTypes = allTypes - commonTypes
+        context.commonSuperTypeOrNull(differentTypes.toList())?.let { commonTypes += it }
+        return commonTypes
+    }
+
+    private fun lowestCommonFlow(left: DelegatingFlowImpl, right: DelegatingFlowImpl): DelegatingFlowImpl {
+        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 DelegatingFlowImpl.collectInfos(untilFlow: DelegatingFlowImpl): ApprovedInfos {
+        val approvedInfos: ApprovedInfos = mutableMapOf()
+
+        // val conditionalInfos: ConditionalInfos = ArrayListMultimap.create()
+
+        var flow = this
+        while (flow != untilFlow) {
+            flow.approvedInfos.forEach { (variable, info) -> approvedInfos.addInfo(variable, info) }
+            /*
+             * we don't join conditional infos yet
+            flow.conditionalInfos.asMap().forEach { (variable, infos) ->
+                conditionalInfos.putAll(variable, infos)
+            }
+             */
+            flow = flow.previousFlow!!
+        }
+        return approvedInfos
+    }
+}

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

@@ -41,9 +41,9 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
     private val receiverStack: ImplicitReceiverStackImpl = transformer.implicitReceiverStack
 
     private val graphBuilder = ControlFlowGraphBuilder()
-    private val logicSystem = LogicSystem(context)
+    private val logicSystem = DelegatingLogicSystemImpl(context)
     private val variableStorage = DataFlowVariableStorage()
-    private val flowOnNodes = mutableMapOf<CFGNode<*>, Flow>()
+    private val flowOnNodes = mutableMapOf<CFGNode<*>, DelegatingFlow>()
 
     private val variablesForWhenConditions = mutableMapOf<WhenBranchConditionExitNode, DataFlowVariable>()
 
@@ -58,7 +58,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
          * If there is no useful information there is no data flow variable also
          */
         val variable = qualifiedAccessExpression.realVariable?.variableUnderAlias ?: return null
-        return graphBuilder.lastNode.flow.approvedFacts(variable)?.exactType ?: return null
+        return graphBuilder.lastNode.flow.getApprovedInfo(variable)?.exactType ?: return null
     }
 
     // ----------------------------------- Named function -----------------------------------
@@ -111,7 +111,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
         val type = typeOperatorCall.conversionTypeRef.coneTypeSafe<ConeKotlinType>() ?: return
         val operandVariable = getOrCreateRealVariable(typeOperatorCall.argument)?.variableUnderAlias ?: return
 
-        var flow = node.flow
+        val flow = node.flow
         when (typeOperatorCall.operation) {
             FirOperation.IS, FirOperation.NOT_IS -> {
                 val expressionVariable = getOrCreateSyntheticVariable(typeOperatorCall)
@@ -122,14 +122,16 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
                 fun chooseInfo(trueBranch: Boolean) =
                     if ((typeOperatorCall.operation == FirOperation.IS) == trueBranch) trueInfo else falseInfo
 
-                flow = flow.addNotApprovedFact(
+                logicSystem.addConditionalInfo(
+                    flow,
                     expressionVariable,
                     ConditionalFirDataFlowInfo(
                         EqTrue, operandVariable, chooseInfo(true)
                     )
                 )
 
-                flow = flow.addNotApprovedFact(
+                logicSystem.addConditionalInfo(
+                    flow,
                     expressionVariable,
                     ConditionalFirDataFlowInfo(
                         EqFalse, operandVariable, chooseInfo(false)
@@ -138,17 +140,20 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
             }
 
             FirOperation.AS -> {
-                flow = addApprovedFact(flow, operandVariable, FirDataFlowInfo(setOf(type), emptySet()))
+                logicSystem.addApprovedInfo(flow, operandVariable, FirDataFlowInfo(setOf(type), emptySet()))
             }
 
             FirOperation.SAFE_AS -> {
                 val expressionVariable = getOrCreateSyntheticVariable(typeOperatorCall)
-                flow = flow.addNotApprovedFact(
+                logicSystem.addConditionalInfo(
+                    flow,
                     expressionVariable,
                     ConditionalFirDataFlowInfo(
                         NotEqNull, operandVariable, FirDataFlowInfo(setOf(type), emptySet())
                     )
-                ).addNotApprovedFact(
+                )
+                logicSystem.addConditionalInfo(
+                    flow,
                     expressionVariable,
                     ConditionalFirDataFlowInfo(
                         EqNull, operandVariable, FirDataFlowInfo(emptySet(), setOf(type))
@@ -192,12 +197,13 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
         }
 
         val expressionVariable = getOrCreateVariable(node.fir)
-        var flow = node.flow
+        val flow = node.flow
 
         // not null for comparisons with constants
         getRealVariablesForSafeCallChain(operand).takeIf { it.isNotEmpty() }?.let { operandVariables ->
             operandVariables.forEach { operandVariable ->
-                flow = flow.addNotApprovedFact(
+                logicSystem.addConditionalInfo(
+                    flow,
                     expressionVariable, ConditionalFirDataFlowInfo(
                         isEq.toEqBoolean(),
                         operandVariable,
@@ -214,12 +220,10 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
             val constValue = (const.value as Boolean)
             val shouldInvert = isEq xor constValue
 
-            flow.conditionalInfos[operandVariable].forEach { info ->
+            flow.getConditionalInfos(operandVariable).forEach { info ->
-                flow = flow.addNotApprovedFact(expressionVariable, info.let { if (shouldInvert) it.invert() else it })
+                logicSystem.addConditionalInfo(flow, expressionVariable, info.let { if (shouldInvert) it.invert() else it })
             }
         }
-
-        node.flow = flow
     }
 
     private fun processEq(node: OperatorCallNode, leftOperand: FirExpression, rightOperand: FirExpression, operation: FirOperation) {
@@ -234,7 +238,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
     }
 
     private fun processEqNull(node: OperatorCallNode, operand: FirExpression, operation: FirOperation) {
-        var flow = node.flow
+        val flow = node.flow
         val expressionVariable = getOrCreateVariable(node.fir)
 
         variableStorage[operand]?.let { operandVariable ->
@@ -245,12 +249,15 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
             }
             val facts = logicSystem.approveFact(operandVariable, condition, flow)
             facts.forEach { (variable, info) ->
-                flow = flow.addNotApprovedFact(
+                logicSystem.addConditionalInfo(
+                    flow,
                     expressionVariable,
                     ConditionalFirDataFlowInfo(
                         EqTrue, variable, info
                     )
-                ).addNotApprovedFact(
+                )
+                logicSystem.addConditionalInfo(
+                    flow,
                     expressionVariable,
                     ConditionalFirDataFlowInfo(
                         EqFalse, variable, info.invert()
@@ -270,7 +277,8 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
         }
 
         operandVariables.forEach { operandVariable ->
-            flow = flow.addNotApprovedFact(
+            logicSystem.addConditionalInfo(
+                flow,
                 expressionVariable, ConditionalFirDataFlowInfo(
                     condition,
                     operandVariable,
@@ -279,7 +287,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
             )
             // TODO: design do we need casts to Nothing?
             /*
-            flow = addNotApprovedFact(
+            flow = addConditionalInfo(
                 expressionVariable, UnapprovedFirDataFlowInfo(
                     eq(conditionValue.invert()!!),
                     operandVariable,
@@ -289,7 +297,6 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
             */
 
         }
-        node.flow = flow
     }
 
     // ----------------------------------- Jump -----------------------------------
@@ -301,43 +308,54 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
     // ----------------------------------- When -----------------------------------
 
     fun enterWhenExpression(whenExpression: FirWhenExpression) {
-        graphBuilder.enterWhenExpression(whenExpression).mergeIncomingFlow()
+        val node = graphBuilder.enterWhenExpression(whenExpression).mergeIncomingFlow()
+        node.flow = node.flow.createNewFlow()
     }
 
     fun enterWhenBranchCondition(whenBranch: FirWhenBranch) {
         val node = graphBuilder.enterWhenBranchCondition(whenBranch).mergeIncomingFlow()
         val previousNode = node.previousNodes.single()
         if (previousNode is WhenBranchConditionExitNode) {
-            node.flow = approveFactsAndUpdateImplicitReceivers(variablesForWhenConditions.remove(previousNode)!!, EqFalse, node.flow)
+            node.flow = logicSystem.approveFactsInsideFlow(
+                variablesForWhenConditions.remove(previousNode)!!,
+                EqFalse,
+                node.flow,
+                shouldCreateNewFlow = true,
+                shouldRemoveSynthetics = true
+            )
         }
-        node.flow.freeze()
     }
 
     fun exitWhenBranchCondition(whenBranch: FirWhenBranch) {
         val (conditionExitNode, branchEnterNode) = graphBuilder.exitWhenBranchCondition(whenBranch)
         conditionExitNode.mergeIncomingFlow()
-        conditionExitNode.flow.freeze()
 
         val conditionVariable = getOrCreateVariable(whenBranch.condition)
         variablesForWhenConditions[conditionExitNode] = conditionVariable
-        branchEnterNode.mergeIncomingFlow()
+        branchEnterNode.flow = logicSystem.approveFactsInsideFlow(
-        branchEnterNode.flow = approveFactsAndUpdateImplicitReceivers(conditionVariable, EqTrue, branchEnterNode.flow)
+            conditionVariable,
+            EqTrue,
+            conditionExitNode.flow,
+            shouldCreateNewFlow = true,
+            shouldRemoveSynthetics = false
+        )
     }
 
     fun exitWhenBranchResult(whenBranch: FirWhenBranch) {
-        val node = graphBuilder.exitWhenBranchResult(whenBranch).mergeIncomingFlow()
+        graphBuilder.exitWhenBranchResult(whenBranch).mergeIncomingFlow()
-        val conditionVariable = getOrCreateVariable(whenBranch.condition)
-        node.flow = node.flow.removeSyntheticVariable(conditionVariable).apply { freeze() }
     }
 
     fun exitWhenExpression(whenExpression: FirWhenExpression) {
-        val node = graphBuilder.exitWhenExpression(whenExpression).mergeIncomingFlow()
+        val node = graphBuilder.exitWhenExpression(whenExpression)
-        var flow = node.flow
+        val previousFlows = node.alivePreviousNodes.map { it.flow }
-        val subjectSymbol = whenExpression.subjectVariable?.symbol
+        val flow = logicSystem.mergeFlows(previousFlows)
-        if (subjectSymbol != null) {
-            variableStorage[subjectSymbol]?.let { flow = flow.removeVariable(it) }
-        }
         node.flow = flow
+        // TODO
+        // val subjectSymbol = whenExpression.subjectVariable?.symbol
+        // if (subjectSymbol != null) {
+        //     variableStorage[subjectSymbol]?.let { flow = flow.removeVariable(it) }
+        // }
+        // node.flow = flow
     }
 
     // ----------------------------------- While Loop -----------------------------------
@@ -463,7 +481,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
         variableStorage[initializer]?.let { initializerVariable ->
             assert(initializerVariable.isSynthetic())
             val realVariable = getOrCreateRealVariable(variable.symbol)
-            node.flow = node.flow.copyNotApprovedFacts(initializerVariable, realVariable)
+            logicSystem.changeVariableForConditionFlow(node.flow, initializerVariable, realVariable)
         }
 
         initializer.resolvedSymbol?.let { initializerSymbol: FirBasedSymbol<*> ->
@@ -491,93 +509,93 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
 
     fun exitLeftBinaryAndArgument(binaryLogicExpression: FirBinaryLogicExpression) {
         val (leftNode, rightNode) = graphBuilder.exitLeftBinaryAndArgument(binaryLogicExpression)
-        leftNode.mergeIncomingFlow()
+        exitLeftArgumentOfBinaryBooleanOperator(leftNode, rightNode, isAnd = true)
-        leftNode.flow.freeze()
-        rightNode.mergeIncomingFlow()
-        val leftOperandVariable = getOrCreateVariable(leftNode.previousNodes.first().fir)
-        val flow = approveFactsAndUpdateImplicitReceivers(leftOperandVariable, EqTrue, rightNode.flow)
-        rightNode.flow = flow.also { it.freeze() }
     }
 
     fun exitBinaryAnd(binaryLogicExpression: FirBinaryLogicExpression) {
-        val node = graphBuilder.exitBinaryAnd(binaryLogicExpression).mergeIncomingFlow()
+        val node = graphBuilder.exitBinaryAnd(binaryLogicExpression)
-        val (leftVariable, rightVariable) = binaryLogicExpression.getVariables()
+        exitBinaryBooleanOperator(binaryLogicExpression, node, isAnd = true)
-
-        val flowFromLeft = node.leftOperandNode.flow
-        val flowFromRight = node.rightOperandNode.flow
-        val flow = node.flow
-
-        val andVariable = getOrCreateVariable(binaryLogicExpression)
-
-        val leftIsTrue = approveFact(leftVariable, EqTrue, flowFromRight)
-        val leftIsFalse = approveFact(leftVariable, EqFalse, flowFromRight)
-        val rightIsTrue = approveFact(rightVariable, EqTrue, flowFromRight) ?: mutableMapOf()
-        val rightIsFalse = approveFact(rightVariable, EqFalse, flowFromRight)
-
-        flowFromRight.approvedInfos.forEach { (variable, info) ->
-            val actualInfo = flowFromLeft.approvedInfos[variable]?.let { info - it } ?: info
-            if (actualInfo.isNotEmpty) rightIsTrue.compute(variable) { _, existingInfo -> info + existingInfo }
-        }
-
-        logicSystem.andForVerifiedFacts(leftIsTrue, rightIsTrue)?.let {
-            for ((variable, info) in it) {
-                flow.addNotApprovedFact(andVariable, ConditionalFirDataFlowInfo(EqTrue, variable, info))
-            }
-        }
-        logicSystem.orForVerifiedFacts(leftIsFalse, rightIsFalse)?.let {
-            for ((variable, info) in it) {
-                flow.addNotApprovedFact(andVariable, ConditionalFirDataFlowInfo(EqFalse, variable, info))
-            }
-        }
-        node.flow = flow.removeSyntheticVariable(leftVariable).removeSyntheticVariable(rightVariable)
     }
 
     fun enterBinaryOr(binaryLogicExpression: FirBinaryLogicExpression) {
-        graphBuilder.enterBinaryOr(binaryLogicExpression)
+        graphBuilder.enterBinaryOr(binaryLogicExpression).mergeIncomingFlow()
     }
 
     fun exitLeftBinaryOrArgument(binaryLogicExpression: FirBinaryLogicExpression) {
         val (leftNode, rightNode) = graphBuilder.exitLeftBinaryOrArgument(binaryLogicExpression)
-        leftNode.mergeIncomingFlow()
+        exitLeftArgumentOfBinaryBooleanOperator(leftNode, rightNode, isAnd = false)
-        leftNode.flow.freeze()
-        rightNode.mergeIncomingFlow()
-        val leftOperandVariable = getOrCreateVariable(leftNode.previousNodes.first().fir)
-        val flow = approveFactsAndUpdateImplicitReceivers(leftOperandVariable, EqFalse, rightNode.flow)
-        rightNode.flow = flow.also { it.freeze() }
     }
 
     fun exitBinaryOr(binaryLogicExpression: FirBinaryLogicExpression) {
-        val node = graphBuilder.exitBinaryOr(binaryLogicExpression).mergeIncomingFlow()
+        val node = graphBuilder.exitBinaryOr(binaryLogicExpression)
-        val (leftVariable, rightVariable) = binaryLogicExpression.getVariables()
+        exitBinaryBooleanOperator(binaryLogicExpression, node, isAnd = false)
-
-        val flowFromLeft = node.leftOperandNode.flow
-        val flowFromRight = node.rightOperandNode.flow
-        val flow = node.flow
-
-        val orVariable = getOrCreateVariable(binaryLogicExpression)
-
-        val leftIsTrue = approveFact(leftVariable, EqTrue, flowFromLeft)
-        val leftIsFalse = approveFact(leftVariable, EqFalse, flowFromLeft)
-        val rightIsTrue = approveFact(rightVariable, EqTrue, flowFromRight)
-        val rightIsFalse = approveFact(rightVariable, EqFalse, flowFromRight)
-
-        logicSystem.orForVerifiedFacts(leftIsTrue, rightIsTrue)?.let {
-            for ((variable, info) in it) {
-                flow.addNotApprovedFact(orVariable, ConditionalFirDataFlowInfo(EqTrue, variable, info))
-            }
-        }
-        logicSystem.andForVerifiedFacts(leftIsFalse, rightIsFalse)?.let {
-            for ((variable, info) in it) {
-                flow.addNotApprovedFact(orVariable, ConditionalFirDataFlowInfo(EqFalse, variable, info))
-            }
-        }
-        node.flow = flow.removeSyntheticVariable(leftVariable).removeSyntheticVariable(rightVariable)
     }
 
+    private fun exitLeftArgumentOfBinaryBooleanOperator(leftNode: CFGNode<*>, rightNode: CFGNode<*>, isAnd: Boolean) {
+        val parentFlow = leftNode.alivePreviousNodes.first().flow
+        leftNode.flow = parentFlow.createNewFlow()
+        val leftOperandVariable = getOrCreateVariable(leftNode.previousNodes.first().fir)
+        rightNode.flow = logicSystem.approveFactsInsideFlow(
+            leftOperandVariable,
+            if (isAnd) EqTrue else EqFalse,
+            parentFlow,
+            shouldCreateNewFlow = true,
+            shouldRemoveSynthetics = false
+        )
+    }
+
+    private fun exitBinaryBooleanOperator(
+        binaryLogicExpression: FirBinaryLogicExpression,
+        node: AbstractBinaryExitNode<*>,
+        isAnd: Boolean
+    ) {
+        val bothEvaluated = if (isAnd) EqTrue else EqFalse
+        val onlyLeftEvaluated = bothEvaluated.invert()
+
+        // Naming for all variables was chosen in assumption that we processing && expression
+        val flowFromRight = node.rightOperandNode.flow
+
+        val flow = node.mergeIncomingFlow().flow
+
+        val (leftVariable, rightVariable) = binaryLogicExpression.getVariables()
+        val andVariable = getOrCreateVariable(binaryLogicExpression)
+
+        val conditionalFromLeft = flow.conditionalInfosFromTopFlow()[leftVariable]
+        val conditionalFromRight = flowFromRight.conditionalInfosFromTopFlow()[rightVariable]
+
+        // left && right == True
+        // left || right == False
+        val approvedIfTrue: ApprovedInfos = mutableMapOf()
+        logicSystem.approveFactTo(approvedIfTrue, bothEvaluated, conditionalFromLeft)
+        logicSystem.approveFactTo(approvedIfTrue, bothEvaluated, conditionalFromRight)
+        flowFromRight.approvedInfosFromTopFlow().forEach { (variable, info) ->
+            approvedIfTrue.addInfo(variable, info)
+        }
+        approvedIfTrue.forEach { (variable, info) ->
+            logicSystem.addConditionalInfo(flow, andVariable, info.toConditional(bothEvaluated, variable))
+        }
+
+        // left && right == False
+        // left || right == True
+        val approvedIfFalse: ApprovedInfos = mutableMapOf()
+        val leftIsFalse = logicSystem.approveFact(onlyLeftEvaluated, conditionalFromLeft)
+        val rightIsFalse = logicSystem.approveFact(onlyLeftEvaluated, conditionalFromRight)
+        approvedIfFalse.mergeInfo(logicSystem.orForVerifiedFacts(leftIsFalse, rightIsFalse))
+        approvedIfFalse.forEach { (variable, info) ->
+            logicSystem.addConditionalInfo(flow, andVariable, info.toConditional(onlyLeftEvaluated, variable))
+        }
+
+        node.flow = flow
+
+        variableStorage.removeVariableIfSynthetic(leftVariable)
+        variableStorage.removeVariableIfSynthetic(rightVariable)
+    }
+
+
     private fun exitBooleanNot(functionCall: FirFunctionCall, node: FunctionCallNode) {
         val booleanExpressionVariable = getOrCreateVariable(node.previousNodes.first().fir)
         val variable = getOrCreateVariable(functionCall)
-        node.flow = node.flow.copyNotApprovedFacts(booleanExpressionVariable, variable) { it.invert() }
+        logicSystem.changeVariableForConditionFlow(node.flow, booleanExpressionVariable, variable) { it.invert() }
     }
 
     // ----------------------------------- Annotations -----------------------------------
@@ -602,37 +620,26 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
 
     // -------------------------------------------------------------------------------------------------------------------------
 
-    private fun approveFact(variable: DataFlowVariable, condition: Condition, flow: Flow): MutableMap<RealDataFlowVariable, FirDataFlowInfo>? =
-        logicSystem.approveFact(variable, condition, flow)
-
     private fun FirBinaryLogicExpression.getVariables(): Pair<DataFlowVariable, DataFlowVariable> =
         getOrCreateVariable(leftOperand) to getOrCreateVariable(rightOperand)
 
-    private var CFGNode<*>.flow: Flow
+    private var CFGNode<*>.flow: DelegatingFlow
-        get() = flowOnNodes[this] ?: Flow.EMPTY
+        get() = flowOnNodes[this.origin] ?: DelegatingFlow.createDefaultFlow()
         set(value) {
-            flowOnNodes[this] = value
+            flowOnNodes[this.origin] = value
         }
 
+    private val CFGNode<*>.origin: CFGNode<*> get() = if (this is StubNode) previousNodes.first() else this
+
     private fun <T : CFGNode<*>> T.mergeIncomingFlow(): T = this.also { node ->
         val previousFlows = node.alivePreviousNodes.map { it.flow }
-        val flow = logicSystem.or(previousFlows)
+        node.flow = logicSystem.mergeFlows(previousFlows)
-        if (previousFlows.size > 1) {
-            receiverStack.forEachIndexed { index, receiver ->
-                val variable = variableStorage[receiver.boundSymbol] ?: return@forEachIndexed
-                val approvedFacts = flow.approvedFacts(variable)
-                if (approvedFacts == null) {
-                    receiver.replaceType(receiverStack.getOriginalType(index))
-                } else {
-                    updateReceiverType(index, approvedFacts)
-                }
-            }
-        }
-        node.flow = flow
     }
 
     // -------------------------------- get or create variable --------------------------------
-    private fun getOrCreateSyntheticVariable(fir: FirElement): SyntheticDataFlowVariable = variableStorage.getOrCreateNewSyntheticVariable(fir)
+
+    private fun getOrCreateSyntheticVariable(fir: FirElement): SyntheticDataFlowVariable =
+        variableStorage.getOrCreateNewSyntheticVariable(fir)
 
     private fun getOrCreateRealVariable(fir: FirElement): RealDataFlowVariable? {
         if (fir is FirThisReceiverExpressionImpl) {
@@ -694,44 +701,25 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
         return result
     }
 
-    private fun approveFactsAndUpdateImplicitReceivers(variable: DataFlowVariable, condition: Condition, flow: Flow): Flow {
+    private inner class DelegatingLogicSystemImpl(context: DataFlowInferenceContext) : DelegatingLogicSystem(context) {
-        val (newFlow, variables) = logicSystem.approveFactsInsideFlow(variable, condition, flow)
+        override fun processUpdatedReceiverVariable(flow: DelegatingFlow, variable: RealDataFlowVariable) {
-        for (variable in variables) {
+            val symbol = (variable.fir as? FirSymbolOwner<*>)?.symbol ?: return
-            updateReceiverType(newFlow, variable)
-        }
-        return newFlow
-    }
 
-    private fun updateReceiverType(flow: Flow, variable: DataFlowVariable) {
+            val index = receiverStack.getReceiverIndex(symbol) ?: return
-        val symbol = (variable.fir as? FirSymbolOwner<*>)?.symbol ?: return
+            val info = flow.getApprovedInfo(variable)
-        val index = receiverStack.getReceiverIndex(symbol) ?: return
-        val info = flow.approvedInfos[variable] ?: return
-        updateReceiverType(index, info)
-    }
 
-    private fun updateReceiverType(index: Int, info: FirDataFlowInfo) {
+            if (info == null) {
-        val types = info.exactType.toMutableList().also {
+                receiverStack.replaceReceiverType(index, receiverStack.getOriginalType(index))
-            it += receiverStack.getOriginalType(index)
+            } else {
-        }
+                val types = info.exactType.toMutableList().also {
-        receiverStack.replaceReceiverType(index, context.intersectTypesOrNull(types)!!)
+                    it += receiverStack.getOriginalType(index)
-
+                }
-    }
+                receiverStack.replaceReceiverType(index, context.intersectTypesOrNull(types)!!)
-
-    private fun addApprovedFact(flow: Flow, variable: RealDataFlowVariable, info: FirDataFlowInfo): Flow {
-        return flow.addApprovedFact(variable, info).also {
-            if (variable.isThisReference) {
-                updateReceiverType(flow, variable)
             }
         }
-    }
 
-    private fun Flow.removeVariable(variable: DataFlowVariable): Flow {
+        override fun updateAllReceivers(flow: DelegatingFlow) {
-        variableStorage.removeVariable(variable)
+            receiverStack.mapNotNull { variableStorage[it.boundSymbol] }.forEach { processUpdatedReceiverVariable(flow, it) }
-        return removeVariableFromFlow(variable)
+        }
-    }
-
-    private fun Flow.removeSyntheticVariable(variable: DataFlowVariable): Flow {
-        if (!variable.isSynthetic()) return this
-        return removeVariable(variable)
     }
 }

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

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.fir.resolve.dfa
 
+import com.google.common.collect.ArrayListMultimap
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.render
 
@@ -24,23 +25,58 @@ data class ConditionalFirDataFlowInfo(
     private fun Set<ConeKotlinType>.render(): String = joinToString { it.render() }
 }
 
-data class FirDataFlowInfo(
+typealias ApprovedInfos = MutableMap<RealDataFlowVariable, FirDataFlowInfo>
-    val exactType: Set<ConeKotlinType>,
+typealias ConditionalInfos = ArrayListMultimap<DataFlowVariable, ConditionalFirDataFlowInfo>
+
+interface FirDataFlowInfo {
+    companion object {
+        // TODO: temporary
+        operator fun invoke(exactType: Set<ConeKotlinType>, exactNotType: Set<ConeKotlinType>): FirDataFlowInfo =
+            MutableFirDataFlowInfo(exactType.toMutableSet(), exactNotType.toMutableSet())
+    }
+
+    val exactType: Set<ConeKotlinType>
     val exactNotType: Set<ConeKotlinType>
-) {
+    operator fun plus(other: FirDataFlowInfo): FirDataFlowInfo
-    operator fun plus(other: FirDataFlowInfo): FirDataFlowInfo = FirDataFlowInfo(
+    operator fun minus(other: FirDataFlowInfo): FirDataFlowInfo
-        exactType + other.exactType,
+    val isNotEmpty: Boolean
-        exactNotType + other.exactNotType
+    val isEmpty: Boolean get() = !isNotEmpty
-    )
+    fun invert(): FirDataFlowInfo
-
-    operator fun minus(other: FirDataFlowInfo): FirDataFlowInfo = FirDataFlowInfo(
-        exactType - other.exactType,
-        exactNotType - other.exactNotType
-    )
-
-    val isNotEmpty: Boolean get() = exactType.isNotEmpty() || exactNotType.isNotEmpty()
-
-    fun invert(): FirDataFlowInfo = FirDataFlowInfo(exactNotType, exactType)
 }
 
-operator fun FirDataFlowInfo.plus(other: FirDataFlowInfo?): FirDataFlowInfo = other?.let { this + other } ?: this
+data class MutableFirDataFlowInfo(
+    override val exactType: MutableSet<ConeKotlinType> = mutableSetOf(),
+    override val exactNotType: MutableSet<ConeKotlinType> = mutableSetOf()
+) : FirDataFlowInfo {
+    override operator fun plus(other: FirDataFlowInfo): MutableFirDataFlowInfo = MutableFirDataFlowInfo(
+        HashSet(exactType).apply { addAll(other.exactType) },
+        HashSet(exactNotType).apply { addAll(other.exactNotType) }
+    )
+
+    override operator fun minus(other: FirDataFlowInfo): MutableFirDataFlowInfo = MutableFirDataFlowInfo(
+        HashSet(exactType).apply { removeAll(other.exactType) },
+        HashSet(exactNotType).apply { removeAll(other.exactNotType) }
+    )
+
+    override val isNotEmpty: Boolean get() = exactType.isNotEmpty() || exactNotType.isNotEmpty()
+
+    override fun invert(): FirDataFlowInfo = MutableFirDataFlowInfo(exactNotType, exactType)
+
+    operator fun plusAssign(info: FirDataFlowInfo) {
+        exactType += info.exactType
+        exactNotType += info.exactNotType
+    }
+
+}
+
+operator fun FirDataFlowInfo.plus(other: FirDataFlowInfo?): FirDataFlowInfo = other?.let { this + other } ?: this
+
+fun FirDataFlowInfo.toConditional(condition: Condition, variable: RealDataFlowVariable): ConditionalFirDataFlowInfo =
+    ConditionalFirDataFlowInfo(condition, variable, this)
+
+fun ApprovedInfos.addInfo(variable: RealDataFlowVariable, info: FirDataFlowInfo) {
+    compute(variable) { _, existingInfo ->
+        if (existingInfo != null) (existingInfo as MutableFirDataFlowInfo).apply { this += info }
+        else info
+    }
+}