[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 {
+    fun FirLoopBuilder.prepareTarget(): FirLoopTarget = prepareTarget(context.firLabels.pop())
-        label = 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

compiler/testData/diagnostics/tests/controlStructures/breakToLabel.fir.kt

@@ -0,0 +1,20 @@
+//KT-48116
+//WITH_STDLIB
+
+fun foo() {
+    var tokenType: String? = null
+    while (true) {
+        FindTagEnd@ while (tokenType.let { it != null && it !== "XML_END_TAG_START" }) {
+            if (tokenType === "XML_COMMENT_CHARACTERS") {
+                // we should terminate on first occurence of </
+                val end = tokenType
+                for (i in tokenType) {
+                    if (i == ' ') {
+                        break@FindTagEnd
+                    }
+                }
+            }
+            tokenType = "abc"
+        }
+    }
+}

compiler/testData/diagnostics/tests/controlStructures/breakToLabel.kt

@@ -0,0 +1,20 @@
+//KT-48116
+//WITH_STDLIB
+
+fun foo() {
+    var tokenType: String? = null
+    while (true) {
+        FindTagEnd@ while (tokenType.let { it != null && it !== "XML_END_TAG_START" }) {
+            if (tokenType === "XML_COMMENT_CHARACTERS") {
+                // we should terminate on first occurence of </
+                val end = tokenType
+                for (i in <!DEBUG_INFO_SMARTCAST!>tokenType<!>) {
+                    if (i == ' ') {
+                        break@FindTagEnd
+                    }
+                }
+            }
+            tokenType = "abc"
+        }
+    }
+}

compiler/testData/diagnostics/tests/controlStructures/breakToLabel.txt

@@ -0,0 +1,3 @@
+package
+
+public fun foo(): kotlin.Unit

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

@@ -6025,6 +6025,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 KtTestUtil.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("compiler/testData/diagnostics/tests/controlStructures"), Pattern.compile("^(.*)\\.kts?$"), 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 {

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

@@ -6019,6 +6019,12 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                 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 {