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 e38965be5d8..051025f0c1b 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 @@ -35,7 +35,8 @@ abstract class AbstractSuspendFunctionsLowering(val co protected abstract fun buildStateMachine( stateMachineFunction: IrFunction, transformingFunction: IrFunction, - argumentToPropertiesMap: Map + argumentToPropertiesMap: Map, + tailSuspendCalls: Set ) protected abstract fun IrBlockBodyBuilder.generateCoroutineStart(invokeSuspendFunction: IrFunction, receiver: IrExpression) @@ -126,11 +127,12 @@ abstract class AbstractSuspendFunctionsLowering(val co 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)) + // 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(buildCoroutine(irFunction, functionReference), irFunction) + hasNotTailSuspendCalls -> + listOf(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(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(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): IrClass { + val coroutine = CoroutineBuilder(irFunction, functionReference, tailSuspendCalls).build() builtCoroutines[irFunction] = coroutine if (functionReference == null) { @@ -204,7 +206,11 @@ abstract class AbstractSuspendFunctionsLowering(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 + ) { 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(val co suspendFunctionClassTypeArguments = unboundParameterTypes + irFunction.returnType superTypes += suspendFunctionClass.typeWith(suspendFunctionClassTypeArguments) } + coroutineClass.superTypes = superTypes val coroutineConstructor = buildConstructor() @@ -285,7 +292,6 @@ abstract class AbstractSuspendFunctionsLowering(val co ) } - coroutineClass.superTypes += superTypes coroutineClass.addFakeOverrides( context.typeSystem, ignoredParentSymbols = ignoredParentSymbols @@ -515,7 +521,7 @@ abstract class AbstractSuspendFunctionsLowering(val co overriddenSymbols += stateMachineFunction.symbol } - buildStateMachine(function, irFunction, argumentToPropertiesMap) + buildStateMachine(function, irFunction, argumentToPropertiesMap, tailSuspendCalls) return function } } diff --git a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/ir/Ir.kt b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/ir/Ir.kt index f54ff2c35b9..d169f4b1bd6 100644 --- a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/ir/Ir.kt +++ b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/ir/Ir.kt @@ -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) diff --git a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/NativeSuspendFunctionLowering.kt b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/NativeSuspendFunctionLowering.kt index 3d9581a793d..7f47c72ffa0 100644 --- a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/NativeSuspendFunctionLowering.kt +++ b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/NativeSuspendFunctionLowering.kt @@ -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) { + argumentToPropertiesMap: Map, + tailSuspendCalls: Set) { 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 { + 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)