JS: refactor coroutine transformer

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

@@ -18,29 +18,25 @@ package org.jetbrains.kotlin.js.coroutine
 
 import com.google.dart.compiler.backend.js.ast.*
 import com.google.dart.compiler.backend.js.ast.metadata.*
+import org.jetbrains.kotlin.js.inline.util.collectBreakContinueTargets
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
 import org.jetbrains.kotlin.utils.DFS
 import org.jetbrains.kotlin.utils.singletonOrEmptyList
 
-class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val throwFunctionName: JsName?) : RecursiveJsVisitor() {
+class CoroutineBodyTransformer(
-    val entryBlock = CoroutineBlock()
+        private val program: JsProgram,
-    val globalCatchBlock = CoroutineBlock()
+        private val context: CoroutineTransformationContext,
+        private val throwFunctionName: JsName?
+) : RecursiveJsVisitor() {
+    private val entryBlock = context.entryBlock
+    private val globalCatchBlock = context.globalCatchBlock
     private var currentBlock = entryBlock
     private val currentStatements: MutableList<JsStatement>
         get() = currentBlock.statements
-    val resultFieldName by lazy { scope.declareFreshName("\$result") }
+    private val breakContinueTargetStatements = mutableMapOf<JsContinue, JsStatement>()
-    val exceptionFieldName by lazy { scope.declareFreshName("\$exception") }
+    private val breakTargets = mutableMapOf<JsStatement, JumpTarget>()
-    val stateFieldName by lazy { scope.declareFreshName("\$state") }
+    private val continueTargets = mutableMapOf<JsStatement, JumpTarget>()
-    val controllerFieldName by lazy { scope.declareFreshName("\$controller") }
-    val exceptionStateName by lazy { scope.declareFreshName("\$exceptionState") }
-    val finallyPathFieldName by lazy { scope.declareFreshName("\$finallyPath") }
-    private val breakTargets = mutableMapOf<JsName, JumpTarget>()
-    private val continueTargets = mutableMapOf<JsName, JumpTarget>()
-    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))
@@ -54,140 +50,19 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
         get() = tryStack.lastIndex
 
     fun preProcess(node: JsNode) {
-        val nodes = mutableSetOf<JsNode>()
+        breakContinueTargetStatements += node.collectBreakContinueTargets()
-
+        nodesToSplit = node.collectNodesToSplit()
-        val visitor = object : RecursiveJsVisitor() {
-            var childrenInSet = false
-
-            override fun visitInvocation(invocation: JsInvocation) {
-                super.visitInvocation(invocation)
-                if (invocation.isSuspend) {
-                    nodes += invocation
-                    childrenInSet = true
-                }
-            }
-
-            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
-                }
-                else {
-                    childrenInSet = oldChildrenInSet
-                }
-            }
-        }
-
-        visitor.accept(node)
-        visitor.accept(node)
-
-        nodesToSplit = nodes
-
         currentStatements += exceptionState(currentCatchBlock)
     }
 
     fun postProcess(): List<CoroutineBlock> {
         currentBlock.statements += JsReturn()
-        val graph = buildBlockGraph()
+        val graph = entryBlock.buildGraph()
         val orderedBlocks = DFS.topologicalOrder(listOf(entryBlock)) { graph[it].orEmpty() }
-        val blockIndexes = orderedBlocks.withIndex().associate { (index, block) -> Pair(block, index) }
+        orderedBlocks.replaceCoroutineFlowStatements(context, program)
-
-        val blockReplacementVisitor = object : JsVisitorWithContextImpl() {
-            override fun endVisit(x: JsDebugger, ctx: JsContext<in JsStatement>) {
-                val target = x.targetBlock
-                if (target != null) {
-                    val lhs = JsNameRef(stateFieldName, JsLiteral.THIS)
-                    val rhs = program.getNumberLiteral(blockIndexes[target]!!)
-                    ctx.replaceMe(JsExpressionStatement(JsAstUtils.assignment(lhs, rhs)).apply {
-                        targetBlock = true
-                    })
-                }
-
-                val exceptionTarget = x.targetExceptionBlock
-                if (exceptionTarget != null) {
-                    val lhs = JsNameRef(exceptionStateName, JsLiteral.THIS)
-                    val rhs = program.getNumberLiteral(blockIndexes[exceptionTarget]!!)
-                    ctx.replaceMe(JsExpressionStatement(JsAstUtils.assignment(lhs, rhs)).apply {
-                        targetExceptionBlock = true
-                    })
-                }
-
-                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(JsExpressionStatement(JsAstUtils.assignment(lhs, rhs)).apply {
-                            this.finallyPath = true
-                        })
-                    }
-                    else {
-                        ctx.removeMe()
-                    }
-                }
-            }
-        }
-        for (block in orderedBlocks) {
-            blockReplacementVisitor.accept(block.jsBlock)
-        }
-
         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)
