FIR CFG: refactor handling of boolean operators a bit

To reduce the number of functions in FirDataFlowAnalyzer.

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

@@ -1038,46 +1038,22 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
 
     // ----------------------------------- Boolean operators -----------------------------------
 
-    fun enterBinaryAnd(binaryLogicExpression: FirBinaryLogicExpression) {
+    fun enterBinaryLogicExpression(binaryLogicExpression: FirBinaryLogicExpression) {
-        graphBuilder.enterBinaryAnd(binaryLogicExpression).mergeIncomingFlow()
+        graphBuilder.enterBinaryLogicExpression(binaryLogicExpression).mergeIncomingFlow()
     }
 
-    fun exitLeftBinaryAndArgument(binaryLogicExpression: FirBinaryLogicExpression) {
+    fun exitLeftBinaryLogicExpressionArgument(binaryLogicExpression: FirBinaryLogicExpression) {
-        val (leftExitNode, rightEnterNode) = graphBuilder.exitLeftBinaryAndArgument(binaryLogicExpression)
+        val (leftExitNode, rightEnterNode) = graphBuilder.exitLeftBinaryLogicExpressionArgument(binaryLogicExpression)
-        exitLeftArgumentOfBinaryBooleanOperator(leftExitNode, rightEnterNode, isAnd = true)
-    }
-
-    fun exitBinaryAnd(binaryLogicExpression: FirBinaryLogicExpression) {
-        val node = graphBuilder.exitBinaryAnd(binaryLogicExpression)
-        exitBinaryBooleanOperator(binaryLogicExpression, node, isAnd = true)
-    }
-
-    fun enterBinaryOr(binaryLogicExpression: FirBinaryLogicExpression) {
-        graphBuilder.enterBinaryOr(binaryLogicExpression).mergeIncomingFlow()
-    }
-
-    fun exitLeftBinaryOrArgument(binaryLogicExpression: FirBinaryLogicExpression) {
-        val (leftExitNode, rightEnterNode) = graphBuilder.exitLeftBinaryOrArgument(binaryLogicExpression)
-        exitLeftArgumentOfBinaryBooleanOperator(leftExitNode, rightEnterNode, isAnd = false)
-    }
-
-    fun exitBinaryOr(binaryLogicExpression: FirBinaryLogicExpression) {
-        val node = graphBuilder.exitBinaryOr(binaryLogicExpression)
-        exitBinaryBooleanOperator(binaryLogicExpression, node, isAnd = false)
-    }
-
-    private fun exitLeftArgumentOfBinaryBooleanOperator(leftExitNode: CFGNode<*>, rightEnterNode: CFGNode<*>, isAnd: Boolean) {
         val leftExitFlow = leftExitNode.mergeIncomingFlow()
         val rightEnterFlow = rightEnterNode.mergeIncomingFlow()
         val leftOperandVariable = variableStorage.get(leftExitFlow, leftExitNode.firstPreviousNode.fir) ?: return
+        val isAnd = binaryLogicExpression.kind == LogicOperationKind.AND
         rightEnterFlow.commitOperationStatement(leftOperandVariable eq isAnd)
     }
 
-    private fun exitBinaryBooleanOperator(
+    fun exitBinaryLogicExpression(binaryLogicExpression: FirBinaryLogicExpression) {
-        binaryLogicExpression: FirBinaryLogicExpression,
+        val node = graphBuilder.exitBinaryLogicExpression()
-        node: AbstractBinaryExitNode<*>,
+        val isAnd = binaryLogicExpression.kind == LogicOperationKind.AND
-        isAnd: Boolean
-    ) {
         val flowFromLeft = node.leftOperandNode.flow
         val flowFromRight = node.rightOperandNode.flow
         val flow = node.mergeIncomingFlow()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/cfg/ControlFlowGraphBuilder.kt

@@ -93,8 +93,7 @@ class ControlFlowGraphBuilder {
     private val whenExitNodes: NodeStorage<FirWhenExpression, WhenExitNode> = NodeStorage()
     private val whenBranchIndices: Stack<Map<FirWhenBranch, Int>> = stackOf()
 
-    private val binaryAndExitNodes: Stack<BinaryAndExitNode> = stackOf()
+    private val binaryLogicExpressionExitNodes: Stack<AbstractBinaryExitNode<FirBinaryLogicExpression>> = stackOf()
-    private val binaryOrExitNodes: Stack<BinaryOrExitNode> = stackOf()
 
     private val tryExitNodes: NodeStorage<FirTryExpression, TryExpressionExitNode> = NodeStorage()
     private val tryMainExitNodes: NodeStorage<FirTryExpression, TryMainBlockExitNode> = NodeStorage()
@@ -778,7 +777,7 @@ class ControlFlowGraphBuilder {
         levelCounter--
         val conditionExitNode = createLoopConditionExitNode(loop.condition)
         addNewSimpleNode(conditionExitNode)
-        val conditionConstBooleanValue = conditionExitNode.booleanConstValue
+        val conditionConstBooleanValue = loop.condition.booleanConstValue
         addEdge(conditionExitNode, loopExitNodes.top(), propagateDeadness = false, isDead = conditionConstBooleanValue == true)
         val loopBlockEnterNode = createLoopBlockEnterNode(loop)
         addNewSimpleNode(loopBlockEnterNode, conditionConstBooleanValue == false)
@@ -829,7 +828,7 @@ class ControlFlowGraphBuilder {
         loopConditionEnterNodes.pop()
         levelCounter--
         val conditionExitNode = createLoopConditionExitNode(loop.condition)
-        val conditionBooleanValue = conditionExitNode.booleanConstValue
+        val conditionBooleanValue = loop.condition.booleanConstValue
         popAndAddEdge(conditionExitNode)
         val blockEnterNode = lastNodes.pop()
         require(blockEnterNode is LoopBlockEnterNode)
@@ -844,79 +843,49 @@ class ControlFlowGraphBuilder {
 
     // ----------------------------------- Boolean operators -----------------------------------
 
-    fun enterBinaryAnd(binaryLogicExpression: FirBinaryLogicExpression): BinaryAndEnterNode {
+    fun enterBinaryLogicExpression(binaryLogicExpression: FirBinaryLogicExpression): CFGNode<FirBinaryLogicExpression> {
-        assert(binaryLogicExpression.kind == LogicOperationKind.AND)
+        val (enterNode, exitNode) = when (binaryLogicExpression.kind) {
-        binaryAndExitNodes.push(createBinaryAndExitNode(binaryLogicExpression))
+            LogicOperationKind.AND -> createBinaryAndEnterNode(binaryLogicExpression) to createBinaryAndExitNode(binaryLogicExpression)
-        return createBinaryAndEnterNode(binaryLogicExpression).also { addNewSimpleNode(it) }.also { levelCounter++ }
+            LogicOperationKind.OR -> createBinaryOrEnterNode(binaryLogicExpression) to createBinaryOrExitNode(binaryLogicExpression)
+        }
+        addNewSimpleNode(enterNode)
+        binaryLogicExpressionExitNodes.push(exitNode)
+        levelCounter++
+        return enterNode
     }
 
-    fun exitLeftBinaryAndArgument(binaryLogicExpression: FirBinaryLogicExpression): Pair<BinaryAndExitLeftOperandNode, BinaryAndEnterRightOperandNode> {
+    fun exitLeftBinaryLogicExpressionArgument(
-        assert(binaryLogicExpression.kind == LogicOperationKind.AND)
+        binaryLogicExpression: FirBinaryLogicExpression
+    ): Pair<CFGNode<FirBinaryLogicExpression>, CFGNode<FirBinaryLogicExpression>> {
+        levelCounter--
+        val leftBooleanConstValue = binaryLogicExpression.leftOperand.booleanConstValue
+        val (leftExitNode, rightEnterNode) = when (binaryLogicExpression.kind) {
+            LogicOperationKind.AND ->
+                createBinaryAndExitLeftOperandNode(binaryLogicExpression) to createBinaryAndEnterRightOperandNode(binaryLogicExpression)
+            LogicOperationKind.OR ->
+                createBinaryOrExitLeftOperandNode(binaryLogicExpression) to createBinaryOrEnterRightOperandNode(binaryLogicExpression)
+        }
+        val isAnd = binaryLogicExpression.kind == LogicOperationKind.AND
         val lastNode = lastNodes.pop()
-        val leftBooleanConstValue = lastNode.booleanConstValue
+        addEdge(lastNode, leftExitNode)
-
+        addEdge(leftExitNode, binaryLogicExpressionExitNodes.top(), propagateDeadness = false, isDead = leftBooleanConstValue == isAnd)
-        val leftExitNode = createBinaryAndExitLeftOperandNode(binaryLogicExpression).also {
+        addEdge(leftExitNode, rightEnterNode, isDead = leftBooleanConstValue == !isAnd)
-            addEdge(lastNode, it)
+        lastNodes.push(rightEnterNode)
-            addEdge(it, binaryAndExitNodes.top(), propagateDeadness = false, isDead = leftBooleanConstValue == true)
+        levelCounter++
-        }
-
-        val rightEnterNode = createBinaryAndEnterRightOperandNode(binaryLogicExpression).also {
-            addEdge(leftExitNode, it, isDead = leftBooleanConstValue == false)
-            lastNodes.push(it)
-        }
         return leftExitNode to rightEnterNode
     }
 
-    fun exitBinaryAnd(binaryLogicExpression: FirBinaryLogicExpression): BinaryAndExitNode {
+    fun exitBinaryLogicExpression(): AbstractBinaryExitNode<FirBinaryLogicExpression> {
         levelCounter--
-        assert(binaryLogicExpression.kind == LogicOperationKind.AND)
+        val exitNode = binaryLogicExpressionExitNodes.pop()
-        return binaryAndExitNodes.pop().also {
+        val rightNode = lastNodes.pop()
-            val rightNode = lastNodes.pop()
+        addEdge(rightNode, exitNode, propagateDeadness = false)
-            addEdge(rightNode, it, propagateDeadness = false, isDead = it.leftOperandNode.booleanConstValue == false)
+        exitNode.updateDeadStatus()
-            it.updateDeadStatus()
+        lastNodes.push(exitNode)
-            lastNodes.push(it)
+        return exitNode
-        }
     }
 
-    fun enterBinaryOr(binaryLogicExpression: FirBinaryLogicExpression): BinaryOrEnterNode {
+    private val FirExpression.booleanConstValue: Boolean? get() = (this as? FirConstExpression<*>)?.value as? Boolean?
-        assert(binaryLogicExpression.kind == LogicOperationKind.OR)
-        binaryOrExitNodes.push(createBinaryOrExitNode(binaryLogicExpression))
-        return createBinaryOrEnterNode(binaryLogicExpression).also {
-            addNewSimpleNode(it)
-        }.also { levelCounter++ }
-    }
-
-    fun exitLeftBinaryOrArgument(binaryLogicExpression: FirBinaryLogicExpression): Pair<BinaryOrExitLeftOperandNode, BinaryOrEnterRightOperandNode> {
-        levelCounter--
-        assert(binaryLogicExpression.kind == LogicOperationKind.OR)
-        val previousNode = lastNodes.pop()
-        val leftBooleanValue = previousNode.booleanConstValue
-
-        val leftExitNode = createBinaryOrExitLeftOperandNode(binaryLogicExpression).also {
-            addEdge(previousNode, it)
-            addEdge(it, binaryOrExitNodes.top(), propagateDeadness = false, isDead = leftBooleanValue == false)
-        }
-
-        val rightExitNode = createBinaryOrEnterRightOperandNode(binaryLogicExpression).also {
-            addEdge(leftExitNode, it, propagateDeadness = true, isDead = leftBooleanValue == true)
-            lastNodes.push(it)
-            levelCounter++
-        }
-        return leftExitNode to rightExitNode
-    }
-
-    fun exitBinaryOr(binaryLogicExpression: FirBinaryLogicExpression): BinaryOrExitNode {
-        assert(binaryLogicExpression.kind == LogicOperationKind.OR)
-        levelCounter--
-        return binaryOrExitNodes.pop().also {
-            val rightNode = lastNodes.pop()
-            addEdge(rightNode, it, propagateDeadness = false)
-            it.updateDeadStatus()
-            lastNodes.push(it)
-        }
-    }
-
-    private val CFGNode<*>.booleanConstValue: Boolean? get() = (fir as? FirConstExpression<*>)?.value as? Boolean?
 
     // ----------------------------------- Try-catch-finally -----------------------------------
 

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

@@ -825,21 +825,13 @@ open class FirExpressionsResolveTransformer(transformer: FirAbstractBodyResolveT
         data: ResolutionMode,
     ): FirStatement = whileAnalysing(binaryLogicExpression) {
         val booleanType = binaryLogicExpression.typeRef.resolvedTypeFromPrototype(builtinTypes.booleanType.type)
-        return when (binaryLogicExpression.kind) {
+        return binaryLogicExpression.also(dataFlowAnalyzer::enterBinaryLogicExpression)
-            LogicOperationKind.AND ->
+            .transformLeftOperand(this, ResolutionMode.WithExpectedType(booleanType))
-                binaryLogicExpression.also(dataFlowAnalyzer::enterBinaryAnd)
+            .also(dataFlowAnalyzer::exitLeftBinaryLogicExpressionArgument)
-                    .transformLeftOperand(this, ResolutionMode.WithExpectedType(booleanType))
+            .transformRightOperand(this, ResolutionMode.WithExpectedType(booleanType))
-                    .also(dataFlowAnalyzer::exitLeftBinaryAndArgument)
+            .also(dataFlowAnalyzer::exitBinaryLogicExpression)
-                    .transformRightOperand(this, ResolutionMode.WithExpectedType(booleanType)).also(dataFlowAnalyzer::exitBinaryAnd)
+            .transformOtherChildren(transformer, ResolutionMode.WithExpectedType(booleanType))
-
+            .also { it.resultType = booleanType }
-            LogicOperationKind.OR ->
-                binaryLogicExpression.also(dataFlowAnalyzer::enterBinaryOr)
-                    .transformLeftOperand(this, ResolutionMode.WithExpectedType(booleanType))
-                    .also(dataFlowAnalyzer::exitLeftBinaryOrArgument)
-                    .transformRightOperand(this, ResolutionMode.WithExpectedType(booleanType)).also(dataFlowAnalyzer::exitBinaryOr)
-        }.transformOtherChildren(transformer, ResolutionMode.WithExpectedType(booleanType)).also {
-            it.resultType = booleanType
-        }
     }
 
     override fun transformVariableAssignment(