Handle break/continue in loop header (condition for 'while', range for 'for').

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

@@ -39,8 +39,8 @@ class LoopExpressionGenerator(statementGenerator: StatementGenerator) : Statemen
                                                       context.builtIns.unitType, IrOperator.DO_WHILE_LOOP))
 
     private fun generateConditionalLoop(ktLoop: KtWhileExpressionBase, irLoop: IrLoopBase): IrLoop {
-        statementGenerator.bodyGenerator.putLoop(ktLoop, irLoop)
         irLoop.condition = statementGenerator.generateExpression(ktLoop.condition!!)
+        statementGenerator.bodyGenerator.putLoop(ktLoop, irLoop)
         irLoop.body = ktLoop.body?.let { statementGenerator.generateExpression(ktLoop.body!!) }
         irLoop.label = getLoopLabel(ktLoop)
         return irLoop
@@ -74,14 +74,14 @@ class LoopExpressionGenerator(statementGenerator: StatementGenerator) : Statemen
                 break
             }
             if (targetLabel == null) {
-                return getLoop(finger)
+                return getLoop(finger) ?: continue
             }
             else {
                 val parent = finger.parent
                 if (parent is KtLabeledExpression) {
                     val label = parent.getLabelName()!!
                     if (targetLabel == label) {
-                        return getLoop(finger)
+                        return getLoop(finger) ?: continue
                     }
                 }
             }
@@ -89,9 +89,8 @@ class LoopExpressionGenerator(statementGenerator: StatementGenerator) : Statemen
         throw AssertionError("No parent loop for break/continue @$targetLabel")
     }
 
