diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/TailSuspendCallsCollector.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/TailSuspendCallsCollector.kt new file mode 100644 index 00000000000..2b98855f353 --- /dev/null +++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/TailSuspendCallsCollector.kt @@ -0,0 +1,110 @@ +/* + * Copyright 2010-2022 JetBrains s.r.o. and Kotlin Programming Language contributors. + * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file. + */ + +package org.jetbrains.kotlin.backend.common + +import org.jetbrains.kotlin.ir.IrElement +import org.jetbrains.kotlin.ir.declarations.IrSimpleFunction +import org.jetbrains.kotlin.ir.expressions.* +import org.jetbrains.kotlin.ir.symbols.IrReturnableBlockSymbol +import org.jetbrains.kotlin.ir.types.isUnit +import org.jetbrains.kotlin.ir.util.isSuspend +import org.jetbrains.kotlin.ir.util.render +import org.jetbrains.kotlin.ir.visitors.IrElementVisitor + +data class TailSuspendCalls(val callSites: Set, val hasNotTailSuspendCalls: Boolean) + +/** + * Collects calls to be treated as tail calls: "last" expressions which are either direct return statement with a call + * to other suspend function, or, for a Unit-returning function, a call to other suspend function, also returning Unit. + */ +fun collectTailSuspendCalls(context: CommonBackendContext, irFunction: IrSimpleFunction): TailSuspendCalls { + require(irFunction.isSuspend) { "A suspend function expected: ${irFunction.render()}" } + val body = irFunction.body ?: return TailSuspendCalls(emptySet(), false) + + class VisitorState(val insideTryBlock: Boolean, val isTailExpression: Boolean) + + val isUnitReturn = irFunction.returnType.isUnit() + var hasNotTailSuspendCall = false + val tailSuspendCalls = mutableSetOf() + val tailReturnableBlocks = mutableSetOf() + + val visitor = object : IrElementVisitor { + override fun visitElement(element: IrElement, data: VisitorState) { + element.acceptChildren(this, VisitorState(data.insideTryBlock, isTailExpression = false)) + } + + override fun visitTry(aTry: IrTry, data: VisitorState) { + aTry.tryResult.accept(this, VisitorState(insideTryBlock = true, isTailExpression = false)) + aTry.catches.forEach { it.result.accept(this, data) } + require(aTry.finallyExpression == null) { "All finally clauses should've been lowered out" } + } + + private fun isTailReturn(expression: IrReturn) = + expression.returnTargetSymbol == irFunction.symbol || expression.returnTargetSymbol in tailReturnableBlocks + + override fun visitReturn(expression: IrReturn, data: VisitorState) { + expression.value.accept(this, VisitorState(data.insideTryBlock, isTailReturn(expression))) + } + + override fun visitExpressionBody(body: IrExpressionBody, data: VisitorState) = + body.acceptChildren(this, data) + + override fun visitBlockBody(body: IrBlockBody, data: VisitorState) = + visitStatementContainer(body, data) + + override fun visitContainerExpression(expression: IrContainerExpression, data: VisitorState) { + if (expression is IrReturnableBlock && data.isTailExpression) + tailReturnableBlocks.add(expression.symbol) + visitStatementContainer(expression, data) + } + + private fun visitStatementContainer(expression: IrStatementContainer, data: VisitorState) { + expression.statements.forEachIndexed { index, irStatement -> + val isTailStatement = if (index == expression.statements.lastIndex) { + // The last statement defines the result of the container expression, so it has the same kind. + // Note: this is even true for returnable blocks: if it is a Unit-returning block, this is exactly + // like a usual block; if it is a non-Unit block, then it must end with an explicit return statement. + data.isTailExpression + } else { + // In a Unit-returning function, any statement directly followed by a `return` is a tail statement. + isUnitReturn && expression.statements[index + 1].let { + it is IrReturn && isTailReturn(it) && it.value.isUnitRead() + } + } + irStatement.accept(this, VisitorState(data.insideTryBlock, isTailStatement)) + } + } + + override fun visitWhen(expression: IrWhen, data: VisitorState) { + expression.branches.forEach { + it.condition.accept(this, VisitorState(data.insideTryBlock, isTailExpression = false)) + it.result.accept(this, data) + } + } + + override fun visitCall(expression: IrCall, data: VisitorState) { + if (expression.isSuspend) { + if (!data.insideTryBlock && data.isTailExpression) + tailSuspendCalls.add(expression) + else + hasNotTailSuspendCall = true + } + + // For a tail call, [returnIfSuspended] can be optimized away. + val isTailExpression = data.isTailExpression && expression.isReturnIfSuspendedCall() + expression.acceptChildren(this, VisitorState(data.insideTryBlock, isTailExpression)) + } + + private fun IrExpression.isUnitRead(): Boolean = + this is IrGetObjectValue && symbol == context.irBuiltIns.unitClass + + private fun IrCall.isReturnIfSuspendedCall() = + symbol == context.ir.symbols.returnIfSuspended + } + + body.accept(visitor, VisitorState(insideTryBlock = false, isTailExpression = true)) + return TailSuspendCalls(tailSuspendCalls, hasNotTailSuspendCall) +} \ No newline at end of file diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/AbstractSuspendFunctionsLowering.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/AbstractSuspendFunctionsLowering.kt index 5c7208593d0..e38965be5d8 100644 --- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/AbstractSuspendFunctionsLowering.kt +++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/AbstractSuspendFunctionsLowering.kt @@ -63,19 +63,16 @@ abstract class AbstractSuspendFunctionsLowering(val co override fun visitClass(declaration: IrClass) { declaration.acceptChildrenVoid(this) - declaration.transformDeclarationsFlat(::tryTransformSuspendFunction) + if (declaration.origin != DECLARATION_ORIGIN_COROUTINE_IMPL) + declaration.transformDeclarationsFlat(::tryTransformSuspendFunction) } }) } - - // Suppress since it is used in native - @Suppress("MemberVisibilityCanBePrivate") protected fun IrCall.isReturnIfSuspendedCall() = - symbol.signature == context.ir.symbols.returnIfSuspended.signature + symbol == context.ir.symbols.returnIfSuspended private fun tryTransformSuspendFunction(element: IrElement) = - if (element is IrSimpleFunction && element.isSuspend && element.modality != Modality.ABSTRACT) transformSuspendFunction(element, suspendLambdas[element]) else null @@ -123,103 +120,52 @@ abstract class AbstractSuspendFunctionsLowering(val co }) } - private sealed class SuspendFunctionKind { - object NO_SUSPEND_CALLS : SuspendFunctionKind() - class DELEGATING(val delegatingCall: IrCall) : SuspendFunctionKind() - object NEEDS_STATE_MACHINE : SuspendFunctionKind() + private fun transformSuspendFunction(irFunction: IrSimpleFunction, functionReference: IrFunctionReference?): List? { + val (tailSuspendCalls, hasNotTailSuspendCalls) = collectTailSuspendCalls(context, irFunction) + return when { + irFunction in suspendLambdas -> { + // Suspend lambdas always need coroutine implementation. + // They are called through factory method , thus we can eliminate original body. + listOf(buildCoroutine(irFunction, functionReference)) + } + + // TODO: May be optimize tail suspend calls even in case of a state machine? + hasNotTailSuspendCalls -> listOf(buildCoroutine(irFunction, functionReference), irFunction) + + else -> { + // Otherwise, no suspend calls at all or all of them are tail calls - no need in a state machine. + // Have to simplify them though (convert them to proper return statements). + simplifyTailSuspendCalls(irFunction, tailSuspendCalls) + null + } + } } - private fun transformSuspendFunction(irFunction: IrSimpleFunction, functionReference: IrFunctionReference?) = - when (val suspendFunctionKind = getSuspendFunctionKind(irFunction)) { - is SuspendFunctionKind.NO_SUSPEND_CALLS -> { - null // No suspend function calls - just an ordinary function. - } + private fun simplifyTailSuspendCalls(irFunction: IrSimpleFunction, tailSuspendCalls: Set) { + if (tailSuspendCalls.isEmpty()) return - is SuspendFunctionKind.DELEGATING -> { // Calls another suspend function at the end. - removeReturnIfSuspendedCallAndSimplifyDelegatingCall(irFunction, suspendFunctionKind.delegatingCall) - null // No need in state machine. - } + val irBuilder = context.createIrBuilder(irFunction.symbol) + irFunction.body!!.transformChildrenVoid(object : IrElementTransformerVoid() { + override fun visitCall(expression: IrCall): IrExpression { + val shortCut = if (expression.isReturnIfSuspendedCall()) + expression.getValueArgument(0)!! + else expression - is SuspendFunctionKind.NEEDS_STATE_MACHINE -> { - val coroutine = buildCoroutine(irFunction, functionReference) // Coroutine implementation. - if (irFunction in suspendLambdas) // Suspend lambdas are called through factory method , - listOf(coroutine) // thus we can eliminate original body. - else - listOf(coroutine, irFunction) - } - } + shortCut.transformChildrenVoid(this) - private fun getSuspendFunctionKind(irFunction: IrSimpleFunction): SuspendFunctionKind { - if (irFunction in suspendLambdas) - return SuspendFunctionKind.NEEDS_STATE_MACHINE // Suspend lambdas always need coroutine implementation. - - val body = irFunction.body ?: return SuspendFunctionKind.NO_SUSPEND_CALLS - - var numberOfSuspendCalls = 0 - body.acceptVoid(object : IrElementVisitorVoid { - override fun visitElement(element: IrElement) { - element.acceptChildrenVoid(this) - } - - override fun visitCall(expression: IrCall) { - expression.acceptChildrenVoid(this) - if (expression.isSuspend) - ++numberOfSuspendCalls + return if (!expression.isSuspend) + shortCut + else irBuilder.at(expression).irReturn( + irBuilder.generateDelegatedCall(irFunction.returnType, shortCut) + ) } }) - // It is important to optimize the case where there is only one suspend call and it is the last statement - // because we don't need to build a fat coroutine class in that case. - // This happens a lot in practise because of suspend functions with default arguments. - // TODO: use TailRecursionCallsCollector. - val lastCall = when (val lastStatement = (body as IrBlockBody).statements.lastOrNull()) { - is IrCall -> lastStatement - is IrReturn -> { - var value: IrElement = lastStatement - /* - * Check if matches this pattern: - * block/return { - * block/return { - * .. suspendCall() - * } - * } - */ - loop@ while (true) { - value = when { - value is IrBlock && value.statements.size == 1 -> value.statements.first() - value is IrReturn -> value.value - else -> break@loop - } - } - value as? IrCall - } - else -> null - } - val suspendCallAtEnd = lastCall != null && lastCall.isSuspend // Suspend call. - return when { - numberOfSuspendCalls == 0 -> SuspendFunctionKind.NO_SUSPEND_CALLS - numberOfSuspendCalls == 1 - && suspendCallAtEnd -> SuspendFunctionKind.DELEGATING(lastCall!!) - else -> SuspendFunctionKind.NEEDS_STATE_MACHINE - } } private val symbols = context.ir.symbols private val getContinuationSymbol = symbols.getContinuation private val continuationClassSymbol = getContinuationSymbol.owner.returnType.classifierOrFail as IrClassSymbol - private fun removeReturnIfSuspendedCallAndSimplifyDelegatingCall(irFunction: IrFunction, delegatingCall: IrCall) { - val returnValue = - if (delegatingCall.isReturnIfSuspendedCall()) - delegatingCall.getValueArgument(0)!! - else delegatingCall - context.createIrBuilder(irFunction.symbol).run { - val statements = (irFunction.body as IrBlockBody).statements - val lastStatement = statements.last() - assert(lastStatement == delegatingCall || lastStatement is IrReturn) { "Unexpected statement $lastStatement" } - statements[statements.lastIndex] = irReturn(generateDelegatedCall(irFunction.returnType, returnValue)) - } - } - private fun buildCoroutine(irFunction: IrSimpleFunction, functionReference: IrFunctionReference?): IrClass { val coroutine = CoroutineBuilder(irFunction, functionReference).build() builtCoroutines[irFunction] = coroutine