[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.
2019-10-21 12:34:21 +02:00
parent eebb071ae9
commit 8b97819c04
9 changed files with 101 additions and 26 deletions
@@ -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)
 }
@@ -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)
 }
@@ -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
@@ -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<IrFunction>): Set<IrFunction> {
        val result = hashSetOf<IrFunction>()
@@ -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
@@ -1,5 +1,4 @@
 // TARGET_BACKEND: JVM
-// IGNORE_BACKEND: JVM_IR
 // WITH_RUNTIME
 // WITH_COROUTINES

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // WITH_RUNTIME
 // WITH_COROUTINES
 // FILE: main.kt
@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // WITH_RUNTIME
 // WITH_COROUTINES
 // COMMON_COROUTINES_TEST
@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // WITH_RUNTIME
 // WITH_COROUTINES
 // COMMON_COROUTINES_TEST