JVM IR: Avoid getTopLevelClass in AddContinuationLowering

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/JvmSymbols.kt

@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.backend.jvm
 
 import org.jetbrains.kotlin.backend.common.ir.Symbols
-import org.jetbrains.kotlin.backend.common.ir.copyTo
 import org.jetbrains.kotlin.backend.common.ir.createImplicitParameterDeclarationWithWrappedDescriptor
 import org.jetbrains.kotlin.backend.jvm.intrinsics.IrIntrinsicMethods
 import org.jetbrains.kotlin.descriptors.ClassDescriptor
@@ -22,7 +21,6 @@ import org.jetbrains.kotlin.ir.declarations.impl.IrExternalPackageFragmentImpl
 import org.jetbrains.kotlin.ir.symbols.*
 import org.jetbrains.kotlin.ir.symbols.impl.IrExternalPackageFragmentSymbolImpl
 import org.jetbrains.kotlin.ir.types.*
-import org.jetbrains.kotlin.ir.types.defaultType
 import org.jetbrains.kotlin.ir.util.*
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
@@ -72,10 +70,18 @@ class JvmSymbols(
     private fun createPackage(fqName: FqName): IrPackageFragment =
         IrExternalPackageFragmentImpl(IrExternalPackageFragmentSymbolImpl(EmptyPackageFragmentDescriptor(context.state.module, fqName)))
 
-    private fun createClass(fqName: FqName, classKind: ClassKind = ClassKind.CLASS, block: (IrClass) -> Unit = {}): IrClassSymbol =
+    private fun createClass(
+        fqName: FqName,
+        classKind: ClassKind = ClassKind.CLASS,
+        classModality: Modality = Modality.FINAL,
+        classIsInline: Boolean = false,
+        block: (IrClass) -> Unit = {}
+    ): IrClassSymbol =
         buildClass {
             name = fqName.shortName()
             kind = classKind
+            modality = classModality
+            isInline = classIsInline
         }.apply {
             parent = when (fqName.parent().asString()) {
                 "kotlin" -> kotlinPackage
@@ -183,7 +189,28 @@ class JvmSymbols(
             klass.addTypeParameter("T", irBuiltIns.anyNType, Variance.IN_VARIANCE)
         }
 
-    val suspendFunctionInterface: IrClassSymbol = createClass(FqName("kotlin.coroutines.jvm.internal.SuspendFunction"), ClassKind.INTERFACE)
+    private val resultClassStub: IrClassSymbol =
+        createClass(DescriptorUtils.RESULT_FQ_NAME, classIsInline = true) { klass ->
+            klass.addTypeParameter("T", irBuiltIns.anyNType, Variance.OUT_VARIANCE)
+            klass.addConstructor { isPrimary = true }.apply {
+                addValueParameter("value", irBuiltIns.anyNType)
+            }
+        }
+
+    val continuationImplClass: IrClassSymbol =
+        createClass(FqName("kotlin.coroutines.jvm.internal.ContinuationImpl"), classModality = Modality.ABSTRACT) { klass ->
+            val continuationType = continuationClass.typeWith(irBuiltIns.anyNType)
+            klass.superTypes += continuationType
+            klass.addConstructor().apply {
+                addValueParameter("completion", continuationType.makeNullable())
+            }
+            klass.addFunction("invokeSuspend", irBuiltIns.anyNType, Modality.ABSTRACT).apply {
+                addValueParameter("result", resultClassStub.typeWith(irBuiltIns.anyNType))
+            }
+        }
+
+    val suspendFunctionInterface: IrClassSymbol =
+        createClass(FqName("kotlin.coroutines.jvm.internal.SuspendFunction"), ClassKind.INTERFACE)
 
     val lambdaClass: IrClassSymbol = createClass(FqName("kotlin.jvm.internal.Lambda")) { klass ->
         klass.addConstructor().apply {
@@ -191,6 +218,24 @@ class JvmSymbols(
         }
     }
 
+    val suspendLambdaClass: IrClassSymbol = createClass(FqName("kotlin.coroutines.jvm.internal.SuspendLambda")) { klass ->
+        klass.superTypes += suspendFunctionInterface.typeWith()
+        klass.addConstructor().apply {
+            addValueParameter("arity", irBuiltIns.intType)
+            addValueParameter("completion", continuationClass.typeWith(irBuiltIns.anyNType).makeNullable())
+        }
+        klass.addFunction("invokeSuspend", irBuiltIns.anyNType, Modality.ABSTRACT).apply {
+            addValueParameter("result", resultClassStub.typeWith(irBuiltIns.anyNType))
+        }
+        klass.addFunction("create", continuationClass.typeWith(irBuiltIns.unitType)).apply {
+            addValueParameter("completion", continuationClass.typeWith(irBuiltIns.nothingType))
+        }
+        klass.addFunction("create", continuationClass.typeWith(irBuiltIns.unitType)).apply {
+            addValueParameter("value", irBuiltIns.anyNType)
+            addValueParameter("completion", continuationClass.typeWith(irBuiltIns.nothingType))
+        }
+    }
+
     private fun generateCallableReferenceMethods(klass: IrClass) {
         klass.addFunction("getSignature", irBuiltIns.stringType, Modality.OPEN)
         klass.addFunction("getName", irBuiltIns.stringType, Modality.OPEN)

compiler/ir/backend.jvm/src/org/jetbrains/kotlin/backend/jvm/lower/AddContinuationLowering.kt

@@ -18,7 +18,6 @@ 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.*
-import org.jetbrains.kotlin.config.coroutinesPackageFqName
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.ir.IrElement
 import org.jetbrains.kotlin.ir.IrStatement
@@ -55,16 +54,6 @@ internal val addContinuationPhase = makeIrFilePhase(
 private object CONTINUATION_CLASS : IrDeclarationOriginImpl("CONTINUATION_CLASS")
 
 private class AddContinuationLowering(private val context: JvmBackendContext) : FileLoweringPass {
-    private val continuation by lazy {
-        context.getTopLevelClass(context.state.languageVersionSettings.coroutinesPackageFqName().child(Name.identifier("Continuation")))
-    }
-    private val continuationImpl by lazy {
-        context.getTopLevelClass(context.state.languageVersionSettings.coroutinesJvmInternalPackageFqName().child(Name.identifier("ContinuationImpl")))
-    }
-    private val suspendLambda by lazy {
-        context.getTopLevelClass(context.state.languageVersionSettings.coroutinesJvmInternalPackageFqName().child(Name.identifier("SuspendLambda")))
-    }
-
     override fun lower(irFile: IrFile) {
         val suspendLambdas = markSuspendLambdas(irFile)
         val suspendFunctions = markSuspendFunctions(irFile, suspendLambdas.map { it.function }.toSet())
@@ -104,7 +93,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
     }
 
     private fun generateContinuationClassForLambda(info: SuspendLambdaInfo) {
-        val suspendLambda = suspendLambda.owner
+        val suspendLambda = context.ir.symbols.suspendLambdaClass.owner
         suspendLambda.createContinuationClassFor(info.function).apply {
             copyAttributes(info.reference)
             val functionNClass = context.ir.symbols.getJvmFunctionClass(info.arity + 1)
@@ -155,7 +144,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         fields: List<IrField>,
         receiverField: IrField?
     ): IrFunction {
-        val superMethod = suspendLambda.functions.single {
+        val superMethod = context.ir.symbols.suspendLambdaClass.functions.single {
             it.owner.name.asString() == INVOKE_SUSPEND_METHOD_NAME && it.owner.valueParameters.size == 1 &&
                     it.owner.valueParameters[0].type.isKotlinResult()
         }.owner
@@ -312,7 +301,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             }
             val completionParameterSymbol = constructor.addCompletionValueParameter()
 
-            val superClassConstructor = suspendLambda.owner.constructors.single {
+            val superClassConstructor = context.ir.symbols.suspendLambdaClass.owner.constructors.single {
                 it.valueParameters.size == 2 && it.valueParameters[0].type.isInt() && it.valueParameters[1].type.isNullableContinuation()
             }
             constructor.body = context.createIrBuilder(constructor.symbol).irBlockBody {
@@ -349,11 +338,11 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
     private fun IrFunction.addCompletionValueParameter(): IrValueParameter =
         addValueParameter(SUSPEND_FUNCTION_COMPLETION_PARAMETER_NAME, continuationType())
 
-    private fun IrFunction.continuationType(): IrSimpleType =
+    private fun IrFunction.continuationType(): IrType =
-        continuation.createType(true, listOf(makeTypeProjection(returnType, Variance.INVARIANT)))
+        context.ir.symbols.continuationClass.typeWith(returnType).makeNullable()
 
     private fun generateContinuationClassForNamedFunction(irFunction: IrFunction) {
-        continuationImpl.owner.createContinuationClassFor(irFunction).apply {
+        context.ir.symbols.continuationImplClass.owner.createContinuationClassFor(irFunction).apply {
             val resultField = addField(
                 context.state.languageVersionSettings.dataFieldName(),
                 context.irBuiltIns.anyType,
@@ -381,7 +370,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         val capturedThisParameter = capturedThisField?.let { constructor.addValueParameter(it.name.asString(), it.type) }
         val completionParameterSymbol = constructor.addCompletionValueParameter()
 
-        val superClassConstructor = continuationImpl.owner.constructors.single { it.valueParameters.size == 1 }
+        val superClassConstructor = context.ir.symbols.continuationImplClass.owner.constructors.single { it.valueParameters.size == 1 }
         constructor.body = context.createIrBuilder(constructor.symbol).irBlockBody {
             if (capturedThisField != null) {
                 +irSetField(irGet(thisReceiver!!), capturedThisField, irGet(capturedThisParameter!!))
@@ -406,7 +395,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         capturedThisField: IrField?,
         isStaticSuspendImpl: Boolean
     ) {
-        val invokeSuspend = continuationImpl.owner.functions.single { it.name == Name.identifier(INVOKE_SUSPEND_METHOD_NAME) }
+        val invokeSuspend = context.ir.symbols.continuationImplClass.owner.functions.single { it.name == Name.identifier(INVOKE_SUSPEND_METHOD_NAME) }
         addFunctionOverride(invokeSuspend).also { function ->
             function.body = context.createIrBuilder(function.symbol).irBlockBody {
                 +irSetField(irGet(function.dispatchReceiverParameter!!), resultField, irGet(function.valueParameters[0]))

compiler/testData/loadJava/compiledKotlin/coroutines/Basic.kt

@@ -1,5 +1,4 @@
-// IGNORE_BACKEND: JVM_IR
+// ALLOW_AST_ACCESS
-//ALLOW_AST_ACCESS
 
 package test
 class Controller {

compiler/testData/writeSignature/callableReference/suspendFunctionReference.kt

@@ -1,5 +1,3 @@
-// IGNORE_BACKEND: JVM_IR
-
 interface I
 
 abstract class AbstractTest {