azalea/kotlin-fork
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[K/N][IR] Suspend tail call optimization for coroutines

Browse Source 
Even if a coroutine class has to be built for a suspend function, we
still can perform tail call optimization for it.
This commit is contained in:
Igor Chevdar
2022-11-10 15:17:38 +02:00
committed by Space Team
parent b100dfe8be
commit e16cfd5485
3 changed files with 94 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/lower/AbstractSuspendFunctionsLowering.kt
+16 -10
View File
@@ -35,7 +35,8 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
protected abstract fun buildStateMachine(
stateMachineFunction: IrFunction,
transformingFunction: IrFunction,
argumentToPropertiesMap: Map<IrValueParameter, IrField>
argumentToPropertiesMap: Map<IrValueParameter, IrField>,
tailSuspendCalls: Set<IrCall>
)
protected abstract fun IrBlockBodyBuilder.generateCoroutineStart(invokeSuspendFunction: IrFunction, receiver: IrExpression)
@@ -126,11 +127,12 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
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))
// Tail call optimization for suspend lambdas isn't yet supported by stdlib.
listOf(buildCoroutine(irFunction, functionReference, emptySet()))
}
// TODO: May be optimize tail suspend calls even in case of a state machine?
hasNotTailSuspendCalls -> listOf<IrDeclaration>(buildCoroutine(irFunction, functionReference), irFunction)
hasNotTailSuspendCalls ->
listOf<IrDeclaration>(buildCoroutine(irFunction, functionReference, tailSuspendCalls), irFunction)
else -> {
// Otherwise, no suspend calls at all or all of them are tail calls - no need in a state machine.
@@ -153,7 +155,7 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
shortCut.transformChildrenVoid(this)
return if (!expression.isSuspend)
return if (!expression.isSuspend || expression !in tailSuspendCalls)
shortCut
else irBuilder.at(expression).irReturn(
irBuilder.generateDelegatedCall(irFunction.returnType, shortCut)
@@ -166,8 +168,8 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
private val getContinuationSymbol = symbols.getContinuation
private val continuationClassSymbol = getContinuationSymbol.owner.returnType.classifierOrFail as IrClassSymbol
private fun buildCoroutine(irFunction: IrSimpleFunction, functionReference: IrFunctionReference?): IrClass {
val coroutine = CoroutineBuilder(irFunction, functionReference).build()
private fun buildCoroutine(irFunction: IrSimpleFunction, functionReference: IrFunctionReference?, tailSuspendCalls: Set<IrCall>): IrClass {
val coroutine = CoroutineBuilder(irFunction, functionReference, tailSuspendCalls).build()
builtCoroutines[irFunction] = coroutine
if (functionReference == null) {
@@ -204,7 +206,11 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
val stateMachineFunction: IrFunction
)
private inner class CoroutineBuilder(val irFunction: IrFunction, val functionReference: IrFunctionReference?) {
private inner class CoroutineBuilder(
val irFunction: IrFunction,
val functionReference: IrFunctionReference?,
val tailSuspendCalls: Set<IrCall>
) {
private val functionParameters = irFunction.explicitParameters
private val boundFunctionParameters = functionReference?.getArgumentsWithIr()?.map { it.first }
private val unboundFunctionParameters = boundFunctionParameters?.let { functionParameters - it }
@@ -251,6 +257,7 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
suspendFunctionClassTypeArguments = unboundParameterTypes + irFunction.returnType
superTypes += suspendFunctionClass.typeWith(suspendFunctionClassTypeArguments)
}
coroutineClass.superTypes = superTypes
val coroutineConstructor = buildConstructor()
@@ -285,7 +292,6 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
)
}
coroutineClass.superTypes += superTypes
coroutineClass.addFakeOverrides(
context.typeSystem,
ignoredParentSymbols = ignoredParentSymbols
@@ -515,7 +521,7 @@ abstract class AbstractSuspendFunctionsLowering<C : CommonBackendContext>(val co
overriddenSymbols += stateMachineFunction.symbol
}
buildStateMachine(function, irFunction, argumentToPropertiesMap)
buildStateMachine(function, irFunction, argumentToPropertiesMap, tailSuspendCalls)
return function
}
}
kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/ir/Ir.kt
+6
View File
@@ -351,6 +351,12 @@ internal class KonanSymbols(
val invokeSuspendFunction =
irBuiltIns.findBuiltInClassMemberFunctions(baseContinuationImpl, Name.identifier("invokeSuspend")).single()
val completionGetter = symbolTable.referenceSimpleFunction(
baseContinuationImpl.descriptor.unsubstitutedMemberScope
.getContributedVariables(Name.identifier("completion"), NoLookupLocation.FROM_BACKEND).single()
.getter!!
)
override val coroutineSuspendedGetter = symbolTable.referenceSimpleFunction(
coroutinesIntrinsicsPackage
.getContributedVariables(COROUTINE_SUSPENDED_NAME, NoLookupLocation.FROM_BACKEND)
kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/NativeSuspendFunctionLowering.kt
+72 -3
View File
@@ -2,6 +2,7 @@ package org.jetbrains.kotlin.backend.konan.lower
import org.jetbrains.kotlin.backend.common.descriptors.synthesizedName
import org.jetbrains.kotlin.backend.common.lower.*
import org.jetbrains.kotlin.backend.common.lower.coroutines.getOrCreateFunctionWithContinuationStub
import org.jetbrains.kotlin.backend.konan.Context
import org.jetbrains.kotlin.descriptors.Modality
import org.jetbrains.kotlin.descriptors.DescriptorVisibilities
@@ -58,15 +59,20 @@ internal class NativeSuspendFunctionsLowering(ctx: Context) : AbstractSuspendFun
)
}
private val getContinuation = context.ir.symbols.getContinuation
private val completionGetter = context.ir.symbols.completionGetter
override fun buildStateMachine(stateMachineFunction: IrFunction,
transformingFunction: IrFunction,
argumentToPropertiesMap: Map<IrValueParameter, IrField>) {
argumentToPropertiesMap: Map<IrValueParameter, IrField>,
tailSuspendCalls: Set<IrCall>) {
val originalBody = transformingFunction.body!!
val resultArgument = stateMachineFunction.valueParameters.single()
val coroutineClass = stateMachineFunction.parentAsClass
val thisReceiver = stateMachineFunction.dispatchReceiverParameter!!
val labelField = coroutineClass.addField(Name.identifier("label"), symbols.nativePtrType, true)
val startOffset = transformingFunction.startOffset
@@ -74,6 +80,71 @@ internal class NativeSuspendFunctionsLowering(ctx: Context) : AbstractSuspendFun
val irBuilder = context.createIrBuilder(stateMachineFunction.symbol, startOffset, endOffset)
stateMachineFunction.body = irBuilder.irBlockBody(startOffset, endOffset) {
if (tailSuspendCalls.isNotEmpty()) {
/*
* Usual suspend call of, say, function foo will be transformed to something like this:
* val result = foo(.., continuation = this)
* if (result == COROUTINE_SUSPENDED) return COROUTINE_SUSPENDED
* if (result.failed)
* this.completion.resumeWithException(result.exception)
* else this.completion.resumeWith(result.value)
*
* But, if a call to foo is a tail call, we can replace the above with just:
* return foo(.., continuation = this.completion)
*
* The visitor below does exactly this.
*/
originalBody.transformChildren(object : IrElementTransformer<Boolean> {
fun IrBuilderWithScope.irGetCompletion() = irCall(completionGetter).apply { dispatchReceiver = irGet(thisReceiver) }
override fun visitCall(expression: IrCall, /* substituteContinuation */ data: Boolean): IrExpression {
val isTailSuspendCall = expression in tailSuspendCalls
return when {
expression.symbol == getContinuation -> {
if (data)
irBuilder.at(expression).irGetCompletion()
else
expression
}
expression.isReturnIfSuspendedCall() -> {
if (!isTailSuspendCall)
expression
else {
expression.transformChildren(this, /* substituteContinuation = */ true)
irBuilder.at(expression).irReturn(expression.getValueArgument(0)!!)
}
}
else -> {
// Only the top-most call should have its continuation substituted.
val substituteContinuation = data && !expression.isSuspendCall
expression.transformChildren(this, substituteContinuation)
if (!isTailSuspendCall)
expression
else {
val oldFun = expression.symbol.owner
val newFun: IrSimpleFunction = oldFun.getOrCreateFunctionWithContinuationStub(this@NativeSuspendFunctionsLowering.context)
irBuilder.at(expression).irReturn(
irCall(
expression,
newFun.symbol,
newReturnType = newFun.returnType,
newSuperQualifierSymbol = expression.superQualifierSymbol
).also {
it.putValueArgument(it.valueArgumentsCount - 1, irGetCompletion())
}
)
}
}
}
}
}, /* substituteContinuation = */ false)
}
val suspendResult = irVar("suspendResult".synthesizedName, context.irBuiltIns.anyNType, true)
@@ -100,8 +171,6 @@ internal class NativeSuspendFunctionsLowering(ctx: Context) : AbstractSuspendFun
originalBody.transformChildrenVoid(object : IrElementTransformerVoid() {
private val thisReceiver = stateMachineFunction.dispatchReceiverParameter!!
// Replace returns to refer to the new function.
override fun visitReturn(expression: IrReturn): IrExpression {
expression.transformChildrenVoid(this)