From 716853dd007c89b536f0b29e64494111aec17ed5 Mon Sep 17 00:00:00 2001 From: Igor Chevdar Date: Thu, 2 Feb 2023 20:45:08 +0200 Subject: [PATCH] [K/N][IR] Tail call optimization for suspend lambdas --- .../lower/AbstractSuspendFunctionsLowering.kt | 323 ++---------------- .../konan/lower/FunctionReferenceLowering.kt | 5 - 2 files changed, 23 insertions(+), 305 deletions(-) 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 1fae70c5f66..8e2229a0807 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 @@ -15,13 +15,10 @@ import org.jetbrains.kotlin.ir.declarations.* import org.jetbrains.kotlin.ir.expressions.* import org.jetbrains.kotlin.ir.expressions.impl.IrInstanceInitializerCallImpl import org.jetbrains.kotlin.ir.symbols.IrClassSymbol -import org.jetbrains.kotlin.ir.symbols.IrSimpleFunctionSymbol -import org.jetbrains.kotlin.ir.symbols.IrSymbol import org.jetbrains.kotlin.ir.types.* import org.jetbrains.kotlin.ir.util.* import org.jetbrains.kotlin.ir.visitors.* import org.jetbrains.kotlin.name.Name -import org.jetbrains.kotlin.util.OperatorNameConventions abstract class AbstractSuspendFunctionsLowering(val context: C) : FileLoweringPass { @@ -46,16 +43,7 @@ abstract class AbstractSuspendFunctionsLowering(val co protected open fun IrBuilderWithScope.generateDelegatedCall(expectedType: IrType, delegatingCall: IrExpression): IrExpression = delegatingCall - private val builtCoroutines = mutableMapOf() - private val suspendLambdas = mutableMapOf() - override fun lower(irFile: IrFile) { - markSuspendLambdas(irFile) - buildCoroutines(irFile) - transformCallableReferencesToSuspendLambdas(irFile) - } - - private fun buildCoroutines(irFile: IrFile) { irFile.transformDeclarationsFlat(::tryTransformSuspendFunction) irFile.acceptVoid(object : IrElementVisitorVoid { override fun visitElement(element: IrElement) { @@ -70,79 +58,24 @@ abstract class AbstractSuspendFunctionsLowering(val co }) } - protected fun IrCall.isReturnIfSuspendedCall() = - symbol == context.ir.symbols.returnIfSuspended + private fun tryTransformSuspendFunction(element: IrElement): List? { + val function = (element as? IrSimpleFunction) ?: return null + if (!function.isSuspend || function.modality == Modality.ABSTRACT) return null - private fun tryTransformSuspendFunction(element: IrElement) = - if (element is IrSimpleFunction && element.isSuspend && element.modality != Modality.ABSTRACT) - transformSuspendFunction(element, suspendLambdas[element]) - else null - - private fun markSuspendLambdas(irElement: IrElement) { - irElement.acceptChildrenVoid(object : IrElementVisitorVoid { - override fun visitElement(element: IrElement) { - element.acceptChildrenVoid(this) - } - - override fun visitFunctionReference(expression: IrFunctionReference) { - expression.acceptChildrenVoid(this) - - if (expression.isSuspend) { - suspendLambdas[expression.symbol.owner] = expression - } - } - }) - } - - private fun transformCallableReferencesToSuspendLambdas(irElement: IrElement) { - irElement.transformChildrenVoid(object : IrElementTransformerVoid() { - - override fun visitFunctionReference(expression: IrFunctionReference): IrExpression { - expression.transformChildrenVoid(this) - - if (!expression.isSuspend) - return expression - val coroutine = builtCoroutines[expression.symbol.owner] - ?: throw Error("Non-local callable reference to suspend lambda: $expression") - val constructorParameters = coroutine.coroutineConstructor.valueParameters - val expressionArguments = expression.getArgumentsWithIr().map { it.second } - assert(constructorParameters.size == expressionArguments.size) { - "Inconsistency between callable reference to suspend lambda and the corresponding coroutine" - } - val irBuilder = context.createIrBuilder(expression.symbol, expression.startOffset, expression.endOffset) - irBuilder.run { - return irCall(coroutine.coroutineConstructor.symbol).apply { - expressionArguments.forEachIndexed { index, argument -> - putValueArgument(index, argument) - } - } - } - } - }) - } - - 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. - // Tail call optimization for suspend lambdas isn't yet supported by stdlib. - listOf(buildCoroutine(irFunction, functionReference, emptySet())) - } - - 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. - // Have to simplify them though (convert them to proper return statements). - simplifyTailSuspendCalls(irFunction, tailSuspendCalls) - null - } + val (tailSuspendCalls, hasNotTailSuspendCalls) = collectTailSuspendCalls(context, function) + return if (hasNotTailSuspendCalls) { + listOf(buildCoroutine(function, tailSuspendCalls).clazz, function) + } 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(function, tailSuspendCalls) + null } } + protected fun IrCall.isReturnIfSuspendedCall() = + symbol == context.ir.symbols.returnIfSuspended + private fun simplifyTailSuspendCalls(irFunction: IrSimpleFunction, tailSuspendCalls: Set) { if (tailSuspendCalls.isEmpty()) return @@ -168,15 +101,12 @@ 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?, tailSuspendCalls: Set): IrClass { - val coroutine = CoroutineBuilder(irFunction, functionReference, tailSuspendCalls).build() - builtCoroutines[irFunction] = coroutine - - if (functionReference == null) { - // It is not a lambda - replace original function with a call to constructor of the built coroutine. + private fun buildCoroutine(irFunction: IrSimpleFunction, tailSuspendCalls: Set) = + CoroutineBuilder(irFunction, tailSuspendCalls).build().also { coroutine -> + // Replace original function with a call to constructor of the built coroutine. val irBuilder = context.createIrBuilder(irFunction.symbol, irFunction.startOffset, irFunction.endOffset) irFunction.body = irBuilder.irBlockBody(irFunction) { - val constructor = coroutine.coroutineConstructor + val constructor = coroutine.constructor generateCoroutineStart(coroutine.stateMachineFunction, irCallConstructor(constructor.symbol, irFunction.typeParameters.map { it.defaultType.makeNullable() @@ -197,23 +127,10 @@ abstract class AbstractSuspendFunctionsLowering(val co } } - return coroutine.coroutineClass - } + private class BuiltCoroutine(val clazz: IrClass, val constructor: IrConstructor, val stateMachineFunction: IrFunction) - private class BuiltCoroutine( - val coroutineClass: IrClass, - val coroutineConstructor: IrConstructor, - val stateMachineFunction: IrFunction - ) - - private inner class CoroutineBuilder( - val irFunction: IrFunction, - val functionReference: IrFunctionReference?, - val tailSuspendCalls: Set - ) { + private inner class CoroutineBuilder(val irFunction: IrFunction, val tailSuspendCalls: Set) { private val functionParameters = irFunction.explicitParameters - private val boundFunctionParameters = functionReference?.getArgumentsWithIr()?.map { it.first } - private val unboundFunctionParameters = boundFunctionParameters?.let { functionParameters - it } private val coroutineClass: IrClass = context.irFactory.buildClass { @@ -239,71 +156,22 @@ abstract class AbstractSuspendFunctionsLowering(val co private val coroutineBaseClass = getCoroutineBaseClass(irFunction) private val coroutineBaseClassConstructor = coroutineBaseClass.owner.constructors.single { it.valueParameters.size == 1 } - private val create1Function = coroutineBaseClass.owner.simpleFunctions() - .single { it.name == CREATE_IDENTIFIER && it.valueParameters.size == 1 } - private val create1CompletionParameter = create1Function.valueParameters[0] private val coroutineConstructors = mutableListOf() fun build(): BuiltCoroutine { - val superTypes = mutableListOf(coroutineBaseClass.defaultType) - var suspendFunctionClass: IrClass? = null - val suspendFunctionClassTypeArguments: List? - if (unboundFunctionParameters != null) { - // Suspend lambda inherits SuspendFunction. - val numberOfParameters = unboundFunctionParameters.size - suspendFunctionClass = symbols.suspendFunctionN(numberOfParameters).owner - val unboundParameterTypes = unboundFunctionParameters.map { it.type } - suspendFunctionClassTypeArguments = unboundParameterTypes + irFunction.returnType - superTypes += suspendFunctionClass.typeWith(suspendFunctionClassTypeArguments) - } - coroutineClass.superTypes = superTypes + coroutineClass.superTypes = mutableListOf(coroutineBaseClass.defaultType) val coroutineConstructor = buildConstructor() val superInvokeSuspendFunction = coroutineBaseClass.owner.simpleFunctions().single { it.name == stateMachineMethodName } val invokeSuspendMethod = buildInvokeSuspendMethod(superInvokeSuspendFunction, coroutineClass) - var coroutineFactoryConstructor: IrConstructor? = null - val ignoredParentSymbols = mutableListOf() - if (functionReference != null) { - // Suspend lambda - create factory methods. - coroutineFactoryConstructor = buildFactoryConstructor(boundFunctionParameters!!) - - val createFunctionSymbol = coroutineBaseClass.owner.simpleFunctions() - .atMostOne { it.name == CREATE_IDENTIFIER && it.valueParameters.size == unboundFunctionParameters!!.size + 1 } - ?.symbol - - val createMethod = buildCreateMethod( - unboundArgs = unboundFunctionParameters!!, - superFunctionSymbol = createFunctionSymbol, - coroutineConstructor = coroutineConstructor - ) - - val transformedSuspendInvokeFunction = suspendFunctionClass!!.simpleFunctions().single { it.name == OperatorNameConventions.INVOKE } - val originalSuspendInvokeFunction = context.mapping.functionWithContinuationsToSuspendFunctions[transformedSuspendInvokeFunction] ?: transformedSuspendInvokeFunction - - ignoredParentSymbols.add(transformedSuspendInvokeFunction.symbol) - - buildInvokeMethod( - functionInvokeFunction = originalSuspendInvokeFunction, - createFunction = createMethod, - stateMachineFunction = invokeSuspendMethod, - ) - } - - coroutineClass.addFakeOverrides( - context.typeSystem, - ignoredParentSymbols = ignoredParentSymbols - ) + coroutineClass.addFakeOverrides(context.typeSystem) initializeStateMachine(coroutineConstructors, coroutineClassThis) - return BuiltCoroutine( - coroutineClass = coroutineClass, - coroutineConstructor = coroutineFactoryConstructor ?: coroutineConstructor, - stateMachineFunction = invokeSuspendMethod - ) + return BuiltCoroutine(coroutineClass, coroutineConstructor, invokeSuspendMethod) } fun buildConstructor(): IrConstructor = @@ -346,145 +214,6 @@ abstract class AbstractSuspendFunctionsLowering(val co } } - private fun buildFactoryConstructor(boundParams: List): IrConstructor = - context.irFactory.buildConstructor { - startOffset = irFunction.startOffset - endOffset = irFunction.endOffset - origin = DECLARATION_ORIGIN_COROUTINE_IMPL - visibility = irFunction.visibility - returnType = coroutineClass.defaultType - }.apply { - parent = coroutineClass - coroutineClass.declarations += this - coroutineConstructors += this - - valueParameters = boundParams.mapIndexed { index, parameter -> - parameter.copyTo(this, DECLARATION_ORIGIN_COROUTINE_IMPL, index) - } - - val irBuilder = context.createIrBuilder(symbol, startOffset, endOffset) - body = irBuilder.irBlockBody { - +irDelegatingConstructorCall(coroutineBaseClassConstructor).apply { - putValueArgument(0, irNull()) // Completion. - } - +IrInstanceInitializerCallImpl( - startOffset, endOffset, coroutineClass.symbol, - context.irBuiltIns.unitType - ) - // Save all arguments to fields. - boundParams.forEachIndexed { index, parameter -> - +irSetField( - irGet(coroutineClassThis), argumentToPropertiesMap.getValue(parameter), - irGet(valueParameters[index]) - ) - } - } - } - - private fun buildCreateMethod( - unboundArgs: List, - superFunctionSymbol: IrSimpleFunctionSymbol?, - coroutineConstructor: IrConstructor - ): IrSimpleFunction = context.irFactory.buildFun { - startOffset = irFunction.startOffset - endOffset = irFunction.endOffset - origin = DECLARATION_ORIGIN_COROUTINE_IMPL - name = CREATE_IDENTIFIER - visibility = DescriptorVisibilities.PROTECTED - returnType = coroutineClass.defaultType - }.apply { - parent = coroutineClass - coroutineClass.declarations += this - - typeParameters = irFunction.typeParameters.map { parameter -> - parameter.copyToWithoutSuperTypes(this, origin = DECLARATION_ORIGIN_COROUTINE_IMPL) - .apply { superTypes += parameter.superTypes } - } - - valueParameters = (unboundArgs + create1CompletionParameter).mapIndexed { index, parameter -> - parameter.copyTo(this, DECLARATION_ORIGIN_COROUTINE_IMPL, index) - } - - this.createDispatchReceiverParameter() - - superFunctionSymbol?.let { - overriddenSymbols += it.owner.overriddenSymbols - overriddenSymbols += it - } - - val thisReceiver = this.dispatchReceiverParameter!! - - val irBuilder = context.createIrBuilder(symbol, startOffset, endOffset) - body = irBuilder.irBlockBody(startOffset, endOffset) { - +irReturn( - irCall(coroutineConstructor).apply { - var unboundIndex = 0 - val unboundArgsSet = unboundArgs.toSet() - functionParameters.map { - if (unboundArgsSet.contains(it)) - irGet(valueParameters[unboundIndex++]) - else - irGetField(irGet(thisReceiver), argumentToPropertiesMap.getValue(it)) - }.forEachIndexed { index, argument -> - putValueArgument(index, argument) - } - putValueArgument(functionParameters.size, irGet(valueParameters[unboundIndex])) - assert(unboundIndex == valueParameters.size - 1) { - "Not all arguments of are used" - } - }) - } - } - - private fun buildInvokeMethod( - functionInvokeFunction: IrSimpleFunction, - createFunction: IrFunction, - stateMachineFunction: IrFunction - ): IrSimpleFunction = context.irFactory.buildFun { - startOffset = irFunction.startOffset - endOffset = irFunction.endOffset - origin = DECLARATION_ORIGIN_COROUTINE_IMPL - name = OperatorNameConventions.INVOKE - visibility = DescriptorVisibilities.PROTECTED - returnType = context.irBuiltIns.anyNType - isSuspend = true - }.apply { - parent = coroutineClass - coroutineClass.declarations += this - - typeParameters = irFunction.typeParameters.map { parameter -> - parameter.copyToWithoutSuperTypes(this, origin = DECLARATION_ORIGIN_COROUTINE_IMPL) - .apply { superTypes += parameter.superTypes } - } - - valueParameters = createFunction.valueParameters - // Skip completion - invoke() already has it implicitly as a suspend function. - .take(createFunction.valueParameters.size - 1) - .mapIndexed { index, parameter -> - parameter.copyTo(this, DECLARATION_ORIGIN_COROUTINE_IMPL, index) - } - - this.createDispatchReceiverParameter() - - overriddenSymbols += functionInvokeFunction.symbol - - val thisReceiver = dispatchReceiverParameter!! - - val irBuilder = context.createIrBuilder(symbol, startOffset, endOffset) - body = irBuilder.irBlockBody(startOffset, endOffset) { - generateCoroutineStart(stateMachineFunction, irCall(createFunction).apply { - dispatchReceiver = irGet(thisReceiver) - valueParameters.forEachIndexed { index, parameter -> - putValueArgument(index, irGet(parameter)) - } - putValueArgument( - valueParameters.size, - irCall(getContinuationSymbol, getContinuationSymbol.owner.returnType, listOf(returnType)) - ) - }) - } - } - private fun buildInvokeSuspendMethod( stateMachineFunction: IrSimpleFunction, coroutineClass: IrClass @@ -569,10 +298,4 @@ abstract class AbstractSuspendFunctionsLowering(val co addChild(it) } } - - companion object { - private val CREATE_IDENTIFIER = Name.identifier("create") - } } - - diff --git a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/FunctionReferenceLowering.kt b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/FunctionReferenceLowering.kt index 701291db9dc..3ac8cc8de67 100644 --- a/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/FunctionReferenceLowering.kt +++ b/kotlin-native/backend.native/compiler/ir/backend.native/src/org/jetbrains/kotlin/backend/konan/lower/FunctionReferenceLowering.kt @@ -161,11 +161,6 @@ internal class FunctionReferenceLowering(val generationState: NativeGenerationSt return expression } - if (expression.origin.isLambda && expression.type.isSuspendFunction()) { - // Handled later by coroutines lowerings - return expression - } - return transformFunctionReference(expression) }