Raw FIR: correct loop target for break/continue in do-while loop condition

As shown in KT-44412 (or KT-45319 or KT-17728):
```
fun test5() {
    var i = 0
    Outer@while (true) {
        ++i
        var j = 0
        Inner@do {
            ++j
        } while (if (j >= 3) false else break) // break@Inner
        if (i == 3) break
    }
}
```

To properly set the loop target for `break` in do-while loop condition,
the loop target for that do-while loop should be ready before parsing
the loop condition.

Previously, Raw FIR loop building configures loop target after visiting
loop conditions. This commit splits the configuration and lets the
builders prepare the loop target for do-while loop only.

compiler/fir/raw-fir/light-tree2fir/src/org/jetbrains/kotlin/fir/lightTree/converter/ExpressionsConverter.kt

@@ -888,16 +888,19 @@ class ExpressionsConverter(
      */
     private fun convertDoWhile(doWhileLoop: LighterASTNode): FirElement {
         var block: LighterASTNode? = null
-        var firCondition: FirExpression = buildErrorExpression(null, ConeSimpleDiagnostic("No condition in do-while loop", DiagnosticKind.Syntax))
-        doWhileLoop.forEachChildren {
-            when (it.tokenType) {
-                BODY -> block = it
-                CONDITION -> firCondition = getAsFirExpression(it, "No condition in do-while loop")
-            }
-        }
+        var firCondition: FirExpression =
+            buildErrorExpression(null, ConeSimpleDiagnostic("No condition in do-while loop", DiagnosticKind.Syntax))
 
         return FirDoWhileLoopBuilder().apply {
             source = doWhileLoop.toFirSourceElement()
+            // For break/continue in the do-while loop condition, prepare the loop target first so that it can refer to the same loop.
+            prepareTarget()
+            doWhileLoop.forEachChildren {
+                when (it.tokenType) {
+                    BODY -> block = it
+                    CONDITION -> firCondition = getAsFirExpression(it, "No condition in do-while loop")
+                }
+            }
             condition = firCondition
         }.configure { convertLoopBody(block) }
     }
@@ -919,6 +922,9 @@ class ExpressionsConverter(
         return FirWhileLoopBuilder().apply {
             source = whileLoop.toFirSourceElement()
             condition = firCondition
+            // break/continue in the while loop condition will refer to an outer loop if any.
+            // So, prepare the loop target after building the condition.
+            prepareTarget()
         }.configure { convertLoopBody(block) }
     }
 
@@ -954,6 +960,9 @@ class ExpressionsConverter(
                     calleeReference = buildSimpleNamedReference { name = Name.identifier("hasNext") }
                     explicitReceiver = generateResolvedAccessExpression(null, iteratorVal)
                 }
+                // break/continue in the for loop condition will refer to an outer loop if any.
+                // So, prepare the loop target after building the condition.
+                prepareTarget()
             }.configure {
                 // NB: just body.toFirBlock() isn't acceptable here because we need to add some statements
                 buildBlock block@{

compiler/fir/raw-fir/psi2fir/src/org/jetbrains/kotlin/fir/builder/RawFirBuilder.kt

@@ -1652,6 +1652,8 @@ open class RawFirBuilder(
         override fun visitDoWhileExpression(expression: KtDoWhileExpression, data: Unit): FirElement {
             return FirDoWhileLoopBuilder().apply {
                 source = expression.toFirSourceElement()
+                // For break/continue in the do-while loop condition, prepare the loop target first so that it can refer to the same loop.
+                prepareTarget()
                 condition = expression.condition.toFirExpression("No condition in do-while loop")
             }.configure { expression.body.toFirBlock() }
         }
@@ -1660,6 +1662,9 @@ open class RawFirBuilder(
             return FirWhileLoopBuilder().apply {
                 source = expression.toFirSourceElement()
                 condition = expression.condition.toFirExpression("No condition in while loop")
+                // break/continue in the while loop condition will refer to an outer loop if any.
+                // So, prepare the loop target after building the condition.
+                prepareTarget()
             }.configure { expression.body.toFirBlock() }
         }
 
@@ -1692,6 +1697,9 @@ open class RawFirBuilder(
                         }
                         explicitReceiver = generateResolvedAccessExpression(fakeSource, iteratorVal)
                     }
+                    // break/continue in the for loop condition will refer to an outer loop if any.
+                    // So, prepare the loop target after building the condition.
+                    prepareTarget()
                 }.configure {
                     // NB: just body.toFirBlock() isn't acceptable here because we need to add some statements
                     val blockBuilder = when (val body = expression.body) {

compiler/fir/raw-fir/raw-fir.common/src/org/jetbrains/kotlin/fir/builder/BaseFirBuilder.kt

@@ -135,8 +135,8 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
 
 
     /**** Function utils ****/
-    fun <T> MutableList<T>.removeLast() {
-        removeAt(size - 1)
+    fun <T> MutableList<T>.removeLast(): T {
+        return removeAt(size - 1)
     }
 
     fun <T> MutableList<T>.pop(): T? {
@@ -216,13 +216,16 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
         }
     }
 
-    fun FirLoopBuilder.configure(generateBlock: () -> FirBlock): FirLoop {
+    fun FirLoopBuilder.prepareTarget() {
         label = context.firLabels.pop()
         val target = FirLoopTarget(label?.name)
         context.firLoopTargets += target
+    }
+
+    fun FirLoopBuilder.configure(generateBlock: () -> FirBlock): FirLoop {
         block = generateBlock()
         val loop = build()
-        context.firLoopTargets.removeLast()
+        val target = context.firLoopTargets.removeLast()
         target.bind(loop)
         return loop
     }

compiler/testData/codegen/box/controlStructures/breakContinueInExpressions/breakInLoopConditions.kt

@@ -1,6 +1,5 @@
 // See: https://youtrack.jetbrains.com/issue/KT-45319
 // IGNORE_BACKEND: JVM
-// IGNORE_BACKEND_FIR: JVM_IR
 // IGNORE_BACKEND: JS
 
 fun breakInDoWhileCondition(): String {

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

@@ -0,0 +1,92 @@
+// !LANGUAGE: +AllowBreakAndContinueInsideWhen
+
+fun breakContinueInWhen(i: Int) {
+    for (y in 0..10) {
+        when(i) {
+            0 -> continue
+            1 -> break
+            2 -> {
+                for(z in 0..10) {
+                    break
+                }
+                for(w in 0..10) {
+                    continue
+                }
+            }
+        }
+    }
+}
+
+
+fun breakContinueInWhenWithWhile(i: Int, j: Int) {
+    while (i > 0) {
+        when (i) {
+            0 -> continue
+            1 -> break
+            2 -> {
+                while (j > 0) {
+                    break
+                }
+            }
+        }
+    }
+}
+
+fun breakContinueInWhenWithDoWhile(i: Int, j: Int) {
+    do {
+        when (i) {
+            0 -> continue
+            1 -> break
+            2 -> {
+                do {
+                    if (j == 5) break
+                    if (j == 10) continue
+                } while (j > 0)
+            }
+        }
+    } while (i > 0)
+}
+
+fun labeledBreakContinue(i: Int) {
+    outer@ for (y in 0..10) {
+        when (i) {
+            0 -> continue@outer
+            1 -> break@outer
+        }
+    }
+}
+
+fun testBreakContinueInWhenInWhileCondition() {
+    var i = 0
+    while (
+        when (i) {
+            1 -> <!BREAK_OR_CONTINUE_OUTSIDE_A_LOOP!>break<!>
+            2 -> <!BREAK_OR_CONTINUE_OUTSIDE_A_LOOP!>continue<!>
+            else -> true
+        }
+    ) {
+        ++i
+    }
+}
+
+fun testBreakContinueInWhenInDoWhileCondition() {
+    var i = 0
+    do {
+        ++i
+    } while (
+        when (i) {
+            1 -> break
+            2 -> continue
+            else -> true
+        }
+    )
+}
+
+fun testBreakContinueInWhenInForIteratorExpression(xs: List<Any>, i: Int) {
+    for (x in when (i) {
+        1 -> <!BREAK_OR_CONTINUE_OUTSIDE_A_LOOP!>break<!>
+        2 -> <!BREAK_OR_CONTINUE_OUTSIDE_A_LOOP!>continue<!>
+        else -> xs
+    }) {
+    }
+}

compiler/testData/diagnostics/tests/BreakContinueInWhen_after.kt

@@ -1,4 +1,3 @@
-// FIR_IDENTICAL
 // !LANGUAGE: +AllowBreakAndContinueInsideWhen
 
 fun breakContinueInWhen(i: Int) {
@@ -90,4 +89,4 @@ fun testBreakContinueInWhenInForIteratorExpression(xs: List<Any>, i: Int) {
         else -> xs
     }) {
     }
-}
+}

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

@@ -6,8 +6,8 @@ fun test() {
     while (<!BREAK_OR_CONTINUE_OUTSIDE_A_LOOP!>break<!>) {}
     l@ while (<!BREAK_OR_CONTINUE_OUTSIDE_A_LOOP!>break@l<!>) {}
 
-    do {} while (<!BREAK_OR_CONTINUE_OUTSIDE_A_LOOP!>continue<!>)
-    l@ do {} while (<!BREAK_OR_CONTINUE_OUTSIDE_A_LOOP!>continue@l<!>)
+    do {} while (continue)
+    l@ do {} while (continue@l)
 
     //KT-5704
     var i = 0
@@ -31,4 +31,4 @@ fun test2(b: Boolean) {
         for (j in if (true) 1..10 else continue) {
         }
     }
-}
+}

compiler/testData/ir/irText/expressions/breakContinueInLoopHeader.fir.kt.txt

@@ -70,7 +70,7 @@ fun test5() {
       j
       // }      while (when {
         greaterOrEqual(arg0 = j, arg1 = 3) -> false
-        else -> break@Outer
+        else -> break@Inner
       })
     }
     when {

compiler/testData/ir/irText/expressions/breakContinueInLoopHeader.fir.txt

@@ -124,7 +124,7 @@ FILE fqName:<root> fileName:/breakContinueInLoopHeader.kt
                   then: CONST Boolean type=kotlin.Boolean value=false
                 BRANCH
                   if: CONST Boolean type=kotlin.Boolean value=true
-                  then: BREAK label=Outer loop.label=Outer
+                  then: BREAK label=Inner loop.label=Inner
           WHEN type=kotlin.Unit origin=IF
             BRANCH
               if: CALL 'public final fun EQEQ (arg0: kotlin.Any?, arg1: kotlin.Any?): kotlin.Boolean declared in kotlin.internal.ir' type=kotlin.Boolean origin=EQEQ