[FIR][CFG] Support multiple levels of finally blocks in jumps

Jump inside try expressions now goes through all finally blocks between
it and the target

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

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve.dfa
 
 import org.jetbrains.kotlin.fir.FirElement
+import org.jetbrains.kotlin.fir.asReversedFrozen
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.CFGNode
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
@@ -16,6 +17,11 @@ abstract class Stack<T> {
     abstract fun pop(): T
     abstract fun push(value: T)
     abstract fun reset()
+
+    /**
+     * returns all elements of the stack in order of retrieval
+     */
+    abstract fun all(): List<T>
 }
 
 fun <T> stackOf(vararg values: T): Stack<T> = StackImpl(*values)
@@ -38,6 +44,8 @@ private class StackImpl<T>(vararg values: T) : Stack<T>() {
     override fun reset() {
         stack.clear()
     }
+
+    override fun all(): List<T> = stack.asReversedFrozen()
 }
 
 class NodeStorage<T : FirElement, N : CFGNode<T>> : Stack<N>(){
@@ -65,6 +73,8 @@ class NodeStorage<T : FirElement, N : CFGNode<T>> : Stack<N>(){
         stack.reset()
         map.clear()
     }
+
+    override fun all(): List<N> = stack.all()
 }
 
 class SymbolBasedNodeStorage<T, N : CFGNode<T>> : Stack<N>() where T : FirElement {
@@ -92,4 +102,6 @@ class SymbolBasedNodeStorage<T, N : CFGNode<T>> : Stack<N>() where T : FirElemen
         stack.reset()
         map.clear()
     }
+
+    override fun all(): List<N> = stack.all()
 }

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

@@ -1389,21 +1389,16 @@ class ControlFlowGraphBuilder {
                     addBackEdge(node, targetNode, label = label)
                 }
             } else {
-                //TODO this supports single try-finally block only
+                //go through all final nodes between node and target
-                //need to get all try-finally up to target
                 val finallyNodes = finallyBefore(targetNode)
-                if (finallyNodes != null) {
+                val finalFrom = finallyNodes.fold(node) { from, (finallyEnter, tryExit) ->
-                    val (finallyEnter, finallyExit) = finallyNodes
+                    addEdgeIfNotExist(from, finallyEnter, propagateDeadness = false, label = label)
-                    addEdge(node, finallyEnter, propagateDeadness = false, label = label)
+                    tryExit
-                    if (!finallyExit.followingNodes.contains(targetNode)) {
+                }
-                        addEdge(finallyExit, targetNode, propagateDeadness = false, label = label)
+                addEdgeIfNotExist(finalFrom, targetNode, propagateDeadness = false, label = label)
-                    }
+                if (trackJump && finallyNodes.isNotEmpty()) {
-                    if (trackJump) {
+                    //actually we can store all returns like this, but not sure if it makes anything better
-                        //actually we can store all returns like this, but not sure if it make anything better
+                    nonDirectJumps.put(targetNode, node)
-                        nonDirectJumps.put(targetNode, node)
-                    }
-                } else {
-                    addEdge(node, targetNode, propagateDeadness = false, label = label)
                 }
             }
         }
@@ -1428,14 +1423,29 @@ class ControlFlowGraphBuilder {
         lastNodes.push(stub)
     }
 
-    private fun finallyBefore(target: CFGNode<*>): Pair<FinallyBlockEnterNode, TryExpressionExitNode>? {
+    private fun finallyBefore(target: CFGNode<*>): List<Pair<FinallyBlockEnterNode, TryExpressionExitNode>> {
-        return finallyEnterNodes.topOrNull()?.takeIf { it.level > target.level }?.let { it to tryExitNodes.top() }
+        return finallyEnterNodes.all().takeWhile { it.level > target.level }.map {
+            it to tryExitNodes[it.fir]!!
+        }
     }
 
     private fun popAndAddEdge(to: CFGNode<*>, preferredKind: EdgeKind = EdgeKind.Forward) {
         addEdge(lastNodes.pop(), to, preferredKind = preferredKind)
     }
 
+    private fun addEdgeIfNotExist(
+        from: CFGNode<*>,
+        to: CFGNode<*>,
+        propagateDeadness: Boolean = true,
+        isDead: Boolean = false,
+        preferredKind: EdgeKind = EdgeKind.Forward,
+        label: EdgeLabel = NormalPath
+    ) {
+        if (!from.followingNodes.contains(to)) {
+            addEdge(from, to, propagateDeadness, isDead, preferredKind = preferredKind, label = label)
+        }
+    }
+
     private fun addEdge(
         from: CFGNode<*>,
         to: CFGNode<*>,