[JS IR BE] Fix state machine generation in case of composition of loops, inline functions and finally blocks

* lower finally blocks in any cases * do not optimize exit blocks for if-statements
2019-05-17 19:03:40 +03:00
parent 9c7d47789c
commit dfa38f4a4d
13 changed files with 351 additions and 257 deletions
@@ -36,7 +36,7 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
    protected abstract fun nameForCoroutineClass(function: IrFunction): Name

    protected abstract fun buildStateMachine(
-        originalBody: IrBody, stateMachineFunction: IrFunction,
+        stateMachineFunction: IrFunction,
        transformingFunction: IrFunction,
        argumentToPropertiesMap: Map<IrValueParameter, IrField>
    )
@@ -617,7 +617,7 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
                }
            }

-            buildStateMachine(originalBody, function, irFunction, argumentToPropertiesMap)
+            buildStateMachine(function, irFunction, argumentToPropertiesMap)
            return function
        }
    }
@@ -216,11 +216,6 @@ class FinallyBlocksLowering(val context: CommonBackendContext, private val throw
        val irBuilder = context.createIrBuilder(currentScope!!.scope.scopeOwnerSymbol, startOffset, endOffset)
        val transformer = this
        irBuilder.run {
-            val transformedTry = IrTryImpl(
-                startOffset = startOffset,
-                endOffset = endOffset,
-                type = context.irBuiltIns.unitType
-            )
            val transformedFinallyExpression = finallyExpression.transform(transformer, null)
            val parameter = WrappedVariableDescriptor()
            val catchParameter = IrVariableImpl(
@@ -233,27 +228,42 @@ class FinallyBlocksLowering(val context: CommonBackendContext, private val throw
            val syntheticTry = IrTryImpl(
                startOffset = startOffset,
                endOffset = endOffset,
-                type = context.irBuiltIns.unitType,
-                tryResult = transformedTry,
-                catches = listOf(
-                    irCatch(catchParameter).apply {
-                        result = irComposite {
-                            +finallyExpression.copy()
-                            +irThrow(irGet(catchParameter))
-                        }
-                    }),
-                finallyExpression = null
-            )
+                type = context.irBuiltIns.unitType
+            ).apply {
+                this.catches += irCatch(catchParameter).apply {
+                    result = irComposite {
+                        +finallyExpression.copy()
+                        +irThrow(irGet(catchParameter))
+                    }
+                }
+
+                this.finallyExpression = null
+            }
+
            using(TryScope(syntheticTry, transformedFinallyExpression, this)) {
+
                val fallThroughType = aTry.type
                val fallThroughSymbol = IrReturnableBlockSymbolImpl(WrappedSimpleFunctionDescriptor())
                val transformedResult = aTry.tryResult.transform(transformer, null)
-                transformedTry.tryResult = irReturn(fallThroughSymbol, transformedResult)
-                for (aCatch in aTry.catches) {
-                    val transformedCatch = aCatch.transform(transformer, null)
-                    transformedCatch.result = irReturn(fallThroughSymbol, transformedCatch.result)
-                    transformedTry.catches.add(transformedCatch)
+                val returnedResult = irReturn(fallThroughSymbol, transformedResult)
+
+                if (aTry.catches.isNotEmpty()) {
+                    val transformedTry = IrTryImpl(
+                        startOffset = startOffset,
+                        endOffset = endOffset,
+                        type = context.irBuiltIns.unitType
+                    )
+                    transformedTry.tryResult = returnedResult
+                    for (aCatch in aTry.catches) {
+                        val transformedCatch = aCatch.transform(transformer, null)
+                        transformedCatch.result = irReturn(fallThroughSymbol, transformedCatch.result)
+                        transformedTry.catches.add(transformedCatch)
+                    }
+                    syntheticTry.tryResult = transformedTry
+                } else {
+                    syntheticTry.tryResult = returnedResult
                }
+
                return irInlineFinally(fallThroughSymbol, fallThroughType, it.expression, it.finallyExpression)
            }
        }
@@ -380,8 +380,8 @@ val jsPhases = namedIrModulePhase(
            moveBodilessDeclarationsToSeparatePlacePhase then
            enumClassLoweringPhase then
            enumUsageLoweringPhase then
-            returnableBlockLoweringPhase then
            suspendFunctionsLoweringPhase then
+            returnableBlockLoweringPhase then
            privateMembersLoweringPhase then
            callableReferenceLoweringPhase then
            defaultArgumentStubGeneratorPhase then
@@ -8,7 +8,7 @@ package org.jetbrains.kotlin.ir.backend.js.lower.coroutines
 import org.jetbrains.kotlin.backend.common.descriptors.synthesizedName
 import org.jetbrains.kotlin.backend.common.ir.isSuspend
 import org.jetbrains.kotlin.backend.common.lower.AbstractSuspendFunctionsLowering
-import org.jetbrains.kotlin.ir.IrElement
+import org.jetbrains.kotlin.backend.common.lower.FinallyBlocksLowering
 import org.jetbrains.kotlin.ir.backend.js.JsIrBackendContext
 import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
 import org.jetbrains.kotlin.ir.builders.*
@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.ir.declarations.*
 import org.jetbrains.kotlin.ir.expressions.*
 import org.jetbrains.kotlin.ir.expressions.impl.*
 import org.jetbrains.kotlin.ir.symbols.IrFieldSymbol
+import org.jetbrains.kotlin.ir.symbols.IrValueParameterSymbol
 import org.jetbrains.kotlin.ir.symbols.IrValueSymbol
 import org.jetbrains.kotlin.ir.util.explicitParameters
 import org.jetbrains.kotlin.ir.util.patchDeclarationParents
@@ -43,21 +44,20 @@ class JsSuspendFunctionsLowering(ctx: JsIrBackendContext) : AbstractSuspendFunct
    override fun nameForCoroutineClass(function: IrFunction) = "${function.name}COROUTINE\$${coroutineId++}".synthesizedName

    override fun buildStateMachine(
-        originalBody: IrBody,
        stateMachineFunction: IrFunction,
        transformingFunction: IrFunction,
        argumentToPropertiesMap: Map<IrValueParameter, IrField>
    ) {
-        val body =
-            (originalBody as IrBlockBody).run {
-                IrBlockImpl(
-                    transformingFunction.startOffset,
-                    transformingFunction.endOffset,
-                    context.irBuiltIns.unitType,
-                    STATEMENT_ORIGIN_COROUTINE_IMPL,
-                    statements
-                )
-            }
+        val simplifiedFunction = transformingFunction.transform(FinallyBlocksLowering(context, context.dynamicType), null) as IrFunction
+        val originalBody = simplifiedFunction.body as IrBlockBody
+
+        val body = IrBlockImpl(
+            simplifiedFunction.startOffset,
+            simplifiedFunction.endOffset,
+            context.irBuiltIns.unitType,
+            STATEMENT_ORIGIN_COROUTINE_IMPL,
+            originalBody.statements
+        )

        val coroutineClass = stateMachineFunction.parent as IrClass
        val suspendResult = JsIrBuilder.buildVar(
@@ -83,11 +83,11 @@ class JsSuspendFunctionsLowering(ctx: JsIrBackendContext) : AbstractSuspendFunct
            COROUTINE_ROOT_LOOP,
            rootTry,
            JsIrBuilder.buildBoolean(context.irBuiltIns.booleanType, true)
-        )
+        ).also {
+            it.label = "\$sm"
+        }

-        val suspendableNodes = mutableSetOf<IrElement>()
-        val loweredBody =
-            collectSuspendableNodes(body, suspendableNodes, context, stateMachineFunction, context.dynamicType)
+        val suspendableNodes = collectSuspendableNodes(body)
        val thisReceiver = (stateMachineFunction.dispatchReceiverParameter as IrValueParameter).symbol

        val stateMachineBuilder = StateMachineBuilder(
@@ -104,36 +104,13 @@ class JsSuspendFunctionsLowering(ctx: JsIrBackendContext) : AbstractSuspendFunct
            suspendResult.symbol
        )

-        loweredBody.acceptVoid(stateMachineBuilder)
+        body.acceptVoid(stateMachineBuilder)

        stateMachineBuilder.finalizeStateMachine()

        rootTry.catches += stateMachineBuilder.globalCatch

-        val visited = mutableSetOf<SuspendState>()
-
-        val sortedStates = DFS.topologicalOrder(listOf(stateMachineBuilder.entryState), { it.successors }, { visited.add(it) })
-        sortedStates.withIndex().forEach { it.value.id = it.index }
-
-        fun buildDispatch(target: SuspendState) = target.run {
-            assert(id >= 0)
-            JsIrBuilder.buildInt(context.irBuiltIns.intType, id)
-        }
-
-        val eqeqeqInt = context.irBuiltIns.eqeqeqSymbol
-
-        for (state in sortedStates) {
-            val condition = JsIrBuilder.buildCall(eqeqeqInt).apply {
-                putValueArgument(0, JsIrBuilder.buildCall(coroutineImplLabelPropertyGetter.symbol).also {
-                    it.dispatchReceiver = JsIrBuilder.buildGetValue(thisReceiver)
-                })
-                putValueArgument(1, JsIrBuilder.buildInt(context.irBuiltIns.intType, state.id))
-            }
-
-            switch.branches += IrBranchImpl(state.entryBlock.startOffset, state.entryBlock.endOffset, condition, state.entryBlock)
-        }
-
-        rootLoop.transform(DispatchPointTransformer(::buildDispatch), null)
+        assignStateIds(stateMachineBuilder.entryState, thisReceiver, switch, rootLoop)

        exceptionTrapId = stateMachineBuilder.rootExceptionTrap.id

@@ -149,7 +126,7 @@ class JsSuspendFunctionsLowering(ctx: JsIrBackendContext) : AbstractSuspendFunct
            override fun visitReturn(expression: IrReturn): IrExpression {
                expression.transformChildrenVoid(this)

-                return if (expression.returnTargetSymbol != transformingFunction.symbol)
+                return if (expression.returnTargetSymbol != simplifiedFunction.symbol)
                    expression
                else
                    JsIrBuilder.buildReturn(stateMachineFunction.symbol, expression.value, expression.type)
@@ -169,7 +146,7 @@ class JsSuspendFunctionsLowering(ctx: JsIrBackendContext) : AbstractSuspendFunct
                }
            }
        }
-        transformingFunction.explicitParameters.forEach {
+        simplifiedFunction.explicitParameters.forEach {
            localToPropertyMap.getOrPut(it.symbol) {
                argumentToPropertiesMap.getValue(it).symbol
            }
@@ -178,6 +155,33 @@ class JsSuspendFunctionsLowering(ctx: JsIrBackendContext) : AbstractSuspendFunct
        stateMachineFunction.transform(LiveLocalsTransformer(localToPropertyMap, { JsIrBuilder.buildGetValue(thisReceiver) }, unit), null)
    }

+    private fun assignStateIds(entryState: SuspendState, thisReceiver: IrValueParameterSymbol, switch: IrWhen, rootLoop: IrLoop) {
+        val visited = mutableSetOf<SuspendState>()
+
+        val sortedStates = DFS.topologicalOrder(listOf(entryState), { it.successors }, { visited.add(it) })
+        sortedStates.withIndex().forEach { it.value.id = it.index }
+
+        val eqeqeqInt = context.irBuiltIns.eqeqeqSymbol
+
+        for (state in sortedStates) {
+            val condition = JsIrBuilder.buildCall(eqeqeqInt).apply {
+                putValueArgument(0, JsIrBuilder.buildCall(coroutineImplLabelPropertyGetter.symbol).also {
+                    it.dispatchReceiver = JsIrBuilder.buildGetValue(thisReceiver)
+                })
+                putValueArgument(1, JsIrBuilder.buildInt(context.irBuiltIns.intType, state.id))
+            }
+
+            switch.branches += IrBranchImpl(state.entryBlock.startOffset, state.entryBlock.endOffset, condition, state.entryBlock)
+        }
+
+        val dispatchPointTransformer = DispatchPointTransformer {
+            assert(it.id >= 0)
+            JsIrBuilder.buildInt(context.irBuiltIns.intType, it.id)
+        }
+
+        rootLoop.transformChildrenVoid(dispatchPointTransformer)
+    }
+
    private fun computeLivenessAtSuspensionPoints(body: IrBody): Map<IrCall, List<IrValueDeclaration>> {
        // TODO: data flow analysis.
        // Just save all visible for now.
@@ -201,6 +205,7 @@ class JsSuspendFunctionsLowering(ctx: JsIrBackendContext) : AbstractSuspendFunct
        for (it in coroutineConstructors) {
            (it.body as? IrBlockBody)?.run {
                val receiver = JsIrBuilder.buildGetValue(coroutineClassThis.symbol)
+                assert(exceptionTrapId >= 0)
                val id = JsIrBuilder.buildInt(context.irBuiltIns.intType, exceptionTrapId)
                statements += JsIrBuilder.buildCall(coroutineImplExceptionStatePropertySetter.symbol).also { call ->
                    call.dispatchReceiver = receiver
@@ -23,10 +23,7 @@ import org.jetbrains.kotlin.ir.symbols.IrFunctionSymbol
 import org.jetbrains.kotlin.ir.symbols.IrReturnableBlockSymbol
 import org.jetbrains.kotlin.ir.symbols.IrValueParameterSymbol
 import org.jetbrains.kotlin.ir.symbols.IrVariableSymbol
-import org.jetbrains.kotlin.ir.types.IrDynamicType
-import org.jetbrains.kotlin.ir.types.IrType
-import org.jetbrains.kotlin.ir.types.classifierOrNull
-import org.jetbrains.kotlin.ir.types.isNothing
+import org.jetbrains.kotlin.ir.types.*
 import org.jetbrains.kotlin.ir.util.deepCopyWithSymbols
 import org.jetbrains.kotlin.ir.visitors.*

@@ -38,9 +35,7 @@ class SuspendState(type: IrType) {

 data class LoopBounds(val headState: SuspendState, val exitState: SuspendState)

-data class FinallyTargets(val normal: SuspendState, val fromThrow: SuspendState)
-
-data class TryState(val tryState: SuspendState, val catchState: SuspendState, val finallyState: FinallyTargets?)
+data class TryState(val tryState: SuspendState, val catchState: SuspendState)

 class IrDispatchPoint(val target: SuspendState) : IrExpressionBase(UNDEFINED_OFFSET, UNDEFINED_OFFSET, target.entryBlock.type) {
    override fun <R, D> accept(visitor: IrElementVisitor<R, D>, data: D) = visitor.visitExpression(this, data)
@@ -50,7 +45,6 @@ class IrDispatchPoint(val target: SuspendState) : IrExpressionBase(UNDEFINED_OFF
    override fun <D> transformChildren(transformer: IrElementTransformer<D>, data: D) {}
 }

-
 class DispatchPointTransformer(val action: (SuspendState) -> IrExpression) : IrElementTransformerVoid() {
    override fun visitExpression(expression: IrExpression): IrExpression {
        val dispatchPoint = expression as? IrDispatchPoint
@@ -76,7 +70,6 @@ class StateMachineBuilder(
    private val loopMap = mutableMapOf<IrLoop, LoopBounds>()
    private val unit = context.irBuiltIns.unitType
    private val nothing = context.irBuiltIns.nothingType
-    private val int = context.irBuiltIns.intType
    private val booleanNotSymbol = context.irBuiltIns.booleanNotSymbol
    private val eqeqeqSymbol = context.irBuiltIns.eqeqeqSymbol

@@ -166,9 +159,13 @@ class StateMachineBuilder(

    private fun addStatement(statement: IrStatement) = currentBlock.addStatement(statement)

+    private fun isBlockEnded(): Boolean {
+        val lastExpression = currentBlock.statements.lastOrNull() as? IrExpression ?: return false
+        return lastExpression.type.isNothing()
+    }
+
    private fun maybeDoDispatch(target: SuspendState) {
-        val lastStatement = currentBlock.statements.lastOrNull()
-        if (lastStatement !is IrReturn && lastStatement !is IrContinue && lastStatement !is IrThrow) {
+        if (!isBlockEnded()) {
            doDispatch(target)
        }
    }
@@ -360,9 +357,6 @@ class StateMachineBuilder(
            branches = expression.branches
        }

-        val rootState = currentState
-        val rootBlock = currentBlock
-
        for (branch in branches) {
            if (!isElseBranch(branch)) {
                branch.condition.acceptVoid(this)
@@ -378,35 +372,24 @@ class StateMachineBuilder(
                currentBlock = branchBlock
                branch.result.acceptVoid(this)

-                // TODO: block should not be empty
-                val lastStatement = currentBlock.statements.lastOrNull()
-                if (lastStatement != null && lastStatement !is IrContinue) {
-                    if (currentState !== rootState) {
-                        doDispatch(exitState)
-                    }
+                if (!isBlockEnded()) {
+                    doDispatch(exitState)
                }

                currentState = dispatchState
                currentBlock = elseBlock
            } else {
                branch.result.acceptVoid(this)
-
-                // TODO: block should not be empty
-                val lastStatement = currentBlock.statements.lastOrNull()
-                if (lastStatement != null && lastStatement !is IrContinue) {
-                    if (currentState !== rootState) {
-                        doDispatch(exitState)
-                    }
+                if (!isBlockEnded()) {
+                    doDispatch(exitState)
                }
                break
            }
        }

-        currentState = rootState
-        currentBlock = rootBlock
        maybeDoDispatch(exitState)
-
        updateState(exitState)
+
        if (varSymbol != null) {
            addStatement(JsIrBuilder.buildGetValue(varSymbol))
        }
@@ -542,7 +525,7 @@ class StateMachineBuilder(
            if (varSymbol != null) {
                transformLastExpression { JsIrBuilder.buildSetVariable(varSymbol, it, it.type) }
            }
-            doDispatch(exitState)
+            maybeDoDispatch(exitState)
        } else {
            transformLastExpression { expression.apply { value = it } }
        }
@@ -553,7 +536,7 @@ class StateMachineBuilder(
        currentState.successors += catchBlockStack.peek()!!
    }

-    private fun hasResultingValue(expression: IrExpression) = !expression.type.isNothing()
+    private fun hasResultingValue(expression: IrExpression) = !expression.type.run { isNothing() || isUnit() }

    override fun visitThrow(expression: IrThrow) {
        expression.acceptChildrenVoid(this)
@@ -562,20 +545,18 @@ class StateMachineBuilder(
    }

    override fun visitTry(aTry: IrTry) {
-        val tryState = buildTryState(aTry)
+
+        require(aTry.finallyExpression == null)
+
+        val tryState = buildTryState()
        val enclosingCatch = catchBlockStack.peek()!!

        catchBlockStack.push(tryState.catchState)

-        val finallyStateVar = tempVar(int, "FINALLY_STATE")
        val exitState = SuspendState(unit)

        val varSymbol = if (hasResultingValue(aTry)) tempVar(aTry.type, "TRY_RESULT") else null

-        if (aTry.finallyExpression != null) {
-            finallyStateVar.initializer = IrDispatchPoint(exitState)
-            addStatement(finallyStateVar)
-        }
        if (varSymbol != null) {
            addStatement(varSymbol)
        }
@@ -591,23 +572,16 @@ class StateMachineBuilder(

        tryResult.acceptVoid(this)

-        if (tryState.finallyState != null) {
-            doDispatch(tryState.finallyState.normal)
-        } else {
+        if (!isBlockEnded()) {
            setupExceptionState(enclosingCatch)
            doDispatch(exitState)
        }
-
        addExceptionEdge()
+
        catchBlockStack.pop()
        updateState(tryState.catchState)

-        if (tryState.finallyState != null) {
-            setupExceptionState(tryState.finallyState.fromThrow)
-        } else {
-            setupExceptionState(enclosingCatch)
-        }
-
+        setupExceptionState(enclosingCatch)

        var rethrowNeeded = true

@@ -628,9 +602,7 @@ class StateMachineBuilder(

                addStatement(irVar)
                catchResult.acceptVoid(this)
-                val exitDispatch = tryState.finallyState?.run { normal } ?: exitState
-                maybeDoDispatch(exitDispatch)
-
+                maybeDoDispatch(exitState)
            } else {
                val check = buildIsCheck(pendingException(), type)

@@ -644,8 +616,7 @@ class StateMachineBuilder(

                addStatement(irVar)
                catchResult.acceptVoid(this)
-                val exitDispatch = tryState.finallyState?.run { normal } ?: exitState
-                maybeDoDispatch(exitDispatch)
+                maybeDoDispatch(exitState)

                currentBlock = ifBlock
                addStatement(irIf)
@@ -658,42 +629,10 @@ class StateMachineBuilder(
            addStatement(JsIrBuilder.buildThrow(nothing, pendingException()))
        }

-        if (tryState.finallyState == null) {
-            currentState.successors += enclosingCatch
-        }
-
-        val finallyState = tryState.finallyState
-        if (finallyState != null) {
-            val throwExitState = SuspendState(unit)
-            updateState(finallyState.fromThrow)
-            tryState.tryState.successors += finallyState.fromThrow
-            addStatement(
-                JsIrBuilder.buildSetVariable(
-                    finallyStateVar.symbol,
-                    IrDispatchPoint(throwExitState), int
-                )
-            )
-            doDispatch(finallyState.normal)
-
-            updateState(finallyState.normal)
-            tryState.tryState.successors += finallyState.normal
-            setupExceptionState(enclosingCatch)
-            aTry.finallyExpression?.acceptVoid(this)
-            currentState.successors += listOf(throwExitState, exitState)
-            addStatement(
-                JsIrBuilder.buildCall(stateSymbolSetter.symbol, unit).also {
-                    it.dispatchReceiver = thisReceiver
-                    it.putValueArgument(0, JsIrBuilder.buildGetValue(finallyStateVar.symbol))
-                }
-            )
-            doContinue()
-
-            updateState(throwExitState)
-            addStatement(JsIrBuilder.buildThrow(nothing, pendingException()))
-            addExceptionEdge()
-        }
+        currentState.successors += enclosingCatch

        updateState(exitState)
+
        if (varSymbol != null) {
            addStatement(JsIrBuilder.buildGetValue(varSymbol.symbol))
        }
@@ -711,19 +650,7 @@ class StateMachineBuilder(
    private fun exceptionState() = JsIrBuilder.buildCall(exStateSymbolGetter.symbol).also { it.dispatchReceiver = thisReceiver }
    private fun pendingException() = JsIrBuilder.buildCall(exceptionSymbolGetter.symbol).also { it.dispatchReceiver = thisReceiver }

-    private fun buildTryState(aTry: IrTry) =
-        TryState(
-            currentState,
-            SuspendState(unit),
-            aTry.finallyExpression?.run {
-                FinallyTargets(
-                    SuspendState(
-                        unit
-                    ), SuspendState(unit)
-                )
-            }
-        )
-
+    private fun buildTryState() = TryState(currentState, SuspendState(unit))

    private fun buildIsCheck(value: IrExpression, toType: IrType) =
        JsIrBuilder.buildTypeOperator(
@@ -736,4 +663,4 @@ class StateMachineBuilder(

    private fun tempVar(type: IrType, name: String = "tmp") =
        JsIrBuilder.buildVar(type, function.owner, name)
-}
+}
@@ -5,18 +5,14 @@

 package org.jetbrains.kotlin.ir.backend.js.lower.coroutines

-import org.jetbrains.kotlin.backend.common.CommonBackendContext
 import org.jetbrains.kotlin.backend.common.ir.isSuspend
-import org.jetbrains.kotlin.backend.common.lower.FinallyBlocksLowering
 import org.jetbrains.kotlin.backend.common.pop
 import org.jetbrains.kotlin.backend.common.push
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
 import org.jetbrains.kotlin.ir.backend.js.ir.JsIrBuilder
-import org.jetbrains.kotlin.ir.declarations.IrFunction
 import org.jetbrains.kotlin.ir.declarations.IrVariable
 import org.jetbrains.kotlin.ir.expressions.*
-import org.jetbrains.kotlin.ir.expressions.impl.IrBlockBodyImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrGetFieldImpl
 import org.jetbrains.kotlin.ir.expressions.impl.IrSetFieldImpl
 import org.jetbrains.kotlin.ir.symbols.IrFieldSymbol
@@ -28,16 +24,23 @@ import org.jetbrains.kotlin.ir.visitors.*
 object COROUTINE_ROOT_LOOP : IrStatementOriginImpl("COROUTINE_ROOT_LOOP")
 object COROUTINE_SWITCH : IrStatementOriginImpl("COROUTINE_SWITCH")

-open class SuspendableNodesCollector(protected val suspendableNodes: MutableSet<IrElement>) : IrElementVisitorVoid {
+open class SuspendableNodesCollector(private val suspendableNodes: MutableSet<IrElement>) : IrElementVisitorVoid {

-    protected var hasSuspendableChildren = false
+    private var hasSuspendableChildren = false
+
+    protected fun markNode(node: IrElement) {
+        suspendableNodes += node
+        hasSuspendableChildren = true
+    }
+
+    protected fun isSuspendableNode(node: IrElement) = node in suspendableNodes

    override fun visitElement(element: IrElement) {
        val current = hasSuspendableChildren
        hasSuspendableChildren = false
        element.acceptChildrenVoid(this)
        if (hasSuspendableChildren) {
-            suspendableNodes += element
+            markNode(element)
        }
        hasSuspendableChildren = hasSuspendableChildren || current
    }
@@ -45,105 +48,38 @@ open class SuspendableNodesCollector(protected val suspendableNodes: MutableSet<
    override fun visitCall(expression: IrCall) {
        super.visitCall(expression)
        if (expression.isSuspend) {
-            suspendableNodes += expression
-            hasSuspendableChildren = true
+            markNode(expression)
        }
    }
 }

-fun collectSuspendableNodes(
-    body: IrBlock,
-    suspendableNodes: MutableSet<IrElement>,
-    context: CommonBackendContext,
-    function: IrFunction,
-    throwableType: IrType
-): IrBlock {
-
-    // 1st: mark suspendable loops and tries
-    body.acceptVoid(SuspendableNodesCollector(suspendableNodes))
-    // 2nd: mark inner terminators
-    val terminatorsCollector = SuspendedTerminatorsCollector(suspendableNodes)
-    body.acceptVoid(terminatorsCollector)
-
-    if (terminatorsCollector.shouldFinalliesBeLowered) {
-        val finallyLower = FinallyBlocksLowering(context, throwableType)
-
-        function.body = IrBlockBodyImpl(body.startOffset, body.endOffset, body.statements)
-        function.transform(finallyLower, null)
-
-        val newBody = function.body as IrBlockBody
-        function.body = null
-        suspendableNodes.clear()
-        val newBlock = JsIrBuilder.buildBlock(body.type, newBody.statements)
-
-        return collectSuspendableNodes(newBlock, suspendableNodes, context, function, throwableType)
-    }
-
-    return body
-}
-
 class SuspendedTerminatorsCollector(suspendableNodes: MutableSet<IrElement>) : SuspendableNodesCollector(suspendableNodes) {

-    var shouldFinalliesBeLowered = false
-
    override fun visitBreakContinue(jump: IrBreakContinue) {
-        if (jump.loop in suspendableNodes) {
-            suspendableNodes.add(jump)
-            hasSuspendableChildren = true
+        if (isSuspendableNode(jump.loop)) {
+            markNode(jump)
        }
-
-        shouldFinalliesBeLowered = shouldFinalliesBeLowered || tryStack.any { it.finallyExpression != null && it in suspendableNodes }
-    }
-
-    private val tryStack = mutableListOf<IrTry>()
-    private val tryLoopStack = mutableListOf<IrStatement>()
-
-    private fun pushTry(aTry: IrTry) {
-        tryStack.push(aTry)
-        tryLoopStack.push(aTry)
-    }
-
-    private fun popTry() {
-        tryLoopStack.pop()
-        tryStack.pop()
-    }
-
-    private fun pushLoop(loop: IrLoop) {
-        tryLoopStack.push(loop)
-    }
-
-    private fun popLoop() {
-        tryLoopStack.pop()
-    }
-
-    override fun visitLoop(loop: IrLoop) {
-        pushLoop(loop)
-
-        super.visitLoop(loop)
-
-        popLoop()
-    }
-
-    override fun visitTry(aTry: IrTry) {
-        pushTry(aTry)
-
-        super.visitTry(aTry)
-
-        popTry()
    }

    override fun visitReturn(expression: IrReturn) {
-        shouldFinalliesBeLowered = shouldFinalliesBeLowered || tryStack.any { it.finallyExpression != null && it in suspendableNodes }
-
        super.visitReturn(expression)

-        if (expression.returnTargetSymbol is IrReturnableBlockSymbol) {
-            suspendableNodes.add(expression)
-            hasSuspendableChildren = true
+        if (expression.returnTargetSymbol is IrReturnableBlockSymbol && isSuspendableNode(expression.returnTargetSymbol.owner)) {
+            markNode(expression)
        }
    }
 }

+fun collectSuspendableNodes(function: IrBlock): MutableSet<IrElement> {
+
+    val suspendableNodes = mutableSetOf<IrElement>()
+    // 1st: mark suspendable loops and tries
+    function.acceptVoid(SuspendableNodesCollector(suspendableNodes))
+    // 2nd: mark inner terminators
+    function.acceptVoid(SuspendedTerminatorsCollector(suspendableNodes))
+
+    return suspendableNodes
+}

 class LiveLocalsTransformer(
    private val localMap: Map<IrValueSymbol, IrFieldSymbol>,
@@ -0,0 +1,53 @@
+// KJS_WITH_FULL_RUNTIME
+// IGNORE_BACKEND: JVM_IR
+// WITH_RUNTIME
+// WITH_COROUTINES
+// COMMON_COROUTINES_TEST
+import helpers.*
+import COROUTINES_PACKAGE.*
+import COROUTINES_PACKAGE.intrinsics.*
+
+
+class Controller {
+    var result = ""
+
+    suspend fun <T> suspendWithResult(value: T): T = suspendCoroutineUninterceptedOrReturn { c ->
+        c.resume(value)
+        COROUTINE_SUSPENDED
+    }
+}
+
+suspend fun lock(owner: Controller) {
+    owner.result += "L"
+}
+
+fun unlock(owner: Controller) {
+    owner.result += "U"
+}
+
+public suspend inline fun doInline(owner: Controller, action: () -> Unit): Unit {
+    lock(owner)
+    try {
+        return action()
+    } finally {
+        unlock(owner)
+    }
+}
+
+
+fun builder(c: suspend Controller.() -> Unit): String {
+    val controller = Controller()
+    c.startCoroutine(controller, EmptyContinuation)
+    return controller.result
+}
+
+fun box(): String {
+    val value = builder {
+        doInline(this) {
+            result += "X"
+        }
+    }
+    if (value != "LXU") return "fail: $value"
+
+    return "OK"
+}
@@ -0,0 +1,73 @@
+// KJS_WITH_FULL_RUNTIME
+// IGNORE_BACKEND: JVM_IR
+// WITH_RUNTIME
+// WITH_COROUTINES
+// COMMON_COROUTINES_TEST
+import helpers.*
+import COROUTINES_PACKAGE.*
+import COROUTINES_PACKAGE.intrinsics.*
+
+
+class Controller {
+    var result = ""
+
+    suspend fun <T> suspendWithResult(value: T): T = suspendCoroutineUninterceptedOrReturn { c ->
+        c.resume(value)
+        COROUTINE_SUSPENDED
+    }
+}
+
+fun Controller.consumeCancel(c: Throwable?) {
+    result += if (c == null) "?" else "!"
+}
+
+fun newIterator() = Iterator()
+
+class Iterator() {
+    var hasNextX = true
+    public suspend fun hasNext(): Boolean {
+        val tmp = hasNextX
+        hasNextX = false
+        return tmp
+    }
+    public suspend fun next(): String = "OK"
+}
+
+public inline fun Controller.consume(action: Controller.() -> String?): String? {
+    var cause: Throwable? = null
+    try {
+        return action()
+    } catch(x: Exception) {
+        cause = x
+        throw x
+    } finally {
+        consumeCancel(cause)
+    }
+}
+
+public suspend fun Controller.doTest(): String? {
+    return consume {
+        val iterator = newIterator()
+        if (!iterator.hasNext())
+            return null
+        return iterator.next()
+    }
+}
+
+
+fun builder(c: suspend Controller.() -> Unit): String {
+    val controller = Controller()
+    c.startCoroutine(controller, EmptyContinuation)
+    return controller.result
+}
+
+fun Controller.add(s: String?) {
+    result += s
+}
+
+fun box(): String {
+    val value = builder {
+        add(doTest())
+    }
+    return if (value != "?OK") return "Fail: $value" else "OK"
+}
@@ -6825,6 +6825,16 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/destructuringInLambdas.kt", "kotlin.coroutines");
            }

+            @TestMetadata("inlineSuspendFinally.kt")
+            public void testInlineSuspendFinally_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/inlineSuspendFinally.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("inlineSuspendFinally.kt")
+            public void testInlineSuspendFinally_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/inlineSuspendFinally.kt", "kotlin.coroutines");
+            }
+
            @TestMetadata("safeCallOnTwoReceiversLong.kt")
            public void testSafeCallOnTwoReceiversLong_1_2() throws Exception {
                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/safeCallOnTwoReceiversLong.kt", "kotlin.coroutines.experimental");
@@ -6865,6 +6875,16 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                runTest("compiler/testData/codegen/box/coroutines/featureIntersection/suspendFunctionIsAs.kt");
            }

+            @TestMetadata("suspendInlineSuspendFinally.kt")
+            public void testSuspendInlineSuspendFinally_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/suspendInlineSuspendFinally.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("suspendInlineSuspendFinally.kt")
+            public void testSuspendInlineSuspendFinally_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/suspendInlineSuspendFinally.kt", "kotlin.coroutines");
+            }
+
            @TestMetadata("suspendOperatorPlusAssign.kt")
            public void testSuspendOperatorPlusAssign_1_2() throws Exception {
                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/suspendOperatorPlusAssign.kt", "kotlin.coroutines.experimental");
@@ -6825,6 +6825,16 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/destructuringInLambdas.kt", "kotlin.coroutines");
            }

+            @TestMetadata("inlineSuspendFinally.kt")
+            public void testInlineSuspendFinally_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/inlineSuspendFinally.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("inlineSuspendFinally.kt")
+            public void testInlineSuspendFinally_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/inlineSuspendFinally.kt", "kotlin.coroutines");
+            }
+
            @TestMetadata("safeCallOnTwoReceiversLong.kt")
            public void testSafeCallOnTwoReceiversLong_1_2() throws Exception {
                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/safeCallOnTwoReceiversLong.kt", "kotlin.coroutines.experimental");
@@ -6865,6 +6875,16 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                runTest("compiler/testData/codegen/box/coroutines/featureIntersection/suspendFunctionIsAs.kt");
            }

+            @TestMetadata("suspendInlineSuspendFinally.kt")
+            public void testSuspendInlineSuspendFinally_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/suspendInlineSuspendFinally.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("suspendInlineSuspendFinally.kt")
+            public void testSuspendInlineSuspendFinally_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/suspendInlineSuspendFinally.kt", "kotlin.coroutines");
+            }
+
            @TestMetadata("suspendOperatorPlusAssign.kt")
            public void testSuspendOperatorPlusAssign_1_2() throws Exception {
                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/suspendOperatorPlusAssign.kt", "kotlin.coroutines.experimental");
@@ -6825,6 +6825,16 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/destructuringInLambdas.kt", "kotlin.coroutines");
            }

+            @TestMetadata("inlineSuspendFinally.kt")
+            public void testInlineSuspendFinally_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/inlineSuspendFinally.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("inlineSuspendFinally.kt")
+            public void testInlineSuspendFinally_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/inlineSuspendFinally.kt", "kotlin.coroutines");
+            }
+
            @TestMetadata("safeCallOnTwoReceiversLong.kt")
            public void testSafeCallOnTwoReceiversLong_1_2() throws Exception {
                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/safeCallOnTwoReceiversLong.kt", "kotlin.coroutines.experimental");
@@ -6865,6 +6875,16 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                runTest("compiler/testData/codegen/box/coroutines/featureIntersection/suspendFunctionIsAs.kt");
            }

+            @TestMetadata("suspendInlineSuspendFinally.kt")
+            public void testSuspendInlineSuspendFinally_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/suspendInlineSuspendFinally.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("suspendInlineSuspendFinally.kt")
+            public void testSuspendInlineSuspendFinally_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/suspendInlineSuspendFinally.kt", "kotlin.coroutines");
+            }
+
            @TestMetadata("suspendOperatorPlusAssign.kt")
            public void testSuspendOperatorPlusAssign_1_2() throws Exception {
                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/suspendOperatorPlusAssign.kt", "kotlin.coroutines.experimental");
@@ -5360,6 +5360,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/destructuringInLambdas.kt", "kotlin.coroutines");
            }

+            @TestMetadata("inlineSuspendFinally.kt")
+            public void testInlineSuspendFinally_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/inlineSuspendFinally.kt", "kotlin.coroutines");
+            }
+
            @TestMetadata("safeCallOnTwoReceiversLong.kt")
            public void testSafeCallOnTwoReceiversLong_1_3() throws Exception {
                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/safeCallOnTwoReceiversLong.kt", "kotlin.coroutines");
@@ -5385,6 +5390,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
                runTest("compiler/testData/codegen/box/coroutines/featureIntersection/suspendFunctionIsAs.kt");
            }

+            @TestMetadata("suspendInlineSuspendFinally.kt")
+            public void testSuspendInlineSuspendFinally_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/suspendInlineSuspendFinally.kt", "kotlin.coroutines");
+            }
+
            @TestMetadata("suspendOperatorPlusAssign.kt")
            public void testSuspendOperatorPlusAssign_1_3() throws Exception {
                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/suspendOperatorPlusAssign.kt", "kotlin.coroutines");
@@ -5915,6 +5915,16 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/destructuringInLambdas.kt", "kotlin.coroutines");
            }

+            @TestMetadata("inlineSuspendFinally.kt")
+            public void testInlineSuspendFinally_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/inlineSuspendFinally.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("inlineSuspendFinally.kt")
+            public void testInlineSuspendFinally_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/inlineSuspendFinally.kt", "kotlin.coroutines");
+            }
+
            @TestMetadata("safeCallOnTwoReceiversLong.kt")
            public void testSafeCallOnTwoReceiversLong_1_2() throws Exception {
                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/safeCallOnTwoReceiversLong.kt", "kotlin.coroutines.experimental");
@@ -5955,6 +5965,16 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                runTest("compiler/testData/codegen/box/coroutines/featureIntersection/suspendFunctionIsAs.kt");
            }

+            @TestMetadata("suspendInlineSuspendFinally.kt")
+            public void testSuspendInlineSuspendFinally_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/suspendInlineSuspendFinally.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("suspendInlineSuspendFinally.kt")
+            public void testSuspendInlineSuspendFinally_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/suspendInlineSuspendFinally.kt", "kotlin.coroutines");
+            }
+
            @TestMetadata("suspendOperatorPlusAssign.kt")
            public void testSuspendOperatorPlusAssign_1_2() throws Exception {
                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/featureIntersection/suspendOperatorPlusAssign.kt", "kotlin.coroutines.experimental");