FIR: fix label stealling and crash with multiple labels

Consider the code below

```
fun test() {
  a@ b@ {
    {}
  }
}
```

Currently when the code is converted to FIR, label `b` is bound to the
outer lambda and `a` gets bound to the inner lambda because it's not
consumed. This is wrong and also leads transfromation to fail with
exceptions because of the unexpected consumption of `a`.

This change fixes the above issue by designating a specific node in the
AST as the allowed user of a label when the label is added.

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

@@ -6271,6 +6271,12 @@ public class DiagnosisCompilerTestFE10TestdataTestGenerated extends AbstractDiag
                 runTest("compiler/testData/diagnostics/tests/controlStructures/localReturnInsidePropertyAccessor.kt");
             }
 
+            @Test
+            @TestMetadata("nestedLoopsWithMultipleLabels.kt")
+            public void testNestedLoopsWithMultipleLabels() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlStructures/nestedLoopsWithMultipleLabels.kt");
+            }
+
             @Test
             @TestMetadata("notAFunctionLabel_after.kt")
             public void testNotAFunctionLabel_after() throws Exception {

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

@@ -6271,6 +6271,12 @@ public class FirOldFrontendDiagnosticsTestGenerated extends AbstractFirDiagnosti
                 runTest("compiler/testData/diagnostics/tests/controlStructures/localReturnInsidePropertyAccessor.kt");
             }
 
+            @Test
+            @TestMetadata("nestedLoopsWithMultipleLabels.kt")
+            public void testNestedLoopsWithMultipleLabels() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlStructures/nestedLoopsWithMultipleLabels.kt");
+            }
+
             @Test
             @TestMetadata("notAFunctionLabel_after.kt")
             public void testNotAFunctionLabel_after() throws Exception {

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

@@ -6271,6 +6271,12 @@ public class FirOldFrontendDiagnosticsWithLightTreeTestGenerated extends Abstrac
                 runTest("compiler/testData/diagnostics/tests/controlStructures/localReturnInsidePropertyAccessor.kt");
             }
 
