[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.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
 }