JS: refactor coroutine transformer

2016-11-11 15:57:02 +03:00
parent d41d09ffc4
commit bb61fb0a91
6 changed files with 456 additions and 339 deletions
@@ -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() {
-    val entryBlock = CoroutineBlock()
-    val globalCatchBlock = CoroutineBlock()
+class CoroutineBodyTransformer(
+        private val program: JsProgram,
+        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") }
-    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") }
-    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 breakContinueTargetStatements = mutableMapOf<JsContinue, JsStatement>()
+    private val breakTargets = mutableMapOf<JsStatement, JumpTarget>()
+    private val continueTargets = mutableMapOf<JsStatement, JumpTarget>()
    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>()
-
-        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
-
+        breakContinueTargetStatements += node.collectBreakContinueTargets()
+        nodesToSplit = node.collectNodesToSplit()
        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) }
-
-        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)
-        }
-
+        orderedBlocks.replaceCoroutineFlowStatements(context, program)
        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 -> handleWhile(inner, x.name)
-
-            is JsFor -> handleFor(inner, x.name)
+            is JsWhile,
+            is JsDoWhile,
+            is JsFor -> inner.accept(this)

            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)
-
-    private fun handleWhile(x: JsWhile, label: JsName?) = splitIfNecessary(x) {
+    override fun visitWhile(x: JsWhile) = 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)
-
-    private fun handleDoWhile(x: JsDoWhile, label: JsName?) = splitIfNecessary(x) {
+    override fun visitDoWhile(x: JsDoWhile) = 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)
-
-    private fun handleFor(x: JsFor, label: JsName?) = splitIfNecessary(x) {
+    override fun visitFor(x: JsFor) = 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, targetTryDepth) = if (targetLabel == null) {
-            defaultBreakTarget
-        }
-        else {
-            breakTargets[targetLabel]!!
-        }
-
-        val targetBlock = breakBlocks[targetStatement]!!
+        val targetStatement = breakContinueTargetStatements[x]!!
+        val (targetBlock, targetTryDepth) = breakTargets[targetStatement]!!
        referencedBlocks += targetBlock
        jumpWithFinally(targetTryDepth + 1, targetBlock)
        currentStatements += jump()
    }

    override fun visitContinue(x: JsContinue) {
-        val targetLabel = x.label?.name
-        val (targetStatement, targetTryDepth) = if (targetLabel == null) {
-            defaultContinueTarget
-        }
-        else {
-            continueTargets[targetLabel]!!
-        }
-
-        val targetBlock = continueBlocks[targetStatement]!!
+        val targetStatement = breakContinueTargetStatements[x]!!
+        val (targetBlock, targetTryDepth) = continueTargets[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 stateRef = JsNameRef(stateFieldName, JsLiteral.THIS)
+        val finallyPathRef = JsNameRef(context.finallyPathFieldName, 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
-        }
-
-        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)
+            continueTargets[statement] = JumpTarget(continueBlock, currentTryDepth)
        }

        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
-            }
-        }
+        breakTargets.keys -= statement
+        continueTargets.keys -= statement
    }

    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 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() {
-        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
+
+    private data class JumpTarget(val block: CoroutineBlock, val tryDepth: Int)
+
+    private class TryBlock(val catchBlock: CoroutineBlock, val finallyBlock: CoroutineBlock?)
 }
@@ -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) {
-            coroutineBlock.jsBlock.replaceLocalVariables(function.scope, bodyTransformer)
-        }
+        coroutineBlocks.forEach { it.jsBlock.collectAdditionalLocalVariables() }
+        coroutineBlocks.forEach { it.jsBlock.replaceLocalVariables(function.scope, context, localVariables) }

        val additionalStatements = mutableListOf<JsStatement>()
-        val resumeName = generateDoResume(coroutineBlocks, bodyTransformer, additionalStatements, throwName)
-        generateContinuationConstructor(bodyTransformer, additionalStatements)
+        val resumeName = generateDoResume(coroutineBlocks, context, additionalStatements, throwName)
+        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(bodyTransformer.exceptionStateName, program.getNumberLiteral(0))
-            if (bodyTransformer.hasFinallyBlocks) {
-                assign(bodyTransformer.finallyPathFieldName, JsLiteral.NULL)
+            assign(context.stateFieldName, program.getNumberLiteral(0))
+            assign(context.exceptionStateName, program.getNumberLiteral(0))
+            if (hasFinallyBlocks) {
+                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 stateRef = JsNameRef(transformer.stateFieldName, JsLiteral.THIS)
+        val indexOfGlobalCatch = blocks.indexOf(context.globalCatchBlock)
+        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(JsNameRef(transformer.exceptionFieldName, JsLiteral.THIS), catch.parameter.name.makeRef()).makeStmt()
+                JsAstUtils.assignment(stateRef.deepCopy(), JsNameRef(context.exceptionStateName, JsLiteral.THIS)).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))
-            transformer.globalCatchBlock.statements += JsReturn(JsInvocation(throwMethodRef.deepCopy(), throwResultRef))
+            val throwMethodRef = JsNameRef(throwName, JsNameRef(context.controllerFieldName, JsLiteral.THIS))
+            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(transformer.exceptionStateName, JsLiteral.THIS))
-        val exceptionToResult = JsAstUtils.assignment(JsNameRef(transformer.exceptionFieldName, JsLiteral.THIS), exceptionName.makeRef())
+                JsNameRef(context.stateFieldName, JsLiteral.THIS),
+                JsNameRef(context.exceptionStateName, JsLiteral.THIS))
+        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()
    }
@@ -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}")
@@ -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") }
+}
@@ -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
+}
@@ -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");