JS: coroutines: fixes after code review

compiler/testData/codegen/box/coroutines/controlFlow/breakFinally.kt

@@ -38,13 +38,14 @@ fun box(): String {
                 }
                 result += "ignore"
             }
+            result += "*"
         }
         finally {
             result += "finally"
         }
         result += "."
     }
-    if (value != "AC!ED!@finally.") return "fail: $value"
+    if (value != "AC!ED!@*finally.") return "fail: $value"
 
     return "OK"
 }

compiler/testData/codegen/box/coroutines/controlFlow/forContinue.kt

@@ -0,0 +1,28 @@
+// WITH_RUNTIME
+
+class Controller {
+    var result = ""
+
+    suspend fun <T> suspendWithResult(value: T, c: Continuation<T>) {
+        c.resume(value)
+    }
+}
+
+fun builder(coroutine c: Controller.() -> Continuation<Unit>): String {
+    val controller = Controller()
+    c(controller).resume(Unit)
+    return controller.result
+}
+
+fun box(): String {
+    val value = builder {
+        for (x in listOf("O", "$", "K")) {
+            if (x == "$") continue
+            result += suspendWithResult(x)
+        }
+        result += "."
+    }
+    if (value != "OK.") return "fail: suspend in for body: $value"
+
+    return "OK"
+}

compiler/testData/codegen/box/coroutines/controlFlow/switchLikeWhen.kt

@@ -0,0 +1,34 @@
+// WITH_RUNTIME
+
+// TODO: fix bug in JVM backend and remove this directive
+// TARGET_BACKEND: JS
+
+class Controller {
+    var result = ""
+
+    suspend fun <T> suspendWithResult(value: T, c: Continuation<T>) {
+        result += "["
+        c.resume(value)
+    }
+}
+
+fun builder(coroutine c: Controller.() -> Continuation<Unit>): String {
+    val controller = Controller()
+    c(controller).resume(Unit)
+    return controller.result
+}
+
+fun box(): String {
+    var value = builder {
+        for (v in listOf("A", "B", "C")) {
+            when (v) {
+                "A" -> result += "A;"
+                "B" -> result += suspendWithResult(v) + "]"
+                else -> result += suspendWithResult(v) + "]!"
+            }
+        }
+    }
+    if (value != "A;[B][C]!") return "fail: suspend as if condition: $value"
+
+    return "OK"
+}

compiler/tests-ir-jvm/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -4744,6 +4744,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 doTest(fileName);
             }
 
+            @TestMetadata("forContinue.kt")
+            public void testForContinue() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/controlFlow/forContinue.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("forStatement.kt")
             public void testForStatement() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/controlFlow/forStatement.kt");

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -4744,6 +4744,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("forContinue.kt")
+            public void testForContinue() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/controlFlow/forContinue.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("forStatement.kt")
             public void testForStatement() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/controlFlow/forStatement.kt");

js/js.inliner/src/org/jetbrains/kotlin/js/coroutine/CoroutineBodyTransformer.kt

