JVM_IR: Support suspend lambdas with multiple parameters

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

@@ -128,7 +128,7 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             }
 
             val parametersFields = info.function.valueParameters.map { addField(it.name.asString(), it.type) }
-            val constructor = addPrimaryConstructorForLambda(info.arity, info.function, parametersFields)
+            val constructor = addPrimaryConstructorForLambda(info.arity, info.reference, parametersFields)
             val secondaryConstructor = addSecondaryConstructorForLambda(constructor)
             val invokeToOverride = functionNClass.functions.single {
                 it.owner.valueParameters.size == info.arity + 1 && it.owner.name.asString() == "invoke"
@@ -136,9 +136,9 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             val invokeSuspend = addInvokeSuspendForLambda(info.function, parametersFields, receiverField)
             if (info.arity <= 1) {
                 val create = addCreate(constructor, suspendLambda, info.arity, parametersFields, receiverField)
-                addInvoke(create, invokeSuspend, invokeToOverride)
+                addInvoke(create, invokeSuspend, invokeToOverride, parametersFields, receiverField, isConstructorCall = false)
             } else {
-                addInvoke(constructor, invokeSuspend, invokeToOverride)
+                addInvoke(constructor, invokeSuspend, invokeToOverride, parametersFields, receiverField, isConstructorCall = true)
             }
 
             context.suspendLambdaToOriginalFunctionMap[this] = info.function
@@ -203,19 +203,34 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
     private fun IrClass.addInvoke(
         create: IrFunction,
         invokeSuspend: IrFunction,
-        invokeToOverride: IrSimpleFunctionSymbol
+        invokeToOverride: IrSimpleFunctionSymbol,
+        parametersFields: List<IrField>,
+        receiverField: IrField?,
+        isConstructorCall: Boolean
     ) {
         val unitClass = context.irBuiltIns.unitClass
         val unitField = context.declarationFactory.getFieldForObjectInstance(unitClass.owner)
         addFunctionOverride(invokeToOverride.owner).also { function ->
             function.body = context.createIrBuilder(function.symbol).irBlockBody {
-                +irReturn(irCall(invokeSuspend).also { invokeSuspendCall ->
+                val newlyCreatedObject = irTemporary(irCall(create).also { createCall ->
-                    invokeSuspendCall.dispatchReceiver = irCall(create).also {
+                    if (!isConstructorCall) {
-                        it.dispatchReceiver = irGet(function.dispatchReceiverParameter!!)
+                        createCall.dispatchReceiver = irGet(function.dispatchReceiverParameter!!)
-                        for ((index, param) in function.valueParameters.withIndex()) {
-                            it.putValueArgument(index, irGet(param))
-                        }
                     }
+                    if (receiverField != null) {
+                        createCall.putValueArgument(0, irGet(function.valueParameters[0]))
+                    }
+                    createCall.putValueArgument(if (receiverField != null) 1 else 0, irGet(function.valueParameters.last()))
+                }, "create")
+                // In old BE 'create' function was responsible for putting arguments into fields, but in IR_BE I do not generate create,
+                // unless suspend lambda has no parameters (extension receiver is allowed, however)
+                // Thus, we put arguments into fields
+                if (function.valueParameters.size > create.valueParameters.size) {
+                    for ((index, param) in function.valueParameters.drop(if (receiverField != null) 1 else 0).dropLast(1).withIndex()) {
+                        +irSetField(irGet(newlyCreatedObject), parametersFields[index], irGet(param))
+                    }
+                }
+                +irReturn(irCall(invokeSuspend).also { invokeSuspendCall ->
+                    invokeSuspendCall.dispatchReceiver = irGet(newlyCreatedObject)
                     invokeSuspendCall.putValueArgument(0, irGetField(null, unitField))
                 })
             }
@@ -268,16 +283,19 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
         parent.declarations.add(irClass)
     }
 
+    // Primary constructor accepts parameters equal to function reference arguments + continuation and sets the fields.
     private fun IrClass.addPrimaryConstructorForLambda(
         arity: Int,
-        irFunction: IrFunction,
+        reference: IrFunctionReference,
         fields: List<IrField>
     ): IrConstructor =
         addConstructor {
             isPrimary = true
             returnType = defaultType
         }.also { constructor ->
-            irFunction.valueParameters.mapTo(constructor.valueParameters) { it.copyTo(constructor) }
+            for ((param, arg) in reference.getArguments()) {
+                constructor.addValueParameter(name = param.name.asString(), type = arg.type)
+            }
             val completionParameterSymbol = constructor.addCompletionValueParameter()
 
             val superClassConstructor = suspendLambda.owner.constructors.single {
@@ -295,6 +313,8 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             }
         }
 
+    // Secondary constructor accepts parameters equals to arguments of function reference and is used for callable references
+    // TODO: get rid of it and use primary constructor only
     private fun IrClass.addSecondaryConstructorForLambda(primary: IrConstructor): IrConstructor =
         addConstructor {
             isPrimary = false

compiler/testData/codegen/box/coroutines/bridges/lambdaWithMultipleParameters_ir.kt

@@ -0,0 +1,28 @@
+// WITH_RUNTIME
+// WITH_COROUTINES
+// COMMON_COROUTINES_TEST
+// CHECK_BYTECODE_LISTING
+// IGNORE_BACKEND: JVM
+// TARGET_BACKEND: JVM
+
+import helpers.*
+import COROUTINES_PACKAGE.*
+
+fun builder(c: suspend() -> Unit) {
+    c.startCoroutine(EmptyContinuation)
+}
+
+suspend fun call(c: suspend(String, Long, Long, Long, Long, Long) -> String): String {
+    return c("OK", 1, 2, 3, 4, 5)
+}
+
+fun box(): String {
+    var res = ""
+    builder {
+        res = call { a, b, c, d, e, f ->
+            "$a$b$c$d$e$f"
+        }
+    }
+    if (res != "OK12345") return res
+    return "OK"
+}

compiler/testData/codegen/box/coroutines/bridges/lambdaWithMultipleParameters_ir.txt

@@ -0,0 +1,47 @@
+@kotlin.Metadata
+final class LambdaWithMultipleParameters_irKt$$call$Continuation {
+    field label: int
+    @org.jetbrains.annotations.NotNull field result: java.lang.Object
+    public method <init>(@org.jetbrains.annotations.Nullable p0: kotlin.coroutines.Continuation): void
+    public final @org.jetbrains.annotations.Nullable method invokeSuspend(@org.jetbrains.annotations.NotNull p0: java.lang.Object): java.lang.Object
+}
+
+@kotlin.coroutines.jvm.internal.DebugMetadata
+@kotlin.Metadata
+final class LambdaWithMultipleParameters_irKt$box$1$1 {
+    private field a: java.lang.String
+    private field b: long
+    private field c: long
+    private field d: long
+    private field e: long
+    private field f: long
+    private field label: int
+    inner class LambdaWithMultipleParameters_irKt$box$1$1
+    public method <init>(): void
+    public method <init>(@org.jetbrains.annotations.Nullable p0: kotlin.coroutines.Continuation): void
+    public final method invoke(p0: java.lang.Object, p1: java.lang.Object, p2: java.lang.Object, p3: java.lang.Object, p4: java.lang.Object, p5: java.lang.Object, p6: java.lang.Object): java.lang.Object
+    public final @org.jetbrains.annotations.Nullable method invokeSuspend(@org.jetbrains.annotations.NotNull p0: java.lang.Object): java.lang.Object
+}
+
+@kotlin.coroutines.jvm.internal.DebugMetadata
+@kotlin.Metadata
+final class LambdaWithMultipleParameters_irKt$box$1 {
+    private field $res: kotlin.jvm.internal.Ref$ObjectRef
+    field L$0: java.lang.Object
+    private field label: int
+    inner class LambdaWithMultipleParameters_irKt$box$1
+    public method <init>(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.internal.Ref$ObjectRef): void
+    public method <init>(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.internal.Ref$ObjectRef, @org.jetbrains.annotations.Nullable p1: kotlin.coroutines.Continuation): void
+    public final @org.jetbrains.annotations.NotNull method create(@org.jetbrains.annotations.NotNull p0: kotlin.coroutines.Continuation): kotlin.coroutines.Continuation
+    public final method invoke(p0: java.lang.Object): java.lang.Object
+    public final @org.jetbrains.annotations.Nullable method invokeSuspend(@org.jetbrains.annotations.NotNull p0: java.lang.Object): java.lang.Object
+}
+
+@kotlin.Metadata
+public final class LambdaWithMultipleParameters_irKt {
+    inner class LambdaWithMultipleParameters_irKt$box$1
+    inner class LambdaWithMultipleParameters_irKt$box$1$1
+    public final static @org.jetbrains.annotations.NotNull method box(): java.lang.String
+    public final static method builder(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function1): void
+    public final static @org.jetbrains.annotations.Nullable method call(@org.jetbrains.annotations.NotNull p0: kotlin.jvm.functions.Function7, @org.jetbrains.annotations.NotNull p1: kotlin.coroutines.Continuation): java.lang.Object
+}

compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java

@@ -6654,6 +6654,16 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             public void testLambdaWithMultipleParameters_1_3() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/bridges/lambdaWithMultipleParameters.kt", "kotlin.coroutines");
             }
+
+            @TestMetadata("lambdaWithMultipleParameters_ir.kt")
+            public void testLambdaWithMultipleParameters_ir_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/bridges/lambdaWithMultipleParameters_ir.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("lambdaWithMultipleParameters_ir.kt")
+            public void testLambdaWithMultipleParameters_ir_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/bridges/lambdaWithMultipleParameters_ir.kt", "kotlin.coroutines");
+            }
         }
 
         @TestMetadata("compiler/testData/codegen/box/coroutines/controlFlow")

compiler/tests/org/jetbrains/kotlin/codegen/LightAnalysisModeTestGenerated.java

@@ -6623,6 +6623,16 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTest("compiler/testData/codegen/box/coroutines/bridges/lambdaWithLongReceiver_ir.kt");
             }
 
+            @TestMetadata("lambdaWithMultipleParameters_ir.kt")
+            public void ignoreLambdaWithMultipleParameters_ir_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/bridges/lambdaWithMultipleParameters_ir.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("lambdaWithMultipleParameters_ir.kt")
+            public void ignoreLambdaWithMultipleParameters_ir_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/bridges/lambdaWithMultipleParameters_ir.kt", "kotlin.coroutines");
+            }
+
             private void runTest(String testDataFilePath) throws Exception {
                 KotlinTestUtils.runTest(this::doTest, TargetBackend.JVM, testDataFilePath);
             }

compiler/tests/org/jetbrains/kotlin/codegen/ir/IrBlackBoxCodegenTestGenerated.java

@@ -6199,6 +6199,11 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
             public void testLambdaWithMultipleParameters_1_3() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/bridges/lambdaWithMultipleParameters.kt", "kotlin.coroutines");
             }
+
+            @TestMetadata("lambdaWithMultipleParameters_ir.kt")
+            public void testLambdaWithMultipleParameters_ir_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/bridges/lambdaWithMultipleParameters_ir.kt", "kotlin.coroutines");
+            }
         }
 
         @TestMetadata("compiler/testData/codegen/box/coroutines/controlFlow")