JVM_IR: rework type parameters in AddContinuationLowering

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

@@ -78,6 +78,9 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
                             "Inconsistency between callable reference to suspend lambda and the corresponding continuation"
                         }
                         +irCall(constructor.symbol).apply {
+                            for (tp in info.constructor.parentAsClass.typeParameters) {
+                                putTypeArgument(tp.index, expression.getTypeArgument(tp.index))
+                            }
                             expressionArguments.forEachIndexed { index, argument ->
                                 putValueArgument(index, argument)
                             }
@@ -94,6 +97,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         val suspendLambda = context.ir.symbols.suspendLambdaClass.owner
         return suspendLambda.createContinuationClassFor(parent, JvmLoweredDeclarationOrigin.SUSPEND_LAMBDA).apply {
             copyAttributes(info.reference)
+            copyTypeParametersFrom(info.function)
             val functionNClass = context.ir.symbols.getJvmFunctionClass(info.arity + 1)
             superTypes.add(
                 IrSimpleTypeImpl(
@@ -130,7 +134,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             if (info.arity <= 1) {
                 val singleParameterField = receiverField ?: parametersWithoutArguments.singleOrNull()
                 val create = addCreate(constructor, suspendLambda, info, parametersWithArguments, singleParameterField)
-                addInvokeCallingCreate(create, invokeSuspend, invokeToOverride, singleParameterField)
+                addInvokeCallingCreate(info.function, create, invokeSuspend, invokeToOverride, singleParameterField)
             } else {
                 addInvokeCallingConstructor(
                     constructor,
@@ -180,7 +184,6 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         receiverField: IrField?,
         fields: List<IrField>
     ) {
-        copyTypeParametersFrom(irFunction)
         body = irFunction.body?.deepCopyWithSymbols(this)
         body?.transformChildrenVoid(object : IrElementTransformerVoid() {
             override fun visitGetValue(expression: IrGetValue): IrExpression {
@@ -239,6 +242,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
     //   2) starting newly created coroutine by calling `invokeSuspend`.
     // Thus, it creates a clone of suspend lambda and starts it.
     private fun IrClass.addInvokeCallingCreate(
+        suspendLambdaFunction: IrFunction,
         create: IrFunction,
         invokeSuspend: IrFunction,
         invokeToOverride: IrSimpleFunctionSymbol,
@@ -285,6 +289,9 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             function.body = context.createIrBuilder(function.symbol).irBlockBody {
                 // Create a copy
                 val newlyCreatedObject = irTemporary(irCall(constructor).also { constructorCall ->
+                    for (tp in typeParameters) {
+                        constructorCall.putTypeArgument(tp.index, tp.defaultType)
+                    }
                     for ((index, field) in parametersWithArguments.withIndex()) {
                         constructorCall.putValueArgument(index, irGetField(irGet(function.dispatchReceiverParameter!!), field))
                     }
@@ -321,6 +328,9 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             function.body = context.createIrBuilder(function.symbol).irBlockBody {
                 var index = 0
                 val constructorCall = irCall(constructor).also {
+                    for (tp in typeParameters) {
+                        it.putTypeArgument(tp.index, tp.defaultType)
+                    }
                     for ((i, field) in parametersWithArguments.withIndex()) {
                         if (info.reference.getValueArgument(i) == null) continue
                         it.putValueArgument(index++, irGetField(irGet(function.dispatchReceiverParameter!!), field))
@@ -395,6 +405,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         return context.ir.symbols.continuationImplClass.owner
             .createContinuationClassFor(irFunction, JvmLoweredDeclarationOrigin.CONTINUATION_CLASS)
             .apply {
+                copyTypeParametersFrom(irFunction)
                 val resultField = addField(
                     context.state.languageVersionSettings.dataFieldName(),
                     context.irBuiltIns.anyType,
@@ -473,7 +484,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
 
                 +irReturn(irCall(irFunction).also {
                     for (i in irFunction.typeParameters.indices) {
-                        it.putTypeArgument(i, context.irBuiltIns.anyNType)
+                        it.putTypeArgument(i, typeParameters[i].defaultType)
                     }
                     val capturedThisValue = capturedThisField?.let { irField ->
                         irGetField(irGet(function.dispatchReceiverParameter!!), irField)