@@ -112,29 +112,28 @@ class CoroutineBodyTransformer(
     }
 
     override fun visitWhile(x: JsWhile) = splitIfNecessary(x) {
-        val predecessor = currentBlock
         val successor = CoroutineBlock()
-
         val bodyEntryBlock = CoroutineBlock()
+        currentStatements += stateAndJump(bodyEntryBlock)
+
         currentBlock = bodyEntryBlock
+        if (x.condition != JsLiteral.TRUE) {
+            currentStatements += JsIf(JsAstUtils.notOptimized(x.condition), JsBlock(stateAndJump(successor))).apply { source = x.source }
+        }
+
         withBreakAndContinue(x, successor, bodyEntryBlock) {
             x.body.accept(this)
         }
 
-        if (x.condition != JsLiteral.TRUE) {
+        currentStatements += stateAndJump(bodyEntryBlock)
-            val jsIf = JsIf(JsAstUtils.notOptimized(x.condition), JsBlock(stateAndJump(successor))).apply { source = x.source }
-            bodyEntryBlock.statements.add(0, jsIf)
-        }
-        currentBlock.statements += stateAndJump(bodyEntryBlock)
-        predecessor.statements += stateAndJump(bodyEntryBlock)
         currentBlock = successor
     }
 
     override fun visitDoWhile(x: JsDoWhile) = splitIfNecessary(x) {
-        val predecessor = currentBlock
         val successor = CoroutineBlock()
-
         val bodyEntryBlock = CoroutineBlock()
+        currentStatements += stateAndJump(bodyEntryBlock)
+
         currentBlock = bodyEntryBlock
         withBreakAndContinue(x, successor, bodyEntryBlock) {
             x.body.accept(this)
@@ -142,10 +141,10 @@ class CoroutineBodyTransformer(
 
         if (x.condition != JsLiteral.TRUE) {
             val jsIf = JsIf(JsAstUtils.notOptimized(x.condition), JsBlock(stateAndJump(successor))).apply { source = x.source }
-            currentBlock.statements.add(jsIf)
+            currentStatements.add(jsIf)
         }
         currentBlock.statements += stateAndJump(bodyEntryBlock)
-        predecessor.statements += stateAndJump(bodyEntryBlock)
+
         currentBlock = successor
     }
 
@@ -159,29 +158,26 @@ class CoroutineBodyTransformer(
             }
         }
 
-        val predecessor = currentBlock
         val increment = CoroutineBlock()
         val successor = CoroutineBlock()
-
         val bodyEntryBlock = CoroutineBlock()
+        currentStatements += stateAndJump(bodyEntryBlock)
+
         currentBlock = bodyEntryBlock
-        withBreakAndContinue(x, successor, predecessor) {
+        if (x.condition != null && x.condition != JsLiteral.TRUE) {
+            currentStatements += JsIf(JsAstUtils.notOptimized(x.condition), JsBlock(stateAndJump(successor))).apply { source = x.source }
+        }
+
+        withBreakAndContinue(x, successor, increment) {
             x.body.accept(this)
         }
-        val bodyExitBlock = currentBlock
 
-        if (x.condition != null && x.condition != JsLiteral.TRUE) {
+        currentStatements += stateAndJump(increment)
-            val jsIf = JsIf(JsAstUtils.notOptimized(x.condition), JsBlock(stateAndJump(successor))).apply { source = x.source }
-            bodyEntryBlock.statements.add(0, jsIf)
-        }
-
-        bodyExitBlock.statements += stateAndJump(increment)
         currentBlock = increment
 
         x.incrementExpression?.let { JsExpressionStatement(it).accept(this) }
         currentStatements += stateAndJump(bodyEntryBlock)
 
-        predecessor.statements += stateAndJump(bodyEntryBlock)
         currentBlock = successor
     }
 
@@ -201,6 +197,10 @@ class CoroutineBodyTransformer(
         currentStatements += jump()
     }
 
+    /**
+     * When we perform break, continue or return, we can leave try blocks, so we should update $exceptionHandler correspondingly.
+     * Also, these try blocks can contain finally clauses, therefore we need to update $finallyPath as well.
+     */
     private fun jumpWithFinally(targetTryDepth: Int, successor: CoroutineBlock) {
         if (targetTryDepth < tryStack.size) {
             val tryBlock = tryStack[targetTryDepth]
@@ -223,15 +223,17 @@ class CoroutineBodyTransformer(
         val catchBlock = CoroutineBlock()
         val finallyBlock = CoroutineBlock()
 
-        val tryBlock = TryBlock(catchBlock, if (finallyNode != null) finallyBlock else null)
+        tryStack += TryBlock(catchBlock, if (finallyNode != null) finallyBlock else null)
-        tryStack += tryBlock
 
         val oldCatchBlock = currentCatchBlock
         currentCatchBlock = catchBlock
         currentStatements += exceptionState(catchBlock)
 
         x.tryBlock.statements.forEach { it.accept(this) }
+
         currentStatements += exceptionState(oldCatchBlock)
+        currentCatchBlock = oldCatchBlock
+
         if (finallyNode != null) {
             currentStatements += updateFinallyPath(listOf(successor))
             currentStatements += stateAndJump(finallyBlock)
@@ -240,8 +242,6 @@ class CoroutineBodyTransformer(
             currentStatements += stateAndJump(successor)
         }
 
-        currentCatchBlock = oldCatchBlock
-
         // Handle catch node
         currentBlock = catchBlock
 
@@ -279,6 +279,9 @@ class CoroutineBodyTransformer(
         currentBlock = successor
     }
 
+    // There's no implementation for JsSwitch, since we don't generate it. However, when we implement optimization
+    // for simple `when` statement, we will need to support JsSwitch here
+
     private fun generateFinallyExit() {
         val finallyPathRef = JsNameRef(context.finallyPathFieldName, JsLiteral.THIS)
         val stateRef = JsNameRef(context.stateFieldName, JsLiteral.THIS)
@@ -320,6 +323,7 @@ class CoroutineBodyTransformer(
 
     override fun visitThrow(x: JsThrow) {
         if (throwFunctionName != null) {
+            // TODO: what if we catch exception in coroutine?
             val methodRef = JsNameRef(throwFunctionName, JsNameRef(context.controllerFieldName, JsLiteral.THIS))
             val invocation = JsInvocation(methodRef, x.expression).apply {
                 source = x.source

js/js.inliner/src/org/jetbrains/kotlin/js/coroutine/CoroutinePasses.kt

@@ -46,6 +46,10 @@ fun JsNode.collectNodesToSplit(breakContinueTargets: Map<JsContinue, JsStatement
             }
         }
 
+        // We don't handle JsThrow case here the same way as we do for JsReturn.
+        // Exception will be caught by the surrounding catch and then dispatched to a corresponding $exceptionState.
+        // Even if there's no `catch` clause, we generate a fake one that dispatches to a finally block.
+
         override fun visitBreak(x: JsBreak) {
             super.visitBreak(x)
 
@@ -218,43 +222,24 @@ fun JsBlock.replaceLocalVariables(scope: JsScope, context: CoroutineTransformati
         }
 
         override fun endVisit(x: JsVars, ctx: JsContext<in JsStatement>) {
-            val declaredNames = x.vars.map { it.name }
+            val assignments = x.vars.mapNotNull {
-            val totalCount = declaredNames.size
+                val fieldName = scope.getFieldName(it.name)
-            val localVarCount = declaredNames.count()
+                val initExpression = it.initExpression
-
+                if (initExpression != null) {
-            when {
+                    JsAstUtils.assignment(JsNameRef(fieldName, JsLiteral.THIS), it.initExpression)
-                totalCount == localVarCount -> {
-                    val assignments = x.vars.mapNotNull {
-                        val fieldName = scope.getFieldName(it.name)
-                        val initExpression = it.initExpression
-                        if (initExpression != null) {
-                            JsAstUtils.assignment(JsNameRef(fieldName, JsLiteral.THIS), it.initExpression)
-                        }
-                        else {
-                            null
-                        }
-                    }
-                    if (assignments.isNotEmpty()) {
-                        ctx.replaceMe(JsExpressionStatement(JsAstUtils.newSequence(assignments)))
-                    }
-                    else {
-                        ctx.removeMe()
-                    }
                 }
-                localVarCount > 0 -> {
+                else {
-                    for (declaration in x.vars) {
+                    null
-                        if (declaration.name in localVariables) {
-                            val fieldName = scope.getFieldName(declaration.name)
-                            val assignment = JsAstUtils.assignment(JsNameRef(fieldName, JsLiteral.THIS), declaration.initExpression)
-                            ctx.addPrevious(assignment.makeStmt())
-                        }
-                        else {
-                            ctx.addPrevious(JsVars(declaration))
-                        }
-                    }
-                    ctx.removeMe()
                 }
             }
+
+            if (assignments.isNotEmpty()) {
+                ctx.replaceMe(JsExpressionStatement(JsAstUtils.newSequence(assignments)))
+            }
+            else {
+                ctx.removeMe()
+            }
+
             super.endVisit(x, ctx)
         }
     }

js/js.inliner/src/org/jetbrains/kotlin/js/inline/util/collectUtils.kt

@@ -191,9 +191,9 @@ fun JsNode.collectBreakContinueTargets(): Map<JsContinue, JsStatement> {
 
     accept(object : RecursiveJsVisitor() {
         var defaultBreakTarget: JsStatement? = null
-        var breakTargets = mutableMapOf<JsName, JsStatement>()
+        var breakTargets = mutableMapOf<JsName, JsStatement?>()
         var defaultContinueTarget: JsStatement? = null
-        var continueTargets = mutableMapOf<JsName, JsStatement>()
+        var continueTargets = mutableMapOf<JsName, JsStatement?>()
 
         override fun visitLabel(x: JsLabel) {
             val inner = x.statement
@@ -204,6 +204,8 @@ fun JsNode.collectBreakContinueTargets(): Map<JsContinue, JsStatement> {
 
                 is JsFor -> handleLoop(inner, inner.body, x.name)
 
+                is JsSwitch -> handleSwitch(inner, x.name)
+
                 else -> {
                     withBreakAndContinue(x.name, x.statement, null) {
                         accept(inner)
@@ -218,6 +220,14 @@ fun JsNode.collectBreakContinueTargets(): Map<JsContinue, JsStatement> {
 
         override fun visitFor(x: JsFor) = handleLoop(x, x.body, null)
 
+        override fun visit(x: JsSwitch) = handleSwitch(x, null)
+
+        private fun handleSwitch(statement: JsSwitch, label: JsName?) {
+            withBreakAndContinue(label, statement) {
+                statement.cases.forEach { accept(it) }
+            }
+        }
+
         private fun handleLoop(loop: JsStatement, body: JsStatement, label: JsName?) {
             withBreakAndContinue(label, loop, loop) {
                 body.accept(this)
@@ -262,9 +272,7 @@ fun JsNode.collectBreakContinueTargets(): Map<JsContinue, JsStatement> {
             defaultBreakTarget = breakTargetStatement
             if (label != null) {
                 breakTargets[label] = breakTargetStatement
-                if (continueTargetStatement != null) {
+                continueTargets[label] = continueTargetStatement
-                    continueTargets[label] = continueTargetStatement
-                }
             }
             if (continueTargetStatement != null) {
                 defaultContinueTarget = continueTargetStatement
@@ -275,18 +283,8 @@ fun JsNode.collectBreakContinueTargets(): Map<JsContinue, JsStatement> {
             defaultBreakTarget = oldDefaultBreakTarget
             defaultContinueTarget = oldDefaultContinueTarget
             if (label != null) {
-                if (oldBreakTarget == null) {
+                breakTargets[label] = oldBreakTarget
-                    breakTargets.keys -= label
+                continueTargets[label] = oldContinueTarget
-                }
-                else {
-                    breakTargets[label] = oldBreakTarget
-                }
-                if (oldContinueTarget == null) {
-                    continueTargets.keys -= label
-                }
-                else {
-                    continueTargets[label] = oldContinueTarget
-                }
             }
         }
     })

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

@@ -5920,6 +5920,12 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("forContinue.kt")
+            public void testForContinue() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/controlFlow/forContinue.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("forStatement.kt")
             public void testForStatement() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/controlFlow/forStatement.kt");
@@ -5938,6 +5944,12 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 doTest(fileName);
             }
 
+            @TestMetadata("switchLikeWhen.kt")
+            public void testSwitchLikeWhen() throws Exception {
+                String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/controlFlow/switchLikeWhen.kt");
+                doTest(fileName);
+            }
+
             @TestMetadata("throwFromCatch.kt")
             public void testThrowFromCatch() throws Exception {
                 String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/controlFlow/throwFromCatch.kt");

js/js.translator/src/org/jetbrains/kotlin/js/translate/context/UsageTracker.kt

@@ -49,13 +49,10 @@ class UsageTracker(
 
         // local named function
         if (descriptor is FunctionDescriptor && descriptor.visibility == Visibilities.LOCAL) {
-            if (descriptor.isCoroutineLambda) {
+            assert(descriptor.isCoroutineLambda || !descriptor.getName().isSpecial) {
-                captureIfNeed(descriptor)
+                "Function with special name can not be captured, descriptor: $descriptor"
-            }
-            else {
-                assert(!descriptor.getName().isSpecial) { "Function with special name can not be captured, descriptor: $descriptor" }
-                captureIfNeed(descriptor)
             }
+            captureIfNeed(descriptor)
         }
         // local variable
         else if (descriptor is VariableDescriptor && descriptor !is PropertyDescriptor) {