-    private fun getLoop(ktLoop: KtLoopExpression): IrLoop {
+    private fun getLoop(ktLoop: KtLoopExpression): IrLoop? {
-        return statementGenerator.bodyGenerator.getLoop(ktLoop) ?:
+        return statementGenerator.bodyGenerator.getLoop(ktLoop)
-               throw AssertionError("Loop was not visited:\n${ktLoop.text}")
     }
 
     fun generateForLoop(ktFor: KtForExpression): IrExpression {

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/util/RenderIrElement.kt

@@ -162,10 +162,10 @@ class RenderIrElementVisitor : IrElementVisitor<String, Nothing?> {
             "DO_WHILE label=${loop.label} operator=${loop.operator}"
 
     override fun visitBreak(jump: IrBreak, data: Nothing?): String =
-            "BREAK label=${jump.label} depth=${jump.getDepth()}"
+            "BREAK label=${jump.label} loop.label=${jump.loop.label} depth=${jump.getDepth()}"
 
     override fun visitContinue(jump: IrContinue, data: Nothing?): String =
-            "CONTINUE label=${jump.label} depth=${jump.getDepth()}"
+            "CONTINUE label=${jump.label} loop.label=${jump.loop.label} depth=${jump.getDepth()}"
 
     override fun visitThrow(expression: IrThrow, data: Nothing?): String =
             "THROW type=${expression.type.render()}"

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

@@ -4,18 +4,18 @@ FILE /breakContinue.kt
       WHILE label=null operator=WHILE_LOOP
         condition: CONST Boolean type=kotlin.Boolean value='true'
         body: BLOCK type=kotlin.Nothing operator=null
-          BREAK label=null depth=0
+          BREAK label=null loop.label=null depth=0
       DO_WHILE label=null operator=DO_WHILE_LOOP
         body: BLOCK type=kotlin.Unit operator=null
-          BREAK label=null depth=0
+          BREAK label=null loop.label=null depth=0
         condition: CONST Boolean type=kotlin.Boolean value='true'
       WHILE label=null operator=WHILE_LOOP
         condition: CONST Boolean type=kotlin.Boolean value='true'
         body: BLOCK type=kotlin.Nothing operator=null
-          CONTINUE label=null depth=0
+          CONTINUE label=null loop.label=null depth=0
       DO_WHILE label=null operator=DO_WHILE_LOOP
         body: BLOCK type=kotlin.Unit operator=null
-          CONTINUE label=null depth=0
+          CONTINUE label=null loop.label=null depth=0
         condition: CONST Boolean type=kotlin.Boolean value='true'
   FUN public fun test2(): kotlin.Unit
     BLOCK_BODY
@@ -25,18 +25,18 @@ FILE /breakContinue.kt
           WHILE label=INNER operator=WHILE_LOOP
             condition: CONST Boolean type=kotlin.Boolean value='true'
             body: BLOCK type=kotlin.Nothing operator=null
-              BREAK label=INNER depth=0
+              BREAK label=INNER loop.label=INNER depth=0
-              BREAK label=OUTER depth=1
+              BREAK label=OUTER loop.label=OUTER depth=1
-          BREAK label=OUTER depth=0
+          BREAK label=OUTER loop.label=OUTER depth=0
       WHILE label=OUTER operator=WHILE_LOOP
         condition: CONST Boolean type=kotlin.Boolean value='true'
         body: BLOCK type=kotlin.Nothing operator=null
           WHILE label=INNER operator=WHILE_LOOP
             condition: CONST Boolean type=kotlin.Boolean value='true'
             body: BLOCK type=kotlin.Nothing operator=null
-              CONTINUE label=INNER depth=0
+              CONTINUE label=INNER loop.label=INNER depth=0
-              CONTINUE label=OUTER depth=1
+              CONTINUE label=OUTER loop.label=OUTER depth=1
-          CONTINUE label=OUTER depth=0
+          CONTINUE label=OUTER loop.label=OUTER depth=0
   FUN public fun test3(): kotlin.Unit
     BLOCK_BODY
       WHILE label=L operator=WHILE_LOOP
@@ -45,13 +45,13 @@ FILE /breakContinue.kt
           WHILE label=L operator=WHILE_LOOP
             condition: CONST Boolean type=kotlin.Boolean value='true'
             body: BLOCK type=kotlin.Nothing operator=null
-              BREAK label=L depth=0
+              BREAK label=L loop.label=L depth=0
-          BREAK label=L depth=0
+          BREAK label=L loop.label=L depth=0
       WHILE label=L operator=WHILE_LOOP
         condition: CONST Boolean type=kotlin.Boolean value='true'
         body: BLOCK type=kotlin.Nothing operator=null
           WHILE label=L operator=WHILE_LOOP
             condition: CONST Boolean type=kotlin.Boolean value='true'
             body: BLOCK type=kotlin.Nothing operator=null
-              CONTINUE label=L depth=0
+              CONTINUE label=L loop.label=L depth=0
-          CONTINUE label=L depth=0
+          CONTINUE label=L loop.label=L depth=0

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

@@ -0,0 +1,23 @@
+fun test1(c: Boolean?) {
+    L@ while (true) {
+        while (c ?: break)
+    }
+}
+
+fun test2(c: Boolean?) {
+    L@ while (true) {
+        while (c ?: continue)
+    }
+}
+
+fun test3(ss: List<String>?) {
+    L@ while (true) {
+        for (s in ss ?: continue)
+    }
+}
+
+fun test4(ss: List<String>?) {
+    L@ while (true) {
+        for (s in ss ?: break)
+    }
+}

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

@@ -0,0 +1,79 @@
+FILE /breakContinueInLoopHeader.kt
+  FUN public fun test1(/*0*/ c: kotlin.Boolean?): kotlin.Unit
+    BLOCK_BODY
+      WHILE label=L operator=WHILE_LOOP
+        condition: CONST Boolean type=kotlin.Boolean value='true'
+        body: BLOCK type=kotlin.Unit operator=null
+          WHILE label=null operator=WHILE_LOOP
+            condition: BLOCK type=kotlin.Boolean operator=ELVIS
+              VAR val tmp0_elvis_lhs: kotlin.Boolean?
+                GET_VAR c type=kotlin.Boolean? operator=null
+              WHEN type=kotlin.Boolean operator=null
+                if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
+                  arg0: GET_VAR tmp0_elvis_lhs type=kotlin.Boolean? operator=null
+                  arg1: CONST Null type=kotlin.Nothing? value='null'
+                then: BREAK label=null loop.label=L depth=1
+                else: GET_VAR tmp0_elvis_lhs type=kotlin.Boolean? operator=null
+  FUN public fun test2(/*0*/ c: kotlin.Boolean?): kotlin.Unit
+    BLOCK_BODY
+      WHILE label=L operator=WHILE_LOOP
+        condition: CONST Boolean type=kotlin.Boolean value='true'
+        body: BLOCK type=kotlin.Unit operator=null
+          WHILE label=null operator=WHILE_LOOP
+            condition: BLOCK type=kotlin.Boolean operator=ELVIS
+              VAR val tmp0_elvis_lhs: kotlin.Boolean?
+                GET_VAR c type=kotlin.Boolean? operator=null
+              WHEN type=kotlin.Boolean operator=null
+                if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
+                  arg0: GET_VAR tmp0_elvis_lhs type=kotlin.Boolean? operator=null
+                  arg1: CONST Null type=kotlin.Nothing? value='null'
+                then: CONTINUE label=null loop.label=L depth=1
+                else: GET_VAR tmp0_elvis_lhs type=kotlin.Boolean? operator=null
+  FUN public fun test3(/*0*/ ss: kotlin.collections.List<kotlin.String>?): kotlin.Unit
+    BLOCK_BODY
+      WHILE label=L operator=WHILE_LOOP
+        condition: CONST Boolean type=kotlin.Boolean value='true'
+        body: BLOCK type=kotlin.Unit operator=null
+          BLOCK type=kotlin.Unit operator=FOR_LOOP
+            VAR val tmp1_iterator: kotlin.collections.Iterator<kotlin.String>
+              CALL .iterator type=kotlin.collections.Iterator<kotlin.String> operator=FOR_LOOP_ITERATOR
+                $this: BLOCK type=kotlin.collections.List<kotlin.String> operator=ELVIS
+                  VAR val tmp0_elvis_lhs: kotlin.collections.List<kotlin.String>?
+                    GET_VAR ss type=kotlin.collections.List<kotlin.String>? operator=null
+                  WHEN type=kotlin.collections.List<kotlin.String> operator=null
+                    if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
+                      arg0: GET_VAR tmp0_elvis_lhs type=kotlin.collections.List<kotlin.String>? operator=null
+                      arg1: CONST Null type=kotlin.Nothing? value='null'
+                    then: CONTINUE label=null loop.label=L depth=0
+                    else: GET_VAR tmp0_elvis_lhs type=kotlin.collections.List<kotlin.String>? operator=null
+            WHILE label=null operator=FOR_LOOP_INNER_WHILE
+              condition: CALL .hasNext type=kotlin.Boolean operator=FOR_LOOP_HAS_NEXT
+                $this: GET_VAR tmp1_iterator type=kotlin.collections.Iterator<kotlin.String> operator=null
+              body: BLOCK type=kotlin.Unit operator=FOR_LOOP_INNER_WHILE
+                VAR val s: kotlin.String
+                  CALL .next type=kotlin.String operator=FOR_LOOP_NEXT
+                    $this: GET_VAR tmp1_iterator type=kotlin.collections.Iterator<kotlin.String> operator=null
+  FUN public fun test4(/*0*/ ss: kotlin.collections.List<kotlin.String>?): kotlin.Unit
+    BLOCK_BODY
+      WHILE label=L operator=WHILE_LOOP
+        condition: CONST Boolean type=kotlin.Boolean value='true'
+        body: BLOCK type=kotlin.Unit operator=null
+          BLOCK type=kotlin.Unit operator=FOR_LOOP
+            VAR val tmp1_iterator: kotlin.collections.Iterator<kotlin.String>
+              CALL .iterator type=kotlin.collections.Iterator<kotlin.String> operator=FOR_LOOP_ITERATOR
+                $this: BLOCK type=kotlin.collections.List<kotlin.String> operator=ELVIS
+                  VAR val tmp0_elvis_lhs: kotlin.collections.List<kotlin.String>?
+                    GET_VAR ss type=kotlin.collections.List<kotlin.String>? operator=null
+                  WHEN type=kotlin.collections.List<kotlin.String> operator=null
+                    if: CALL .EQEQ type=kotlin.Boolean operator=EQEQ
+                      arg0: GET_VAR tmp0_elvis_lhs type=kotlin.collections.List<kotlin.String>? operator=null
+                      arg1: CONST Null type=kotlin.Nothing? value='null'
+                    then: BREAK label=null loop.label=L depth=0
+                    else: GET_VAR tmp0_elvis_lhs type=kotlin.collections.List<kotlin.String>? operator=null
+            WHILE label=null operator=FOR_LOOP_INNER_WHILE
+              condition: CALL .hasNext type=kotlin.Boolean operator=FOR_LOOP_HAS_NEXT
+                $this: GET_VAR tmp1_iterator type=kotlin.collections.Iterator<kotlin.String> operator=null
+              body: BLOCK type=kotlin.Unit operator=FOR_LOOP_INNER_WHILE
+                VAR val s: kotlin.String
+                  CALL .next type=kotlin.String operator=FOR_LOOP_NEXT
+                    $this: GET_VAR tmp1_iterator type=kotlin.collections.Iterator<kotlin.String> operator=null

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

@@ -13,7 +13,7 @@ FILE /forWithBreakContinue.kt
               CALL .next type=kotlin.String operator=FOR_LOOP_NEXT
                 $this: GET_VAR tmp0_iterator type=kotlin.collections.Iterator<kotlin.String> operator=null
             BLOCK type=kotlin.Nothing operator=null
-              BREAK label=null depth=0
+              BREAK label=null loop.label=null depth=0
   FUN public fun testForBreak2(/*0*/ ss: kotlin.collections.List<kotlin.String>): kotlin.Unit
     BLOCK_BODY
       BLOCK type=kotlin.Unit operator=FOR_LOOP
@@ -40,10 +40,10 @@ FILE /forWithBreakContinue.kt
                       CALL .next type=kotlin.String operator=FOR_LOOP_NEXT
                         $this: GET_VAR tmp1_iterator type=kotlin.collections.Iterator<kotlin.String> operator=null
                     BLOCK type=kotlin.Nothing operator=null
-                      BREAK label=OUTER depth=1
+                      BREAK label=OUTER loop.label=OUTER depth=1
-                      BREAK label=INNER depth=0
+                      BREAK label=INNER loop.label=INNER depth=0
-                      BREAK label=null depth=0
+                      BREAK label=null loop.label=INNER depth=0
-              BREAK label=OUTER depth=0
+              BREAK label=OUTER loop.label=OUTER depth=0
   FUN public fun testForContinue1(/*0*/ ss: kotlin.collections.List<kotlin.String>): kotlin.Unit
     BLOCK_BODY
       BLOCK type=kotlin.Unit operator=FOR_LOOP
@@ -58,7 +58,7 @@ FILE /forWithBreakContinue.kt
               CALL .next type=kotlin.String operator=FOR_LOOP_NEXT
                 $this: GET_VAR tmp0_iterator type=kotlin.collections.Iterator<kotlin.String> operator=null
             BLOCK type=kotlin.Nothing operator=null
-              CONTINUE label=null depth=0
+              CONTINUE label=null loop.label=null depth=0
   FUN public fun testForContinue2(/*0*/ ss: kotlin.collections.List<kotlin.String>): kotlin.Unit
     BLOCK_BODY
       BLOCK type=kotlin.Unit operator=FOR_LOOP
@@ -85,7 +85,7 @@ FILE /forWithBreakContinue.kt
                       CALL .next type=kotlin.String operator=FOR_LOOP_NEXT
                         $this: GET_VAR tmp1_iterator type=kotlin.collections.Iterator<kotlin.String> operator=null
                     BLOCK type=kotlin.Nothing operator=null
-                      CONTINUE label=OUTER depth=1
+                      CONTINUE label=OUTER loop.label=OUTER depth=1
-                      CONTINUE label=INNER depth=0
+                      CONTINUE label=INNER loop.label=INNER depth=0
-                      CONTINUE label=null depth=0
+                      CONTINUE label=null loop.label=INNER depth=0
-              CONTINUE label=OUTER depth=0
+              CONTINUE label=OUTER loop.label=OUTER depth=0

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

@@ -304,6 +304,12 @@ public class IrTextTestCaseGenerated extends AbstractIrTextTestCase {
             doTest(fileName);
         }
 
+        @TestMetadata("breakContinueInLoopHeader.kt")
+        public void testBreakContinueInLoopHeader() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/ir/irText/expressions/breakContinueInLoopHeader.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("callWithReorderedArguments.kt")
         public void testCallWithReorderedArguments() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/ir/irText/expressions/callWithReorderedArguments.kt");