[FIR] Fix break/continue in try-finally in loop, ^KT-51759

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/DiagnosisCompilerTestFE10TestdataTestGenerated.java

@@ -5434,6 +5434,12 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                 runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/breakContinueInTryFinally.kt");
             }
 
+            @Test
+            @TestMetadata("breakContinueInTryFinallyInLoop.kt")
+            public void testBreakContinueInTryFinallyInLoop() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/breakContinueInTryFinallyInLoop.kt");
+            }
+
             @Test
             @TestMetadata("breakInsideLocal.kt")
             public void testBreakInsideLocal() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsTestGenerated.java

@@ -5434,6 +5434,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                 runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/breakContinueInTryFinally.kt");
             }
 
+            @Test
+            @TestMetadata("breakContinueInTryFinallyInLoop.kt")
+            public void testBreakContinueInTryFinallyInLoop() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/breakContinueInTryFinallyInLoop.kt");
+            }
+
             @Test
             @TestMetadata("breakInsideLocal.kt")
             public void testBreakInsideLocal() throws Exception {

compiler/fir/analysis-tests/tests-gen/org/jetbrains/kotlin/test/runners/FirOldFrontendDiagnosticsWithLightTreeTestGenerated.java

@@ -5434,6 +5434,12 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                 runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/breakContinueInTryFinally.kt");
             }
 
+            @Test
+            @TestMetadata("breakContinueInTryFinallyInLoop.kt")
+            public void testBreakContinueInTryFinallyInLoop() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/breakContinueInTryFinallyInLoop.kt");
+            }
+
             @Test
             @TestMetadata("breakInsideLocal.kt")
             public void testBreakInsideLocal() throws Exception {

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

@@ -16,9 +16,6 @@ abstract class PathAwareControlFlowInfo<P : PathAwareControlFlowInfo<P, S>, S :
     internal val infoAtNormalPath: S
         get() = map.getValue(NormalPath)
 
-    private val hasNormalPath: Boolean
-        get() = map.containsKey(NormalPath)
-
     fun applyLabel(node: CFGNode<*>, label: EdgeLabel): P {
         if (label.isNormal) {
             // Special case: when we exit the try expression, null label means a normal path.
@@ -26,7 +23,8 @@ abstract class PathAwareControlFlowInfo<P : PathAwareControlFlowInfo<P, S>, S :
             // One day, if we allow multiple edges between nodes with different labels, e.g., labeling all paths in try/catch/finally,
             // instead of this kind of special handling, proxy enter/exit nodes per label are preferred.
             if (node is TryExpressionExitNode) {
-                return if (hasNormalPath) {
+                val infoAtNormalPath = map[NormalPath]
+                return if (infoAtNormalPath != null) {
                     constructor(persistentMapOf(NormalPath to infoAtNormalPath))
                 } else {
                     /* This means no info for normal path. */

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

@@ -640,10 +640,10 @@ class ControlFlowGraphBuilder {
             else -> throw IllegalArgumentException("Unknown jump type: ${jump.render()}")
         }
 
-        val labelForFinallyBLock = if (jump is FirReturnExpression) {
-            ReturnPath(jump.target.labeledElement.symbol)
-        } else {
-            NormalPath
+        val labelForFinallyBLock = when (jump) {
+            is FirReturnExpression -> ReturnPath(jump.target.labeledElement.symbol)
+            is FirLoopJump -> LoopPath(jump)
+            else -> NormalPath
         }
 
         addNodeWithJump(

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

@@ -6,6 +6,8 @@
 package org.jetbrains.kotlin.fir.resolve.dfa.cfg
 
 import org.jetbrains.kotlin.fir.declarations.FirDeclaration
+import org.jetbrains.kotlin.fir.expressions.FirBreakExpression
+import org.jetbrains.kotlin.fir.expressions.FirLoopJump
 import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 
 class ControlFlowGraph(val declaration: FirDeclaration?, val name: String, val kind: Kind) {
@@ -144,6 +146,11 @@ object LoopBackPath : EdgeLabel(label = null) {
 
 object UncaughtExceptionPath : EdgeLabel(label = "onUncaughtException")
 
+class LoopPath(
+    firLoopJump: FirLoopJump
+) : EdgeLabel("\"" + (if (firLoopJump is FirBreakExpression) "break" else "continue") +
+                      (firLoopJump.target.labeledElement.label?.let { "@${it.name}" } ?: "") + "\"")
+
 // TODO: Label `return`ing edge with this.
 class ReturnPath(
     returnTargetSymbol: FirFunctionSymbol<*>

compiler/testData/diagnostics/tests/controlFlowAnalysis/breakContinueInTryFinallyInLoop.kt

@@ -0,0 +1,24 @@
+// FIR_IDENTICAL
+// ISSUE: KT-51759
+
+fun testBreak(b: Boolean, s: String?) {
+    while (b) {
+        val x: String?
+        try {
+            x = s ?: break
+        } finally {
+        }
+        x!!.length
+    }
+}
+
+fun testContinue(b: Boolean, s: String?) {
+    while (b) {
+        val x: String?
+        try {
+            x = s ?: continue
+        } finally {
+        }
+        x!!.length
+    }
+}

compiler/testData/diagnostics/tests/controlFlowAnalysis/breakContinueInTryFinallyInLoop.txt

@@ -0,0 +1,4 @@
+package
+
+public fun testBreak(/*0*/ b: kotlin.Boolean, /*1*/ s: kotlin.String?): kotlin.Unit
+public fun testContinue(/*0*/ b: kotlin.Boolean, /*1*/ s: kotlin.String?): kotlin.Unit

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/DiagnosticTestGenerated.java

@@ -5440,6 +5440,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/breakContinueInTryFinally.kt");
             }
 
+            @Test
+            @TestMetadata("breakContinueInTryFinallyInLoop.kt")
+            public void testBreakContinueInTryFinallyInLoop() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlFlowAnalysis/breakContinueInTryFinallyInLoop.kt");
+            }
+
             @Test
             @TestMetadata("breakInsideLocal.kt")
             public void testBreakInsideLocal() throws Exception {