JS: implementing try/finally in coroutines

compiler/testData/codegen/box/coroutines/tryFinallyWithHandleResult.kt

@@ -2,7 +2,7 @@
 var globalResult = ""
 var wasCalled = false
 class Controller {
-    val postponedActions = java.util.ArrayList<() -> Unit>()
+    val postponedActions = mutableListOf<() -> Unit>()
 
     suspend fun suspendWithValue(v: String, x: Continuation<String>) {
         postponedActions.add {
@@ -71,7 +71,7 @@ fun box(): String {
 
     builder(expectException = true) {
         try {
-            suspendWithException(java.lang.RuntimeException("OK"))
+            suspendWithException(RuntimeException("OK"))
         } finally {
             if (suspendWithValue("G") != "G") throw RuntimeException("fail 2")
             wasCalled = true

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

@@ -30,23 +30,32 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
     private val currentStatements: MutableList<JsStatement>
         get() = currentBlock.statements
     val resultFieldName by lazy { scope.declareFreshName("\$result") }
+    val exceptionFieldName by lazy { scope.declareFreshName("\$exception") }
     val stateFieldName by lazy { scope.declareFreshName("\$state") }
     val controllerFieldName by lazy { scope.declareFreshName("\$controller") }
     val exceptionStateName by lazy { scope.declareFreshName("\$exceptionState") }
-    private val breakTargets = mutableMapOf<JsName, JsStatement>()
+    val finallyPathFieldName by lazy { scope.declareFreshName("\$finallyPath") }
-    private val continueTargets = mutableMapOf<JsName, JsStatement>()
+    private val breakTargets = mutableMapOf<JsName, JumpTarget>()
-    private var defaultBreakTarget: JsStatement = JsEmpty
+    private val continueTargets = mutableMapOf<JsName, JumpTarget>()
-    private var defaultContinueTarget: JsStatement = JsEmpty
+    private var defaultBreakTarget: JumpTarget = JumpTarget(JsEmpty, 0)
+    private var defaultContinueTarget: JumpTarget = JumpTarget(JsEmpty, 0)
     private val referencedBlocks = mutableSetOf<CoroutineBlock>()
     private val breakBlocks = mutableMapOf<JsStatement, CoroutineBlock>()
     private val continueBlocks = mutableMapOf<JsStatement, CoroutineBlock>()
     private lateinit var nodesToSplit: Set<JsNode>
     private var currentCatchBlock = globalCatchBlock
+    private val tryStack = mutableListOf(TryBlock(globalCatchBlock, null))
+
+    var hasFinallyBlocks = false
+        get
+        private set
+
+    private val currentTryDepth = tryStack.lastIndex
 
     fun preProcess(node: JsNode) {
         val nodes = mutableSetOf<JsNode>()
 
-        node.accept(object : RecursiveJsVisitor() {
+        val visitor = object : RecursiveJsVisitor() {
             var childrenInSet = false
 
             override fun visitInvocation(invocation: JsInvocation) {
@@ -57,10 +66,33 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
                 }
             }
 
+            override fun visitReturn(x: JsReturn) {
+                super.visitReturn(x)
+                nodes += x
+                childrenInSet = true
+            }
+
+            override fun visitBreak(x: JsBreak) {
+                super.visitBreak(x)
+
+                // It's a simplification
+                // TODO: don't split break and continue statements when possible
+                nodes += x
+                childrenInSet = true
+            }
+
+            override fun visitContinue(x: JsContinue) {
+                super.visitContinue(x)
+                nodes += x
+                childrenInSet = true
+            }
+
             override fun visitElement(node: JsNode) {
                 val oldChildrenInSet = childrenInSet
                 childrenInSet = false
+
                 node.acceptChildren(this)
+
                 if (childrenInSet) {
                     nodes += node
                 }
@@ -68,16 +100,20 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
                     childrenInSet = oldChildrenInSet
                 }
             }
-        })
+        }
+
+        visitor.accept(node)
+        visitor.accept(node)
 
         nodesToSplit = nodes
+
+        currentStatements += exceptionState(currentCatchBlock)
     }
 
     fun postProcess(): List<CoroutineBlock> {
-        if (entryBlock == currentBlock) return listOf(entryBlock)
-
         currentBlock.statements += JsReturn()
-        val orderedBlocks = DFS.topologicalOrder(listOf(entryBlock)) { it.collectTargetBlocks() }
+        val graph = buildBlockGraph()
+        val orderedBlocks = DFS.topologicalOrder(listOf(entryBlock)) { graph[it].orEmpty() }
         val blockIndexes = orderedBlocks.withIndex().associate { (index, block) -> Pair(block, index) }
 
         val blockReplacementVisitor = object : JsVisitorWithContextImpl() {
@@ -95,6 +131,18 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
                     val rhs = program.getNumberLiteral(blockIndexes[exceptionTarget]!!)
                     ctx.replaceMe(JsAstUtils.assignment(lhs, rhs).makeStmt())
                 }
+
+                val finallyPath = x.finallyPath
+                if (finallyPath != null) {
+                    if (finallyPath.isNotEmpty()) {
+                        val lhs = JsNameRef(finallyPathFieldName, JsLiteral.THIS)
+                        val rhs = JsArrayLiteral(finallyPath.map { program.getNumberLiteral(blockIndexes[it]!!) })
+                        ctx.replaceMe(JsAstUtils.assignment(lhs, rhs).makeStmt())
+                    }
+                    else {
+                        ctx.removeMe()
+                    }
+                }
             }
         }
         for (block in orderedBlocks) {
@@ -104,6 +152,35 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
         return orderedBlocks
     }
 
+    private fun buildBlockGraph(): Map<CoroutineBlock, Set<CoroutineBlock>> {
+        // That's a little more than DFS due to need of tracking finally paths
+
+        val visitedBlocks = mutableSetOf<CoroutineBlock>()
+        val graph = mutableMapOf<CoroutineBlock, MutableSet<CoroutineBlock>>()
+
+        fun visitBlock(block: CoroutineBlock) {
+            if (block in visitedBlocks) return
+
+            for (finallyPath in block.collectFinallyPaths()) {
+                for ((finallySource, finallyTarget) in (listOf(block) + finallyPath).zip(finallyPath)) {
+                    if (graph.getOrPut(finallySource) { mutableSetOf() }.add(finallyTarget)) {
+                        visitedBlocks -= finallySource
+                    }
+                }
+            }
+
+            visitedBlocks += block
+
+            val successors = graph.getOrPut(block) { mutableSetOf() }
+            successors += block.collectTargetBlocks()
+            successors.forEach(::visitBlock)
+        }
+
+        visitBlock(entryBlock)
+
+        return graph
+    }
+
     override fun visitBlock(x: JsBlock) = splitIfNecessary(x) {
         for (statement in x.statements) {
             statement.accept(this)
@@ -146,7 +223,9 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
 
             else -> splitIfNecessary(x) {
                 val successor = CoroutineBlock()
-                accept(inner)
+                withBreakAndContinue(x.name, x.statement, successor, null) {
+                    accept(inner)
+                }
                 if (successor in referencedBlocks) {
                     currentBlock.statements += stateAndJump(successor)
                     currentBlock = successor
@@ -237,68 +316,121 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
 
     override fun visitBreak(x: JsBreak) {
         val targetLabel = x.label?.name
-        val targetStatement = if (targetLabel == null) {
+        val (targetStatement, targetTryDepth) = if (targetLabel == null) {
             defaultBreakTarget
         }
         else {
             breakTargets[targetLabel]!!
         }
 
-        if (targetStatement in nodesToSplit) {
+        val targetBlock = breakBlocks[targetStatement]!!
-            val targetBlock = breakBlocks[targetStatement]!!
+        referencedBlocks += targetBlock
-            referencedBlocks += targetBlock
+        jumpWithFinally(targetTryDepth, targetBlock)
-            currentStatements += stateAndJump(targetBlock)
+        currentStatements += jump()
-        }
-        else {
-            currentStatements += x
-        }
     }
 
     override fun visitContinue(x: JsContinue) {
         val targetLabel = x.label?.name
-        val targetStatement = if (targetLabel == null) {
+        val (targetStatement, targetTryDepth) = if (targetLabel == null) {
             defaultContinueTarget
         }
         else {
             continueTargets[targetLabel]!!
         }
 
-        if (targetStatement in nodesToSplit) {
+        val targetBlock = continueBlocks[targetStatement]!!
-            val targetBlock = continueBlocks[targetStatement]!!
+        referencedBlocks += targetBlock
-            referencedBlocks += targetBlock
+
-            currentStatements += stateAndJump(targetBlock)
+        jumpWithFinally(targetTryDepth, targetBlock)
-        }
+        currentStatements += jump()
-        else {
+    }
-            currentStatements += x
+
+    private fun jumpWithFinally(targetTryDepth: Int, successor: CoroutineBlock) {
+        if (targetTryDepth == tryStack.size) return
+
+        val tryBlock = tryStack[targetTryDepth]
+        currentStatements += exceptionState(tryBlock.catchBlock)
+
+        val relativeFinallyPath = relativeFinallyPath(targetTryDepth)
+        val fullPath = relativeFinallyPath + successor
+        if (fullPath.size > 1) {
+            currentStatements += updateFinallyPath(fullPath.drop(1))
         }
+        currentStatements += state(fullPath[0])
     }
 
     override fun visitTry(x: JsTry) = splitIfNecessary(x) {
-        val catch = x.catches.firstOrNull()
+        val catchNode = x.catches.firstOrNull()
+        val finallyNode = x.finallyBlock
         val successor = CoroutineBlock()
 
         val catchBlock = CoroutineBlock()
+        val finallyBlock = CoroutineBlock()
+
+        val tryBlock = TryBlock(catchBlock, if (finallyNode != null) finallyBlock else null)
+        tryStack += tryBlock
+
         val oldCatchBlock = currentCatchBlock
-        if (catch != null) {
+        currentCatchBlock = catchBlock
-            currentCatchBlock = catchBlock
+        currentStatements += exceptionState(catchBlock)
-        }
 
         x.tryBlock.statements.forEach { it.accept(this) }
-        currentBlock.statements += stateAndJump(successor)
+        currentStatements += exceptionState(oldCatchBlock)
+        if (finallyNode != null) {
+            currentStatements += updateFinallyPath(listOf(successor))
+            currentStatements += stateAndJump(finallyBlock)
+        }
+        else {
+            currentStatements += stateAndJump(successor)
+        }
 
         currentCatchBlock = oldCatchBlock
 
-        if (catch != null) {
+        // Handle catch node
-            currentBlock = catchBlock
+        currentBlock = catchBlock
-            currentBlock.statements += JsAstUtils.newVar(catch.parameter.name, JsNameRef(resultFieldName, JsLiteral.THIS))
 
-            catch.body.statements.forEach { it.accept(this) }
+        if (finallyNode != null) {
-            currentBlock.statements += stateAndJump(successor)
+            currentStatements += updateFinallyPath(listOf(oldCatchBlock))
+            currentStatements += exceptionState(finallyBlock)
+        }
+        else {
+            currentStatements += exceptionState(oldCatchBlock)
         }
 
+        if (catchNode != null) {
+            currentStatements += JsAstUtils.newVar(catchNode.parameter.name, JsNameRef(exceptionFieldName, JsLiteral.THIS))
+            catchNode.body.statements.forEach { it.accept(this) }
+
+            if (finallyNode == null) {
+                currentStatements += stateAndJump(successor)
+            }
+            else {
+                currentStatements += updateFinallyPath(listOf(successor))
+                currentStatements += stateAndJump(finallyBlock)
+            }
+        }
+
+        // Handle finally node
+        if (finallyNode != null) {
+            currentBlock = finallyBlock
+            finallyNode.statements.forEach { it.accept(this) }
+            generateFinallyExit()
+            hasFinallyBlocks = true
+        }
+
+        tryStack.removeAt(tryStack.lastIndex)
+
         currentBlock = successor
     }
 
+    private fun generateFinallyExit() {
+        val finallyPathRef = JsNameRef(finallyPathFieldName, JsLiteral.THIS)
+        val stateRef = JsNameRef(stateFieldName, JsLiteral.THIS)
+        val nextState = JsInvocation(JsNameRef("shift", finallyPathRef))
+        currentStatements += JsAstUtils.assignment(stateRef, nextState).makeStmt()
+        currentStatements += jump()
+    }
+
     override fun visitExpressionStatement(x: JsExpressionStatement) {
         val expression = x.expression
         if (expression is JsInvocation && expression.isSuspend) {
@@ -323,11 +455,26 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
     }
 
     override fun visitReturn(x: JsReturn) {
+        val returnBlock = CoroutineBlock()
+        val isInFinally = hasEnclosingFinallyBlock()
+        if (isInFinally) {
+            jumpWithFinally(0, returnBlock)
+        }
+
         val returnExpression = x.expression
         if (returnExpression != null) {
             x.expression = handleExpression(returnExpression)
         }
-        currentStatements += x
+
+        if (isInFinally) {
+            currentStatements += x.expression?.makeStmt().singletonOrEmptyList()
+            currentStatements += jump()
+            currentBlock = returnBlock
+            currentStatements += JsReturn()
+        }
+        else {
+            currentStatements += x
+        }
     }
 
     override fun visitThrow(x: JsThrow) {
@@ -373,13 +520,10 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
     }
 
     private fun state(target: CoroutineBlock): List<JsStatement> {
-        val nextPlaceholder = JsDebugger()
+        val placeholder = JsDebugger()
-        nextPlaceholder.targetBlock = target
+        placeholder.targetBlock = target
 
-        val exceptionPlaceholder = JsDebugger()
+        return listOf(placeholder)
-        exceptionPlaceholder.targetExceptionBlock = currentCatchBlock
-
-        return listOf(nextPlaceholder, exceptionPlaceholder)
     }
 
     private fun jump() = JsContinue()
@@ -388,6 +532,19 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
         return state(target) + jump()
     }
 
+    private fun exceptionState(target: CoroutineBlock): List<JsStatement> {
+        val placeholder = JsDebugger()
+        placeholder.targetExceptionBlock = target
+
+        return listOf(placeholder)
+    }
+
+    private fun updateFinallyPath(path: List<CoroutineBlock>): List<JsStatement> {
+        val placeholder = JsDebugger()
+        placeholder.finallyPath = path
+        return listOf(placeholder)
+    }
+
     private inline fun splitIfNecessary(statement: JsStatement, action: () -> Unit) {
         if (statement in nodesToSplit) {
             action()
@@ -418,15 +575,15 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
             Pair(null, null)
         }
 
-        defaultBreakTarget = statement
+        defaultBreakTarget = JumpTarget(statement, currentTryDepth)
         if (label != null) {
-            breakTargets[label] = statement
+            breakTargets[label] = JumpTarget(statement, currentTryDepth)
             if (continueBlock != null) {
-                continueTargets[label] = statement
+                continueTargets[label] = JumpTarget(statement, currentTryDepth)
             }
         }
         if (continueBlock != null) {
-            defaultContinueTarget = statement
+            defaultContinueTarget = JumpTarget(statement, currentTryDepth)
         }
 
         action()
@@ -448,8 +605,16 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
             }
         }
     }
+
+    private fun relativeFinallyPath(targetTryDepth: Int) = tryStack.subList(targetTryDepth, tryStack.size).mapNotNull { it.finallyBlock }
+
+    private fun hasEnclosingFinallyBlock() = tryStack.any { it.finallyBlock != null }
 }
 
+private data class JumpTarget(val statement: JsStatement, val tryDepth: Int)
+
+private class TryBlock(val catchBlock: CoroutineBlock, val finallyBlock: CoroutineBlock?)
+
 private fun CoroutineBlock.collectTargetBlocks(): Set<CoroutineBlock> {
     val targetBlocks = mutableSetOf<CoroutineBlock>()
     jsBlock.accept(object : RecursiveJsVisitor() {
@@ -460,5 +625,16 @@ private fun CoroutineBlock.collectTargetBlocks(): Set<CoroutineBlock> {
     return targetBlocks
 }
 
+private fun CoroutineBlock.collectFinallyPaths(): List<List<CoroutineBlock>> {
+    val finallyPaths = mutableListOf<List<CoroutineBlock>>()
+    jsBlock.accept(object : RecursiveJsVisitor() {
+        override fun visitDebugger(x: JsDebugger) {
+            x.finallyPath?.let { finallyPaths += it }
+        }
+    })
+    return finallyPaths
+}
+
 private var JsDebugger.targetBlock: CoroutineBlock? by MetadataProperty(default = null)
 private var JsDebugger.targetExceptionBlock: CoroutineBlock? by MetadataProperty(default = null)
+private var JsDebugger.finallyPath: List<CoroutineBlock>? by MetadataProperty(default = null)

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

@@ -98,6 +98,10 @@ class CoroutineFunctionTransformer(
 
         constructor.body.statements.run {
             assign(bodyTransformer.stateFieldName, program.getNumberLiteral(0))
+            assign(bodyTransformer.exceptionStateName, program.getNumberLiteral(0))
+            if (bodyTransformer.hasFinallyBlocks) {
+                assign(bodyTransformer.finallyPathFieldName, JsLiteral.NULL)
+            }
             assign(bodyTransformer.controllerFieldName, controllerName.makeRef())
             for (localVariable in localVariables) {
                 val value = if (localVariable !in parameterNames) JsLiteral.NULL else localVariable.makeRef()
@@ -185,13 +189,24 @@ class CoroutineFunctionTransformer(
             throwName: JsName?,
             exceptionName: JsName
     ): List<JsStatement> {
+        val indexOfGlobalCatch = blocks.indexOf(transformer.globalCatchBlock)
+        val stateRef = JsNameRef(transformer.stateFieldName, JsLiteral.THIS)
+
+        val isFromGlobalCatch = JsAstUtils.equality(stateRef, program.getNumberLiteral(indexOfGlobalCatch))
         val catch = JsCatch(functionWithBody.scope, "e")
+        val continueWithException = JsBlock(
+                JsAstUtils.assignment(stateRef.deepCopy(), JsNameRef(transformer.exceptionStateName, JsLiteral.THIS)).makeStmt(),
+                JsAstUtils.assignment(JsNameRef(transformer.exceptionFieldName, JsLiteral.THIS), catch.parameter.name.makeRef()).makeStmt()
+        )
+        catch.body = JsBlock(JsIf(isFromGlobalCatch, JsThrow(catch.parameter.name.makeRef()), continueWithException))
+
+        val throwResultRef = JsNameRef(transformer.exceptionFieldName, JsLiteral.THIS)
         if (throwName != null) {
             val throwMethodRef = JsNameRef(throwName, JsNameRef(transformer.controllerFieldName, JsLiteral.THIS))
-            catch.body = JsBlock(JsReturn(JsInvocation(throwMethodRef, catch.parameter.name.makeRef())))
+            transformer.globalCatchBlock.statements += JsReturn(JsInvocation(throwMethodRef.deepCopy(), throwResultRef))
-
+        }
-            val resultRef = JsNameRef(transformer.resultFieldName, JsLiteral.THIS)
+        else {
-            transformer.globalCatchBlock.statements += JsReturn(JsInvocation(throwMethodRef.deepCopy(), resultRef))
+            transformer.globalCatchBlock.statements += JsThrow(throwResultRef)
         }
 
         val cases = blocks.withIndex().map { (index, block) ->
@@ -200,20 +215,18 @@ class CoroutineFunctionTransformer(
                 statements += block.statements
             }
         }
-        val stateRef = JsNameRef(transformer.stateFieldName, JsLiteral.THIS)
+        val switchStatement = JsSwitch(stateRef.deepCopy(), cases)
-        val switchStatement = JsSwitch(stateRef, cases)
+        val loop = JsDoWhile(JsLiteral.TRUE, JsTry(JsBlock(switchStatement), catch, null))
-        val loop = JsDoWhile(JsLiteral.TRUE, switchStatement)
 
         val testExceptionPassed = JsAstUtils.notOptimized(
                 JsAstUtils.typeOfIs(exceptionName.makeRef(), program.getStringLiteral("undefined")))
         val stateToException = JsAstUtils.assignment(
                 JsNameRef(transformer.stateFieldName, JsLiteral.THIS),
                 JsNameRef(transformer.exceptionStateName, JsLiteral.THIS))
-        val exceptionToResult = JsAstUtils.assignment(JsNameRef(transformer.resultFieldName, JsLiteral.THIS), exceptionName.makeRef())
+        val exceptionToResult = JsAstUtils.assignment(JsNameRef(transformer.exceptionFieldName, JsLiteral.THIS), exceptionName.makeRef())
         val throwExceptionIfNeeded = JsIf(testExceptionPassed, JsBlock(stateToException.makeStmt(), exceptionToResult.makeStmt()))
 
-        val resultBlock = JsBlock(throwExceptionIfNeeded, loop)
+        return listOf(throwExceptionIfNeeded, loop)
-        return if (throwName != null) listOf(JsTry(resultBlock, catch, null)) else resultBlock.statements
     }
 
     private fun JsBlock.replaceHandleResult(transformer: CoroutineBodyTransformer) {