FIR CFG: reconfigure exception throwing paths in try expression

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

@@ -153,12 +153,12 @@ digraph complex_kt {
     22 -> {23};
     23 -> {24};
     24 -> {25};
-    25 -> {26};
-    26 -> {54 31 27};
+    25 -> {26 31};
+    26 -> {27};
     27 -> {28};
     28 -> {29};
     29 -> {30};
-    30 -> {38};
+    30 -> {38 31};
     31 -> {54 32};
     32 -> {33};
     33 -> {34};

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

@@ -167,8 +167,8 @@ digraph inplaceLambdaInControlFlowExpressions_kt {
     }
     37 -> {38};
     38 -> {39};
-    39 -> {40};
-    40 -> {56 47 41};
+    39 -> {40 47};
+    40 -> {41};
     41 -> {42};
     42 -> {57};
     42 -> {43} [color=red];
@@ -176,7 +176,7 @@ digraph inplaceLambdaInControlFlowExpressions_kt {
     43 -> {44};
     44 -> {45};
     45 -> {46};
-    46 -> {52};
+    46 -> {52 47};
     47 -> {56 48};
     48 -> {49};
     49 -> {50};

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

@@ -308,13 +308,13 @@ digraph propertiesAndInitBlocks_kt {
         102 [label="Exit property" style="filled" fillcolor=red];
     }
     84 -> {85};
-    85 -> {86};
-    86 -> {102 91 87};
-    86 -> {96}    [label=onUncaughtException];
+    85 -> {86 91};
+    85 -> {96}    [label=onUncaughtException];
+    86 -> {87};
     87 -> {88};
     88 -> {89};
     89 -> {90};
-    90 -> {96};
+    90 -> {96 91};
     91 -> {102 92};
     92 -> {93};
     93 -> {94};

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

@@ -56,13 +56,13 @@ digraph tryCatch_kt {
     }
     0 -> {1};
     1 -> {2};
-    2 -> {3};
-    3 -> {23 15 9 4};
+    2 -> {3 15 9};
+    3 -> {4};
     4 -> {5};
     5 -> {6};
     6 -> {7};
     7 -> {8};
-    8 -> {21};
+    8 -> {21 15 9};
     9 -> {23 10};
     10 -> {11};
     11 -> {12};
@@ -118,12 +118,12 @@ digraph tryCatch_kt {
     }
     24 -> {25};
     25 -> {26};
-    26 -> {27};
-    27 -> {40 32 28};
+    26 -> {27 32};
+    27 -> {28};
     28 -> {29};
     29 -> {30};
     30 -> {31};
-    31 -> {37};
+    31 -> {37 32};
     32 -> {40 33};
     33 -> {34};
     34 -> {35};
@@ -261,8 +261,8 @@ digraph tryCatch_kt {
     46 -> {87} [style=dotted];
     47 -> {48};
     48 -> {49};
-    49 -> {50};
-    50 -> {102 91 81 51};
+    49 -> {50 91 81};
+    50 -> {51};
     51 -> {52};
     52 -> {53};
     53 -> {54};
@@ -294,7 +294,7 @@ digraph tryCatch_kt {
     77 -> {78} [style=dotted];
     78 -> {79};
     79 -> {80};
-    80 -> {97};
+    80 -> {97 91 81};
     81 -> {102 82};
     82 -> {83};
     83 -> {87};

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

@@ -101,8 +101,8 @@ finally {
     11 -> {12};
     12 -> {13};
     13 -> {14};
-    14 -> {30 15};
-    14 -> {22}    [label=onUncaughtException];
+    13 -> {22}    [label=onUncaughtException];
+    14 -> {15};
     15 -> {16};
     16 -> {17};
     17 -> {31};
@@ -192,8 +192,8 @@ finally {
     36 -> {37};
     37 -> {38};
     38 -> {39};
-    39 -> {55 40};
-    39 -> {47}    [label=onUncaughtException];
+    38 -> {47}    [label=onUncaughtException];
+    39 -> {40};
     40 -> {41};
     41 -> {42};
     42 -> {43 56};

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

@@ -84,6 +84,7 @@ class ControlFlowGraphBuilder {
     private val binaryOrExitNodes: Stack<BinaryOrExitNode> = stackOf()
 
     private val tryExitNodes: NodeStorage<FirTryExpression, TryExpressionExitNode> = NodeStorage()
+    private val tryMainExitNodes: NodeStorage<FirTryExpression, TryMainBlockExitNode> = NodeStorage()
     private val catchNodeStorages: Stack<NodeStorage<FirCatch, CatchClauseEnterNode>> = stackOf()
     private val catchNodeStorage: NodeStorage<FirCatch, CatchClauseEnterNode> get() = catchNodeStorages.top()
     private val finallyEnterNodes: Stack<FinallyBlockEnterNode> = stackOf()
@@ -786,19 +787,21 @@ class ControlFlowGraphBuilder {
         levelCounter++
         val enterTryNodeBlock = createTryMainBlockEnterNode(tryExpression)
         addNewSimpleNode(enterTryNodeBlock)
-        addEdge(enterTryNodeBlock, exitTargetsForTry.top())
 
         for (catch in tryExpression.catches) {
             val catchNode = createCatchClauseEnterNode(catch)
             catchNodeStorage.push(catchNode)
-            addEdge(enterTryNodeBlock, catchNode)
+            // a flow where an exception of interest is thrown and caught before executing any of try-main block.
+            addEdge(enterTryExpressionNode, catchNode)
+            // TODO: Should go to the finally block w/ a label, and exit to the exit target w/ the same label.
             addEdge(catchNode, exitTargetsForTry.top())
         }
         levelCounter++
 
         if (tryExpression.finallyBlock != null) {
             val finallyEnterNode = createFinallyBlockEnterNode(tryExpression)
-            addEdge(enterTryNodeBlock, finallyEnterNode, label = UncaughtExceptionPath)
+            // a flow where an uncaught exception is thrown before executing any of try-main block.
+            addEdge(enterTryExpressionNode, finallyEnterNode, label = UncaughtExceptionPath)
             finallyEnterNodes.push(finallyEnterNode)
         }
 
@@ -808,6 +811,7 @@ class ControlFlowGraphBuilder {
     fun exitTryMainBlock(tryExpression: FirTryExpression): TryMainBlockExitNode {
         levelCounter--
         val node = createTryMainBlockExitNode(tryExpression)
+        tryMainExitNodes.push(node)
         popAndAddEdge(node)
         val finallyEnterNode = finallyEnterNodes.topOrNull()
         // NB: Check the level to avoid adding an edge to the finally block at an upper level.
@@ -820,7 +824,13 @@ class ControlFlowGraphBuilder {
     }
 
     fun enterCatchClause(catch: FirCatch): CatchClauseEnterNode {
-        return catchNodeStorage[catch]!!.also { lastNodes.push(it) }.also { levelCounter++ }
+        return catchNodeStorage[catch]!!.also {
+            val tryMainExitNode = tryMainExitNodes.top()
+            // a flow where an exception of interest is thrown and caught after executing all of try-main block.
+            addEdge(tryMainExitNode, it)
+            lastNodes.push(it)
+            levelCounter++
+        }
     }
 
     fun exitCatchClause(catch: FirCatch): CatchClauseExitNode {
@@ -847,7 +857,9 @@ class ControlFlowGraphBuilder {
         return createFinallyBlockExitNode(tryExpression).also {
             popAndAddEdge(it)
             val tryExitNode = tryExitNodes.top()
+            // a flow where either there wasn't any exception or caught if any.
             addEdge(it, tryExitNode)
+            // a flow that exits to the exit target while there was an uncaught exception.
             addEdge(it, exitTargetsForTry.top(), label = UncaughtExceptionPath)
         }
     }
@@ -857,6 +869,7 @@ class ControlFlowGraphBuilder {
     ): Pair<TryExpressionExitNode, UnionFunctionCallArgumentsNode?> {
         levelCounter--
         catchNodeStorages.pop()
+        tryMainExitNodes.pop()
         val node = tryExitNodes.pop()
         node.updateDeadStatus()
         lastNodes.push(node)