From 8b97819c046faa9163e7c2a2c39b62b87489fd24 Mon Sep 17 00:00:00 2001 From: Mads Ager Date: Mon, 21 Oct 2019 12:34:21 +0200 Subject: [PATCH] [JVM_IR] Support suspend functions that can be overriden. The challenge for overridable suspend functions is that the calling the state machine method to resume after suspension would be virtually dispatced to the wrong method. To avoid that a static suspend implementation method is generated which becomes the state machine method used to resume. --- .../kotlin/backend/common/ir/IrUtils.kt | 38 ++++++++-- .../kotlin/backend/jvm/DeclarationOrigins.kt | 1 + .../backend/jvm/codegen/CoroutineCodegen.kt | 5 +- .../jvm/lower/AddContinuationLowering.kt | 75 ++++++++++++++++--- .../kotlin/ir/expressions/impl/IrCallImpl.kt | 4 +- .../suspendCovariantJavaOverrides.kt | 1 - .../openFunWithJava.kt | 1 - .../suspendFunctionAsCoroutine/superCall.kt | 1 - .../superCallAbstractClass.kt | 1 - 9 files changed, 101 insertions(+), 26 deletions(-) diff --git a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/ir/IrUtils.kt b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/ir/IrUtils.kt index 263a982844f..de672bccc20 100644 --- a/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/ir/IrUtils.kt +++ b/compiler/ir/backend.common/src/org/jetbrains/kotlin/backend/common/ir/IrUtils.kt @@ -24,10 +24,7 @@ import org.jetbrains.kotlin.ir.declarations.impl.IrTypeParameterImpl import org.jetbrains.kotlin.ir.declarations.impl.IrValueParameterImpl import org.jetbrains.kotlin.ir.descriptors.IrBuiltIns import org.jetbrains.kotlin.ir.expressions.* -import org.jetbrains.kotlin.ir.expressions.impl.IrBlockBodyImpl -import org.jetbrains.kotlin.ir.expressions.impl.IrDelegatingConstructorCallImpl -import org.jetbrains.kotlin.ir.expressions.impl.IrGetValueImpl -import org.jetbrains.kotlin.ir.expressions.impl.IrInstanceInitializerCallImpl +import org.jetbrains.kotlin.ir.expressions.impl.* import org.jetbrains.kotlin.ir.symbols.impl.IrConstructorSymbolImpl import org.jetbrains.kotlin.ir.symbols.impl.IrSimpleFunctionSymbolImpl import org.jetbrains.kotlin.ir.symbols.impl.IrTypeParameterSymbolImpl @@ -36,7 +33,9 @@ import org.jetbrains.kotlin.ir.types.* import org.jetbrains.kotlin.ir.types.impl.IrSimpleTypeImpl import org.jetbrains.kotlin.ir.types.impl.makeTypeProjection import org.jetbrains.kotlin.ir.util.* +import org.jetbrains.kotlin.ir.visitors.IrElementTransformerVoid import org.jetbrains.kotlin.ir.visitors.IrElementVisitor +import org.jetbrains.kotlin.ir.visitors.transformChildrenVoid import org.jetbrains.kotlin.name.Name import org.jetbrains.kotlin.serialization.deserialization.descriptors.DescriptorWithContainerSource import java.io.StringWriter @@ -509,7 +508,8 @@ fun createStaticFunctionWithReceivers( oldFunction: IrFunction, dispatchReceiverType: IrType? = oldFunction.dispatchReceiverParameter?.type, origin: IrDeclarationOrigin = oldFunction.origin, - modality: Modality = Modality.FINAL + modality: Modality = Modality.FINAL, + copyMetadata: Boolean = true ): IrSimpleFunction { val descriptor = (oldFunction.descriptor as? DescriptorWithContainerSource)?.let { WrappedFunctionDescriptorWithContainerSource(it.containerSource) @@ -551,7 +551,7 @@ fun createStaticFunctionWithReceivers( oldFunction.valueParameters.map { it.copyTo(this, index = it.index + offset) } ) - metadata = oldFunction.metadata + if (copyMetadata) metadata = oldFunction.metadata } } @@ -560,6 +560,28 @@ fun copyBodyToStatic(oldFunction: IrFunction, staticFunction: IrFunction) { (listOfNotNull(oldFunction.dispatchReceiverParameter, oldFunction.extensionReceiverParameter) + oldFunction.valueParameters) .zip(staticFunction.valueParameters).toMap() staticFunction.body = oldFunction.body - ?.transform(VariableRemapper(mapping), null) - ?.patchDeclarationParents(staticFunction) + ?.transform( + object: IrElementTransformerVoid() { + // Remap return targets to the static method so they do not appear to be + // non-local returns. + override fun visitReturn(expression: IrReturn): IrExpression { + expression.transformChildrenVoid(this); + return if (expression.returnTargetSymbol == oldFunction.symbol) { + IrReturnImpl( + expression.startOffset, + expression.endOffset, + expression.type, + staticFunction.symbol, + expression.value) + } else expression + } + + // Remap argument values. + override fun visitGetValue(expression: IrGetValue): IrExpression = + mapping[expression.symbol.owner]?.let { + IrGetValueImpl(expression.startOffset, expression.endOffset, it.type, it.symbol, expression.origin) + } ?: expression + + }, null) + ?.patchDeclarationParents(staticFunction) } diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/DeclarationOrigins.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/DeclarationOrigins.kt index 0d224f6fca8..c99c0a1014b 100644 --- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/DeclarationOrigins.kt +++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/DeclarationOrigins.kt @@ -44,4 +44,5 @@ interface JvmLoweredDeclarationOrigin : IrDeclarationOrigin { object INLINE_CLASS_GENERATED_IMPL_METHOD : IrDeclarationOriginImpl("INLINE_CLASS_GENERATED_IMPL_METHOD") object GENERATED_ASSERTION_ENABLED_FIELD : IrDeclarationOriginImpl("GENERATED_ASSERTION_ENABLED_FIELD", isSynthetic = true) object GENERATED_MAIN_FOR_PARAMETERLESS_MAIN_METHOD : IrDeclarationOriginImpl("GENERATED_MAIN_FOR_PARAMETERLESS_MAIN_METHOD", isSynthetic = true) + object SUSPEND_IMPL_STATIC_FUNCTION : IrDeclarationOriginImpl("SUSPEND_IMPL_STATIC_FUNCTION", isSynthetic = true) } diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/CoroutineCodegen.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/CoroutineCodegen.kt index e358cb227fa..e909be88719 100644 --- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/CoroutineCodegen.kt +++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/codegen/CoroutineCodegen.kt @@ -66,7 +66,8 @@ internal fun generateStateMachineForNamedFunction( shouldPreserveClassInitialization = state.constructorCallNormalizationMode.shouldPreserveClassInitialization, containingClassInternalName = classCodegen.visitor.thisName, isForNamedFunction = true, - needDispatchReceiver = irFunction.dispatchReceiverParameter != null, + needDispatchReceiver = irFunction.dispatchReceiverParameter != null + || irFunction.origin == JvmLoweredDeclarationOrigin.SUSPEND_IMPL_STATIC_FUNCTION, internalNameForDispatchReceiver = classCodegen.visitor.thisName, putContinuationParameterToLvt = false ) @@ -186,7 +187,7 @@ internal fun IrCall.createSuspendFunctionCallViewIfNeeded( if (!isSuspend) return this val view = (symbol.owner as IrSimpleFunction).getOrCreateSuspendFunctionViewIfNeeded(context) if (view == symbol.owner) return this - return IrCallImpl(startOffset, endOffset, view.returnType, view.symbol).also { + return IrCallImpl(startOffset, endOffset, view.returnType, view.symbol, superQualifierSymbol = superQualifierSymbol).also { it.copyTypeArgumentsFrom(this) it.dispatchReceiver = dispatchReceiver it.extensionReceiver = extensionReceiver diff --git a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/AddContinuationLowering.kt b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/AddContinuationLowering.kt index 14e2f7e657d..4585f51bf6d 100644 --- a/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/AddContinuationLowering.kt +++ b/compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/AddContinuationLowering.kt @@ -7,16 +7,14 @@ package org.jetbrains.kotlin.backend.jvm.lower import org.jetbrains.kotlin.backend.common.FileLoweringPass import org.jetbrains.kotlin.backend.common.descriptors.synthesizedName -import org.jetbrains.kotlin.backend.common.ir.copyTo -import org.jetbrains.kotlin.backend.common.ir.copyTypeParametersFrom -import org.jetbrains.kotlin.backend.common.ir.createImplicitParameterDeclarationWithWrappedDescriptor -import org.jetbrains.kotlin.backend.common.ir.isSuspend +import org.jetbrains.kotlin.backend.common.ir.* import org.jetbrains.kotlin.backend.common.lower.createIrBuilder import org.jetbrains.kotlin.backend.common.lower.parents import org.jetbrains.kotlin.backend.common.phaser.makeIrFilePhase import org.jetbrains.kotlin.backend.common.pop import org.jetbrains.kotlin.backend.common.push import org.jetbrains.kotlin.backend.jvm.JvmBackendContext +import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin import org.jetbrains.kotlin.backend.jvm.codegen.isInlineIrBlock import org.jetbrains.kotlin.backend.jvm.codegen.isInvokeOfSuspendCallableReference import org.jetbrains.kotlin.codegen.coroutines.* @@ -364,9 +362,15 @@ private class AddContinuationLowering(private val context: JvmBackendContext) : val capturedThisField = irFunction.dispatchReceiverParameter?.let { addField("this\$0", it.type) } val labelField = addField(COROUTINE_LABEL_FIELD_NAME, context.irBuiltIns.intType, JavaVisibilities.PACKAGE_VISIBILITY) addConstructorForNamedFunction(capturedThisField) - addInvokeSuspendForNamedFunction(irFunction, resultField, labelField, capturedThisField) - - context.suspendFunctionContinuations[irFunction] = this + var function = irFunction + if (function is IrSimpleFunction && function.isOverridable && function.body != null) { + // Create static method for the suspend state machine method so that reentering the method + // does not lead to virtual dispatch to the wrong method. + context.suspendFunctionContinuations[function] = this + function = createStaticSuspendImpl(function) + } + addInvokeSuspendForNamedFunction(function, resultField, labelField, capturedThisField, function != irFunction) + context.suspendFunctionContinuations[function] = this } } @@ -388,11 +392,19 @@ private class AddContinuationLowering(private val context: JvmBackendContext) : } } + + private fun Name.toSuspendImplementationName() = when { + isSpecial -> Name.special(asString() + SUSPEND_IMPL_NAME_SUFFIX) + else -> Name.identifier(asString() + SUSPEND_IMPL_NAME_SUFFIX) + } + + private fun IrClass.addInvokeSuspendForNamedFunction( irFunction: IrFunction, resultField: IrField, labelField: IrField, - capturedThisField: IrField? + capturedThisField: IrField?, + isStaticSuspendImpl: Boolean ) { val invokeSuspend = continuationImpl.owner.functions.single { it.name == Name.identifier(INVOKE_SUSPEND_METHOD_NAME) } addFunctionOverride(invokeSuspend).also { function -> @@ -416,26 +428,69 @@ private class AddContinuationLowering(private val context: JvmBackendContext) : irInt(signBit) ) ) + +irReturn(irCall(irFunction).also { for (i in irFunction.typeParameters.indices) { it.putTypeArgument(i, context.irBuiltIns.anyNType) } - it.dispatchReceiver = capturedThisField?.let { irField -> + val capturedThisValue = capturedThisField?.let { irField -> irGetField(irGet(function.dispatchReceiverParameter!!), irField) } + if (irFunction.dispatchReceiverParameter != null) { + it.dispatchReceiver = capturedThisValue + } if (irFunction.extensionReceiverParameter != null) { it.extensionReceiver = irNull() } for ((i, parameter) in irFunction.valueParameters.withIndex()) { val defaultValueForParameter = IrConstImpl.defaultValueForType( - UNDEFINED_OFFSET, UNDEFINED_OFFSET, parameter.type) + UNDEFINED_OFFSET, UNDEFINED_OFFSET, parameter.type + ) it.putValueArgument(i, defaultValueForParameter) } + if (isStaticSuspendImpl) { + it.putValueArgument(0, capturedThisValue) + } }) } } } + private fun createStaticSuspendImpl(irFunction: IrSimpleFunction): IrSimpleFunction { + // Create static suspend impl method. + val static = createStaticFunctionWithReceivers( + irFunction.parent, + irFunction.name.toSuspendImplementationName(), + irFunction, + origin = JvmLoweredDeclarationOrigin.SUSPEND_IMPL_STATIC_FUNCTION, + copyMetadata = false + ) + copyBodyToStatic(irFunction, static) + (irFunction.parent as IrClass).declarations.add(static) + // Rewrite the body of the original suspend method to forward to the new static method. + irFunction.body = context.createIrBuilder(irFunction.symbol).irBlockBody { + +irReturn(irCall(static).also { + for (i in irFunction.typeParameters.indices) { + it.putTypeArgument(i, context.irBuiltIns.anyNType) + } + var i = 0 + if (irFunction.dispatchReceiverParameter != null) { + it.putValueArgument(i++, irGet(irFunction.dispatchReceiverParameter!!)) + } + if (irFunction.extensionReceiverParameter != null) { + it.putValueArgument(i++, irNull()) + } + for (parameter in irFunction.valueParameters) { + val defaultValueForParameter = IrConstImpl.defaultValueForType( + UNDEFINED_OFFSET, UNDEFINED_OFFSET, parameter.type + ) + it.putValueArgument(i++, defaultValueForParameter) + } + }) + } + return static + } + // TODO: Generate two copies of inline suspend functions private fun markSuspendFunctions(irFile: IrFile, suspendLambdas: Set): Set { val result = hashSetOf() diff --git a/compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/expressions/impl/IrCallImpl.kt b/compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/expressions/impl/IrCallImpl.kt index 98b94c19f25..a8ddb783a8d 100644 --- a/compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/expressions/impl/IrCallImpl.kt +++ b/compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/expressions/impl/IrCallImpl.kt @@ -79,8 +79,8 @@ class IrCallImpl( descriptor.valueParameters.size, origin, superQualifierSymbol ) - constructor(startOffset: Int, endOffset: Int, type: IrType, symbol: IrFunctionSymbol) : - this(startOffset, endOffset, type, symbol, symbol.descriptor) + constructor(startOffset: Int, endOffset: Int, type: IrType, symbol: IrFunctionSymbol, superQualifierSymbol : IrClassSymbol? = null) : + this(startOffset, endOffset, type, symbol, symbol.descriptor, superQualifierSymbol = superQualifierSymbol) override val superQualifier: ClassDescriptor? = superQualifierSymbol?.descriptor diff --git a/compiler/testData/codegen/box/coroutines/suspendCovariantJavaOverrides.kt b/compiler/testData/codegen/box/coroutines/suspendCovariantJavaOverrides.kt index a694edb4eed..96d237855f9 100644 --- a/compiler/testData/codegen/box/coroutines/suspendCovariantJavaOverrides.kt +++ b/compiler/testData/codegen/box/coroutines/suspendCovariantJavaOverrides.kt @@ -1,5 +1,4 @@ // TARGET_BACKEND: JVM -// IGNORE_BACKEND: JVM_IR // WITH_RUNTIME // WITH_COROUTINES diff --git a/compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/openFunWithJava.kt b/compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/openFunWithJava.kt index 69bb06e176a..e16b000c5f3 100644 --- a/compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/openFunWithJava.kt +++ b/compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/openFunWithJava.kt @@ -1,4 +1,3 @@ -// IGNORE_BACKEND: JVM_IR // WITH_RUNTIME // WITH_COROUTINES // FILE: main.kt diff --git a/compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/superCall.kt b/compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/superCall.kt index 51e35dc64f6..10ef23a3462 100644 --- a/compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/superCall.kt +++ b/compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/superCall.kt @@ -1,4 +1,3 @@ -// IGNORE_BACKEND: JVM_IR // WITH_RUNTIME // WITH_COROUTINES // COMMON_COROUTINES_TEST diff --git a/compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/superCallAbstractClass.kt b/compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/superCallAbstractClass.kt index 2e203f01621..25d9fd9baa2 100644 --- a/compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/superCallAbstractClass.kt +++ b/compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/superCallAbstractClass.kt @@ -1,4 +1,3 @@ -// IGNORE_BACKEND: JVM_IR // WITH_RUNTIME // WITH_COROUTINES // COMMON_COROUTINES_TEST