[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

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/AbstractSuspendFunctionsLowering.kt

@@ -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
         }
     }

compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/FinallyBlocksLowering.kt

@@ -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,
+                type = context.irBuiltIns.unitType
-                tryResult = transformedTry,
+            ).apply {
-                catches = listOf(
+                this.catches += irCatch(catchParameter).apply {
-                    irCatch(catchParameter).apply {
+                    result = irComposite {
-                        result = irComposite {
+                        +finallyExpression.copy()
-                            +finallyExpression.copy()
+                        +irThrow(irGet(catchParameter))
-                            +irThrow(irGet(catchParameter))
+                    }
-                        }
+                }
-                    }),
+
-                finallyExpression = null
+                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)
+                val returnedResult = irReturn(fallThroughSymbol, transformedResult)
-                for (aCatch in aTry.catches) {
+
-                    val transformedCatch = aCatch.transform(transformer, null)
+                if (aTry.catches.isNotEmpty()) {
-                    transformedCatch.result = irReturn(fallThroughSymbol, transformedCatch.result)
+                    val transformedTry = IrTryImpl(
-                    transformedTry.catches.add(transformedCatch)
+                        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)
             }
         }

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/JsLoweringPhases.kt

@@ -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

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/coroutines/JsSuspendFunctionsLowering.kt

@@ -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 =
+        val simplifiedFunction = transformingFunction.transform(FinallyBlocksLowering(context, context.dynamicType), null) as IrFunction
-            (originalBody as IrBlockBody).run {
+        val originalBody = simplifiedFunction.body as IrBlockBody
-                IrBlockImpl(
+
-                    transformingFunction.startOffset,
+        val body = IrBlockImpl(
-                    transformingFunction.endOffset,
+            simplifiedFunction.startOffset,
-                    context.irBuiltIns.unitType,
+            simplifiedFunction.endOffset,
-                    STATEMENT_ORIGIN_COROUTINE_IMPL,
+            context.irBuiltIns.unitType,
-                    statements
+            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 suspendableNodes = collectSuspendableNodes(body)
-        val loweredBody =
-            collectSuspendableNodes(body, suspendableNodes, context, stateMachineFunction, context.dynamicType)
         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>()
+        assignStateIds(stateMachineBuilder.entryState, thisReceiver, switch, rootLoop)
-
-        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)
 
         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

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/coroutines/StateMachineBuilder.kt

@@ -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.*
-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.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)
-
-data class TryState(val tryState: SuspendState, val catchState: SuspendState, val finallyState: FinallyTargets?)
 
 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 (!isBlockEnded()) {
-        if (lastStatement !is IrReturn && lastStatement !is IrContinue && lastStatement !is IrThrow) {
             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
+                if (!isBlockEnded()) {
-                val lastStatement = currentBlock.statements.lastOrNull()
+                    doDispatch(exitState)
-                if (lastStatement != null && lastStatement !is IrContinue) {
-                    if (currentState !== rootState) {
-                        doDispatch(exitState)
-                    }
                 }
 
                 currentState = dispatchState
                 currentBlock = elseBlock
             } else {
                 branch.result.acceptVoid(this)
-
+                if (!isBlockEnded()) {
-                // TODO: block should not be empty
+                    doDispatch(exitState)
-                val lastStatement = currentBlock.statements.lastOrNull()
-                if (lastStatement != null && lastStatement !is IrContinue) {
-                    if (currentState !== rootState) {
-                        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) {
+        if (!isBlockEnded()) {
-            doDispatch(tryState.finallyState.normal)
-        } else {
             setupExceptionState(enclosingCatch)
             doDispatch(exitState)
         }
-
         addExceptionEdge()
+
         catchBlockStack.pop()
         updateState(tryState.catchState)
 
-        if (tryState.finallyState != null) {
+        setupExceptionState(enclosingCatch)
-            setupExceptionState(tryState.finallyState.fromThrow)
-        } else {
-            setupExceptionState(enclosingCatch)
-        }
-
 
         var rethrowNeeded = true
 
@@ -628,9 +602,7 @@ class StateMachineBuilder(
 
                 addStatement(irVar)
                 catchResult.acceptVoid(this)
-                val exitDispatch = tryState.finallyState?.run { normal } ?: exitState
+                maybeDoDispatch(exitState)
-                maybeDoDispatch(exitDispatch)
-
             } 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(exitState)
-                maybeDoDispatch(exitDispatch)
 
                 currentBlock = ifBlock
                 addStatement(irIf)
@@ -658,42 +629,10 @@ class StateMachineBuilder(
             addStatement(JsIrBuilder.buildThrow(nothing, pendingException()))
         }
 
-        if (tryState.finallyState == null) {
+        currentState.successors += enclosingCatch
-            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()
-        }
 
         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) =
+    private fun buildTryState() = TryState(currentState, SuspendState(unit))
-        TryState(
-            currentState,
-            SuspendState(unit),
-            aTry.finallyExpression?.run {
-                FinallyTargets(
-                    SuspendState(
-                        unit
-                    ), SuspendState(unit)
-                )
-            }
-        )
-
 
     private fun buildIsCheck(value: IrExpression, toType: IrType) =
         JsIrBuilder.buildTypeOperator(

compiler/ir/backend.js/src/org/jetbrains/kotlin/ir/backend/js/lower/coroutines/SuspendLoweringUtils.kt

@@ -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
+            markNode(expression)
-            hasSuspendableChildren = true
         }
     }
 }
 
-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) {
+        if (isSuspendableNode(jump.loop)) {
-            suspendableNodes.add(jump)
+            markNode(jump)
-            hasSuspendableChildren = true
         }
-
-        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) {
+        if (expression.returnTargetSymbol is IrReturnableBlockSymbol && isSuspendableNode(expression.returnTargetSymbol.owner)) {
-            suspendableNodes.add(expression)
+            markNode(expression)
-            hasSuspendableChildren = true
         }
     }
 }
 
+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>,

compiler/testData/codegen/box/coroutines/featureIntersection/inlineSuspendFinally.kt

@@ -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"
+}

compiler/testData/codegen/box/coroutines/featureIntersection/suspendInlineSuspendFinally.kt

@@ -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"
+}

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -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");

compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -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");

compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -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");

js/js.tests/test/org/jetbrains/kotlin/js/test/ir/semantics/IrJsCodegenBoxTestGenerated.java

@@ -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");

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

@@ -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");