[FIR] Improve the control flow graph around try expressions

1. throw goes to catches instead of main exist block
2. return goes via finally (single level only supported atm)
3. collect non-direct return to retrieve all return expressions easier

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

@@ -16,7 +16,6 @@ import org.jetbrains.kotlin.fir.declarations.utils.isLateInit
 import org.jetbrains.kotlin.fir.declarations.utils.referredPropertySymbol
 import org.jetbrains.kotlin.fir.expressions.FirQualifiedAccess
 import org.jetbrains.kotlin.fir.expressions.FirVariableAssignment
-import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
 import org.jetbrains.kotlin.fir.symbols.SymbolInternals
 import org.jetbrains.kotlin.fir.symbols.impl.FirPropertySymbol

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

@@ -9,8 +9,10 @@ import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.analysis.diagnostics.reportOn
-import org.jetbrains.kotlin.fir.declarations.*
+import org.jetbrains.kotlin.fir.declarations.FirAnonymousFunction
-import org.jetbrains.kotlin.fir.declarations.utils.hasBody
+import org.jetbrains.kotlin.fir.declarations.FirConstructor
+import org.jetbrains.kotlin.fir.declarations.FirFunction
+import org.jetbrains.kotlin.fir.declarations.FirPropertyAccessor
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.BlockExitNode
 import org.jetbrains.kotlin.fir.resolve.dfa.controlFlowGraph
 import org.jetbrains.kotlin.fir.types.isNothing

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

@@ -12,7 +12,6 @@ import org.jetbrains.kotlin.fir.analysis.checkers.context.CheckerContext
 import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.analysis.diagnostics.reportOn
-import org.jetbrains.kotlin.fir.expressions.FirTryExpression
 import org.jetbrains.kotlin.fir.resolve.dfa.cfg.*
 import org.jetbrains.kotlin.fir.visitors.FirVisitorVoid
 
