[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.

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.*
-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.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)
+        ?.transform(
-            ?.patchDeclarationParents(staticFunction)
+            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)
 }

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)
 }

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

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.*
-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.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)
+            var function = irFunction
-
+            if (function is IrSimpleFunction && function.isOverridable && function.body != null) {
-            context.suspendFunctionContinuations[irFunction] = this
+                // 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>()

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) :
+    constructor(startOffset: Int, endOffset: Int, type: IrType, symbol: IrFunctionSymbol, superQualifierSymbol : IrClassSymbol? = null) :
-            this(startOffset, endOffset, type, symbol, symbol.descriptor)
+            this(startOffset, endOffset, type, symbol, symbol.descriptor, superQualifierSymbol = superQualifierSymbol)
 
 
     override val superQualifier: ClassDescriptor? = superQualifierSymbol?.descriptor

compiler/testData/codegen/box/coroutines/suspendCovariantJavaOverrides.kt

@@ -1,5 +1,4 @@
 // TARGET_BACKEND: JVM
-// IGNORE_BACKEND: JVM_IR
 // WITH_RUNTIME
 // WITH_COROUTINES
 

compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/openFunWithJava.kt

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

compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/superCall.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // WITH_RUNTIME
 // WITH_COROUTINES
 // COMMON_COROUTINES_TEST

compiler/testData/codegen/box/coroutines/suspendFunctionAsCoroutine/superCallAbstractClass.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
 // WITH_RUNTIME
 // WITH_COROUTINES
 // COMMON_COROUTINES_TEST