[JS IR] Fix state machine control flow

- exception loop unwinding: make sure exception state is reset after try block is finished
 - break/continue of suspended loops

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/ReturnableBlockLowering.kt

@@ -75,7 +75,7 @@ class ReturnableBlockLowering(val context: CommonBackendContext) : FileLoweringP
     }
 }
 
-private class ReturnableBlockTransformer(val context: CommonBackendContext) : IrElementTransformerVoidWithContext() {
+class ReturnableBlockTransformer(val context: CommonBackendContext) : IrElementTransformerVoidWithContext() {
     private var labelCnt = 0
     private val returnMap = mutableMapOf<IrReturnableBlockSymbol, (IrReturn) -> IrExpression>()
 

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/coroutines/JsSuspendFunctionsLowering.kt

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.backend.common.descriptors.synthesizedName
 import org.jetbrains.kotlin.backend.common.ir.isSuspend
 import org.jetbrains.kotlin.backend.common.lower.AbstractSuspendFunctionsLowering
 import org.jetbrains.kotlin.backend.common.lower.FinallyBlocksLowering
+import org.jetbrains.kotlin.backend.common.lower.ReturnableBlockTransformer
 import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
 import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
 import org.jetbrains.kotlin.ir.builders.*
@@ -51,7 +52,11 @@ class JsSuspendFunctionsLowering(ctx: JsIrBackendContext) : AbstractSuspendFunct
         transformingFunction: IrFunction,
         argumentToPropertiesMap: Map<IrValueParameter, IrField>
     ) {
-        val simplifiedFunction = transformingFunction.transform(FinallyBlocksLowering(context, context.dynamicType), null) as IrFunction
+        val returnableBlockTransformer = ReturnableBlockTransformer(context)
+        val finallyBlockTransformer = FinallyBlocksLowering(context, context.dynamicType)
+        val simplifiedFunction =
+            transformingFunction.transform(finallyBlockTransformer, null).transform(returnableBlockTransformer, null) as IrFunction
+
         val originalBody = simplifiedFunction.body as IrBlockBody
 
         val body = IrBlockImpl(

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/coroutines/StateMachineBuilder.kt

@@ -135,6 +135,8 @@ class StateMachineBuilder(
     private val returnableBlockMap = mutableMapOf<IrReturnableBlockSymbol, Pair<SuspendState, IrVariableSymbol?>>()
 
     private val catchBlockStack = mutableListOf(rootExceptionTrap)
+    private val tryStateMap = mutableMapOf<IrExpression, TryState>()
+    private val tryLoopStack = mutableListOf<IrExpression>()
 
     private fun buildExceptionTrapState(): SuspendState {
         val state = SuspendState(unit)
@@ -217,8 +219,12 @@ class StateMachineBuilder(
 
         loopMap[loop] = LoopBounds(loopHeadState, loopExitState)
 
+        tryLoopStack.push(loop)
+
         transformer(loop, loopHeadState, loopExitState)
 
+        tryLoopStack.pop().also { assert(it === loop) }
+
         loopMap.remove(loop)
 
         updateState(loopExitState)
@@ -251,35 +257,6 @@ class StateMachineBuilder(
         doDispatch(exit)
     }
 
-    private fun processReturnableBlock(expression: IrReturnableBlock) {
-
-        if (expression !in suspendableNodes) return super.visitBlock(expression)
-
-        val exitState = SuspendState(unit)
-        val resultVariable = if (hasResultingValue(expression)) {
-            val irVar = tempVar(expression.type, "RETURNABLE_BLOCK")
-            addStatement(irVar)
-            irVar.symbol
-        } else null
-
-        returnableBlockMap[expression.symbol] = Pair(exitState, resultVariable)
-
-        super.visitBlock(expression)
-
-        returnableBlockMap.remove(expression.symbol)
-
-        maybeDoDispatch(exitState)
-
-        updateState(exitState)
-
-        if (resultVariable != null) {
-            addStatement(JsIrBuilder.buildGetValue(resultVariable))
-        }
-    }
-
-    override fun visitBlock(expression: IrBlock) =
-        if (expression is IrReturnableBlock) processReturnableBlock(expression) else super.visitBlock(expression)
-
     private fun implicitCast(value: IrExpression, toType: IrType) = JsIrBuilder.buildImplicitCast(value, toType)
     private fun reinterpretCast(value: IrExpression, toType: IrType) = JsIrBuilder.buildReinterpretCast(value, toType)
 
@@ -347,14 +324,58 @@ class StateMachineBuilder(
 
     override fun visitBreak(jump: IrBreak) {
         val exitState = loopMap[jump.loop]!!.exitState
+        resetExceptionStateIfNeeded(jump.loop)
         doDispatch(exitState)
     }
 
     override fun visitContinue(jump: IrContinue) {
         val headState = loopMap[jump.loop]!!.headState
+        resetExceptionStateIfNeeded(jump.loop)
         doDispatch(headState)
     }
 
+    private fun resetExceptionStateIfNeeded(loop: IrLoop) {
+
+        /**
+         * First find the nearest try statement following after terminating circle
+         * In case we have tryLoopStack like this
+         *
+         * [try 1] <- current exception state
+         * [loop] <- terminating loop
+         * [try 2] <- enclosing try-catch
+         *
+         * our goal to find [try 2]
+         *
+         * Second set exception state to either found try's catch block or root catch
+         */
+
+        var nearestTry: IrExpression? = null
+        var found = false
+        var needReset = false
+        for (e in tryLoopStack.asReversed()) {
+
+            if (e is IrTry) {
+                needReset = !found
+            }
+
+            if (e === loop) {
+                found = true
+            }
+
+            if (found) {
+                if (e is IrTry) {
+                    nearestTry = e
+                    break
+                }
+            }
+        }
+
+        if (needReset) {
+            val tryState = tryStateMap[nearestTry]?.catchState ?: rootExceptionTrap
+            setupExceptionState(tryState)
+        }
+    }
+
     private fun wrap(expression: IrExpression, variable: IrVariableSymbol) =
         JsIrBuilder.buildSetVariable(variable, expression, unit)
 
@@ -581,7 +602,10 @@ class StateMachineBuilder(
         val tryState = buildTryState()
         val enclosingCatch = catchBlockStack.peek()!!
 
+        tryStateMap[aTry] = tryState
+
         catchBlockStack.push(tryState.catchState)
+        tryLoopStack.push(aTry)
 
         val exitState = SuspendState(unit)
 
@@ -608,7 +632,11 @@ class StateMachineBuilder(
         }
         addExceptionEdge()
 
+        tryStateMap.remove(aTry)
+        tryLoopStack.pop().also { assert(it === aTry) }
+
         catchBlockStack.pop()
+
         updateState(tryState.catchState)
 
         setupExceptionState(enclosingCatch)
@@ -662,6 +690,7 @@ class StateMachineBuilder(
         currentState.successors += enclosingCatch
 
         updateState(exitState)
+        setupExceptionState(enclosingCatch)
 
         if (varSymbol != null) {
             addStatement(JsIrBuilder.buildGetValue(varSymbol.symbol))

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/coroutines/SuspendLoweringUtils.kt

@@ -28,12 +28,12 @@ open class SuspendableNodesCollector(private val suspendableNodes: MutableSet<Ir
 
     private var hasSuspendableChildren = false
 
-    protected fun markNode(node: IrElement) {
+    private fun markNode(node: IrElement) {
         suspendableNodes += node
         hasSuspendableChildren = true
     }
 
-    protected fun isSuspendableNode(node: IrElement) = node in suspendableNodes
+    private fun isSuspendableNode(node: IrElement) = node in suspendableNodes
 
     override fun visitElement(element: IrElement) {
         val current = hasSuspendableChildren
@@ -51,9 +51,6 @@ open class SuspendableNodesCollector(private val suspendableNodes: MutableSet<Ir
             markNode(expression)
         }
     }
-}
-
-class SuspendedTerminatorsCollector(suspendableNodes: MutableSet<IrElement>) : SuspendableNodesCollector(suspendableNodes) {
 
     override fun visitBreakContinue(jump: IrBreakContinue) {
         if (isSuspendableNode(jump.loop)) {
@@ -71,12 +68,13 @@ class SuspendedTerminatorsCollector(suspendableNodes: MutableSet<IrElement>) : S
 }
 
 fun collectSuspendableNodes(function: IrBlock): MutableSet<IrElement> {
-
     val suspendableNodes = mutableSetOf<IrElement>()
-    // 1st: mark suspendable loops and tries
-    function.acceptVoid(SuspendableNodesCollector(suspendableNodes))
-    // 2nd: mark inner terminators
-    function.acceptVoid(SuspendedTerminatorsCollector(suspendableNodes))
+    var size: Int
+
+    do {
+        size = suspendableNodes.size
+        function.acceptVoid(SuspendableNodesCollector(suspendableNodes))
+    } while (size != suspendableNodes.size)
 
     return suspendableNodes
 }

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

@@ -0,0 +1,70 @@
+// WITH_RUNTIME
+// WITH_COROUTINES
+// COMMON_COROUTINES_TEST
+import helpers.*
+import COROUTINES_PACKAGE.*
+import COROUTINES_PACKAGE.intrinsics.*
+
+
+class Controller {
+    var result = ""
+
+    suspend fun <T> suspendAndLog(value: T): T = suspendCoroutineUninterceptedOrReturn { c ->
+        result += "suspend($value);"
+        c.resume(value)
+        COROUTINE_SUSPENDED
+    }
+
+    var count = 0
+
+    fun <T> log(value: T) {
+        result += "$value"
+    }
+
+    fun check(): Boolean = count > 1
+
+    suspend fun foo() { count++ }
+}
+
+fun builder(c: suspend Controller.() -> Unit): String {
+    val controller = Controller()
+    c.startCoroutine(controller, handleResultContinuation {
+        controller.result += "return;"
+    })
+    return controller.result
+}
+
+suspend fun Controller.test() {
+    var exception: Throwable? = null
+    suspendLoop@do {
+        log("slh;")
+        foo()
+        regularLoop@do {
+            log("rlh;")
+            if (!check()) {
+                log("rlb;")
+                break@regularLoop
+            }
+
+            log("rlc;")
+
+            if (check()) {
+                log("slb;")
+                break@suspendLoop
+            }
+
+            log("fail1;")
+        } while (false)
+
+        log("slt;")
+    } while (true)
+}
+
+fun box(): String {
+    val res = builder {
+        test()
+    }
+
+    if (res != "slh;rlh;rlb;slt;slh;rlh;rlc;slb;return;") return "FAIL: $res"
+    return "OK"
+}

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

@@ -0,0 +1,65 @@
+// WITH_RUNTIME
+// WITH_COROUTINES
+// COMMON_COROUTINES_TEST
+import helpers.*
+import COROUTINES_PACKAGE.*
+import COROUTINES_PACKAGE.intrinsics.*
+
+class Controller {
+    var result = ""
+
+    var count = 0
+    fun expect(i: Int) {
+        if (++count != i) throw Exception("EXPECTED $i")
+    }
+
+    fun <T> log(value: T) {
+        result += "$value"
+    }
+}
+
+fun builder(c: suspend Controller.() -> Unit): String {
+    val controller = Controller()
+    c.startCoroutine(controller, handleResultContinuation {
+        controller.result += "return;"
+    })
+    return controller.result
+}
+
+suspend fun Controller.makeException(i: Int): Throwable? {
+    expect(i)
+    return Error("CHECK")
+}
+suspend fun Controller.rethrowException(i: Int, t: Throwable?) {
+    expect(i)
+    t?.let { throw it }
+}
+
+suspend fun Controller.test() {
+    var exception: Throwable? = null
+    expect(1)
+    try {
+        exception = makeException(2)
+        log("try(t);")
+    } finally {
+        // Separate method because of KT-32220
+        rethrowException(3, exception)
+        log("FAIL2")
+        exception?.let { throw it }
+    }
+    log("FAIL3")
+}
+
+fun box(): String {
+    val res = builder {
+        try {
+            log("try;")
+            test()
+        } catch (e: Error) {
+            log("catch;")
+        }
+    }
+
+    if (res != "try;try(t);catch;return;") return "FAIL: $res"
+    return "OK"
+}

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

@@ -6772,6 +6772,16 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/doWhileStatement.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("doubleBreak.kt")
+            public void testDoubleBreak_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/doubleBreak.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("doubleBreak.kt")
+            public void testDoubleBreak_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/doubleBreak.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("finallyCatch.kt")
             public void testFinallyCatch_1_2() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/finallyCatch.kt", "kotlin.coroutines.experimental");
@@ -6887,6 +6897,16 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/throwFromCatch.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("throwFromFinally.kt")
+            public void testThrowFromFinally_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/throwFromFinally.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("throwFromFinally.kt")
+            public void testThrowFromFinally_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/throwFromFinally.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("throwInTryWithHandleResult.kt")
             public void testThrowInTryWithHandleResult_1_2() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/throwInTryWithHandleResult.kt", "kotlin.coroutines.experimental");

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

@@ -6772,6 +6772,16 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/doWhileStatement.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("doubleBreak.kt")
+            public void testDoubleBreak_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/doubleBreak.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("doubleBreak.kt")
+            public void testDoubleBreak_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/doubleBreak.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("finallyCatch.kt")
             public void testFinallyCatch_1_2() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/finallyCatch.kt", "kotlin.coroutines.experimental");
@@ -6887,6 +6897,16 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/throwFromCatch.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("throwFromFinally.kt")
+            public void testThrowFromFinally_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/throwFromFinally.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("throwFromFinally.kt")
+            public void testThrowFromFinally_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/throwFromFinally.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("throwInTryWithHandleResult.kt")
             public void testThrowInTryWithHandleResult_1_2() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/throwInTryWithHandleResult.kt", "kotlin.coroutines.experimental");

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

@@ -6287,6 +6287,11 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/doWhileStatement.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("doubleBreak.kt")
+            public void testDoubleBreak_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/doubleBreak.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("finallyCatch.kt")
             public void testFinallyCatch_1_3() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/finallyCatch.kt", "kotlin.coroutines");
@@ -6347,6 +6352,11 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/throwFromCatch.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("throwFromFinally.kt")
+            public void testThrowFromFinally_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/throwFromFinally.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("throwInTryWithHandleResult.kt")
             public void testThrowInTryWithHandleResult_1_3() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/throwInTryWithHandleResult.kt", "kotlin.coroutines");

js/js.tests/test/org/jetbrains/kotlin/js/test/ir/semantics/IrJsCodegenBoxTestGenerated.java

@@ -5367,6 +5367,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/doWhileStatement.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("doubleBreak.kt")
+            public void testDoubleBreak_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/doubleBreak.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("finallyCatch.kt")
             public void testFinallyCatch_1_3() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/finallyCatch.kt", "kotlin.coroutines");
@@ -5427,6 +5432,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/throwFromCatch.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("throwFromFinally.kt")
+            public void testThrowFromFinally_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/throwFromFinally.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("throwInTryWithHandleResult.kt")
             public void testThrowInTryWithHandleResult_1_3() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/throwInTryWithHandleResult.kt", "kotlin.coroutines");

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

@@ -5812,6 +5812,16 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/doWhileStatement.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("doubleBreak.kt")
+            public void testDoubleBreak_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/doubleBreak.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("doubleBreak.kt")
+            public void testDoubleBreak_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/doubleBreak.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("finallyCatch.kt")
             public void testFinallyCatch_1_2() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/finallyCatch.kt", "kotlin.coroutines.experimental");
@@ -5927,6 +5937,16 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/throwFromCatch.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("throwFromFinally.kt")
+            public void testThrowFromFinally_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/throwFromFinally.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("throwFromFinally.kt")
+            public void testThrowFromFinally_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/throwFromFinally.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("throwInTryWithHandleResult.kt")
             public void testThrowInTryWithHandleResult_1_2() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/controlFlow/throwInTryWithHandleResult.kt", "kotlin.coroutines.experimental");