[FIR] Fix while's label, when condition has lambda

#KT-48116 Fixed

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

@@ -6019,6 +6019,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/controlStructures"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
             }
 
+            @Test
+            @TestMetadata("breakToLabel.kt")
+            public void testBreakToLabel() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlStructures/breakToLabel.kt");
+            }
+
             @Test
             @TestMetadata("catchGenerics.kt")
             public void testCatchGenerics() throws Exception {

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

@@ -6019,6 +6019,12 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/controlStructures"), Pattern.compile("^(.+)\\.kt$"), Pattern.compile("^(.+)\\.fir\\.kts?$"), true);
             }
 
+            @Test
+            @TestMetadata("breakToLabel.kt")
+            public void testBreakToLabel() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlStructures/breakToLabel.kt");
+            }
+
             @Test
             @TestMetadata("catchGenerics.kt")
             public void testCatchGenerics() throws Exception {

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

@@ -1022,6 +1022,7 @@ class ExpressionsConverter(
     private fun convertWhile(whileLoop: LighterASTNode): FirElement {
         var block: LighterASTNode? = null
         var firCondition: FirExpression? = null
+        val label = stashLabel() //get label of while, otherwise if condition has lambda, it will steal the label
         whileLoop.forEachChildren {
             when (it.tokenType) {
                 BODY -> block = it
@@ -1036,7 +1037,7 @@ class ExpressionsConverter(
                 firCondition ?: buildErrorExpression(null, ConeSimpleDiagnostic("No condition in while loop", DiagnosticKind.Syntax))
             // break/continue in the while loop condition will refer to an outer loop if any.
             // So, prepare the loop target after building the condition.
-            target = prepareTarget()
+            target = prepareTarget(label)
         }.configure(target) { convertLoopBody(block) }
     }
 

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

@@ -1878,10 +1878,11 @@ open class RawFirBuilder(
             val target: FirLoopTarget
             return FirWhileLoopBuilder().apply {
                 source = expression.toFirSourceElement()
+                val label = stashLabel() //get label of while, otherwise if condition has lambda, it will steal the label
                 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.
-                target = prepareTarget()
+                target = prepareTarget(label)
             }.configure(target) { expression.body.toFirBlock() }
         }
 

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

@@ -213,8 +213,12 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
         }
     }
 
-    fun FirLoopBuilder.prepareTarget(): FirLoopTarget {
-        label = context.firLabels.pop()
+    fun FirLoopBuilder.prepareTarget(): FirLoopTarget = prepareTarget(context.firLabels.pop())
+
+    fun stashLabel(): FirLabel? = context.firLabels.pop()
+
+    fun FirLoopBuilder.prepareTarget(label: FirLabel?): FirLoopTarget {
+        this.label = label
         val target = FirLoopTarget(label?.name)
         context.firLoopTargets += target
         return target