FIR CFG: label edges from try-enter through finally block to exit target

compiler/fir/analysis-tests/testData/resolve/cfg/propertiesAndInitBlocks.dot

@@ -309,7 +309,8 @@ digraph propertiesAndInitBlocks_kt {
     }
     84 -> {85};
     85 -> {86};
-    86 -> {102 91 96 87};
+    86 -> {102 91 87};
+    86 -> {96}    [label=onUncaughtException];
     87 -> {88};
     88 -> {89};
     89 -> {90};
@@ -324,6 +325,7 @@ digraph propertiesAndInitBlocks_kt {
     98 -> {99};
     99 -> {100};
     100 -> {101};
+    100 -> {102}    [label=onUncaughtException];
     101 -> {102};
 
 }

compiler/fir/analysis-tests/testData/resolveWithStdlib/smartcasts/tryWithLambdaInside.dot

@@ -101,7 +101,8 @@ finally {
     11 -> {12};
     12 -> {13};
     13 -> {14};
-    14 -> {30 22 15};
+    14 -> {30 15};
+    14 -> {22}    [label=onUncaughtException];
     15 -> {16};
     16 -> {17};
     17 -> {31};
@@ -115,6 +116,7 @@ finally {
     23 -> {24};
     24 -> {25};
     25 -> {26};
+    25 -> {30}    [label=onUncaughtException];
     26 -> {27};
     27 -> {30};
     27 -> {28} [style=dotted];
@@ -190,7 +192,8 @@ finally {
     36 -> {37};
     37 -> {38};
     38 -> {39};
-    39 -> {55 47 40};
+    39 -> {55 40};
+    39 -> {47}    [label=onUncaughtException];
     40 -> {41};
     41 -> {42};
     42 -> {43 56};
@@ -203,6 +206,7 @@ finally {
     48 -> {49};
     49 -> {50};
     50 -> {51};
+    50 -> {55}    [label=onUncaughtException];
     51 -> {52};
     52 -> {55};
     52 -> {53} [style=dotted];

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

@@ -26,7 +26,7 @@ fun CFGNode<*>.isEnterNode(direction: TraverseDirection): Boolean = when (direct
 
 val CFGNode<*>.previousCfgNodes: List<CFGNode<*>>
     get() = previousNodes.filter {
-        val kind = incomingEdges.getValue(it)
+        val kind = incomingEdges.getValue(it).kind
         if (this.isDead) {
             kind.usedInCfa
         } else {
@@ -39,7 +39,7 @@ val CFGNode<*>.followingCfgNodes: List<CFGNode<*>>
         val nodes = mutableListOf<CFGNode<*>>()
 
         followingNodes.filterTo(nodes) {
-            val kind = outgoingEdges.getValue(it)
+            val kind = outgoingEdges.getValue(it).kind
             kind.usedInCfa && !kind.isDead
         }
         (this as? CFGNodeWithCfgOwner<*>)?.subGraphs?.mapTo(nodes) { it.enterNode }

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

@@ -1118,9 +1118,9 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         shouldForkFlow: Boolean = false
     ): T = this.also { node ->
         val previousFlows = if (node.isDead)
-            node.previousNodes.mapNotNull { runIf(!node.incomingEdges.getValue(it).isBack) { it.flow } }
+            node.previousNodes.mapNotNull { runIf(!node.incomingEdges.getValue(it).kind.isBack) { it.flow } }
         else
-            node.previousNodes.mapNotNull { prev -> prev.takeIf { node.incomingEdges.getValue(it).usedInDfa }?.flow }
+            node.previousNodes.mapNotNull { prev -> prev.takeIf { node.incomingEdges.getValue(it).kind.usedInDfa }?.flow }
         var flow = logicSystem.joinFlow(previousFlows)
         if (updateReceivers) {
             logicSystem.updateAllReceivers(flow)

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

@@ -19,25 +19,46 @@ import org.jetbrains.kotlin.utils.addToStdlib.runIf
 
 sealed class CFGNode<out E : FirElement>(val owner: ControlFlowGraph, val level: Int, private val id: Int) {
     companion object {
-        internal fun addEdge(from: CFGNode<*>, to: CFGNode<*>, kind: EdgeKind, propagateDeadness: Boolean) {
+        internal fun addEdge(
+            from: CFGNode<*>,
+            to: CFGNode<*>,
+            kind: EdgeKind,
+            propagateDeadness: Boolean,
+            label: EdgeLabel = NormalPath
+        ) {
             from._followingNodes += to
             to._previousNodes += from
-            addJustKindEdge(from, to, kind, propagateDeadness, edgeExists = false)
+            addJustKindEdge(from, to, kind, propagateDeadness, edgeExists = false, label = label)
         }
 
-        internal fun addJustKindEdge(from: CFGNode<*>, to: CFGNode<*>, kind: EdgeKind, propagateDeadness: Boolean) {
+        internal fun addJustKindEdge(
-            addJustKindEdge(from, to, kind, propagateDeadness, edgeExists = true)
+            from: CFGNode<*>,
+            to: CFGNode<*>,
+            kind: EdgeKind,
+            propagateDeadness: Boolean,
+            label: EdgeLabel = NormalPath
+        ) {
+            addJustKindEdge(from, to, kind, propagateDeadness, edgeExists = true, label = label)
         }
 
-        private fun addJustKindEdge(from: CFGNode<*>, to: CFGNode<*>, kind: EdgeKind, propagateDeadness: Boolean, edgeExists: Boolean) {
+        private fun addJustKindEdge(
-            if (kind != EdgeKind.Forward) {
+            from: CFGNode<*>,
-                val fromToKind = from._outgoingEdges[to] ?: runIf(edgeExists) { EdgeKind.Forward }
+            to: CFGNode<*>,
+            kind: EdgeKind,
+            propagateDeadness: Boolean,
+            edgeExists: Boolean,
+            label: EdgeLabel = NormalPath
+        ) {
+            // It's hard to define label merging, hence overwritten with the latest one.
+            // One day, if we allow multiple edges between nodes with different labels, we won't even need kind merging.
+            if (kind != EdgeKind.Forward || label != NormalPath) {
+                val fromToKind = from._outgoingEdges[to]?.kind ?: runIf(edgeExists) { EdgeKind.Forward }
                 merge(kind, fromToKind)?.let {
-                    from._outgoingEdges[to] = it
+                    from._outgoingEdges[to] = Edge.create(label, it)
                 } ?: from._outgoingEdges.remove(to)
-                val toFromKind = to._incomingEdges[from] ?: runIf(edgeExists) { EdgeKind.Forward }
+                val toFromKind = to._incomingEdges[from]?.kind ?: runIf(edgeExists) { EdgeKind.Forward }
                 merge(kind, toFromKind)?.let {
-                    to._incomingEdges[from] = it
+                    to._incomingEdges[from] = Edge.create(label, it)
                 } ?: to._incomingEdges.remove(from)
             }
             if (propagateDeadness && kind == EdgeKind.DeadForward) {
@@ -87,18 +108,18 @@ sealed class CFGNode<out E : FirElement>(val owner: ControlFlowGraph, val level:
     val previousNodes: List<CFGNode<*>> get() = _previousNodes
     val followingNodes: List<CFGNode<*>> get() = _followingNodes
 
-    private val _incomingEdges = mutableMapOf<CFGNode<*>, EdgeKind>().withDefault { EdgeKind.Forward }
+    private val _incomingEdges = mutableMapOf<CFGNode<*>, Edge>().withDefault { Edge.Normal_Forward }
-    private val _outgoingEdges = mutableMapOf<CFGNode<*>, EdgeKind>().withDefault { EdgeKind.Forward }
+    private val _outgoingEdges = mutableMapOf<CFGNode<*>, Edge>().withDefault { Edge.Normal_Forward }
 
-    val incomingEdges: Map<CFGNode<*>, EdgeKind> get() = _incomingEdges
+    val incomingEdges: Map<CFGNode<*>, Edge> get() = _incomingEdges
-    val outgoingEdges: Map<CFGNode<*>, EdgeKind> get() = _outgoingEdges
+    val outgoingEdges: Map<CFGNode<*>, Edge> get() = _outgoingEdges
 
     abstract val fir: E
     var isDead: Boolean = false
         protected set
 
     internal fun updateDeadStatus() {
-        isDead = incomingEdges.size == previousNodes.size && incomingEdges.values.all { it == EdgeKind.DeadForward }
+        isDead = incomingEdges.size == previousNodes.size && incomingEdges.values.all { it.kind == EdgeKind.DeadForward }
     }
 
     abstract fun <R, D> accept(visitor: ControlFlowGraphVisitor<R, D>, data: D): R

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve.dfa.cfg
 
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 
 class ControlFlowGraph(val declaration: FirDeclaration?, val name: String, val kind: Kind) {
     private var _nodes: MutableList<CFGNode<*>> = mutableListOf()
@@ -86,6 +87,58 @@ class ControlFlowGraph(val declaration: FirDeclaration?, val name: String, val k
     }
 }
 
+data class Edge(
+    val label: EdgeLabel,
+    val kind: EdgeKind,
+) {
+    companion object {
+        val Normal_Forward = Edge(NormalPath, EdgeKind.Forward)
+        private val Normal_DeadForward = Edge(NormalPath, EdgeKind.DeadForward)
+        private val Normal_DfgForward = Edge(NormalPath, EdgeKind.DfgForward)
+        private val Normal_CfgForward = Edge(NormalPath, EdgeKind.CfgForward)
+        private val Normal_CfgBackward = Edge(NormalPath, EdgeKind.CfgBackward)
+        private val Normal_DeadBackward = Edge(NormalPath, EdgeKind.DeadBackward)
+
+        fun create(label: EdgeLabel, kind: EdgeKind): Edge =
+            when (label) {
+                NormalPath -> {
+                    when (kind) {
+                        EdgeKind.Forward -> Normal_Forward
+                        EdgeKind.DeadForward -> Normal_DeadForward
+                        EdgeKind.DfgForward -> Normal_DfgForward
+                        EdgeKind.CfgForward -> Normal_CfgForward
+                        EdgeKind.CfgBackward -> Normal_CfgBackward
+                        EdgeKind.DeadBackward -> Normal_DeadBackward
+                    }
+                }
+                else -> {
+                    Edge(label, kind)
+                }
+            }
+    }
+}
+
+sealed class EdgeLabel(val label: String?) {
+    open val isNormal: Boolean
+        get() = false
+
+    override fun toString(): String {
+        return label ?: ""
+    }
+}
+
+object NormalPath : EdgeLabel(label = null) {
+    override val isNormal: Boolean
+        get() = true
+}
+
+object UncaughtExceptionPath : EdgeLabel(label = "onUncaughtException")
+
+// TODO: Label `return`ing edge with this.
+class ReturnPath(
+    returnTargetSymbol: FirFunctionSymbol<*>
+) : EdgeLabel(label = "return@${returnTargetSymbol.callableId}")
+
 enum class EdgeKind(
     val usedInDfa: Boolean,
     val usedInCfa: Boolean,
@@ -113,7 +166,7 @@ private fun ControlFlowGraph.orderNodes(): LinkedHashSet<CFGNode<*>> {
     while (stack.isNotEmpty()) {
         val node = stack.removeFirst()
         val previousNodes = node.previousNodes
-        if (previousNodes.any { it !in visitedNodes && it.owner == this && !node.incomingEdges.getValue(it).isBack }) {
+        if (previousNodes.any { it !in visitedNodes && it.owner == this && !node.incomingEdges.getValue(it).kind.isBack }) {
             delayedNodes.add(node)
             stack.addLast(node)
             continue

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

@@ -777,7 +777,8 @@ class ControlFlowGraphBuilder {
         catchNodeStorages.push(NodeStorage())
         val enterTryExpressionNode = createTryExpressionEnterNode(tryExpression)
         addNewSimpleNode(enterTryExpressionNode)
-        tryExitNodes.push(createTryExpressionExitNode(tryExpression))
+        val tryExitNode = createTryExpressionExitNode(tryExpression)
+        tryExitNodes.push(tryExitNode)
         levelCounter++
         val enterTryNodeBlock = createTryMainBlockEnterNode(tryExpression)
         addNewSimpleNode(enterTryNodeBlock)
@@ -793,7 +794,7 @@ class ControlFlowGraphBuilder {
 
         if (tryExpression.finallyBlock != null) {
             val finallyEnterNode = createFinallyBlockEnterNode(tryExpression)
-            addEdge(enterTryNodeBlock, finallyEnterNode)
+            addEdge(enterTryNodeBlock, finallyEnterNode, label = UncaughtExceptionPath)
             finallyEnterNodes.push(finallyEnterNode)
         }
 
@@ -841,7 +842,9 @@ class ControlFlowGraphBuilder {
     fun exitFinallyBlock(tryExpression: FirTryExpression): FinallyBlockExitNode {
         return createFinallyBlockExitNode(tryExpression).also {
             popAndAddEdge(it)
-            addEdge(it, tryExitNodes.top())
+            val tryExitNode = tryExitNodes.top()
+            addEdge(it, tryExitNode)
+            addEdge(it, exitTargetsForTry.top(), label = UncaughtExceptionPath)
         }
     }
 
@@ -1193,20 +1196,22 @@ class ControlFlowGraphBuilder {
         propagateDeadness: Boolean = true,
         isDead: Boolean = false,
         isBack: Boolean = false,
-        preferredKind: EdgeKind = EdgeKind.Forward
+        preferredKind: EdgeKind = EdgeKind.Forward,
+        label: EdgeLabel = NormalPath
     ) {
         val kind = if (isDead || from.isDead || to.isDead) {
             if (isBack) EdgeKind.DeadBackward else EdgeKind.DeadForward
         } else preferredKind
-        CFGNode.addEdge(from, to, kind, propagateDeadness)
+        CFGNode.addEdge(from, to, kind, propagateDeadness, label)
     }
 
     private fun addBackEdge(
         from: CFGNode<*>,
         to: CFGNode<*>,
-        isDead: Boolean = false
+        isDead: Boolean = false,
+        label: EdgeLabel = NormalPath
     ) {
-        addEdge(from, to, propagateDeadness = false, isDead = isDead, isBack = true, preferredKind = EdgeKind.CfgBackward)
+        addEdge(from, to, propagateDeadness = false, isDead = isDead, isBack = true, preferredKind = EdgeKind.CfgBackward, label = label)
     }
 
     // ----------------------------------- Utils -----------------------------------

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

@@ -108,15 +108,35 @@ class FirControlFlowGraphRenderVisitor(
             if (node.followingNodes.isEmpty()) continue
 
             fun renderEdges(kind: EdgeKind) {
-                val edges = node.followingNodes.filter { node.outgoingEdges.getValue(it) == kind }
+                val edges = node.followingNodes.filter { node.outgoingEdges.getValue(it).kind == kind }
                 if (edges.isEmpty()) return
-                print(
+
-                    indices.getValue(node),
+                fun renderEdgesWithoutLabel(edges: List<CFGNode<*>>) {
-                    EDGE,
+                    print(
-                    edges.joinToString(prefix = "{", postfix = "}", separator = " ") { indices.getValue(it).toString() }
+                        indices.getValue(node),
-                )
+                        EDGE,
-                EDGE_STYLE.getValue(kind).takeIf { it.isNotBlank() }?.let { printWithNoIndent(" $it") }
+                        edges.joinToString(prefix = "{", postfix = "}", separator = " ") { indices.getValue(it).toString() }
-                printlnWithNoIndent(";")
+                    )
+                    EDGE_STYLE.getValue(kind).takeIf { it.isNotBlank() }?.let { printWithNoIndent(" $it") }
+                    printlnWithNoIndent(";")
+                }
+
+                if (edges.any { node.outgoingEdges[it]?.label?.label != null }) {
+                    val edgeGroups = edges.groupBy { node.outgoingEdges[it]?.label?.label != null }
+                    edgeGroups[false]?.let { renderEdgesWithoutLabel(it) }
+                    for (edge in edgeGroups[true]!!) {
+                        print(
+                            indices.getValue(node),
+                            EDGE,
+                            "{", indices.getValue(edge), "}"
+                        )
+                        EDGE_STYLE.getValue(kind).takeIf { it.isNotBlank() }?.let { printWithNoIndent(" $it") }
+                        print("[label=${node.outgoingEdges[edge]!!.label}]")
+                        printlnWithNoIndent(";")
+                    }
+                } else {
+                    renderEdgesWithoutLabel(edges)
+                }
             }
 
             for (kind in EdgeKind.values()) {

compiler/tests-common/tests/org/jetbrains/kotlin/fir/AbstractFirDiagnosticTest.kt

@@ -326,8 +326,8 @@ abstract class AbstractFirDiagnosticsTest : AbstractFirBaseDiagnosticsTest() {
         private fun checkEdge(from: CFGNode<*>, to: CFGNode<*>) {
             KtUsefulTestCase.assertContainsElements(from.followingNodes, to)
             KtUsefulTestCase.assertContainsElements(to.previousNodes, from)
-            val fromKind = from.outgoingEdges.getValue(to)
+            val fromKind = from.outgoingEdges.getValue(to).kind
-            val toKind = to.incomingEdges.getValue(from)
+            val toKind = to.incomingEdges.getValue(from).kind
             TestCase.assertEquals(fromKind, toKind)
             if (from.isDead && to.isDead) {
                 KtUsefulTestCase.assertContainsElements(cfgKinds, toKind)
@@ -339,7 +339,7 @@ abstract class AbstractFirDiagnosticsTest : AbstractFirBaseDiagnosticsTest() {
             for (node in graph.nodes) {
                 for (previousNode in node.previousNodes) {
                     if (previousNode.owner != graph) continue
-                    if (!node.incomingEdges.getValue(previousNode).isBack) {
+                    if (!node.incomingEdges.getValue(previousNode).kind.isBack) {
                         assertTrue(previousNode in visited)
                     }
                 }