+            @Test
+            @TestMetadata("nestedLoopsWithMultipleLabels.kt")
+            public void testNestedLoopsWithMultipleLabels() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlStructures/nestedLoopsWithMultipleLabels.kt");
+            }
+
             @Test
             @TestMetadata("notAFunctionLabel_after.kt")
             public void testNotAFunctionLabel_after() throws Exception {

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

@@ -67,6 +67,8 @@ class ExpressionsConverter(
         return when (expression.tokenType) {
             LAMBDA_EXPRESSION -> {
                 val lambdaTree = LightTree2Fir.buildLightTreeLambdaExpression(expression.asText)
+                // Pass on label user to the lambda root
+                context.forwardLabelUsagePermission(expression, lambdaTree.root)
                 declarationsConverter.withOffset(offset + expression.startOffset) {
                     ExpressionsConverter(baseSession, lambdaTree, declarationsConverter, context)
                         .convertLambdaExpression(lambdaTree.root)
@@ -100,7 +102,11 @@ class ExpressionsConverter(
             BREAK, CONTINUE -> convertLoopJump(expression)
             RETURN -> convertReturn(expression)
             THROW -> convertThrow(expression)
-            PARENTHESIZED -> getAsFirExpression(expression.getExpressionInParentheses(), "Empty parentheses")
+            PARENTHESIZED -> {
+                val content = expression.getExpressionInParentheses()
+                context.forwardLabelUsagePermission(expression, content)
+                getAsFirExpression(content, "Empty parentheses")
+            }
             PROPERTY_DELEGATE, INDICES, CONDITION, LOOP_RANGE ->
                 getAsFirExpression(expression.getExpressionInParentheses(), errorReason)
             THIS_EXPRESSION -> convertThisExpression(expression)
@@ -136,7 +142,7 @@ class ExpressionsConverter(
             receiverTypeRef = implicitType
             symbol = FirAnonymousFunctionSymbol()
             isLambda = true
-            label = context.firLabels.pop() ?: context.calleeNamesForLambda.lastOrNull()?.let {
+            label = context.getLastLabel(lambdaExpression) ?: context.calleeNamesForLambda.lastOrNull()?.let {
                 buildLabel {
                     source = expressionSource.fakeElement(FirFakeSourceElementKind.GeneratedLambdaLabel)
                     name = it.asString()
@@ -326,15 +332,15 @@ class ExpressionsConverter(
      */
     private fun convertLabeledExpression(labeledExpression: LighterASTNode): FirElement {
         var firExpression: FirElement? = null
-        val previousLabelsSize = context.firLabels.size
         var errorLabelSource: FirSourceElement? = null
 
         labeledExpression.forEachChildren {
+            context.setNewLabelUserNode(it)
             when (it.tokenType) {
                 LABEL_QUALIFIER -> {
                     val rawName = it.toString()
                     val pair = buildLabelAndErrorSource(rawName.substring(0, rawName.length - 1), it.toFirSourceElement())
-                    context.firLabels += pair.first
+                    context.addNewLabel(pair.first)
                     errorLabelSource = pair.second
                 }
                 BLOCK -> firExpression = declarationsConverter.convertBlock(it)
@@ -343,9 +349,7 @@ class ExpressionsConverter(
             }
         }
 
-        if (context.firLabels.size != previousLabelsSize) {
-            context.firLabels.removeLast()
-        }
+        context.dropLastLabel()
 
         return buildExpressionWithErrorLabel(firExpression, errorLabelSource, labeledExpression.toFirSourceElement())
     }
@@ -426,7 +430,10 @@ class ExpressionsConverter(
             when (it.tokenType) {
                 ANNOTATION -> firAnnotationList += declarationsConverter.convertAnnotation(it)
                 ANNOTATION_ENTRY -> firAnnotationList += declarationsConverter.convertAnnotationEntry(it)
-                else -> if (it.isExpression()) firExpression = getAsFirExpression(it)
+                else -> if (it.isExpression()) {
+                    context.forwardLabelUsagePermission(annotatedExpression, it)
+                    firExpression = getAsFirExpression(it)
+                }
             }
         }
 
@@ -998,7 +1005,7 @@ class ExpressionsConverter(
         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.
-            target = prepareTarget()
+            target = prepareTarget(doWhileLoop)
             doWhileLoop.forEachChildren {
                 when (it.tokenType) {
                     BODY -> block = it
@@ -1017,7 +1024,6 @@ 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
@@ -1032,7 +1038,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(label)
+            target = prepareTarget(whileLoop)
         }.configure(target) { convertLoopBody(block) }
     }
 
@@ -1085,7 +1091,7 @@ class ExpressionsConverter(
                 }
                 // break/continue in the for loop condition will refer to an outer loop if any.
                 // So, prepare the loop target after building the condition.
-                target = prepareTarget()
+                target = prepareTarget(forLoop)
             }.configure(target) {
                 // 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

@@ -1323,7 +1323,7 @@ open class RawFirBuilder(
                     }
                 }
                 val expressionSource = expression.toFirSourceElement()
-                label = context.firLabels.pop() ?: context.calleeNamesForLambda.lastOrNull()?.let {
+                label = context.getLastLabel(expression) ?: context.calleeNamesForLambda.lastOrNull()?.let {
                     buildLabel {
                         source = expressionSource.fakeElement(FirFakeSourceElementKind.GeneratedLambdaLabel)
                         name = it.asString()
@@ -1945,7 +1945,7 @@ open class RawFirBuilder(
             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.
-                target = prepareTarget()
+                target = prepareTarget(expression)
                 condition = expression.condition.toFirExpression("No condition in do-while loop")
             }.configure(target) { expression.body.toFirBlock() }
         }
@@ -1954,11 +1954,10 @@ 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(label)
+                target = prepareTarget(expression)
             }.configure(target) { expression.body.toFirBlock() }
         }
 
@@ -1995,7 +1994,7 @@ open class RawFirBuilder(
                     }
                     // break/continue in the for loop condition will refer to an outer loop if any.
                     // So, prepare the loop target after building the condition.
-                    target = prepareTarget()
+                    target = prepareTarget(expression)
                 }.configure(target) {
                     // NB: just body.toFirBlock() isn't acceptable here because we need to add some statements
                     val blockBuilder = when (val body = expression.body) {
@@ -2323,33 +2322,33 @@ open class RawFirBuilder(
         }
 
         override fun visitParenthesizedExpression(expression: KtParenthesizedExpression, data: Unit): FirElement {
+            context.forwardLabelUsagePermission(expression, expression.expression)
             return expression.expression?.accept(this, data)
                 ?: buildErrorExpression(expression.toFirSourceElement(), ConeSimpleDiagnostic("Empty parentheses", DiagnosticKind.Syntax))
         }
 
         override fun visitLabeledExpression(expression: KtLabeledExpression, data: Unit): FirElement {
             val label = expression.getTargetLabel()
-            val previousLabelsSize = context.firLabels.size
             var errorLabelSource: FirSourceElement? = null
 
-            if (label != null) {
+            val result = if (label != null) {
                 val rawName = label.getReferencedNameElement().node!!.text
                 val labelAndErrorSource = buildLabelAndErrorSource(rawName, label.toFirPsiSourceElement())
-                context.firLabels += labelAndErrorSource.first
                 errorLabelSource = labelAndErrorSource.second
-            }
-
-            val result = expression.baseExpression?.accept(this, data)
-
-            if (context.firLabels.size != previousLabelsSize) {
-                context.firLabels.removeLast()
+                context.withNewLabel(labelAndErrorSource.first, expression.baseExpression) {
+                    expression.baseExpression?.accept(this, data)
+                }
+            } else {
+                expression.baseExpression?.accept(this, data)
             }
 
             return buildExpressionWithErrorLabel(result, errorLabelSource, expression.toFirSourceElement())
         }
 
         override fun visitAnnotatedExpression(expression: KtAnnotatedExpression, data: Unit): FirElement {
-            val rawResult = expression.baseExpression?.accept(this, data)
+            val baseExpression = expression.baseExpression
+            context.forwardLabelUsagePermission(expression, baseExpression)
+            val rawResult = baseExpression?.accept(this, data)
             val result = rawResult as? FirAnnotationContainer
                 ?: buildErrorExpression(
                     expression.toFirSourceElement(),

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

@@ -183,7 +183,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
                     }
                 }
                 target = FirFunctionTarget(labelName, false).apply {
-                    if (context.firLabels.any { it.name == labelName } || context.firLoopTargets.any { it.labelName == labelName }) {
+                    if (context.firLabels.any { it.name == labelName }) {
                         bindToErrorFunction("Label $labelName does not target a function", DiagnosticKind.NotAFunctionLabel)
                     } else {
                         bindToErrorFunction("Cannot bind label $labelName to a function", DiagnosticKind.UnresolvedLabel)
@@ -217,9 +217,7 @@ abstract class BaseFirBuilder<T>(val baseSession: FirSession, val context: Conte
         }
     }
 
-    fun FirLoopBuilder.prepareTarget(): FirLoopTarget = prepareTarget(context.firLabels.pop())
-
-    fun stashLabel(): FirLabel? = context.firLabels.pop()
+    fun FirLoopBuilder.prepareTarget(firLabelUser: Any): FirLoopTarget = prepareTarget(context.getLastLabel(firLabelUser))
 
     fun FirLoopBuilder.prepareTarget(label: FirLabel?): FirLoopTarget {
         this.label = label

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

@@ -5,10 +5,7 @@
 
 package org.jetbrains.kotlin.fir.builder
 
-import org.jetbrains.kotlin.fir.FirFunctionTarget
-import org.jetbrains.kotlin.fir.FirLabel
-import org.jetbrains.kotlin.fir.FirLoopTarget
-import org.jetbrains.kotlin.fir.FirSourceElementKind
+import org.jetbrains.kotlin.fir.*
 import org.jetbrains.kotlin.fir.declarations.FirTypeParameterRef
 import org.jetbrains.kotlin.fir.expressions.FirExpression
 import org.jetbrains.kotlin.fir.symbols.impl.FirTypeParameterSymbol
@@ -26,7 +23,19 @@ class Context<T> {
 
     val firFunctionTargets = mutableListOf<FirFunctionTarget>()
     val calleeNamesForLambda = mutableListOf<Name>()
-    val firLabels = mutableListOf<FirLabel>()
+
+    @PrivateForInline
+    val _firLabels = mutableListOf<FirLabel>()
+
+    @OptIn(PrivateForInline::class)
+    val firLabels: List<FirLabel>
+        get() = _firLabels
+
+    /**
+     * A designated `KtElement` or `LighterASTNode` object that is allowed to claim the last label in [firLabels].
+     */
+    @PrivateForInline
+    var firLabelUserNode: Any? = null
     val firLoopTargets = mutableListOf<FirLoopTarget>()
     val capturedTypeParameters = mutableListOf<StatusFirTypeParameterSymbolList>()
     val arraySetArgument = mutableMapOf<T, FirExpression>()
@@ -60,5 +69,50 @@ class Context<T> {
         }
     }
 
+    /**
+     * Gets the last label that was added or null if the current node does not have permission to use the label.
+     */
+    @OptIn(PrivateForInline::class)
+    fun getLastLabel(currentNode: Any): FirLabel? {
+        if (this.firLabelUserNode == currentNode) return firLabels.last()
+        return null
+    }
+
+    @OptIn(PrivateForInline::class)
+    fun addNewLabel(label: FirLabel) {
+        _firLabels += label
+    }
+
+    @OptIn(PrivateForInline::class)
+    fun setNewLabelUserNode(useNode: Any?) {
+        this.firLabelUserNode = useNode
+    }
+
+    @OptIn(PrivateForInline::class)
+    fun dropLastLabel() {
+        _firLabels.removeLast()
+        firLabelUserNode = null
+    }
+
+    inline fun <T> withNewLabel(label: FirLabel, userNode: Any?, block: () -> T): T {
+        addNewLabel(label)
+        setNewLabelUserNode(userNode)
+        try {
+            return block()
+        } finally {
+            dropLastLabel()
+        }
+    }
+
+    /**
+     * Forwards the permission to use the last label to a different node if the current user node has the permission.
+     */
+    @OptIn(PrivateForInline::class)
+    fun forwardLabelUsagePermission(currentUserNode: Any, newUserNode: Any?) {
+        if (currentUserNode == firLabelUserNode) {
+            firLabelUserNode = newUserNode
+        }
+    }
+
     data class StatusFirTypeParameterSymbolList(val notNested: Boolean, val list: List<FirTypeParameterSymbol> = listOf())
 }

compiler/testData/codegen/boxInline/nonLocalReturns/returnFromFunctionExpr.kt

@@ -1,5 +1,4 @@
 // NO_CHECK_LAMBDA_INLINING
-// IGNORE_BACKEND_FIR: JVM_IR
 // FILE: 1.kt
 
 inline fun foo(f: () -> Unit) {

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

@@ -0,0 +1,9 @@
+// FIR_IDENTICAL
+fun test() {
+    a@ b@ while(true) {
+        val f = {
+            <!NOT_A_FUNCTION_LABEL!>return@a<!>
+        }
+        break@b
+    }
+}

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

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

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

@@ -51,12 +51,21 @@ fun testLoopLabelInReturn(xs: List<Int>) {
 }
 
 fun testValLabelInReturn() {
-    L@ val fn = { return@L }
+    L@ val fn = { <!NOT_A_FUNCTION_LABEL!>return@L<!> }
     fn()
 }
 
 fun testHighOrderFunctionCallLabelInReturn() {
     L@ run {
-        return@L
+        <!NOT_A_FUNCTION_LABEL!>return@L<!>
+    }
+}
+
+fun testMultipleLabelsWithNestedLambda() {
+    l1@ l2@{
+        {
+            <!NOT_A_FUNCTION_LABEL!>return@l1<!>
+        }
+        return@l2
     }
 }

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

@@ -59,4 +59,13 @@ fun testHighOrderFunctionCallLabelInReturn() {
     <!REDUNDANT_LABEL_WARNING!>L@<!> run {
         <!NOT_A_FUNCTION_LABEL!>return@L<!>
     }
-}
+}
+
+fun testMultipleLabelsWithNestedLambda() {
+    l1@ l2@{
+        {
+            <!RETURN_NOT_ALLOWED!>return@l1<!>
+        }
+        return@l2
+    }
+}

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

@@ -9,6 +9,7 @@ public fun testLambdaLabel(): () -> kotlin.Unit
 public fun testLambdaMultipleLabels1(): () -> kotlin.Unit
 public fun testLambdaMultipleLabels2(): () -> kotlin.Unit
 public fun testLoopLabelInReturn(/*0*/ xs: kotlin.collections.List<kotlin.Int>): kotlin.Unit
+public fun testMultipleLabelsWithNestedLambda(): kotlin.Unit
 public fun testParenthesizedLambdaLabel(): () -> kotlin.Unit
 public fun testValLabelInReturn(): kotlin.Unit
 
@@ -18,3 +19,4 @@ public fun testValLabelInReturn(): kotlin.Unit
     public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
     public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
 }
+

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

@@ -51,12 +51,21 @@ fun testLoopLabelInReturn(xs: List<Int>) {
 }
 
 fun testValLabelInReturn() {
-    L@ val fn = { return@L }
+    L@ val fn = { <!NOT_A_FUNCTION_LABEL!>return@L<!> }
     fn()
 }
 
 fun testHighOrderFunctionCallLabelInReturn() {
     L@ run {
-        return@L
+        <!NOT_A_FUNCTION_LABEL!>return@L<!>
+    }
+}
+
+fun testMultipleLabelsWithNestedLambda() {
+    l1@ l2@{
+        {
+            <!NOT_A_FUNCTION_LABEL!>return@l1<!>
+        }
+        return@l2
     }
 }

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

@@ -59,4 +59,13 @@ fun testHighOrderFunctionCallLabelInReturn() {
     <!REDUNDANT_LABEL_WARNING!>L@<!> run {
         <!NOT_A_FUNCTION_LABEL_WARNING!>return@L<!>
     }
-}
+}
+
+fun testMultipleLabelsWithNestedLambda() {
+    l1@ l2@{
+        {
+            <!RETURN_NOT_ALLOWED!>return@l1<!>
+        }
+        return@l2
+    }
+}

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

@@ -9,6 +9,7 @@ public fun testLambdaLabel(): () -> kotlin.Unit
 public fun testLambdaMultipleLabels1(): () -> kotlin.Unit
 public fun testLambdaMultipleLabels2(): () -> kotlin.Unit
 public fun testLoopLabelInReturn(/*0*/ xs: kotlin.collections.List<kotlin.Int>): kotlin.Unit
+public fun testMultipleLabelsWithNestedLambda(): kotlin.Unit
 public fun testParenthesizedLambdaLabel(): () -> kotlin.Unit
 public fun testValLabelInReturn(): kotlin.Unit
 
@@ -18,3 +19,4 @@ public fun testValLabelInReturn(): kotlin.Unit
     public open override /*1*/ /*fake_override*/ fun hashCode(): kotlin.Int
     public open override /*1*/ /*fake_override*/ fun toString(): kotlin.String
 }
+

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

@@ -6277,6 +6277,12 @@ public class DiagnosticTestGenerated extends AbstractDiagnosticTest {
                 runTest("compiler/testData/diagnostics/tests/controlStructures/localReturnInsidePropertyAccessor.kt");
             }
 
+            @Test
+            @TestMetadata("nestedLoopsWithMultipleLabels.kt")
+            public void testNestedLoopsWithMultipleLabels() throws Exception {
+                runTest("compiler/testData/diagnostics/tests/controlStructures/nestedLoopsWithMultipleLabels.kt");
+            }
+
             @Test
             @TestMetadata("notAFunctionLabel_after.kt")
             public void testNotAFunctionLabel_after() throws Exception {