KJS: Label loops in LabeledBlockToDoWhileTransformation

Fix KT-24542: insert labes for affected loops, breaks and continues
statements during LabeledBlockToDoWhileTransformation

js/js.inliner/src/org/jetbrains/kotlin/js/inline/clean/LabeledBlockToDoWhileTransformation.kt

@@ -17,6 +17,8 @@
 package org.jetbrains.kotlin.js.inline.clean
 
 import org.jetbrains.kotlin.js.backend.ast.*
+import org.jetbrains.kotlin.utils.addIfNotNull
+import java.util.*
 
 object LabeledBlockToDoWhileTransformation {
     fun apply(fragments: List<JsProgramFragment>) {
@@ -26,15 +28,97 @@ object LabeledBlockToDoWhileTransformation {
         }
     }
 
+    fun JsStatement.isLoop(): Boolean = when (this) {
+        is JsWhile -> true
+        is JsDoWhile -> true
+        is JsFor -> true
+        is JsForIn -> true
+        else -> false
+    }
+
     fun apply(root: JsNode) {
         object : JsVisitorWithContextImpl() {
+            val loopStack = Stack<JsStatement>()
+            val newFakeLoops = HashSet<JsDoWhile>()
+            val loopLabels = HashMap<JsStatement, JsLabel>()
+
+            // If labeled block sits in between loop and corresponding unlabeled breaks/continues
+            // we have to label this loop, breaks and continues in order to preserve their
+            // relationships across new fake do-while loop
+            val loopsToLabel = HashSet<JsStatement>()
+
             override fun endVisit(x: JsLabel, ctx: JsContext<JsNode>) {
                 if (x.statement is JsBlock) {
-                    x.statement = JsDoWhile(JsBooleanLiteral(false), x.statement)
+                    loopsToLabel.addIfNotNull(loopStack.lastOrNull())
+                    val fakeLoop = JsDoWhile(JsBooleanLiteral(false), x.statement)
+                    newFakeLoops.add(fakeLoop)
+                    x.statement = fakeLoop
                 }
-
                 super.endVisit(x, ctx)
             }
+
+            override fun visit(x: JsLabel, ctx: JsContext<JsNode>): Boolean {
+                if (x.statement.isLoop()) {
+                    loopLabels[x.statement] = x
+                }
+                return true
+            }
+
+            override fun visit(x: JsWhile, ctx: JsContext<JsNode>) = visitLoop(x)
+            override fun visit(x: JsDoWhile, ctx: JsContext<JsNode>) = visitLoop(x)
+            override fun visit(x: JsFor, ctx: JsContext<JsNode>) = visitLoop(x)
+            override fun visit(x: JsForIn, ctx: JsContext<JsNode>) = visitLoop(x)
+
+            override fun endVisit(x: JsWhile, ctx: JsContext<JsNode>) = endVisitLoop(x, ctx)
+            override fun endVisit(x: JsDoWhile, ctx: JsContext<JsNode>) = endVisitLoop(x, ctx)
+            override fun endVisit(x: JsFor, ctx: JsContext<JsNode>) = endVisitLoop(x, ctx)
+            override fun endVisit(x: JsForIn, ctx: JsContext<JsNode>) = endVisitLoop(x, ctx)
+
+            private fun visitLoop(x: JsStatement): Boolean {
+                loopStack.push(x)
+                return true
+            }
+
+            private fun endVisitLoop(x: JsStatement, ctx: JsContext<JsNode>) {
+                loopStack.pop()
+                if (loopsToLabel.contains(x)) {
+                    // Reuse loop label if present. Otherwise create new label.
+                    var label = loopLabels[x]
+                    if (label == null) {
+                        val labelName = JsScope.declareTemporaryName("loop_label")
+                        label = JsLabel(labelName, x)
+                        loopLabels[x] = label
+                        ctx.replaceMe(label)
+                    }
+                    labelLoopBreaksAndContinues(x, newFakeLoops, label.name.makeRef())
+                }
+            }
         }.accept(root)
     }
+
+    /*
+    Label unlabeled breaks that correspond to current loop.
+    Skip newly created fake do-while loops.
+     */
+    private fun labelLoopBreaksAndContinues(loop: JsStatement, fakeLoops: Set<JsDoWhile>, label: JsNameRef) {
+        object : JsVisitorWithContextImpl() {
+            override fun visit(x: JsWhile, ctx: JsContext<JsNode>): Boolean = visitLoop(x)
+            override fun visit(x: JsDoWhile, ctx: JsContext<JsNode>): Boolean = visitLoop(x)
+            override fun visit(x: JsFor, ctx: JsContext<JsNode>): Boolean = visitLoop(x)
+            override fun visit(x: JsForIn, ctx: JsContext<JsNode>): Boolean = visitLoop(x)
+            private fun visitLoop(x: JsStatement): Boolean = fakeLoops.contains(x) || x === loop
+
+            override fun endVisit(x: JsBreak, ctx: JsContext<JsNode>) {
+                if (x.label == null)
+                    ctx.replaceMe(JsBreak(label))
+                super.endVisit(x, ctx)
+            }
+
+            override fun endVisit(x: JsContinue, ctx: JsContext<JsNode>) {
+                if (x.label == null)
+                    ctx.replaceMe(JsContinue(label))
+                super.endVisit(x, ctx)
+            }
+        }.accept(loop)
+    }
 }

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/BoxJsTestGenerated.java