@@ -220,15 +95,13 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
     override fun visitLabel(x: JsLabel) {
         val inner = x.statement
         when (inner) {
-            is JsDoWhile -> handleDoWhile(inner, x.name)
+            is JsWhile,
-
+            is JsDoWhile,
-            is JsWhile -> handleWhile(inner, x.name)
+            is JsFor -> inner.accept(this)
-
-            is JsFor -> handleFor(inner, x.name)
 
             else -> splitIfNecessary(x) {
                 val successor = CoroutineBlock()
-                withBreakAndContinue(x.name, x.statement, successor, null) {
+                withBreakAndContinue(x.statement, successor, null) {
                     accept(inner)
                 }
                 if (successor in referencedBlocks) {
@@ -239,15 +112,13 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
         }
     }
 
-    override fun visitWhile(x: JsWhile) = handleWhile(x, null)
+    override fun visitWhile(x: JsWhile) = splitIfNecessary(x) {
-
-    private fun handleWhile(x: JsWhile, label: JsName?) = splitIfNecessary(x) {
         val predecessor = currentBlock
         val successor = CoroutineBlock()
 
         val bodyEntryBlock = CoroutineBlock()
         currentBlock = bodyEntryBlock
-        withBreakAndContinue(label, x, successor, bodyEntryBlock) {
+        withBreakAndContinue(x, successor, bodyEntryBlock) {
             x.body.accept(this)
         }
 
@@ -260,15 +131,13 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
         currentBlock = successor
     }
 
-    override fun visitDoWhile(x: JsDoWhile) = handleDoWhile(x, null)
+    override fun visitDoWhile(x: JsDoWhile) = splitIfNecessary(x) {
-
-    private fun handleDoWhile(x: JsDoWhile, label: JsName?) = splitIfNecessary(x) {
         val predecessor = currentBlock
         val successor = CoroutineBlock()
 
         val bodyEntryBlock = CoroutineBlock()
         currentBlock = bodyEntryBlock
-        withBreakAndContinue(label, x, successor, bodyEntryBlock) {
+        withBreakAndContinue(x, successor, bodyEntryBlock) {
             x.body.accept(this)
         }
 
@@ -281,9 +150,7 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
         currentBlock = successor
     }
 
-    override fun visitFor(x: JsFor) = handleFor(x, null)
+    override fun visitFor(x: JsFor) = splitIfNecessary(x) {
-
-    private fun handleFor(x: JsFor, label: JsName?) = splitIfNecessary(x) {
         x.initExpression?.let {
             JsExpressionStatement(it).accept(this)
         }
@@ -299,7 +166,7 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
 
         val bodyEntryBlock = CoroutineBlock()
         currentBlock = bodyEntryBlock
-        withBreakAndContinue(label, x, successor, predecessor) {
+        withBreakAndContinue(x, successor, predecessor) {
             x.body.accept(this)
         }
         val bodyExitBlock = currentBlock
@@ -320,32 +187,17 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
     }
 
     override fun visitBreak(x: JsBreak) {
-        val targetLabel = x.label?.name
+        val targetStatement = breakContinueTargetStatements[x]!!
-        val (targetStatement, targetTryDepth) = if (targetLabel == null) {
+        val (targetBlock, targetTryDepth) = breakTargets[targetStatement]!!
-            defaultBreakTarget
-        }
-        else {
-            breakTargets[targetLabel]!!
-        }
-
-        val targetBlock = breakBlocks[targetStatement]!!
         referencedBlocks += targetBlock
         jumpWithFinally(targetTryDepth + 1, targetBlock)
         currentStatements += jump()
     }
 
     override fun visitContinue(x: JsContinue) {
-        val targetLabel = x.label?.name
+        val targetStatement = breakContinueTargetStatements[x]!!
-        val (targetStatement, targetTryDepth) = if (targetLabel == null) {
+        val (targetBlock, targetTryDepth) = continueTargets[targetStatement]!!
-            defaultContinueTarget
-        }
-        else {
-            continueTargets[targetLabel]!!
-        }
-
-        val targetBlock = continueBlocks[targetStatement]!!
         referencedBlocks += targetBlock
-
         jumpWithFinally(targetTryDepth + 1, targetBlock)
         currentStatements += jump()
     }
@@ -403,7 +255,7 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
         }
 
         if (catchNode != null) {
-            currentStatements += JsAstUtils.newVar(catchNode.parameter.name, JsNameRef(exceptionFieldName, JsLiteral.THIS))
+            currentStatements += JsAstUtils.newVar(catchNode.parameter.name, JsNameRef(context.exceptionFieldName, JsLiteral.THIS))
             catchNode.body.statements.forEach { it.accept(this) }
 
             if (finallyNode == null) {
@@ -429,8 +281,8 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
     }
 
     private fun generateFinallyExit() {
-        val finallyPathRef = JsNameRef(finallyPathFieldName, JsLiteral.THIS)
+        val finallyPathRef = JsNameRef(context.finallyPathFieldName, JsLiteral.THIS)
-        val stateRef = JsNameRef(stateFieldName, JsLiteral.THIS)
+        val stateRef = JsNameRef(context.stateFieldName, JsLiteral.THIS)
         val nextState = JsInvocation(JsNameRef("shift", finallyPathRef))
         currentStatements += JsAstUtils.assignment(stateRef, nextState).makeStmt()
         currentStatements += jump()
@@ -469,7 +321,7 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
 
     override fun visitThrow(x: JsThrow) {
         if (throwFunctionName != null) {
-            val methodRef = JsNameRef(throwFunctionName, JsNameRef(controllerFieldName, JsLiteral.THIS))
+            val methodRef = JsNameRef(throwFunctionName, JsNameRef(context.controllerFieldName, JsLiteral.THIS))
             val invocation = JsInvocation(methodRef, x.expression).apply {
                 source = x.source
             }
@@ -547,55 +399,20 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
     }
 
     private fun withBreakAndContinue(
-            label: JsName?,
             statement: JsStatement,
             breakBlock: CoroutineBlock,
             continueBlock: CoroutineBlock? = null,
             action: () -> Unit
     ) {
-        breakBlocks[statement] = breakBlock
+        breakTargets[statement] = JumpTarget(breakBlock, currentTryDepth)
         if (continueBlock != null) {
-            continueBlocks[statement] = continueBlock
+            continueTargets[statement] = JumpTarget(continueBlock, currentTryDepth)
-        }
-
-        val oldDefaultBreakTarget = defaultBreakTarget
-        val oldDefaultContinueTarget = defaultContinueTarget
-        val (oldBreakTarget, oldContinueTarget) = if (label != null) {
-            Pair(breakTargets[label], continueTargets[label])
-        }
-        else {
-            Pair(null, null)
-        }
-
-        defaultBreakTarget = JumpTarget(statement, currentTryDepth)
-        if (label != null) {
-            breakTargets[label] = JumpTarget(statement, currentTryDepth)
-            if (continueBlock != null) {
-                continueTargets[label] = JumpTarget(statement, currentTryDepth)
-            }
-        }
-        if (continueBlock != null) {
-            defaultContinueTarget = JumpTarget(statement, currentTryDepth)
         }
 
         action()
 
-        defaultBreakTarget = oldDefaultBreakTarget
+        breakTargets.keys -= statement
-        defaultContinueTarget = oldDefaultContinueTarget
+        continueTargets.keys -= statement
-        if (label != null) {
-            if (oldBreakTarget == null) {
-                breakTargets.keys -= label
-            }
-            else {
-                breakTargets[label] = oldBreakTarget
-            }
-            if (oldContinueTarget == null) {
-                continueTargets.keys -= label
-            }
-            else {
-                continueTargets[label] = oldContinueTarget
-            }
-        }
     }
 
     private fun relativeFinallyPath(targetTryDepth: Int) = tryStack
@@ -604,28 +421,8 @@ class CoroutineBodyTransformer(val program: JsProgram, val scope: JsScope, val t
             .reversed()
 
     private fun hasEnclosingFinallyBlock() = tryStack.any { it.finallyBlock != null }
-}
+
-
+    private data class JumpTarget(val block: CoroutineBlock, val tryDepth: Int)
-private data class JumpTarget(val statement: JsStatement, val tryDepth: Int)
+
-
+    private class TryBlock(val catchBlock: CoroutineBlock, val finallyBlock: CoroutineBlock?)
-private class TryBlock(val catchBlock: CoroutineBlock, val finallyBlock: CoroutineBlock?)
-
-private fun CoroutineBlock.collectTargetBlocks(): Set<CoroutineBlock> {
-    val targetBlocks = mutableSetOf<CoroutineBlock>()
-    jsBlock.accept(object : RecursiveJsVisitor() {
-        override fun visitDebugger(x: JsDebugger) {
-            targetBlocks += x.targetExceptionBlock.singletonOrEmptyList() + x.targetBlock.singletonOrEmptyList()
-        }
-    })
-    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
 }

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

@@ -48,18 +48,18 @@ class CoroutineFunctionTransformer(
             function.scope.declareName(throwId)
         }
 
-        val bodyTransformer = CoroutineBodyTransformer(program, function.scope, throwName)
+        val context = CoroutineTransformationContext(function.scope)
+        val bodyTransformer = CoroutineBodyTransformer(program, context, throwName)
         bodyTransformer.preProcess(body)
         body.statements.forEach { it.accept(bodyTransformer) }
         val coroutineBlocks = bodyTransformer.postProcess()
 
-        for (coroutineBlock in coroutineBlocks) {
+        coroutineBlocks.forEach { it.jsBlock.collectAdditionalLocalVariables() }
-            coroutineBlock.jsBlock.replaceLocalVariables(function.scope, bodyTransformer)
+        coroutineBlocks.forEach { it.jsBlock.replaceLocalVariables(function.scope, context, localVariables) }
-        }
 
         val additionalStatements = mutableListOf<JsStatement>()
-        val resumeName = generateDoResume(coroutineBlocks, bodyTransformer, additionalStatements, throwName)
+        val resumeName = generateDoResume(coroutineBlocks, context, additionalStatements, throwName)
-        generateContinuationConstructor(bodyTransformer, additionalStatements)
+        generateContinuationConstructor(context, additionalStatements, bodyTransformer.hasFinallyBlocks)
         generateContinuationMethods(resumeName, additionalStatements)
 
         generateCoroutineInstantiation()
@@ -67,7 +67,11 @@ class CoroutineFunctionTransformer(
         return additionalStatements
     }
 
-    private fun generateContinuationConstructor(bodyTransformer: CoroutineBodyTransformer, statements: MutableList<JsStatement>) {
+    private fun generateContinuationConstructor(
+            context: CoroutineTransformationContext,
+            statements: MutableList<JsStatement>,
+            hasFinallyBlocks: Boolean
+    ) {
         val constructor = JsFunction(function.scope.parent, JsBlock(), "Continuation")
         constructor.name = className
         constructor.parameters += function.parameters.map { JsParameter(it.name) }
@@ -81,12 +85,12 @@ class CoroutineFunctionTransformer(
         val parameterNames = (function.parameters.map { it.name } + innerFunction?.parameters?.map { it.name }.orEmpty()).toSet()
 
         constructor.body.statements.run {
-            assign(bodyTransformer.stateFieldName, program.getNumberLiteral(0))
+            assign(context.stateFieldName, program.getNumberLiteral(0))
-            assign(bodyTransformer.exceptionStateName, program.getNumberLiteral(0))
+            assign(context.exceptionStateName, program.getNumberLiteral(0))
-            if (bodyTransformer.hasFinallyBlocks) {
+            if (hasFinallyBlocks) {
-                assign(bodyTransformer.finallyPathFieldName, JsLiteral.NULL)
+                assign(context.finallyPathFieldName, JsLiteral.NULL)
             }
-            assign(bodyTransformer.controllerFieldName, controllerName.makeRef())
+            assign(context.controllerFieldName, controllerName.makeRef())
             for (localVariable in localVariables) {
                 val value = if (localVariable !in parameterNames) JsLiteral.NULL else localVariable.makeRef()
                 assign(function.scope.getFieldName(localVariable), value)
@@ -142,7 +146,7 @@ class CoroutineFunctionTransformer(
 
     private fun generateDoResume(
             coroutineBlocks: List<CoroutineBlock>,
-            bodyTransformer: CoroutineBodyTransformer,
+            context: CoroutineTransformationContext,
             statements: MutableList<JsStatement>,
             throwName: JsName?
     ): JsName {
@@ -151,11 +155,11 @@ class CoroutineFunctionTransformer(
         val resumeException = resumeFunction.scope.declareFreshName("exception")
         resumeFunction.parameters += listOf(JsParameter(resumeParameter), JsParameter(resumeException))
 
-        val coroutineBody = generateCoroutineBody(bodyTransformer, coroutineBlocks, throwName, resumeException)
+        val coroutineBody = generateCoroutineBody(context, coroutineBlocks, throwName, resumeException)
         functionWithBody.body.statements.clear()
 
         resumeFunction.body.statements.apply {
-            assign(bodyTransformer.resultFieldName, resumeParameter.makeRef())
+            assign(context.resultFieldName, resumeParameter.makeRef())
             this += coroutineBody
         }
 
@@ -182,29 +186,29 @@ class CoroutineFunctionTransformer(
     }
 
     private fun generateCoroutineBody(
-            transformer: CoroutineBodyTransformer,
+            context: CoroutineTransformationContext,
             blocks: List<CoroutineBlock>,
             throwName: JsName?,
             exceptionName: JsName
     ): List<JsStatement> {
-        val indexOfGlobalCatch = blocks.indexOf(transformer.globalCatchBlock)
+        val indexOfGlobalCatch = blocks.indexOf(context.globalCatchBlock)
-        val stateRef = JsNameRef(transformer.stateFieldName, JsLiteral.THIS)
+        val stateRef = JsNameRef(context.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(stateRef.deepCopy(), JsNameRef(context.exceptionStateName, JsLiteral.THIS)).makeStmt(),
-                JsAstUtils.assignment(JsNameRef(transformer.exceptionFieldName, JsLiteral.THIS), catch.parameter.name.makeRef()).makeStmt()
+                JsAstUtils.assignment(JsNameRef(context.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)
+        val throwResultRef = JsNameRef(context.exceptionFieldName, JsLiteral.THIS)
         if (throwName != null) {
-            val throwMethodRef = JsNameRef(throwName, JsNameRef(transformer.controllerFieldName, JsLiteral.THIS))
+            val throwMethodRef = JsNameRef(throwName, JsNameRef(context.controllerFieldName, JsLiteral.THIS))
-            transformer.globalCatchBlock.statements += JsReturn(JsInvocation(throwMethodRef.deepCopy(), throwResultRef))
+            context.globalCatchBlock.statements += JsReturn(JsInvocation(throwMethodRef.deepCopy(), throwResultRef))
         }
         else {
-            transformer.globalCatchBlock.statements += JsThrow(throwResultRef)
+            context.globalCatchBlock.statements += JsThrow(throwResultRef)
         }
 
         val cases = blocks.withIndex().map { (index, block) ->
@@ -219,96 +223,30 @@ class CoroutineFunctionTransformer(
         val testExceptionPassed = JsAstUtils.notOptimized(
                 JsAstUtils.typeOfIs(exceptionName.makeRef(), program.getStringLiteral("undefined")))
         val stateToException = JsAstUtils.assignment(
-                JsNameRef(transformer.stateFieldName, JsLiteral.THIS),
+                JsNameRef(context.stateFieldName, JsLiteral.THIS),
-                JsNameRef(transformer.exceptionStateName, JsLiteral.THIS))
+                JsNameRef(context.exceptionStateName, JsLiteral.THIS))
-        val exceptionToResult = JsAstUtils.assignment(JsNameRef(transformer.exceptionFieldName, JsLiteral.THIS), exceptionName.makeRef())
+        val exceptionToResult = JsAstUtils.assignment(JsNameRef(context.exceptionFieldName, JsLiteral.THIS), exceptionName.makeRef())
         val throwExceptionIfNeeded = JsIf(testExceptionPassed, JsBlock(stateToException.makeStmt(), exceptionToResult.makeStmt()))
 
         return listOf(throwExceptionIfNeeded, loop)
     }
 
-    private fun JsBlock.replaceLocalVariables(scope: JsScope, transformer: CoroutineBodyTransformer) {
+    private fun JsBlock.collectAdditionalLocalVariables() {
         accept(object : RecursiveJsVisitor() {
             override fun visit(x: JsVars.JsVar) {
                 super.visit(x)
                 localVariables += x.name
             }
         })
-
-        val visitor = object : JsVisitorWithContextImpl() {
-            override fun endVisit(x: JsNameRef, ctx: JsContext<in JsNode>) {
-                when {
-                    x.coroutineController -> {
-                        ctx.replaceMe(JsNameRef(transformer.controllerFieldName, x.qualifier).apply {
-                            sideEffects = SideEffectKind.PURE
-                        })
-                    }
-
-                    x.coroutineResult -> {
-                        ctx.replaceMe(JsNameRef(transformer.resultFieldName, x.qualifier).apply {
-                            sideEffects = SideEffectKind.DEPENDS_ON_STATE
-                        })
-                    }
-
-                    x.qualifier == null && x.name in localVariables -> {
-                        val fieldName = scope.getFieldName(x.name!!)
-                        ctx.replaceMe(JsNameRef(fieldName, JsLiteral.THIS))
-                    }
-                }
-            }
-
-            override fun endVisit(x: JsVars, ctx: JsContext<in JsStatement>) {
-                val declaredNames = x.vars.map { it.name }
-                val totalCount = declaredNames.size
-                val localVarCount = declaredNames.count()
-
-                when {
-                    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 -> {
-                        for (declaration in x.vars) {
-                            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()
-                    }
-                }
-                super.endVisit(x, ctx)
-            }
-        }
-        visitor.accept(this)
     }
 
+
     private fun ClassDescriptor.findFunction(name: String): FunctionDescriptor? {
         val functions = unsubstitutedMemberScope.getContributedDescriptors(DescriptorKindFilter.FUNCTIONS)
                 .filter { it.name.asString() == name }
         return functions.mapNotNull { it as? FunctionDescriptor }.firstOrNull { it.kind.isReal }
     }
 
-    private fun JsScope.getFieldName(variableName: JsName) = declareName("local\$${variableName.ident}")
-
     private fun MutableList<JsStatement>.assign(fieldName: JsName, value: JsExpression) {
         this += JsAstUtils.assignment(JsNameRef(fieldName, JsLiteral.THIS), value).makeStmt()
     }

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

@@ -0,0 +1,237 @@
+/*
+ * Copyright 2010-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.js.coroutine
+
+import com.google.dart.compiler.backend.js.ast.*
+import com.google.dart.compiler.backend.js.ast.metadata.*
+import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
+import org.jetbrains.kotlin.utils.singletonOrEmptyList
+
+fun JsNode.collectNodesToSplit(): Set<JsNode> {
+    val nodes = mutableSetOf<JsNode>()
+
+    val visitor = object : RecursiveJsVisitor() {
+        var childrenInSet = false
+
+        override fun visitInvocation(invocation: JsInvocation) {
+            super.visitInvocation(invocation)
+            if (invocation.isSuspend) {
+                nodes += invocation
+                childrenInSet = true
+            }
+        }
+
+        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
+            }
+            else {
+                childrenInSet = oldChildrenInSet
+            }
+        }
+    }
+
+    visitor.accept(this)
+    visitor.accept(this)
+
+    return nodes
+}
+
+fun List<CoroutineBlock>.replaceCoroutineFlowStatements(context: CoroutineTransformationContext, program: JsProgram) {
+    val blockIndexes = withIndex().associate { (index, block) -> Pair(block, index) }
+
+    val blockReplacementVisitor = object : JsVisitorWithContextImpl() {
+        override fun endVisit(x: JsDebugger, ctx: JsContext<in JsStatement>) {
+            val target = x.targetBlock
+            if (target != null) {
+                val lhs = JsNameRef(context.stateFieldName, JsLiteral.THIS)
+                val rhs = program.getNumberLiteral(blockIndexes[target]!!)
+                ctx.replaceMe(JsExpressionStatement(JsAstUtils.assignment(lhs, rhs)).apply {
+                    targetBlock = true
+                })
+            }
+
+            val exceptionTarget = x.targetExceptionBlock
+            if (exceptionTarget != null) {
+                val lhs = JsNameRef(context.exceptionStateName, JsLiteral.THIS)
+                val rhs = program.getNumberLiteral(blockIndexes[exceptionTarget]!!)
+                ctx.replaceMe(JsExpressionStatement(JsAstUtils.assignment(lhs, rhs)).apply {
+                    targetExceptionBlock = true
+                })
+            }
+
+            val finallyPath = x.finallyPath
+            if (finallyPath != null) {
+                if (finallyPath.isNotEmpty()) {
+                    val lhs = JsNameRef(context.finallyPathFieldName, JsLiteral.THIS)
+                    val rhs = JsArrayLiteral(finallyPath.map { program.getNumberLiteral(blockIndexes[it]!!) })
+                    ctx.replaceMe(JsExpressionStatement(JsAstUtils.assignment(lhs, rhs)).apply {
+                        this.finallyPath = true
+                    })
+                }
+                else {
+                    ctx.removeMe()
+                }
+            }
+        }
+    }
+    return forEach { blockReplacementVisitor.accept(it.jsBlock) }
+}
+
+fun CoroutineBlock.buildGraph(): 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(this)
+
+    return graph
+}
+
+private fun CoroutineBlock.collectTargetBlocks(): Set<CoroutineBlock> {
+    val targetBlocks = mutableSetOf<CoroutineBlock>()
+    jsBlock.accept(object : RecursiveJsVisitor() {
+        override fun visitDebugger(x: JsDebugger) {
+            targetBlocks += x.targetExceptionBlock.singletonOrEmptyList() + x.targetBlock.singletonOrEmptyList()
+        }
+    })
+    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
+}
+
+fun JsBlock.replaceLocalVariables(scope: JsScope, context: CoroutineTransformationContext, localVariables: Set<JsName>) {
+    val visitor = object : JsVisitorWithContextImpl() {
+        override fun endVisit(x: JsNameRef, ctx: JsContext<in JsNode>) {
+            when {
+                x.coroutineController -> {
+                    ctx.replaceMe(JsNameRef(context.controllerFieldName, x.qualifier).apply {
+                        sideEffects = SideEffectKind.PURE
+                    })
+                }
+
+                x.coroutineResult -> {
+                    ctx.replaceMe(JsNameRef(context.resultFieldName, x.qualifier).apply {
+                        sideEffects = SideEffectKind.DEPENDS_ON_STATE
+                    })
+                }
+
+                x.qualifier == null && x.name in localVariables -> {
+                    val fieldName = scope.getFieldName(x.name!!)
+                    ctx.replaceMe(JsNameRef(fieldName, JsLiteral.THIS))
+                }
+            }
+        }
+
+        override fun endVisit(x: JsVars, ctx: JsContext<in JsStatement>) {
+            val declaredNames = x.vars.map { it.name }
+            val totalCount = declaredNames.size
+            val localVarCount = declaredNames.count()
+
+            when {
+                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 -> {
+                    for (declaration in x.vars) {
+                        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()
+                }
+            }
+            super.endVisit(x, ctx)
+        }
+    }
+    visitor.accept(this)
+}
+
+fun JsScope.getFieldName(variableName: JsName) = declareName("local\$${variableName.ident}")

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

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2010-2016 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.js.coroutine
+
+import com.google.dart.compiler.backend.js.ast.JsScope
+
+class CoroutineTransformationContext(private val scope: JsScope) {
+    val entryBlock = CoroutineBlock()
+    val globalCatchBlock = CoroutineBlock()
+    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") }
+    val finallyPathFieldName by lazy { scope.declareFreshName("\$finallyPath") }
+}

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

@@ -18,6 +18,7 @@ package org.jetbrains.kotlin.js.inline.util
 
 import com.google.dart.compiler.backend.js.ast.*
 import com.google.dart.compiler.backend.js.ast.metadata.staticRef
+
 import org.jetbrains.kotlin.js.inline.util.collectors.InstanceCollector
 import org.jetbrains.kotlin.js.translate.expression.InlineMetadata
 import org.jetbrains.kotlin.js.translate.utils.JsAstUtils
@@ -184,3 +185,111 @@ fun <T : JsNode> collectInstances(klass: Class<T>, scope: JsNode): List<T> {
         collected
     }
 }
+
+fun JsNode.collectBreakContinueTargets(): Map<JsContinue, JsStatement> {
+    val targets = mutableMapOf<JsContinue, JsStatement>()
+
+    accept(object : RecursiveJsVisitor() {
+        var defaultBreakTarget: JsStatement? = null
+        var breakTargets = mutableMapOf<JsName, JsStatement>()
+        var defaultContinueTarget: JsStatement? = null
+        var continueTargets = mutableMapOf<JsName, JsStatement>()
+
+        override fun visitLabel(x: JsLabel) {
+            val inner = x.statement
+            when (inner) {
+                is JsDoWhile -> handleLoop(inner, inner.body, x.name)
+
+                is JsWhile -> handleLoop(inner, inner.body, x.name)
+
+                is JsFor -> handleLoop(inner, inner.body, x.name)
+
+                else -> {
+                    withBreakAndContinue(x.name, x.statement, null) {
+                        accept(inner)
+                    }
+                }
+            }
+        }
+
+        override fun visitWhile(x: JsWhile) = handleLoop(x, x.body, null)
+
+        override fun visitDoWhile(x: JsDoWhile) = handleLoop(x, x.body, null)
+
+        override fun visitFor(x: JsFor) = handleLoop(x, x.body, null)
+
+        private fun handleLoop(loop: JsStatement, body: JsStatement, label: JsName?) {
+            withBreakAndContinue(label, loop, loop) {
+                body.accept(this)
+            }
+        }
+
+        override fun visitBreak(x: JsBreak) {
+            val targetLabel = x.label?.name
+            targets[x] = if (targetLabel == null) {
+                defaultBreakTarget!!
+            }
+            else {
+                breakTargets[targetLabel]!!
+            }
+        }
+
+        override fun visitContinue(x: JsContinue) {
+            val targetLabel = x.label?.name
+            targets[x] = if (targetLabel == null) {
+                defaultContinueTarget!!
+            }
+            else {
+                continueTargets[targetLabel]!!
+            }
+        }
+
+        private fun withBreakAndContinue(
+                label: JsName?,
+                breakTargetStatement: JsStatement,
+                continueTargetStatement: JsStatement? = null,
+                action: () -> Unit
+        ) {
+            val oldDefaultBreakTarget = defaultBreakTarget
+            val oldDefaultContinueTarget = defaultContinueTarget
+            val (oldBreakTarget, oldContinueTarget) = if (label != null) {
+                Pair(breakTargets[label], continueTargets[label])
+            }
+            else {
+                Pair(null, null)
+            }
+
+            defaultBreakTarget = breakTargetStatement
+            if (label != null) {
+                breakTargets[label] = breakTargetStatement
+                if (continueTargetStatement != null) {
+                    continueTargets[label] = continueTargetStatement
+                }
+            }
+            if (continueTargetStatement != null) {
+                defaultContinueTarget = continueTargetStatement
+            }
+
+            action()
+
+            defaultBreakTarget = oldDefaultBreakTarget
+            defaultContinueTarget = oldDefaultContinueTarget
+            if (label != null) {
+                if (oldBreakTarget == null) {
+                    breakTargets.keys -= label
+                }
+                else {
+                    breakTargets[label] = oldBreakTarget
+                }
+                if (oldContinueTarget == null) {
+                    continueTargets.keys -= label
+                }
+                else {
+                    continueTargets[label] = oldContinueTarget
+                }
+            }
+        }
+    })
+
+    return targets
+}

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

@@ -5696,6 +5696,12 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
             doTest(fileName);
         }
 
+        @TestMetadata("instanceOfContinuation.kt")
+        public void testInstanceOfContinuation() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/instanceOfContinuation.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("iterateOverArray.kt")
         public void testIterateOverArray() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/coroutines/iterateOverArray.kt");