@@ -27,11 +26,6 @@ object UnreachableCodeChecker : FirControlFlowChecker() {
         val unreachableElements = unreachableNodes.map { it.fir }
         val innerNodes = mutableSetOf<FirElement>()
         unreachableElements.forEach { it.collectInnerNodes(innerNodes) }
-        //todo cfg is broken in catch and finally blocks, so exclude reporting anything
-        nodes.mapNotNull { it.fir as? FirTryExpression }.distinct().forEach { tryNode ->
-            tryNode.finallyBlock?.collectInnerNodes(innerNodes)
-            tryNode.catches.forEach { it.collectInnerNodes(innerNodes) }
-        }
         unreachableElements.distinctBy { it.source }.forEach { element ->
             if (element !in innerNodes) {
                 reporter.reportOn(element.source, FirErrors.UNREACHABLE_CODE, reachableSources, unreachableSources, context)

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

@@ -18,6 +18,8 @@ import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.types.isNothing
+import org.jetbrains.kotlin.fir.util.ListMultimap
+import org.jetbrains.kotlin.fir.util.listMultimapOf
 import org.jetbrains.kotlin.fir.visitors.FirDefaultVisitor
 import org.jetbrains.kotlin.utils.addToStdlib.runIf
 import kotlin.random.Random
@@ -69,6 +71,9 @@ class ControlFlowGraphBuilder {
     private val exitsOfAnonymousFunctions: MutableMap<FirFunctionSymbol<*>, FunctionExitNode> = mutableMapOf()
     private val enterToLocalClassesMembers: MutableMap<FirBasedSymbol<*>, CFGNode<*>?> = mutableMapOf()
 
+    //return jumps via finally blocks, target -> jumps
+    private val nonDirectJumps: ListMultimap<CFGNode<*>, CFGNode<*>> = listMultimapOf()
+
     private val postponedLambdas: MutableSet<FirFunctionSymbol<*>> = mutableSetOf()
     private val entersToPostponedAnonymousFunctions: MutableMap<FirFunctionSymbol<*>, PostponedLambdaEnterNode> = mutableMapOf()
     private val exitsFromPostponedAnonymousFunctions: MutableMap<FirFunctionSymbol<*>, PostponedLambdaExitNode> = mutableMapOf()
@@ -89,6 +94,7 @@ class ControlFlowGraphBuilder {
     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 catchExitNodeStorages: Stack<NodeStorage<FirCatch, CatchClauseExitNode>> = stackOf()
     private val finallyEnterNodes: Stack<FinallyBlockEnterNode> = stackOf()
 
     private val initBlockExitNodes: Stack<InitBlockExitNode> = stackOf()
@@ -143,7 +149,8 @@ class ControlFlowGraphBuilder {
 
         val exitNode = function.controlFlowGraphReference?.controlFlowGraph?.exitNode
             ?: exitsOfAnonymousFunctions.getValue(function.symbol)
-        return exitNode.previousNodes.mapNotNullTo(mutableSetOf()) {
+        val nonDirect = nonDirectJumps[exitNode]
+        return (nonDirect + exitNode.previousNodes).mapNotNullTo(mutableSetOf()) {
             it.extractArgument() as FirStatement?
         }
     }
@@ -597,14 +604,13 @@ class ControlFlowGraphBuilder {
 
     fun exitJump(jump: FirJump<*>): JumpNode {
         val node = createJumpNode(jump)
-        // TODO: if within `try` with `finally`, don't go to the target directly.
         val nextNode = when (jump) {
             is FirReturnExpression -> exitTargetsForReturn[jump.target.labeledElement.symbol]
             is FirContinueExpression -> loopEnterNodes[jump.target.labeledElement]
             is FirBreakExpression -> loopExitNodes[jump.target.labeledElement]
             else -> throw IllegalArgumentException("Unknown jump type: ${jump.render()}")
         }
-        addNodeWithJump(node, nextNode, isBack = jump is FirContinueExpression, isBreak = jump is FirBreakExpression)
+        addNodeWithJump(node, nextNode, isBack = jump is FirContinueExpression, trackJump = jump is FirReturnExpression)
         return node
     }
 
@@ -841,6 +847,7 @@ class ControlFlowGraphBuilder {
 
     fun enterTryExpression(tryExpression: FirTryExpression): Pair<TryExpressionEnterNode, TryMainBlockEnterNode> {
         catchNodeStorages.push(NodeStorage())
+        catchExitNodeStorages.push(NodeStorage())
         val enterTryExpressionNode = createTryExpressionEnterNode(tryExpression)
         addNewSimpleNode(enterTryExpressionNode)
         val tryExitNode = createTryExpressionExitNode(tryExpression)
@@ -874,13 +881,17 @@ class ControlFlowGraphBuilder {
         levelCounter--
         val node = tryMainExitNodes.top()
         popAndAddEdge(node)
+        node.updateDeadStatus()
         val finallyEnterNode = finallyEnterNodes.topOrNull()
         // NB: Check the level to avoid adding an edge to the finally block at an upper level.
         if (finallyEnterNode != null && finallyEnterNode.level == levelCounter + 1) {
             // TODO: in case of return/continue/break in try main block, we need a unique label.
             addEdge(node, finallyEnterNode)
+            //in case try exit is dead, but there is other edges to finally (eg return)
+            // actually finallyEnterNode can't be dead, except for the case when the whole try is dead
+            finallyEnterNode.updateDeadStatus()
         } else {
-            addEdge(node, tryExitNodes.top())
+            addEdge(node, tryExitNodes.top(), propagateDeadness = false)
         }
         return node
     }
@@ -890,6 +901,8 @@ class ControlFlowGraphBuilder {
             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)
+            //tryMainExitNode might be dead (eg main block contains return), but it doesn't mean catch block is also dead
+            it.updateDeadStatus()
             val finallyEnterNode = finallyEnterNodes.topOrNull()
             // a flow where an uncaught exception is thrown before executing any of catch clause.
             // NB: Check the level to avoid adding an edge to the finally block at an upper level.
@@ -915,6 +928,7 @@ class ControlFlowGraphBuilder {
             } else {
                 addEdge(it, tryExitNodes.top(), propagateDeadness = false)
             }
+            catchExitNodeStorages.top().push(it)
         }
     }
 
@@ -947,11 +961,24 @@ class ControlFlowGraphBuilder {
     ): Pair<TryExpressionExitNode, UnionFunctionCallArgumentsNode?> {
         levelCounter--
         catchNodeStorages.pop()
-        tryMainExitNodes.pop()
+        val catchExitNodes = catchExitNodeStorages.pop()
+        val tryMainExitNode = tryMainExitNodes.pop()
         val node = tryExitNodes.pop()
         node.updateDeadStatus()
         lastNodes.push(node)
         val (_, unionNode) = processUnionOfArguments(node, callCompleted)
+
+        val allCatchesAreDead = tryMainExitNode.fir.catches.all { catch -> catchExitNodes[catch]!!.isDead }
+        val tryMainBlockIsDead = tryMainExitNode.previousNodes.all { prev ->
+            prev.isDead || tryMainExitNode.incomingEdges.getValue(prev).label != NormalPath
+        }
+        if (tryMainBlockIsDead && allCatchesAreDead) {
+            //if try expression doesn't have any regular non-dead exits, we add stub to make everything after dead
+            val stub = createStubNode()
+            popAndAddEdge(stub)
+            lastNodes.push(stub)
+        }
+
         return node to unionNode
     }
 
@@ -1331,7 +1358,17 @@ class ControlFlowGraphBuilder {
             // (3)... within finally.
             else -> exitTargetsForTry.top()
         }
-        addNodeWithJump(node, targetNode, preferredKind, label = UncaughtExceptionPath)
+        if (targetNode is TryMainBlockExitNode) {
+            val catches = targetNode.fir.catches
+            if (catches.isEmpty()) {
+                addNodeWithJump(node, exitTargetsForTry.top(), preferredKind, label = UncaughtExceptionPath)
+            } else {
+                //edges to all the catches
+                addNodeWithJumpToCatches(node, catches.map { catchNodeStorage[it]!! }, preferredKind = preferredKind)
+            }
+        } else {
+            addNodeWithJump(node, targetNode, preferredKind, label = UncaughtExceptionPath)
+        }
     }
 
     private fun addNodeWithJump(
@@ -1340,7 +1377,7 @@ class ControlFlowGraphBuilder {
         preferredKind: EdgeKind = EdgeKind.Forward,
         isBack: Boolean = false,
         label: EdgeLabel = NormalPath,
-        isBreak: Boolean = false
+        trackJump: Boolean = false
     ) {
         popAndAddEdge(node, preferredKind)
         if (targetNode != null) {
@@ -1355,10 +1392,16 @@ class ControlFlowGraphBuilder {
                 //TODO this supports single try-finally block only
                 //need to get all try-finally up to target
                 val finallyNodes = finallyBefore(targetNode)
-                if (finallyNodes != null && isBreak) {
+                if (finallyNodes != null) {
                     val (finallyEnter, finallyExit) = finallyNodes
                     addEdge(node, finallyEnter, propagateDeadness = false, label = label)
-                    addEdge(finallyExit, targetNode, propagateDeadness = false, label = label)
+                    if (!finallyExit.followingNodes.contains(targetNode)) {
+                        addEdge(finallyExit, targetNode, propagateDeadness = false, label = label)
+                    }
+                    if (trackJump) {
+                        //actually we can store all returns like this, but not sure if it make anything better
+                        nonDirectJumps.put(targetNode, node)
+                    }
                 } else {
                     addEdge(node, targetNode, propagateDeadness = false, label = label)
                 }
@@ -1369,6 +1412,22 @@ class ControlFlowGraphBuilder {
         lastNodes.push(stub)
     }
 
+    private fun addNodeWithJumpToCatches(
+        node: CFGNode<*>,
+        targets: List<CatchClauseEnterNode>,
+        label: EdgeLabel = UncaughtExceptionPath,
+        preferredKind: EdgeKind = EdgeKind.Forward
+    ) {
+        require(targets.isNotEmpty())
+        popAndAddEdge(node, preferredKind)
+        targets.forEach { target ->
+            addEdge(node, target, propagateDeadness = false, label = label)
+        }
+        val stub = createStubNode()
+        addEdge(node, stub)
+        lastNodes.push(stub)
+    }
+
     private fun finallyBefore(target: CFGNode<*>): Pair<FinallyBlockEnterNode, TryExpressionExitNode>? {
         return finallyEnterNodes.topOrNull()?.takeIf { it.level > target.level }?.let { it to tryExitNodes.top() }
     }

compiler/testData/codegen/boxInline/complexStack/spillConstructorArgumentsAndInlineLambdaParameter.kt

@@ -1,4 +1,3 @@
-// IGNORE_FIR_DIAGNOSTICS
 // FILE: 1.kt
 class A(val s: String)
 

compiler/testData/diagnostics/tests/controlFlowAnalysis/deadCode/deadCodeDifferentExamples.fir.kt

@@ -113,7 +113,7 @@ fun t7() : Int {
   catch (<!THROWABLE_TYPE_MISMATCH!>e : Any<!>) {
     2
   }
-  <!UNREACHABLE_CODE!>return 1<!> // this is OK, like in Java
+  return 1 // this is OK, like in Java
 }
 
 fun t8() : Int {
@@ -123,7 +123,7 @@ fun t8() : Int {
   }
   catch (<!THROWABLE_TYPE_MISMATCH!>e : Any<!>) {
     return 1
-    2
+    <!UNREACHABLE_CODE!>2<!>
   }
   <!UNREACHABLE_CODE!>return 1<!>
 }

compiler/testData/diagnostics/tests/controlFlowAnalysis/reassignmentInTryCatch.fir.kt

@@ -22,7 +22,7 @@ fun f2() {
     finally {
         <!VAL_REASSIGNMENT!>n<!> = 2
     }
-    n.hashCode()
+    <!UNINITIALIZED_VARIABLE!>n<!>.hashCode()
 }
 
 fun g1(flag: Boolean) {

compiler/testData/diagnostics/tests/initializedAfterRethrow.fir.kt

@@ -1,222 +0,0 @@
-fun foo(): Int = 42
-
-object ThrowInTryWithCatch {
-    private val p: String
-
-    init {
-        try {
-            throw Exception()
-        } catch (e: Exception) {
-        }
-        p = "OK"
-    }
-}
-
-object ThrowInTryWithCatchAndFinally {
-    private val p: String
-
-    init {
-        try {
-            throw Exception()
-        } catch (e: Exception) {
-        } finally {
-        }
-        p = "OK"
-    }
-}
-
-object ThrowInFinally {
-    <!MUST_BE_INITIALIZED_OR_BE_ABSTRACT!>private val p: String<!>
-
-    init {
-        try {
-            foo()
-        } catch (e: Exception) {
-        } finally {
-            throw Exception()
-        }
-        p = "OK"
-    }
-}
-
-object RethrowInCatch {
-    private val p: String
-
-    init {
-        try {
-            foo()
-        } catch (e: Exception) {
-            throw e
-        }
-        p = "OK"
-    }
-}
-
-object RethrowInCatchWithFinally {
-    private val p: String
-
-    init {
-        try {
-            foo()
-        } catch (e: Exception) {
-            throw e
-        } finally {
-        }
-        p = "OK"
-    }
-}
-
-object InnerTryWithCatch {
-    <!MUST_BE_INITIALIZED_OR_BE_ABSTRACT!>private val p: String<!>
-
-    init {
-        try {
-            foo()
-        } catch (e: Exception) {
-            try {
-                throw e
-            } catch (ee: Exception) {
-            }
-        }
-        p = "OK"
-    }
-}
-
-object InnerTryWithFinally {
-    private val p: String
-
-    init {
-        try {
-            foo()
-        } catch (e: Exception) {
-            try {
-                throw e
-            } finally {
-            }
-        }
-        p = "OK"
-    }
-}
-
-
-object InnerTryWithCatchAndFinally {
-    private val p: String
-
-    init {
-        try {
-            foo()
-        } catch (e: Exception) {
-            try {
-                throw e
-            } catch (ee: Exception) {
-            } finally {
-            }
-        }
-        p = "OK"
-    }
-}
-
-object InnerCatch {
-    private val p: String
-
-    init {
-        try {
-            foo()
-        } catch (e: Exception) {
-            try {
-                foo()
-            } catch (ee: Exception) {
-                throw ee
-            }
-        }
-        p = "OK"
-    }
-}
-
-object InnerCatchWithFinally {
-    private val p: String
-
-    init {
-        try {
-            foo()
-        } catch (e: Exception) {
-            try {
-                foo()
-            } catch (ee: Exception) {
-                throw ee
-            } finally {
-            }
-        }
-        p = "OK"
-    }
-}
-
-object InnerCatchOuterRethrow {
-    private val p: String
-
-    init {
-        try {
-            foo()
-        } catch (e: Exception) {
-            try {
-                foo()
-            } catch (ee: Exception) {
-                throw e
-            }
-        }
-        p = "OK"
-    }
-}
-
-object InnerCatchOuterRethrowWithFinally {
-    private val p: String
-
-    init {
-        try {
-            foo()
-        } catch (e: Exception) {
-            try {
-                foo()
-            } catch (ee: Exception) {
-                throw e
-            } finally {
-            }
-        }
-        p = "OK"
-    }
-}
-
-object InnerFinally {
-    private val p: String
-
-    init {
-        try {
-            foo()
-        } catch (e: Exception) {
-            try {
-                foo()
-            } finally {
-                throw e
-            }
-        }
-        p = "OK"
-    }
-}
-
-object InnerFinallyWithCatch {
-    private val p: String
-
-    init {
-        try {
-            foo()
-        } catch (e: Exception) {
-            try {
-                foo()
-            } catch (ee: Exception) {
-            } finally {
-                throw e
-            }
-        }
-        p = "OK"
-    }
-}

compiler/testData/diagnostics/tests/initializedAfterRethrow.kt

@@ -1,3 +1,5 @@
+// FIR_IDENTICAL
+
 fun foo(): Int = 42
 
 object ThrowInTryWithCatch {

compiler/testData/diagnostics/tests/regressions/kt58.fir.kt

@@ -43,7 +43,7 @@ fun t3() : Int {
     finally {
         doSmth(3)
     }
-}
+<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 
 fun t4() : Int {
     try {
@@ -52,7 +52,7 @@ fun t4() : Int {
     catch (e: UnsupportedOperationException) {
         doSmth(2)
     }
-}
+<!NO_RETURN_IN_FUNCTION_WITH_BLOCK_BODY!>}<!>
 
 fun t5() : Int {
     try {

compiler/testData/diagnostics/tests/smartCasts/throwInTry.fir.kt

@@ -61,7 +61,7 @@ fun conditionalThrowInTry_rethrow_smartcastAfterTryCatch(a: A) {
     } catch (e: Throwable) {
         throw e
     }
-    takeB(<!ARGUMENT_TYPE_MISMATCH!>a<!>)
+    takeB(a)
 }
 
 fun conditionalThrowInTry_rethrow_smartcastAfterTryCatchFinally(a: A) {