FIR CFA: make node unionness a runtime property

Class exit node should only unify data flow if it is an anonymous
object, in which case it only has one control flow input. Coversely,
any other class exit node has normal merge semantics for its control
flow inputs, but it won't have any data flow inputs.

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

@@ -158,8 +158,6 @@ object FirCallsEffectAnalyzer : FirControlFlowChecker() {
 
         override fun visitNode(node: CFGNode<*>, data: IllegalScopeContext) {}
 
-        override fun <T> visitUnionNode(node: T, data: IllegalScopeContext) where T : CFGNode<*>, T : UnionNodeMarker {}
-
         override fun visitFunctionEnterNode(node: FunctionEnterNode, data: IllegalScopeContext) {
             // TODO: this is not how CFG works, this should be done by FIR tree traversal. Especially considering that
             //  none of these methods use anything from the CFG other than `node.fir`, which should've been a hint.
@@ -237,7 +235,7 @@ object FirCallsEffectAnalyzer : FirControlFlowChecker() {
             node: FunctionCallNode,
             data: PathAwareLambdaInvocationInfo
         ): PathAwareLambdaInvocationInfo {
-            var dataForNode = visitUnionNode(node, data)
+            var dataForNode = visitNode(node, data)
 
             val functionSymbol = node.fir.toResolvedCallableSymbol() as? FirFunctionSymbol<*>?
             val contractDescription = functionSymbol?.resolvedContractDescription

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

@@ -51,10 +51,8 @@ object FirPropertyInitializationAnalyzer : AbstractFirPropertyInitializationChec
         val reporter: DiagnosticReporter,
         val context: CheckerContext
     ) : ControlFlowGraphVisitorVoid() {
-        override fun visitNode(node: CFGNode<*>) {}
+        override fun visitNode(node: CFGNode<*>) {
-
+            // TODO: `node.isUnion` - f({ x = 1 }, { x = 2 }) - which to report?
-        override fun <T> visitUnionNode(node: T) where T : CFGNode<*>, T : UnionNodeMarker {
-            // TODO: f({ x = 1 }, { x = 2 }) - which to report?
             //  Also this is currently indistinguishable from x = 1; f({}, {}).
         }
 

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

@@ -45,7 +45,7 @@ private fun <I : ControlFlowInfo<I, *, *>> ControlFlowGraph.collectDataForNodeIn
             TraverseDirection.Backward -> node.followingCfgNodes
         }
         // TODO: if data for previousNodes hasn't changed, then should be no need to recompute data for this one
-        val union = node is UnionNodeMarker
+        val union = node.isUnion
         val previousData =
             previousNodes.mapNotNull {
                 val k = when (direction) {

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

@@ -34,8 +34,6 @@ class LocalPropertyAndCapturedWriteCollector private constructor() : ControlFlow
 
     override fun visitNode(node: CFGNode<*>) {}
 
-    override fun <T> visitUnionNode(node: T) where T : CFGNode<*>, T : UnionNodeMarker {}
-
     override fun visitVariableDeclarationNode(node: VariableDeclarationNode) {
         symbols[node.fir.symbol] = lambdaOrLocalFunctionStack.lastOrNull() == null
     }

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

@@ -65,9 +65,4 @@ abstract class PathAwareControlFlowGraphVisitor<I : ControlFlowInfo<I, *, *>> :
         node: CFGNode<*>,
         data: PathAwareControlFlowInfo<I>
     ): PathAwareControlFlowInfo<I> = data
-
-    override fun <T> visitUnionNode(
-        node: T,
-        data: PathAwareControlFlowInfo<I>
-    ): PathAwareControlFlowInfo<I> where T : CFGNode<*>, T : UnionNodeMarker = data
 }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/declaration/FirTailrecFunctionChecker.kt

@@ -40,8 +40,6 @@ object FirTailrecFunctionChecker : FirSimpleFunctionChecker() {
         graph.traverse(object : ControlFlowGraphVisitorVoid() {
             override fun visitNode(node: CFGNode<*>) {}
 
-            override fun <T> visitUnionNode(node: T) where T : CFGNode<*>, T : UnionNodeMarker {}
-
             override fun visitTryMainBlockEnterNode(node: TryMainBlockEnterNode) {
                 tryScopeCount++
             }

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/extended/CanBeValChecker.kt

@@ -82,8 +82,6 @@ object CanBeValChecker : AbstractFirPropertyInitializationChecker() {
     ) : ControlFlowGraphVisitorVoid() {
         override fun visitNode(node: CFGNode<*>) {}
 
-        override fun <T> visitUnionNode(node: T) where T : CFGNode<*>, T : UnionNodeMarker {}
-
         override fun visitVariableAssignmentNode(node: VariableAssignmentNode) {
             val symbol = node.fir.calleeReference.toResolvedPropertySymbol() ?: return
             if (symbol !in localProperties) return

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/extended/UnusedChecker.kt

@@ -50,8 +50,6 @@ object UnusedChecker : AbstractFirPropertyInitializationChecker() {
     ) : ControlFlowGraphVisitorVoid() {
         override fun visitNode(node: CFGNode<*>) {}
 
-        override fun <T> visitUnionNode(node: T) where T : CFGNode<*>, T : UnionNodeMarker {}
-
         override fun visitVariableAssignmentNode(node: VariableAssignmentNode) {
             val variableSymbol = node.fir.calleeReference.toResolvedPropertySymbol() ?: return
             val dataPerNode = data[node] ?: return
@@ -182,12 +180,6 @@ object UnusedChecker : AbstractFirPropertyInitializationChecker() {
         ): PathAwareVariableStatusInfo =
             super.visitNode(node, data).withAnnotationsFrom(node)
 
-        override fun <T> visitUnionNode(
-            node: T,
-            data: PathAwareVariableStatusInfo
-        ): PathAwareVariableStatusInfo where T : CFGNode<*>, T : UnionNodeMarker =
-            super.visitUnionNode(node, data).withAnnotationsFrom(node)
-
         override fun visitVariableDeclarationNode(
             node: VariableDeclarationNode,
             data: PathAwareVariableStatusInfo
@@ -289,7 +281,7 @@ object UnusedChecker : AbstractFirPropertyInitializationChecker() {
             node: FunctionCallNode,
             data: PathAwareVariableStatusInfo
         ): PathAwareVariableStatusInfo {
-            val dataForNode = visitUnionNode(node, data)
+            val dataForNode = visitNode(node, data)
             val reference = node.fir.calleeReference.resolved ?: return dataForNode
             val functionSymbol = reference.resolvedSymbol as? FirFunctionSymbol<*> ?: return dataForNode
             val symbol = if (functionSymbol.callableId.callableName.identifier == "invoke") {

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

@@ -1119,7 +1119,7 @@ abstract class FirDataFlowAnalyzer(
             // data flow for any of them until reaching a completed call.
             if (it is MergePostponedLambdaExitsNode) it.mergeIncomingFlow() else it.flow
         }
-        val result = logicSystem.joinFlow(previousFlows, union = this is UnionNodeMarker)
+        val result = logicSystem.joinFlow(previousFlows, isUnion)
         if (graphBuilder.lastNodeOrNull == this) {
             // Here it is, the new `lastNode`. If the previous state is the only predecessor, then there is actually
             // nothing to update; `addTypeStatement` has already ensured we have the correct information.

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

@@ -273,7 +273,7 @@ class ControlFlowGraphBuilder {
         postponedLambdaExits.push(mutableListOf())
     }
 
-    private fun <T> unifyDataFlowFromPostponedLambdas(node: T, callCompleted: Boolean) where T : CFGNode<*>, T : UnionNodeMarker {
+    private fun unifyDataFlowFromPostponedLambdas(node: CFGNode<*>, callCompleted: Boolean) {
         val currentLevelExits = postponedLambdaExits.pop()
         if (currentLevelExits.isEmpty()) return
 
@@ -449,10 +449,6 @@ class ControlFlowGraphBuilder {
             }
         }
 
-        // Classes are not initialized in place so no point in merging data flow - it will not be used.
-        val mergeDataFlow = klass is FirAnonymousObject && klass.classKind != ClassKind.ENUM_ENTRY
-        val exitKind = if (mergeDataFlow) EdgeKind.Forward else EdgeKind.CfgForward
-
         val lastInPlaceExit = calledInPlace.fold<_, CFGNode<*>>(enterNode) { lastNode, graph ->
             // In local classes, we already have control flow (+ data flow) edge from `enterNode`
             // to first in-place initializer.
@@ -461,10 +457,9 @@ class ControlFlowGraphBuilder {
             }
             graph.exitNode
         }
-        if (mergeDataFlow) {
+        if (exitNode.isUnion) {
-            // TODO: this is wrong. Not sure if it's possible to describe the correct data flow graph here. Things
+            // => this is an anonymous object => there's only one constructor => can use `exitNode`
-            //  would be much easier if it was guaranteed that constructors are analyzed after all other members:
+            // to unify data flow from all in-place-called members, including said constructor.
-            //  they have no implicit return type anyway, and we could then pass correct data flow to them.
             for (graph in calledInPlace) {
                 addEdge(graph.exitNode, exitNode, preferredKind = EdgeKind.DfgForward)
             }
@@ -486,7 +481,7 @@ class ControlFlowGraphBuilder {
             if (delegatedConstructorExit !== enterNode || delegatedConstructorExit.previousNodes.isEmpty()) {
                 addEdgeToSubGraph(delegatedConstructorExit, graph.enterNode)
             }
-            addEdge(graph.exitNode, exitNode, preferredKind = exitKind)
+            addEdge(graph.exitNode, exitNode, preferredKind = if (exitNode.isUnion) EdgeKind.Forward else EdgeKind.CfgForward)
         }
 
         if (constructors.isEmpty()) {
@@ -511,7 +506,7 @@ class ControlFlowGraphBuilder {
 
         enterNode.subGraphs = calledInPlace + constructors.values
         exitNode.subGraphs = calledLater
-        return exitNode.takeIf { mergeDataFlow } to popGraph()
+        return exitNode.takeIf { it.isUnion } to popGraph()
     }
 
     fun exitAnonymousObjectExpression(anonymousObjectExpression: FirAnonymousObjectExpression): AnonymousObjectExpressionExitNode? {

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

@@ -70,7 +70,7 @@ class FirControlFlowGraphRenderVisitor(
             when {
                 node.isDead -> "gray"
                 node == node.owner.enterNode || node == node.owner.exitNode -> "red"
-                node is UnionNodeMarker -> "yellow"
+                node.isUnion -> "yellow"
                 else -> null
             }?.let {
                 attributes += "style=\"filled\""

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

@@ -8,6 +8,7 @@
 package org.jetbrains.kotlin.fir.resolve.dfa.cfg
 
 import org.jetbrains.kotlin.KtSourceElement
+import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.*
@@ -31,6 +32,12 @@ sealed class CFGNode<out E : FirElement>(val owner: ControlFlowGraph, val level:
 
     open val canThrow: Boolean get() = false
 
+    //   a ---> b ---> d
+    //      \-> c -/
+    // Normal CFG semantics: a, then either b or c, then d
+    // If d is a union node: a, then *both* b and c in some unknown order, then d
+    open val isUnion: Boolean get() = false
+
     companion object {
         @CfgInternals
         fun addEdge(
@@ -107,7 +114,7 @@ sealed class CFGNode<out E : FirElement>(val owner: ControlFlowGraph, val level:
 
     @CfgInternals
     fun updateDeadStatus() {
-        isDead = if (this is UnionNodeMarker)
+        isDead = if (isUnion)
             _incomingEdges?.let { map -> map.values.any { it.kind.isDead } } == true
         else
             _incomingEdges?.let { map -> map.size == previousNodes.size && map.values.all { it.kind.isDead || !it.kind.usedInCfa } } == true
@@ -128,12 +135,6 @@ interface ExitNodeMarker
 interface GraphEnterNodeMarker : EnterNodeMarker
 interface GraphExitNodeMarker : ExitNodeMarker
 
-//   a ---> b ---> d
-//      \-> c -/
-// Normal CFG semantics: a, then either b or c, then d
-// If d is an instance of this interface: a, then *both* b and c in some unknown order, then d
-interface UnionNodeMarker
-
 // ----------------------------------- EnterNode for declaration with CFG -----------------------------------
 
 sealed class CFGNodeWithSubgraphs<out E : FirElement>(owner: ControlFlowGraph, level: Int) : CFGNode<E>(owner, level) {
@@ -245,7 +246,10 @@ class ClassEnterNode(owner: ControlFlowGraph, override val fir: FirClass, level:
 }
 
 class ClassExitNode(owner: ControlFlowGraph, override val fir: FirClass, level: Int) : CFGNodeWithSubgraphs<FirClass>(owner, level),
-    GraphExitNodeMarker, UnionNodeMarker {
+    GraphExitNodeMarker {
+
+    override val isUnion: Boolean
+        get() = fir is FirAnonymousObject && fir.classKind != ClassKind.ENUM_ENTRY
 
     @set:CfgInternals
     override lateinit var subGraphs: List<ControlFlowGraph>
@@ -305,7 +309,10 @@ class PropertyInitializerExitNode(owner: ControlFlowGraph, override val fir: Fir
 
 
 class DelegateExpressionExitNode(owner: ControlFlowGraph, override val fir: FirExpression, level: Int)
-    : CFGNode<FirExpression>(owner, level), UnionNodeMarker {
+    : CFGNode<FirExpression>(owner, level) {
+
+    override val isUnion: Boolean get() = true
+
     override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
         return visitor.visitDelegateExpressionExitNode(this, data)
     }
@@ -610,10 +617,13 @@ class ConstExpressionNode(owner: ControlFlowGraph, override val fir: FirConstExp
 // ----------------------------------- Check not null call -----------------------------------
 
 class CheckNotNullCallNode(owner: ControlFlowGraph, override val fir: FirCheckNotNullCall, level: Int)
-    : CFGNode<FirCheckNotNullCall>(owner, level), UnionNodeMarker {
+    : CFGNode<FirCheckNotNullCall>(owner, level) {
     override val canThrow: Boolean
         get() = true
 
+    override val isUnion: Boolean
+        get() = true
+
     override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
         return visitor.visitCheckNotNullCallNode(this, data)
     }
@@ -645,10 +655,14 @@ class ResolvedQualifierNode(
 }
 
 class FunctionCallNode(owner: ControlFlowGraph, override val fir: FirFunctionCall, level: Int)
-    : CFGNode<FirFunctionCall>(owner, level), UnionNodeMarker {
+    : CFGNode<FirFunctionCall>(owner, level) {
+
     override val canThrow: Boolean
         get() = true
 
+    override val isUnion: Boolean
+        get() = true
+
     override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
         return visitor.visitFunctionCallNode(this, data)
     }
@@ -671,7 +685,11 @@ class GetClassCallNode(owner: ControlFlowGraph, override val fir: FirGetClassCal
 }
 
 class DelegatedConstructorCallNode(owner: ControlFlowGraph, override val fir: FirDelegatedConstructorCall, level: Int)
-    : CFGNode<FirDelegatedConstructorCall>(owner, level), UnionNodeMarker {
+    : CFGNode<FirDelegatedConstructorCall>(owner, level) {
+
+    override val isUnion: Boolean
+        get() = true
+
     override val canThrow: Boolean
         get() = true // shouldn't matter since delegated constructor calls aren't wrapped in try-finally, but still
 
@@ -681,7 +699,11 @@ class DelegatedConstructorCallNode(owner: ControlFlowGraph, override val fir: Fi
 }
 
 class StringConcatenationCallNode(owner: ControlFlowGraph, override val fir: FirStringConcatenationCall, level: Int)
-    : CFGNode<FirStringConcatenationCall>(owner, level), UnionNodeMarker {
+    : CFGNode<FirStringConcatenationCall>(owner, level) {
+
+    override val isUnion: Boolean
+        get() = true
+
     override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
         return visitor.visitStringConcatenationCallNode(this, data)
     }
@@ -804,4 +826,4 @@ class SmartCastExpressionExitNode(owner: ControlFlowGraph, override val fir: Fir
     override fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R {
         return visitor.visitSmartCastExpressionExitNode(this, data)
     }
-}
+}

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

@@ -8,8 +8,6 @@ package org.jetbrains.kotlin.fir.resolve.dfa.cfg
 abstract class ControlFlowGraphVisitor<out R, in D> {
     abstract fun visitNode(node: CFGNode<*>, data: D): R
 
-    abstract fun <T> visitUnionNode(node: T, data: D): R where T : CFGNode<*>, T : UnionNodeMarker
-
     // ----------------------------------- Simple function -----------------------------------
 
     open fun visitFunctionEnterNode(node: FunctionEnterNode, data: D): R {
@@ -103,7 +101,7 @@ abstract class ControlFlowGraphVisitor<out R, in D> {
     }
 
     open fun visitDelegateExpressionExitNode(node: DelegateExpressionExitNode, data: D): R {
-        return visitUnionNode(node, data)
+        return visitNode(node, data)
     }
 
     // ----------------------------------- Field -----------------------------------
@@ -296,7 +294,7 @@ abstract class ControlFlowGraphVisitor<out R, in D> {
     // ----------------------------------- Check not null call -----------------------------------
 
     open fun visitCheckNotNullCallNode(node: CheckNotNullCallNode, data: D): R {
-        return visitUnionNode(node, data)
+        return visitNode(node, data)
     }
 
     // ----------------------------------- Resolvable call -----------------------------------
@@ -310,7 +308,7 @@ abstract class ControlFlowGraphVisitor<out R, in D> {
     }
 
     open fun visitFunctionCallNode(node: FunctionCallNode, data: D): R {
-        return visitUnionNode(node, data)
+        return visitNode(node, data)
     }
 
     open fun visitCallableReferenceNode(node: CallableReferenceNode, data: D): R {
@@ -322,11 +320,11 @@ abstract class ControlFlowGraphVisitor<out R, in D> {
     }
 
     open fun visitDelegatedConstructorCallNode(node: DelegatedConstructorCallNode, data: D): R {
-        return visitUnionNode(node, data)
+        return visitNode(node, data)
     }
 
     open fun visitStringConcatenationCallNode(node: StringConcatenationCallNode, data: D): R {
-        return visitUnionNode(node, data)
+        return visitNode(node, data)
     }
 
     open fun visitThrowExceptionNode(node: ThrowExceptionNode, data: D): R {

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

@@ -8,8 +8,6 @@ package org.jetbrains.kotlin.fir.resolve.dfa.cfg
 abstract class ControlFlowGraphVisitorVoid : ControlFlowGraphVisitor<Unit, Nothing?>() {
     abstract fun visitNode(node: CFGNode<*>)
 
-    abstract fun <T> visitUnionNode(node: T) where T : CFGNode<*>, T : UnionNodeMarker
-
     // ----------------------------------- Simple function -----------------------------------
 
     open fun visitFunctionEnterNode(node: FunctionEnterNode) {
@@ -53,7 +51,7 @@ abstract class ControlFlowGraphVisitorVoid : ControlFlowGraphVisitor<Unit, Nothi
     }
 
     open fun visitDelegateExpressionExitNode(node: DelegateExpressionExitNode) {
-        visitUnionNode(node)
+        visitNode(node)
     }
 
     // ----------------------------------- Init -----------------------------------
@@ -236,7 +234,7 @@ abstract class ControlFlowGraphVisitorVoid : ControlFlowGraphVisitor<Unit, Nothi
     // ----------------------------------- Check not null call -----------------------------------
 
     open fun visitCheckNotNullCallNode(node: CheckNotNullCallNode) {
-        visitUnionNode(node)
+        visitNode(node)
     }
 
     // ----------------------------------- Resolvable call -----------------------------------
@@ -250,15 +248,15 @@ abstract class ControlFlowGraphVisitorVoid : ControlFlowGraphVisitor<Unit, Nothi
     }
 
     open fun visitFunctionCallNode(node: FunctionCallNode) {
-        visitUnionNode(node)
+        visitNode(node)
     }
 
     open fun visitDelegatedConstructorCallNode(node: DelegatedConstructorCallNode) {
-        visitUnionNode(node)
+        visitNode(node)
     }
 
     open fun visitStringConcatenationCallNode(node: StringConcatenationCallNode) {
-        visitUnionNode(node)
+        visitNode(node)
     }
 
     open fun visitThrowExceptionNode(node: ThrowExceptionNode) {
@@ -291,10 +289,6 @@ abstract class ControlFlowGraphVisitorVoid : ControlFlowGraphVisitor<Unit, Nothi
         visitNode(node)
     }
 
-    final override fun <T> visitUnionNode(node: T, data: Nothing?) where T : CFGNode<*>, T : UnionNodeMarker {
-        visitUnionNode(node)
-    }
-
     // ----------------------------------- Simple function -----------------------------------
 
     final override fun visitFunctionEnterNode(node: FunctionEnterNode, data: Nothing?) {