@@ -3919,6 +3919,11 @@ public class BoxJsTestGenerated extends AbstractBoxJsTest {
             runTest("js/js.translator/testData/box/inline/localInlineFunctionReference.kt");
         }
 
+        @TestMetadata("loopWithInlinableCondition.kt")
+        public void testLoopWithInlinableCondition() throws Exception {
+            runTest("js/js.translator/testData/box/inline/loopWithInlinableCondition.kt");
+        }
+
         @TestMetadata("metadataForPublicFunction.kt")
         public void testMetadataForPublicFunction() throws Exception {
             runTest("js/js.translator/testData/box/inline/metadataForPublicFunction.kt");

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/IrBoxJsTestGenerated.java

@@ -3919,6 +3919,11 @@ public class IrBoxJsTestGenerated extends AbstractIrBoxJsTest {
             runTest("js/js.translator/testData/box/inline/localInlineFunctionReference.kt");
         }
 
+        @TestMetadata("loopWithInlinableCondition.kt")
+        public void testLoopWithInlinableCondition() throws Exception {
+            runTest("js/js.translator/testData/box/inline/loopWithInlinableCondition.kt");
+        }
+
         @TestMetadata("metadataForPublicFunction.kt")
         public void testMetadataForPublicFunction() throws Exception {
             runTest("js/js.translator/testData/box/inline/metadataForPublicFunction.kt");

js/js.translator/testData/box/inline/loopWithInlinableCondition.kt

@@ -0,0 +1,110 @@
+// EXPECTED_REACHABLE_NODES: 1122
+/*
+Modified test case from issue: https://youtrack.jetbrains.com/issue/KT-24542
+ */
+package foo
+
+class Test() {
+    var output: String = ""
+
+    inline fun foo(): Boolean {
+        output += "foo "
+        return false
+    }
+
+    fun run(doBreak: Boolean, doContinue: Boolean) {
+        do {
+            output += "1 "
+            if (doBreak) break
+            output += "2 "
+            if (doContinue) continue
+            output += "3 "
+        }
+        while(foo())
+    }
+
+    fun runNested(doBreak: Boolean, doContinue: Boolean) {
+        do {
+            output += "0_1 "
+
+            do {
+                output += "1_1 "
+                if (doBreak) break
+                output += "1_2 "
+                if (doContinue) continue
+                output += "1_3 "
+            }
+            while(foo())
+
+            output += "0_2 "
+
+            if (doBreak) break
+
+            output += "0_3 "
+
+            do {
+                output += "2_1 "
+                if (doBreak) break
+                output += "2_2 "
+                if (doContinue) continue
+                output += "2_3 "
+            }
+            while(foo())
+
+            output += "0_4 "
+
+            if (doContinue) continue
+
+            output += "0_5 "
+
+            loop_with_label@ do {
+                output += "3_1 "
+                if (doBreak) break
+                output += "3_2 "
+                if (doContinue) continue
+                output += "3_3 "
+            }
+            while(foo())
+
+            output += "0_6 "
+        }
+        while(foo())
+    }
+}
+
+fun test(doBreak: Boolean, doContinue: Boolean): String {
+    var x = Test()
+    x.run(doBreak, doContinue)
+    return x.output
+}
+
+fun testNested(doBreak: Boolean, doContinue: Boolean): String {
+    var x = Test()
+    x.runNested(doBreak, doContinue)
+    return x.output
+}
+
+fun box(): String {
+    val test1 = test(true, true)
+    val test2 = test(true, false)
+    val test3 = test(false, true)
+    val test4 = test(false, false)
+
+    if (test1 != "1 ") return "Test1 output: ${test1}"
+    if (test2 != "1 ") return "Test2 output: ${test2}"
+    if (test3 != "1 2 foo ") return "Test3 output: ${test3}"
+    if (test4 != "1 2 3 foo ") return "Test4 output: ${test4}"
+
+    val testNested1 = testNested(true, true)
+    val testNested2 = testNested(true, false)
+    val testNested3 = testNested(false, true)
+    val testNested4 = testNested(false, false)
+
+    if (testNested1 != "0_1 1_1 0_2 ") return "testNested1 output: ${testNested1}"
+    if (testNested2 != "0_1 1_1 0_2 ") return "testNested2 output: ${testNested2}"
+    if (testNested3 != "0_1 1_1 1_2 foo 0_2 0_3 2_1 2_2 foo 0_4 foo ") return "testNested3 output: ${testNested3}"
+    if (testNested4 != "0_1 1_1 1_2 1_3 foo 0_2 0_3 2_1 2_2 2_3 foo 0_4 0_5 3_1 3_2 3_3 foo 0_6 foo ")
+        return "testNested4 output: ${testNested4}"
+
+    return "OK"
+}