Fix KT-24804 PSI2IR Crashes on loop with two labels

compiler/ir/ir.psi2ir/src/org/jetbrains/kotlin/psi2ir/generators/LoopExpressionGenerator.kt

@@ -40,6 +40,8 @@ class LoopExpressionGenerator(statementGenerator: StatementGenerator) : Statemen
 
         statementGenerator.bodyGenerator.putLoop(ktWhile, irLoop)
 
+        irLoop.label = getLoopLabel(ktWhile)
+
         irLoop.body = ktWhile.body?.let { ktLoopBody ->
             if (ktLoopBody is KtBlockExpression)
                 generateWhileLoopBody(ktLoopBody)
@@ -47,8 +49,6 @@ class LoopExpressionGenerator(statementGenerator: StatementGenerator) : Statemen
                 ktLoopBody.genExpr()
         }
 
-        irLoop.label = getLoopLabel(ktWhile)
-
         return irLoop
     }
 
@@ -60,6 +60,8 @@ class LoopExpressionGenerator(statementGenerator: StatementGenerator) : Statemen
 
         statementGenerator.bodyGenerator.putLoop(ktDoWhile, irLoop)
 
+        irLoop.label = getLoopLabel(ktDoWhile)
+
         irLoop.body = ktDoWhile.body?.let { ktLoopBody ->
             if (ktLoopBody is KtBlockExpression)
                 generateDoWhileLoopBody(ktLoopBody)
@@ -69,8 +71,6 @@ class LoopExpressionGenerator(statementGenerator: StatementGenerator) : Statemen
 
         irLoop.condition = ktDoWhile.condition!!.genExpr()
 
-        irLoop.label = getLoopLabel(ktDoWhile)
-
         return IrBlockImpl(ktDoWhile.startOffsetSkippingComments, ktDoWhile.endOffset, context.irBuiltIns.unitType).apply {
             statements.add(irLoop)
         }
@@ -90,20 +90,26 @@ class LoopExpressionGenerator(statementGenerator: StatementGenerator) : Statemen
         )
 
     fun generateBreak(ktBreak: KtBreakExpression): IrExpression {
-        val parentLoop = findParentLoop(ktBreak) ?: return ErrorExpressionGenerator(statementGenerator).generateErrorExpression(
+        val parentLoop = findParentLoop(ktBreak)
-            ktBreak, RuntimeException("Loop not found for break expression: ${ktBreak.text}")
+            ?: return ErrorExpressionGenerator(statementGenerator).generateErrorExpression(
-        )
+                ktBreak, RuntimeException("Loop not found for break expression: ${ktBreak.text}")
-        return IrBreakImpl(ktBreak.startOffsetSkippingComments, ktBreak.endOffset, context.irBuiltIns.nothingType, parentLoop).apply {
+            )
-            label = ktBreak.getLabelName()
+        return IrBreakImpl(
+            ktBreak.startOffsetSkippingComments, ktBreak.endOffset, context.irBuiltIns.nothingType, parentLoop
+        ).apply {
+            label = parentLoop.label.takeIf { ktBreak.getLabelName() != null }
         }
     }
 
     fun generateContinue(ktContinue: KtContinueExpression): IrExpression {
-        val parentLoop = findParentLoop(ktContinue) ?: return ErrorExpressionGenerator(statementGenerator).generateErrorExpression(
+        val parentLoop = findParentLoop(ktContinue)
-            ktContinue, RuntimeException("Loop not found for continue expression: ${ktContinue.text}")
+            ?: return ErrorExpressionGenerator(statementGenerator).generateErrorExpression(
-        )
+                ktContinue, RuntimeException("Loop not found for continue expression: ${ktContinue.text}")
-        return IrContinueImpl(ktContinue.startOffsetSkippingComments, ktContinue.endOffset, context.irBuiltIns.nothingType, parentLoop).apply {
+            )
-            label = ktContinue.getLabelName()
+        return IrContinueImpl(
+            ktContinue.startOffsetSkippingComments, ktContinue.endOffset, context.irBuiltIns.nothingType, parentLoop
+        ).apply {
+            label = parentLoop.label.takeIf { ktContinue.getLabelName() != null }
         }
     }
 
@@ -115,20 +121,21 @@ class LoopExpressionGenerator(statementGenerator: StatementGenerator) : Statemen
 
     private fun findParentLoop(ktExpression: KtExpression, targetLabel: String?): IrLoop? {
         var finger: KtExpression? = ktExpression
-        while (finger != null) {
+        BY_LOOP_EXPRESSIONS@ while (finger != null) {
             finger = finger.getParentOfType<KtLoopExpression>(true)
             if (finger == null) {
                 break
             }
             if (targetLabel == null) {
-                return getLoop(finger) ?: continue
+                return getLoop(finger) ?: continue@BY_LOOP_EXPRESSIONS
             } else {
-                val parent = finger.parent
+                var parent = finger.parent
-                if (parent is KtLabeledExpression) {
+                while (parent is KtLabeledExpression) {
                     val label = parent.getLabelName()!!
                     if (targetLabel == label) {
-                        return getLoop(finger) ?: continue
+                        return getLoop(finger) ?: continue@BY_LOOP_EXPRESSIONS
                     }
+                    parent = parent.parent
                 }
             }
         }

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

@@ -134,4 +134,3 @@ FILE fqName:<root> fileName:/breakContinueInLoopHeader.kt
                 arg0: GET_VAR 'i: Int' type=kotlin.Int origin=null
                 arg1: CONST Int type=kotlin.Int value=3
               then: BREAK label=null loop.label=Outer
-

compiler/testData/ir/irText/expressions/kt24804.kt

@@ -0,0 +1,17 @@
+inline fun foo() = false
+
+fun run(x: Boolean, y: Boolean): String {
+    var z = 10
+    l1@ l2@ do {
+        z += 1
+        if (z > 100) return "NOT_OK"
+        if (x) continue@l1
+        if (y) continue@l2
+    } while(foo())
+
+    return "OK"
+}
+
+fun box(): String {
+    return run(true, true)
+}

compiler/testData/ir/irText/expressions/kt24804.txt

@@ -0,0 +1,42 @@
+FILE fqName:<root> fileName:/kt24804.kt
+  FUN name:foo visibility:public modality:FINAL <> () returnType:kotlin.Boolean flags:inline
+    BLOCK_BODY
+      RETURN type=kotlin.Nothing from='foo(): Boolean'
+        CONST Boolean type=kotlin.Boolean value=false
+  FUN name:run visibility:public modality:FINAL <> (x:kotlin.Boolean, y:kotlin.Boolean) returnType:kotlin.String flags:
+    VALUE_PARAMETER name:x index:0 type:kotlin.Boolean flags:
+    VALUE_PARAMETER name:y index:1 type:kotlin.Boolean flags:
+    BLOCK_BODY
+      VAR name:z type:kotlin.Int flags:var
+        CONST Int type=kotlin.Int value=10
+      BLOCK type=kotlin.Unit origin=null
+        DO_WHILE label=l2 origin=DO_WHILE_LOOP
+          body: COMPOSITE type=kotlin.Unit origin=null
+            SET_VAR 'z: Int' type=kotlin.Unit origin=PLUSEQ
+              CALL 'plus(Int): Int' type=kotlin.Int origin=PLUSEQ
+                $this: GET_VAR 'z: Int' type=kotlin.Int origin=PLUSEQ
+                other: CONST Int type=kotlin.Int value=1
+            WHEN type=kotlin.Unit origin=null
+              BRANCH
+                if: CALL 'greater(Int, Int): Boolean' type=kotlin.Boolean origin=GT
+                  arg0: GET_VAR 'z: Int' type=kotlin.Int origin=null
+                  arg1: CONST Int type=kotlin.Int value=100
+                then: RETURN type=kotlin.Nothing from='run(Boolean, Boolean): String'
+                  CONST String type=kotlin.String value="NOT_OK"
+            WHEN type=kotlin.Unit origin=null
+              BRANCH
+                if: GET_VAR 'value-parameter x: Boolean' type=kotlin.Boolean origin=null
+                then: CONTINUE label=l2 loop.label=l2
+            WHEN type=kotlin.Unit origin=null
+              BRANCH
+                if: GET_VAR 'value-parameter y: Boolean' type=kotlin.Boolean origin=null
+                then: CONTINUE label=l2 loop.label=l2
+          condition: CALL 'foo(): Boolean' type=kotlin.Boolean origin=null
+      RETURN type=kotlin.Nothing from='run(Boolean, Boolean): String'
+        CONST String type=kotlin.String value="OK"
+  FUN name:box visibility:public modality:FINAL <> () returnType:kotlin.String flags:
+    BLOCK_BODY
+      RETURN type=kotlin.Nothing from='box(): String'
+        CALL 'run(Boolean, Boolean): String' type=kotlin.String origin=null
+          x: CONST Boolean type=kotlin.Boolean value=true
+          y: CONST Boolean type=kotlin.Boolean value=true

compiler/tests/org/jetbrains/kotlin/ir/IrTextTestCaseGenerated.java

@@ -947,6 +947,11 @@ public class IrTextTestCaseGenerated extends AbstractIrTextTestCase {
             runTest("compiler/testData/ir/irText/expressions/kt23030.kt");
         }
 
+        @TestMetadata("kt24804.kt")
+        public void testKt24804() throws Exception {
+            runTest("compiler/testData/ir/irText/expressions/kt24804.kt");
+        }
+
         @TestMetadata("kt27933.kt")
         public void testKt27933() throws Exception {
             runTest("compiler/testData/ir/irText/expressions/kt27933.kt");