[FIR] Make Flow and LogicSystem abstract

It's needed for simple changing of different Flow implementations

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

@@ -7,26 +7,12 @@ 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 {
+class DelegatingFlow(
-    companion object {
+    val previousFlow: DelegatingFlow?,
-        fun createDefaultFlow(): DelegatingFlow = DelegatingFlowImpl(null)
+    val approvedInfos: MutableApprovedInfos = mutableMapOf(),
-    }
-
-    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 {
+) : Flow {
     val level: Int = previousFlow?.level?.plus(1) ?: 0
 
     override fun getApprovedInfo(variable: RealDataFlowVariable): FirDataFlowInfo? {
@@ -64,12 +50,8 @@ private class DelegatingFlowImpl(
         return result
     }
 
-    override fun createNewFlow(): DelegatingFlowImpl {
+    private inline fun collect(block: (DelegatingFlow) -> Boolean) {
-        return DelegatingFlowImpl(this)
+        var flow: DelegatingFlow? = this
-    }
-
-    private inline fun collect(block: (DelegatingFlowImpl) -> Boolean) {
-        var flow: DelegatingFlowImpl? = this
         while (flow != null) {
             val shouldContinue = block(flow)
             if (!shouldContinue) break
@@ -78,25 +60,14 @@ private class DelegatingFlowImpl(
     }
 }
 
-fun DelegatingFlow.approvedInfosFromTopFlow(): Map<RealDataFlowVariable, FirDataFlowInfo> {
+fun DelegatingFlow.approvedInfosFromTopFlow(): ApprovedInfos {
-    require(this is DelegatingFlowImpl)
     return approvedInfos
 }
 
-fun DelegatingFlow.conditionalInfosFromTopFlow(): Multimap<DataFlowVariable, ConditionalFirDataFlowInfo> {
+fun DelegatingFlow.conditionalInfosFromTopFlow(): ConditionalInfos {
-    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()
@@ -123,16 +94,33 @@ private class ImmutableMultimap<K, V>(private val original: Multimap<K, V>) : Mu
     }
 }
 
-abstract class DelegatingLogicSystem(private val context: DataFlowInferenceContext) {
+abstract class DelegatingLogicSystem(context: DataFlowInferenceContext) : LogicSystem(context) {
+    override fun createEmptyFlow(): Flow {
+        return DelegatingFlow(null)
+    }
+
+    override fun forkFlow(flow: Flow): Flow {
+        require(flow is DelegatingFlow)
+        return DelegatingFlow(flow)
+    }
+
+    override fun collectInfoForBooleanOperator(leftFlow: Flow, rightFlow: Flow): InfoForBooleanOperator {
+        require(leftFlow is DelegatingFlow && rightFlow is DelegatingFlow)
+        return InfoForBooleanOperator(
+            leftFlow.previousFlow!!.conditionalInfosFromTopFlow(),
+            rightFlow.conditionalInfosFromTopFlow(),
+            rightFlow.approvedInfosFromTopFlow()
+        )
+    }
 
     // ------------------------------- Flow operations -------------------------------
 
-    fun mergeFlows(flows: Collection<DelegatingFlow>): DelegatingFlow {
+    override fun joinFlow(flows: Collection<Flow>): Flow {
-        if (flows.isEmpty()) return DelegatingFlow.createDefaultFlow()
+        if (flows.isEmpty()) return createEmptyFlow()
         flows.singleOrNull()?.let { return it }
 
         @Suppress("UNCHECKED_CAST", "NAME_SHADOWING")
-        val flows = flows as Collection<DelegatingFlowImpl>
+        val flows = flows as Collection<DelegatingFlow>
         val commonFlow = flows.reduce(this::lowestCommonFlow)
 
         val approvedInfosFromAllFlows = mutableListOf<ApprovedInfos>()
@@ -151,7 +139,7 @@ abstract class DelegatingLogicSystem(private val context: DataFlowInferenceConte
             if (existingInfo == null) {
                 commonFlow.approvedInfos[variable] = intersectedInfo
             } else {
-                (existingInfo as MutableFirDataFlowInfo) += intersectedInfo
+                existingInfo += intersectedInfo
             }
         }
 
@@ -160,19 +148,13 @@ abstract class DelegatingLogicSystem(private val context: DataFlowInferenceConte
         return commonFlow
     }
 
-    /*
+    override fun changeVariableForConditionFlow(
-     *  used for:
+        flow: Flow,
-     *   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
+        transform: ((ConditionalFirDataFlowInfo) -> ConditionalFirDataFlowInfo)?
     ) {
-        require(flow is DelegatingFlowImpl)
+        require(flow is DelegatingFlow)
         var infos = flow.getConditionalInfos(sourceVariable)
         if (transform != null) {
             infos = infos.map(transform)
@@ -183,14 +165,14 @@ abstract class DelegatingLogicSystem(private val context: DataFlowInferenceConte
         }
     }
 
-    fun approveFactsInsideFlow(
+    override fun approveFactsInsideFlow(
         variable: DataFlowVariable,
         condition: Condition,
-        flow: DelegatingFlow,
+        flow: Flow,
-        shouldCreateNewFlow: Boolean,
+        shouldForkFlow: Boolean,
         shouldRemoveSynthetics: Boolean
-    ): DelegatingFlow {
+    ): Flow {
-        require(flow is DelegatingFlowImpl)
+        require(flow is DelegatingFlow)
 
         val notApprovedFacts: Collection<ConditionalFirDataFlowInfo> = if (shouldRemoveSynthetics && variable.isSynthetic()) {
             flow.conditionalInfos.removeAll(variable)
@@ -198,7 +180,7 @@ abstract class DelegatingLogicSystem(private val context: DataFlowInferenceConte
             flow.getConditionalInfos(variable)
         }
 
-        val resultFlow = if (shouldCreateNewFlow) flow.createNewFlow() else flow
+        val resultFlow = if (shouldForkFlow) forkFlow(flow) else flow
         if (notApprovedFacts.isEmpty()) {
             return resultFlow
         }
@@ -229,102 +211,23 @@ abstract class DelegatingLogicSystem(private val context: DataFlowInferenceConte
         return resultFlow
     }
 
-    fun addApprovedInfo(flow: DelegatingFlow, variable: RealDataFlowVariable, info: FirDataFlowInfo) {
+    override fun addApprovedInfo(flow: Flow, variable: RealDataFlowVariable, info: FirDataFlowInfo) {
         assert(info is MutableFirDataFlowInfo)
-        require(flow is DelegatingFlowImpl)
+        require(flow is DelegatingFlow)
         flow.approvedInfos.addInfo(variable, info)
         if (variable.isThisReference) {
             processUpdatedReceiverVariable(flow, variable)
         }
     }
 
-    fun addConditionalInfo(flow: DelegatingFlow, variable: DataFlowVariable, info: ConditionalFirDataFlowInfo) {
+    override fun addConditionalInfo(flow: Flow, variable: DataFlowVariable, info: ConditionalFirDataFlowInfo) {
-        require(flow is DelegatingFlowImpl)
+        require(flow is DelegatingFlow)
         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)
+    private fun lowestCommonFlow(left: DelegatingFlow, right: DelegatingFlow): DelegatingFlow {
-
-    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
@@ -343,8 +246,8 @@ abstract class DelegatingLogicSystem(private val context: DataFlowInferenceConte
         return left
     }
 
-    private fun DelegatingFlowImpl.collectInfos(untilFlow: DelegatingFlowImpl): ApprovedInfos {
+    private fun DelegatingFlow.collectInfos(untilFlow: DelegatingFlow): ApprovedInfos {
-        val approvedInfos: ApprovedInfos = mutableMapOf()
+        val approvedInfos: MutableApprovedInfos = mutableMapOf()
 
         // val conditionalInfos: ConditionalInfos = ArrayListMultimap.create()
 

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

@@ -7,7 +7,6 @@ package org.jetbrains.kotlin.fir.resolve.dfa
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirResolvedCallableReference
-import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
 import org.jetbrains.kotlin.fir.declarations.FirAnonymousInitializer
 import org.jetbrains.kotlin.fir.declarations.FirFunction
 import org.jetbrains.kotlin.fir.declarations.FirProperty
@@ -42,9 +41,9 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
     private val receiverStack: ImplicitReceiverStackImpl = transformer.implicitReceiverStack
 
     private val graphBuilder = ControlFlowGraphBuilder()
-    private val logicSystem = DelegatingLogicSystemImpl(context)
+    private val logicSystem: LogicSystem = LogicSystemImpl(context)
     private val variableStorage = DataFlowVariableStorage()
-    private val flowOnNodes = mutableMapOf<CFGNode<*>, DelegatingFlow>()
+    private val flowOnNodes = mutableMapOf<CFGNode<*>, Flow>()
 
     private val variablesForWhenConditions = mutableMapOf<WhenBranchConditionExitNode, DataFlowVariable>()
 
@@ -309,8 +308,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
     // ----------------------------------- When -----------------------------------
 
     fun enterWhenExpression(whenExpression: FirWhenExpression) {
-        val node = graphBuilder.enterWhenExpression(whenExpression).mergeIncomingFlow()
+        graphBuilder.enterWhenExpression(whenExpression).mergeIncomingFlow()
-        node.flow = node.flow.createNewFlow()
     }
 
     fun enterWhenBranchCondition(whenBranch: FirWhenBranch) {
@@ -321,7 +319,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
                 variablesForWhenConditions.remove(previousNode)!!,
                 EqFalse,
                 node.flow,
-                shouldCreateNewFlow = true,
+                shouldForkFlow = true,
                 shouldRemoveSynthetics = true
             )
         }
@@ -337,7 +335,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
             conditionVariable,
             EqTrue,
             conditionExitNode.flow,
-            shouldCreateNewFlow = true,
+            shouldForkFlow = true,
             shouldRemoveSynthetics = false
         )
     }
@@ -349,7 +347,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
     fun exitWhenExpression(whenExpression: FirWhenExpression) {
         val node = graphBuilder.exitWhenExpression(whenExpression)
         val previousFlows = node.alivePreviousNodes.map { it.flow }
-        val flow = logicSystem.mergeFlows(previousFlows)
+        val flow = logicSystem.joinFlow(previousFlows)
         node.flow = flow
         // TODO
         // val subjectSymbol = whenExpression.subjectVariable?.symbol
@@ -541,13 +539,13 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
 
     private fun exitLeftArgumentOfBinaryBooleanOperator(leftNode: CFGNode<*>, rightNode: CFGNode<*>, isAnd: Boolean) {
         val parentFlow = leftNode.alivePreviousNodes.first().flow
-        leftNode.flow = parentFlow.createNewFlow()
+        leftNode.flow = logicSystem.forkFlow(parentFlow)
         val leftOperandVariable = getOrCreateVariable(leftNode.previousNodes.first().fir)
         rightNode.flow = logicSystem.approveFactsInsideFlow(
             leftOperandVariable,
             if (isAnd) EqTrue else EqFalse,
             parentFlow,
-            shouldCreateNewFlow = true,
+            shouldForkFlow = true,
             shouldRemoveSynthetics = false
         )
     }
@@ -561,6 +559,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
         val onlyLeftEvaluated = bothEvaluated.invert()
 
         // Naming for all variables was chosen in assumption that we processing && expression
+        val flowFromLeft = node.leftOperandNode.flow
         val flowFromRight = node.rightOperandNode.flow
 
         val flow = node.mergeIncomingFlow().flow
@@ -568,15 +567,20 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
         val (leftVariable, rightVariable) = binaryLogicExpression.getVariables()
         val andVariable = getOrCreateVariable(binaryLogicExpression)
 
-        val conditionalFromLeft = flow.conditionalInfosFromTopFlow()[leftVariable]
+        val (conditionalFromLeft, conditionalFromRight, approvedFromRight) = logicSystem.collectInfoForBooleanOperator(
-        val conditionalFromRight = flowFromRight.conditionalInfosFromTopFlow()[rightVariable]
+            flowFromLeft,
+            flowFromRight
+        )
+
+        val conditionalFromLeftArgument = conditionalFromLeft[leftVariable]
+        val conditionalFromRightArgument = conditionalFromRight[rightVariable]
 
         // left && right == True
         // left || right == False
-        val approvedIfTrue: ApprovedInfos = mutableMapOf()
+        val approvedIfTrue: MutableApprovedInfos = mutableMapOf()
-        logicSystem.approveFactTo(approvedIfTrue, bothEvaluated, conditionalFromLeft)
+        logicSystem.approveFactTo(approvedIfTrue, bothEvaluated, conditionalFromLeftArgument)
-        logicSystem.approveFactTo(approvedIfTrue, bothEvaluated, conditionalFromRight)
+        logicSystem.approveFactTo(approvedIfTrue, bothEvaluated, conditionalFromRightArgument)
-        flowFromRight.approvedInfosFromTopFlow().forEach { (variable, info) ->
+        approvedFromRight.forEach { (variable, info) ->
             approvedIfTrue.addInfo(variable, info)
         }
         approvedIfTrue.forEach { (variable, info) ->
@@ -585,9 +589,9 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
 
         // left && right == False
         // left || right == True
-        val approvedIfFalse: ApprovedInfos = mutableMapOf()
+        val approvedIfFalse: MutableApprovedInfos = mutableMapOf()
-        val leftIsFalse = logicSystem.approveFact(onlyLeftEvaluated, conditionalFromLeft)
+        val leftIsFalse = logicSystem.approveFact(onlyLeftEvaluated, conditionalFromLeftArgument)
-        val rightIsFalse = logicSystem.approveFact(onlyLeftEvaluated, conditionalFromRight)
+        val rightIsFalse = logicSystem.approveFact(onlyLeftEvaluated, conditionalFromRightArgument)
         approvedIfFalse.mergeInfo(logicSystem.orForVerifiedFacts(leftIsFalse, rightIsFalse))
         approvedIfFalse.forEach { (variable, info) ->
             logicSystem.addConditionalInfo(flow, andVariable, info.toConditional(onlyLeftEvaluated, variable))
@@ -631,7 +635,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
     private fun FirBinaryLogicExpression.getVariables(): Pair<DataFlowVariable, DataFlowVariable> =
         getOrCreateVariable(leftOperand) to getOrCreateVariable(rightOperand)
 
-    private var CFGNode<*>.flow: DelegatingFlow
+    private var CFGNode<*>.flow: Flow
         get() = flowOnNodes.getValue(this.origin)
         set(value) {
             flowOnNodes[this.origin] = value
@@ -641,7 +645,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
 
     private fun <T : CFGNode<*>> T.mergeIncomingFlow(): T = this.also { node ->
         val previousFlows = node.alivePreviousNodes.map { it.flow }
-        node.flow = logicSystem.mergeFlows(previousFlows)
+        node.flow = logicSystem.joinFlow(previousFlows)
     }
 
     // -------------------------------- get or create variable --------------------------------
@@ -709,8 +713,8 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
         return result
     }
 
-    private inner class DelegatingLogicSystemImpl(context: DataFlowInferenceContext) : DelegatingLogicSystem(context) {
+    private inner class LogicSystemImpl(context: DataFlowInferenceContext) : DelegatingLogicSystem(context) {
-        override fun processUpdatedReceiverVariable(flow: DelegatingFlow, variable: RealDataFlowVariable) {
+        override fun processUpdatedReceiverVariable(flow: Flow, variable: RealDataFlowVariable) {
             val symbol = (variable.fir as? FirSymbolOwner<*>)?.symbol ?: return
 
             val index = receiverStack.getReceiverIndex(symbol) ?: return
@@ -726,7 +730,7 @@ class FirDataFlowAnalyzer(transformer: FirBodyResolveTransformer) : BodyResolveC
             }
         }
 
-        override fun updateAllReceivers(flow: DelegatingFlow) {
+        override fun updateAllReceivers(flow: Flow) {
             receiverStack.mapNotNull { variableStorage[it.boundSymbol] }.forEach { processUpdatedReceiverVariable(flow, it) }
         }
     }

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

@@ -5,7 +5,7 @@
 
 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
 import org.jetbrains.kotlin.fir.types.render
 
@@ -25,8 +25,9 @@ data class ConditionalFirDataFlowInfo(
     private fun Set<ConeKotlinType>.render(): String = joinToString { it.render() }
 }
 
-typealias ApprovedInfos = MutableMap<RealDataFlowVariable, FirDataFlowInfo>
+typealias ApprovedInfos = Map<RealDataFlowVariable, FirDataFlowInfo>
-typealias ConditionalInfos = ArrayListMultimap<DataFlowVariable, ConditionalFirDataFlowInfo>
+typealias MutableApprovedInfos = MutableMap<RealDataFlowVariable, MutableFirDataFlowInfo>
+typealias ConditionalInfos = Multimap<DataFlowVariable, ConditionalFirDataFlowInfo>
 
 interface FirDataFlowInfo {
     companion object {
@@ -74,9 +75,14 @@ operator fun FirDataFlowInfo.plus(other: FirDataFlowInfo?): FirDataFlowInfo = ot
 fun FirDataFlowInfo.toConditional(condition: Condition, variable: RealDataFlowVariable): ConditionalFirDataFlowInfo =
     ConditionalFirDataFlowInfo(condition, variable, this)
 
-fun ApprovedInfos.addInfo(variable: RealDataFlowVariable, info: FirDataFlowInfo) {
+fun MutableApprovedInfos.addInfo(variable: RealDataFlowVariable, info: FirDataFlowInfo) {
-    compute(variable) { _, existingInfo ->
+    merge(variable, info as MutableFirDataFlowInfo) { existingInfo, newInfo ->
-        if (existingInfo != null) (existingInfo as MutableFirDataFlowInfo).apply { this += info }
+        existingInfo.apply { this += newInfo }
-        else info
+    }
+}
+
+fun MutableApprovedInfos.mergeInfo(other: Map<RealDataFlowVariable, FirDataFlowInfo>) {
+    other.forEach { (variable, info) ->
+        addInfo(variable, info)
     }
 }

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

@@ -5,122 +5,119 @@
 
 package org.jetbrains.kotlin.fir.resolve.dfa
 
-import com.google.common.collect.HashMultimap
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
 
-class LogicSystem(private val context: DataFlowInferenceContext) {
+interface Flow {
-    private fun <E> List<Set<E>>.intersectSets(): Set<E> = takeIf { isNotEmpty() }?.reduce { x, y -> x.intersect(y) } ?: emptySet()
+    fun getApprovedInfo(variable: RealDataFlowVariable): FirDataFlowInfo?
+    fun getConditionalInfos(variable: DataFlowVariable): Collection<ConditionalFirDataFlowInfo>
 
-    fun or(storages: Collection<Flow>): Flow {
+    fun getVariablesInApprovedInfos(): Collection<RealDataFlowVariable>
-        storages.singleOrNull()?.let {
+}
-            return it
-        }
-        val approvedFacts = mutableMapOf<RealDataFlowVariable, FirDataFlowInfo>()
-        storages.map { it.approvedInfos.keys }
-            .intersectSets()
-            .forEach { variable ->
-                val infos = storages.map { it.approvedInfos[variable]!! }
-                if (infos.isNotEmpty()) {
-                    approvedFacts[variable] = or(infos)
-                }
-            }
 
+abstract class LogicSystem(private val context: DataFlowInferenceContext) {
+    // ------------------------------- Flow operations -------------------------------
 
-        val notApprovedFacts = HashMultimap.create<DataFlowVariable, ConditionalFirDataFlowInfo>()
+    abstract fun createEmptyFlow(): Flow
-        storages.map { it.conditionalInfos.keySet() }
+    abstract fun forkFlow(flow: Flow): Flow
-            .intersectSets()
+    abstract fun joinFlow(flows: Collection<Flow>): Flow
-            .forEach { variable ->
-                val infos = storages.map { it.conditionalInfos[variable] }.intersectSets()
-                if (infos.isNotEmpty()) {
-                    notApprovedFacts.putAll(variable, infos)
-                }
-            }
 
-        return Flow(approvedFacts, notApprovedFacts)
+    abstract fun addApprovedInfo(flow: Flow, variable: RealDataFlowVariable, info: FirDataFlowInfo)
-    }
+    abstract fun addConditionalInfo(flow: Flow, variable: DataFlowVariable, info: ConditionalFirDataFlowInfo)
 
-    fun andForVerifiedFacts(
+    /*
-        left: Map<RealDataFlowVariable, FirDataFlowInfo>?,
+     *  used for:
-        right: Map<RealDataFlowVariable, FirDataFlowInfo>?
+     *   1. val b = x is String
-    ): Map<RealDataFlowVariable, FirDataFlowInfo>? {
+     *   2. b = x is String
-        if (left.isNullOrEmpty()) return right
+     *   3. !b | b.not()   for Booleans
-        if (right.isNullOrEmpty()) return left
+     */
+    abstract fun changeVariableForConditionFlow(
+        flow: Flow,
+        sourceVariable: DataFlowVariable,
+        newVariable: DataFlowVariable,
+        transform: ((ConditionalFirDataFlowInfo) -> ConditionalFirDataFlowInfo)? = null
+    )
 
-        val map = mutableMapOf<RealDataFlowVariable, FirDataFlowInfo>()
+    abstract fun approveFactsInsideFlow(
-        for (variable in left.keys.union(right.keys)) {
+        variable: DataFlowVariable,
-            val leftInfo = left[variable]
+        condition: Condition,
-            val rightInfo = right[variable]
+        flow: Flow,
-            map[variable] = and(listOfNotNull(leftInfo, rightInfo))
+        shouldForkFlow: Boolean,
-        }
+        shouldRemoveSynthetics: Boolean
-        return map
+    ): Flow
-    }
+
+    // ------------------------------- Callbacks for updating implicit receiver stack -------------------------------
+
+    abstract fun processUpdatedReceiverVariable(flow: Flow, variable: RealDataFlowVariable)
+    abstract fun updateAllReceivers(flow: Flow)
+
+    // ------------------------------- Public DataFlowInfo util functions -------------------------------
+
+    data class InfoForBooleanOperator(
+        val conditionalFromLeft: ConditionalInfos,
+        val conditionalFromRight: ConditionalInfos,
+        val approvedFromRight: ApprovedInfos
+    )
+
+    abstract fun collectInfoForBooleanOperator(leftFlow: Flow, rightFlow: Flow): InfoForBooleanOperator
 
     fun orForVerifiedFacts(
-        left: Map<RealDataFlowVariable, FirDataFlowInfo>?,
+        left: ApprovedInfos,
-        right: Map<RealDataFlowVariable, FirDataFlowInfo>?
+        right: ApprovedInfos
-    ): Map<RealDataFlowVariable, FirDataFlowInfo>? {
+    ): MutableApprovedInfos {
-        if (left.isNullOrEmpty() || right.isNullOrEmpty()) return null
+        if (left.isNullOrEmpty() || right.isNullOrEmpty()) return mutableMapOf()
-        val map = mutableMapOf<RealDataFlowVariable, FirDataFlowInfo>()
+        val map = mutableMapOf<RealDataFlowVariable, MutableFirDataFlowInfo>()
         for (variable in left.keys.intersect(right.keys)) {
-            val leftInfo = left[variable]!!
+            val leftInfo = left.getValue(variable)
-            val rightInfo = right[variable]!!
+            val rightInfo = right.getValue(variable)
             map[variable] = or(listOf(leftInfo, rightInfo))
         }
         return map
     }
 
-    fun approveFactsInsideFlow(variable: DataFlowVariable, condition: Condition, flow: Flow): Pair<Flow, Collection<RealDataFlowVariable>> {
+    fun approveFactTo(destination: MutableApprovedInfos, variable: DataFlowVariable, condition: Condition, flow: Flow) {
-        val notApprovedFacts: Set<ConditionalFirDataFlowInfo> = flow.conditionalInfos[variable]
+        val notApprovedFacts: Collection<ConditionalFirDataFlowInfo> = flow.getConditionalInfos(variable)
-        if (notApprovedFacts.isEmpty()) {
+        approveFactTo(destination, condition, notApprovedFacts)
-            return flow to emptyList()
-        }
-        @Suppress("NAME_SHADOWING")
-        val flow = flow.copyForBuilding()
-        val newFacts = HashMultimap.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 infos = ArrayList(infos)
-            flow.approvedInfos[variable]?.let {
-                infos.add(it)
-            }
-            flow.approvedInfos[variable] = and(infos)
-            if (variable.isThisReference) {
-                updatedReceivers += variable
-            }
-        }
-        return flow to updatedReceivers
     }
 
-    fun approveFact(variable: DataFlowVariable, condition: Condition, flow: Flow): MutableMap<RealDataFlowVariable, FirDataFlowInfo> {
+    fun approveFactTo(destination: MutableApprovedInfos, condition: Condition, notApprovedFacts: Collection<ConditionalFirDataFlowInfo>) {
-        val notApprovedFacts: Set<ConditionalFirDataFlowInfo> = flow.conditionalInfos[variable]
+        if (notApprovedFacts.isEmpty()) return
-        if (notApprovedFacts.isEmpty()) {
-            return mutableMapOf()
-        }
-        val newFacts = HashMultimap.create<RealDataFlowVariable, FirDataFlowInfo>()
         notApprovedFacts.forEach {
             if (it.condition == condition) {
-                newFacts.put(it.variable, it.info)
+                destination.addInfo(it.variable, it.info)
             }
         }
-        return newFacts.asMap().mapValuesTo(mutableMapOf()) { (_, infos) -> and(infos) }
     }
 
-    private fun or(infos: Collection<FirDataFlowInfo>): FirDataFlowInfo {
+    fun approveFact(condition: Condition, notApprovedFacts: Collection<ConditionalFirDataFlowInfo>): MutableApprovedInfos {
-        infos.singleOrNull()?.let { return it }
+        return mutableMapOf<RealDataFlowVariable, MutableFirDataFlowInfo>().apply { approveFactTo(this, condition, notApprovedFacts) }
+    }
+
+    fun approveFact(variable: DataFlowVariable, condition: Condition, flow: Flow): MutableApprovedInfos {
+        return mutableMapOf<RealDataFlowVariable, MutableFirDataFlowInfo>().apply { approveFactTo(this, variable, condition, flow) }
+    }
+
+    // ------------------------------- Util functions -------------------------------
+
+    // TODO
+    protected fun <E> Collection<Collection<E>>.intersectSets(): Set<E> {
+        if (isEmpty()) return emptySet()
+        val iterator = iterator()
+        val result = HashSet<E>(iterator.next())
+        while (iterator.hasNext()) {
+            result.retainAll(iterator.next())
+        }
+        return result
+    }
+
+    protected fun or(infos: Collection<FirDataFlowInfo>): MutableFirDataFlowInfo {
+        infos.singleOrNull()?.let { return it as MutableFirDataFlowInfo }
         val exactType = orTypes(infos.map { it.exactType })
         val exactNotType = orTypes(infos.map { it.exactNotType })
-        return FirDataFlowInfo(exactType, exactNotType)
+        return MutableFirDataFlowInfo(exactType, exactNotType)
     }
 
-    private fun orTypes(types: Collection<Set<ConeKotlinType>>): Set<ConeKotlinType> {
+    private fun orTypes(types: Collection<Set<ConeKotlinType>>): MutableSet<ConeKotlinType> {
-        if (types.any { it.isEmpty() }) return emptySet()
+        if (types.any { it.isEmpty() }) return mutableSetOf()
         val allTypes = types.flatMapTo(mutableSetOf()) { it }
         val commonTypes = allTypes.toMutableSet()
         types.forEach { commonTypes.retainAll(it) }
@@ -128,11 +125,4 @@ class LogicSystem(private val context: DataFlowInferenceContext) {
         context.commonSuperTypeOrNull(differentTypes.toList())?.let { commonTypes += it }
         return commonTypes
     }
-
-    private fun and(infos: Collection<FirDataFlowInfo>): FirDataFlowInfo {
-        infos.singleOrNull()?.let { return it }
-        val exactType = infos.flatMapTo(mutableSetOf()) { it.exactType }
-        val exactNotType = infos.flatMapTo(mutableSetOf()) { it.exactNotType }
-        return FirDataFlowInfo(exactType, exactNotType)
-    }
 }

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

@@ -7,7 +7,7 @@ package org.jetbrains.kotlin.fir.resolve.dfa
 
 import com.google.common.collect.HashMultimap
 
-class Flow private constructor(
+class SimpleFlow private constructor(
     val approvedInfos: MutableMap<RealDataFlowVariable, FirDataFlowInfo>,
     val conditionalInfos: HashMultimap<DataFlowVariable, ConditionalFirDataFlowInfo>,
     private var state: State = State.Building
@@ -18,7 +18,7 @@ class Flow private constructor(
     ) : this(approvedInfos, conditionalInfos, State.Building)
 
     companion object {
-        val EMPTY = Flow(mutableMapOf(), HashMultimap.create(), State.Frozen)
+        val EMPTY = SimpleFlow(mutableMapOf(), HashMultimap.create(), State.Frozen)
     }
 
     private val isFrozen: Boolean get() = state == State.Frozen
@@ -27,7 +27,7 @@ class Flow private constructor(
         state = State.Frozen
     }
 
-    fun addApprovedFact(variable: RealDataFlowVariable, info: FirDataFlowInfo): Flow {
+    fun addApprovedFact(variable: RealDataFlowVariable, info: FirDataFlowInfo): SimpleFlow {
         if (isFrozen) return copyForBuilding().addApprovedFact(variable, info)
         approvedInfos.compute(variable) { _, existingInfo ->
             if (existingInfo == null) info
@@ -36,7 +36,7 @@ class Flow private constructor(
         return this
     }
 
-    fun addNotApprovedFact(variable: DataFlowVariable, info: ConditionalFirDataFlowInfo): Flow {
+    fun addNotApprovedFact(variable: DataFlowVariable, info: ConditionalFirDataFlowInfo): SimpleFlow {
         if (isFrozen) return copyForBuilding().addNotApprovedFact(variable, info)
         conditionalInfos.put(variable, info)
         return this
@@ -46,7 +46,7 @@ class Flow private constructor(
         from: DataFlowVariable,
         to: DataFlowVariable,
         transform: ((ConditionalFirDataFlowInfo) -> ConditionalFirDataFlowInfo)? = null
-    ): Flow {
+    ): SimpleFlow {
         if (isFrozen) {
             return copyForBuilding().copyNotApprovedFacts(from, to, transform)
         }
@@ -67,7 +67,7 @@ class Flow private constructor(
         return approvedInfos[variable]
     }
 
-    fun removeVariableFromFlow(variable: DataFlowVariable): Flow {
+    fun removeVariableFromFlow(variable: DataFlowVariable): SimpleFlow {
         if (isFrozen) return copyForBuilding().removeVariableFromFlow(variable)
         conditionalInfos.removeAll(variable)
         approvedInfos.remove(variable)
@@ -78,15 +78,15 @@ class Flow private constructor(
         Building, Frozen
     }
 
-    fun copy(): Flow {
+    fun copy(): SimpleFlow {
         return when (state) {
             State.Frozen -> this
             State.Building -> copyForBuilding()
         }
     }
 
-    fun copyForBuilding(): Flow {
+    fun copyForBuilding(): SimpleFlow {
-        return Flow(approvedInfos.toMutableMap(), conditionalInfos.copy(), State.Building)
+        return SimpleFlow(approvedInfos.toMutableMap(), conditionalInfos.copy(), State.Building)
     }
 }
 

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

@@ -0,0 +1,138 @@
+/*
+ * 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.HashMultimap
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
+
+class SimpleLogicSystem(private val context: DataFlowInferenceContext) {
+    private fun <E> List<Set<E>>.intersectSets(): Set<E> = takeIf { isNotEmpty() }?.reduce { x, y -> x.intersect(y) } ?: emptySet()
+
+    fun or(storages: Collection<SimpleFlow>): SimpleFlow {
+        storages.singleOrNull()?.let {
+            return it
+        }
+        val approvedFacts = mutableMapOf<RealDataFlowVariable, FirDataFlowInfo>()
+        storages.map { it.approvedInfos.keys }
+            .intersectSets()
+            .forEach { variable ->
+                val infos = storages.map { it.approvedInfos[variable]!! }
+                if (infos.isNotEmpty()) {
+                    approvedFacts[variable] = or(infos)
+                }
+            }
+
+
+        val notApprovedFacts = HashMultimap.create<DataFlowVariable, ConditionalFirDataFlowInfo>()
+        storages.map { it.conditionalInfos.keySet() }
+            .intersectSets()
+            .forEach { variable ->
+                val infos = storages.map { it.conditionalInfos[variable] }.intersectSets()
+                if (infos.isNotEmpty()) {
+                    notApprovedFacts.putAll(variable, infos)
+                }
+            }
+
+        return SimpleFlow(approvedFacts, notApprovedFacts)
+    }
+
+    fun andForVerifiedFacts(
+        left: Map<RealDataFlowVariable, FirDataFlowInfo>?,
+        right: Map<RealDataFlowVariable, FirDataFlowInfo>?
+    ): Map<RealDataFlowVariable, FirDataFlowInfo>? {
+        if (left.isNullOrEmpty()) return right
+        if (right.isNullOrEmpty()) return left
+
+        val map = mutableMapOf<RealDataFlowVariable, FirDataFlowInfo>()
+        for (variable in left.keys.union(right.keys)) {
+            val leftInfo = left[variable]
+            val rightInfo = right[variable]
+            map[variable] = and(listOfNotNull(leftInfo, rightInfo))
+        }
+        return map
+    }
+
+    fun orForVerifiedFacts(
+        left: Map<RealDataFlowVariable, FirDataFlowInfo>?,
+        right: Map<RealDataFlowVariable, FirDataFlowInfo>?
+    ): Map<RealDataFlowVariable, FirDataFlowInfo>? {
+        if (left.isNullOrEmpty() || right.isNullOrEmpty()) return null
+        val map = mutableMapOf<RealDataFlowVariable, FirDataFlowInfo>()
+        for (variable in left.keys.intersect(right.keys)) {
+            val leftInfo = left[variable]!!
+            val rightInfo = right[variable]!!
+            map[variable] = or(listOf(leftInfo, rightInfo))
+        }
+        return map
+    }
+
+    fun approveFactsInsideFlow(variable: DataFlowVariable, condition: Condition, flow: SimpleFlow): Pair<SimpleFlow, Collection<RealDataFlowVariable>> {
+        val notApprovedFacts: Set<ConditionalFirDataFlowInfo> = flow.conditionalInfos[variable]
+        if (notApprovedFacts.isEmpty()) {
+            return flow to emptyList()
+        }
+        @Suppress("NAME_SHADOWING")
+        val flow = flow.copyForBuilding()
+        val newFacts = HashMultimap.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 infos = ArrayList(infos)
+            flow.approvedInfos[variable]?.let {
+                infos.add(it)
+            }
+            flow.approvedInfos[variable] = and(infos)
+            if (variable.isThisReference) {
+                updatedReceivers += variable
+            }
+        }
+        return flow to updatedReceivers
+    }
+
+    fun approveFact(variable: DataFlowVariable, condition: Condition, flow: SimpleFlow): MutableMap<RealDataFlowVariable, FirDataFlowInfo> {
+        val notApprovedFacts: Set<ConditionalFirDataFlowInfo> = flow.conditionalInfos[variable]
+        if (notApprovedFacts.isEmpty()) {
+            return mutableMapOf()
+        }
+        val newFacts = HashMultimap.create<RealDataFlowVariable, FirDataFlowInfo>()
+        notApprovedFacts.forEach {
+            if (it.condition == condition) {
+                newFacts.put(it.variable, it.info)
+            }
+        }
+        return newFacts.asMap().mapValuesTo(mutableMapOf()) { (_, infos) -> and(infos) }
+    }
+
+    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 and(infos: Collection<FirDataFlowInfo>): FirDataFlowInfo {
+        infos.singleOrNull()?.let { return it }
+        val exactType = infos.flatMapTo(mutableSetOf()) { it.exactType }
+        val exactNotType = infos.flatMapTo(mutableSetOf()) { it.exactNotType }
+        return FirDataFlowInfo(exactType, exactNotType)
+    }
+}