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