[K/N][IR] Implemented suspend tail call optimization
If for a suspend function all calls to other suspend functions a tail calls, it is not needed to build the coroutine implementation.
This commit is contained in:
Igor Chevdar
+110
@@ -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<IrCall>, 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<IrCall>()
|
||||
val tailReturnableBlocks = mutableSetOf<IrReturnableBlockSymbol>()
|
||||
|
||||
val visitor = object : IrElementVisitor<Unit, VisitorState> {
|
||||
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)
|
||||
}
|
||||
+36
-90
@@ -63,19 +63,16 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(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<C : CommonBackendContext>(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<IrDeclaration>? {
|
||||
val (tailSuspendCalls, hasNotTailSuspendCalls) = collectTailSuspendCalls(context, irFunction)
|
||||
return when {
|
||||
irFunction in suspendLambdas -> {
|
||||
// Suspend lambdas always need coroutine implementation.
|
||||
// They are called through factory method <create>, 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<IrDeclaration>(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<IrCall>) {
|
||||
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 <create>,
|
||||
listOf(coroutine) // thus we can eliminate original body.
|
||||
else
|
||||
listOf<IrDeclaration>(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
|
||||
|
||||
Reference in New Issue
Block a user