JVM_IR: Support suspend lambdas with receivers and multiple parameters

In old BE we generate `create` for this kind of suspend lambdas, which,
like in simple suspend lambdas is responsible for putting arguments
(including extension receiver) to fields.
But, if number of parameters of suspend lambda (including receiver) >= 1,
there is no need to generate `create`, since `create` is called only by
`createCoroutine` and friends from stdlib, and there is no version of
`createCoroutine` for lambdas with multiple parameters.
Thus, in old BE we generate a redundant method, which affects method
count.
In JVM_BE we decided to 'inline' create into `invoke` for suspend lambdas
with multiple parameters.

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

@@ -128,9 +128,15 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             if (info.arity <= 1) {
                 val singleParameterField = receiverField ?: parametersWithoutArguments.singleOrNull()
                 val create = addCreate(constructor, suspendLambda, info, parametersWithArguments, singleParameterField)
-                addInvoke(create, invokeSuspend, invokeToOverride, singleParameterField, emptyList(), isConstructorCall = false)
+                addInvokeCallingCreate(create, invokeSuspend, invokeToOverride, singleParameterField)
             } else {
-                addInvoke(constructor, invokeSuspend, invokeToOverride, receiverField, parametersWithoutArguments, isConstructorCall = true)
+                addInvokeCallingConstructor(
+                    constructor,
+                    invokeSuspend,
+                    invokeToOverride,
+                    parametersWithArguments,
+                    listOfNotNull(receiverField) + parametersWithoutArguments
+                )
             }
 
             context.suspendLambdaToOriginalFunctionMap[this] = info.function
@@ -199,37 +205,69 @@ private class AddContinuationLowering(private val context: JvmBackendContext) :
             function.valueParameters.mapTo(valueParameters) { it.copyTo(this) }
         }
 
-    private fun IrClass.addInvoke(
+    // Invoke function in lambdas is responsible for
+    //   1) calling `create`
+    //   2) starting newly created coroutine by calling `invokeSuspend`.
+    // Thus, it creates a clone of suspend lambda and starts it.
+    private fun IrClass.addInvokeCallingCreate(
         create: IrFunction,
         invokeSuspend: IrFunction,
         invokeToOverride: IrSimpleFunctionSymbol,
-        receiverField: IrField?,
-        parametersWithoutArguments: List<IrField>,
-        isConstructorCall: Boolean
+        receiverField: IrField?
     ) {
         val unitClass = context.irBuiltIns.unitClass
         val unitField = context.declarationFactory.getFieldForObjectInstance(unitClass.owner)
         addFunctionOverride(invokeToOverride.owner).also { function ->
             function.body = context.createIrBuilder(function.symbol).irBlockBody {
+                // Call `create`
                 val newlyCreatedObject = irTemporary(irCall(create).also { createCall ->
-                    if (!isConstructorCall) {
-                        createCall.dispatchReceiver = irGet(function.dispatchReceiverParameter!!)
-                    }
+                    createCall.dispatchReceiver = irGet(function.dispatchReceiverParameter!!)
                     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 we do not generate create,
-                // unless suspend lambda has no or one parameter, including extension receiver
-                // This is because 'create' is called only from 'createCoroutineUnintercepted' from stdlib. And we have only versions
-                // for suspend lambdas without parameters and for ones with exactly one parameter (or extension receiver).
-                // Thus, we put arguments into fields in 'invoke'.
+                // Start coroutine
+                +irReturn(irCall(invokeSuspend).also { invokeSuspendCall ->
+                    invokeSuspendCall.dispatchReceiver = irGet(newlyCreatedObject)
+                    invokeSuspendCall.putValueArgument(0, irGetField(null, unitField))
+                })
+            }
+        }
+    }
+
+    // The same as @see addInvokeCallingCreate, but without `create`.
+    // In old BE 'create' function was responsible for putting arguments into fields, but in IR_BE we do not generate create,
+    // unless suspend lambda has no or one parameter, including extension receiver.
+    // This is because 'create' is called only from 'createCoroutineUnintercepted' from stdlib. And we have only versions
+    // for suspend lambdas without parameters and for ones with exactly one parameter (or extension receiver).
+    // And since we still value method count, we do not want to inflate it with unnecessary functions.
+    // Thus, we put arguments into fields in 'invoke', instead of `create`.
+    private fun IrClass.addInvokeCallingConstructor(
+        constructor: IrFunction,
+        invokeSuspend: IrFunction,
+        invokeToOverride: IrSimpleFunctionSymbol,
+        parametersWithArguments: List<IrField>,
+        parametersWithoutArguments: List<IrField>
+    ) {
+        val unitClass = context.irBuiltIns.unitClass
+        val unitField = context.declarationFactory.getFieldForObjectInstance(unitClass.owner)
+        addFunctionOverride(invokeToOverride.owner).also { function ->
+            function.body = context.createIrBuilder(function.symbol).irBlockBody {
+                // Create a copy
+                val newlyCreatedObject = irTemporary(irCall(constructor).also { constructorCall ->
+                    for ((index, field) in parametersWithArguments.withIndex()) {
+                        constructorCall.putValueArgument(index, irGetField(irGet(function.dispatchReceiverParameter!!), field))
+                    }
+                    constructorCall.putValueArgument(parametersWithArguments.size, irGet(function.valueParameters.last()))
+                }, "constructor")
+                // Move parameters into fields (instead of `create`)
                 if (parametersWithoutArguments.isNotEmpty()) {
-                    for ((index, param) in function.valueParameters.drop(if (receiverField != null) 1 else 0).dropLast(1).withIndex()) {
+                    for ((index, param) in function.valueParameters.dropLast(1).withIndex()) {
                         +irSetField(irGet(newlyCreatedObject), parametersWithoutArguments[index], irGet(param))
                     }
                 }
+                // Start coroutine
                 +irReturn(irCall(invokeSuspend).also { invokeSuspendCall ->
                     invokeSuspendCall.dispatchReceiver = irGet(newlyCreatedObject)
                     invokeSuspendCall.putValueArgument(0, irGetField(null, unitField))

compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/manyParameters.kt

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

compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/manyParametersNoCapture.kt

@@ -0,0 +1,42 @@
+// IGNORE_BACKEND_FIR: JVM_IR
+// WITH_RUNTIME
+// WITH_COROUTINES
+// COMMON_COROUTINES_TEST
+import helpers.*
+import COROUTINES_PACKAGE.*
+import COROUTINES_PACKAGE.intrinsics.*
+
+suspend fun suspendHere(v: String): String = suspendCoroutineUninterceptedOrReturn { x ->
+    x.resume(v)
+    COROUTINE_SUSPENDED
+}
+
+fun builder(c: suspend () -> Unit) {
+    c.startCoroutine(EmptyContinuation)
+}
+
+class Controller {
+    var result = ""
+
+    override fun toString() = "Controller"
+}
+
+val controller = Controller()
+
+suspend fun foo(c: suspend Controller.(Long, Int, String) -> String) = controller.c(56L, 55, "abc")
+
+fun box(): String {
+    var final = ""
+
+    builder {
+        final = foo { l, i, s ->
+            result = suspendHere("$this#$l#$i#$s")
+            "OK"
+        }
+    }
+
+    if (controller.result != "Controller#56#55#abc") return "fail: ${controller.result}"
+
+    return final
+}
+

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

@@ -8402,6 +8402,16 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/localVal.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("manyParametersNoCapture.kt")
+            public void testManyParametersNoCapture_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/manyParametersNoCapture.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("manyParametersNoCapture.kt")
+            public void testManyParametersNoCapture_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/manyParametersNoCapture.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("manyParameters.kt")
             public void testManyParameters_1_2() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/manyParameters.kt", "kotlin.coroutines.experimental");

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

@@ -8402,6 +8402,16 @@ public class LightAnalysisModeTestGenerated extends AbstractLightAnalysisModeTes
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/localVal.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("manyParametersNoCapture.kt")
+            public void testManyParametersNoCapture_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/manyParametersNoCapture.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("manyParametersNoCapture.kt")
+            public void testManyParametersNoCapture_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/manyParametersNoCapture.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("manyParameters.kt")
             public void testManyParameters_1_2() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/manyParameters.kt", "kotlin.coroutines.experimental");

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

@@ -7347,6 +7347,11 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/localVal.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("manyParametersNoCapture.kt")
+            public void testManyParametersNoCapture_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/manyParametersNoCapture.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("manyParameters.kt")
             public void testManyParameters_1_3() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/manyParameters.kt", "kotlin.coroutines");

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

@@ -7347,6 +7347,11 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/localVal.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("manyParametersNoCapture.kt")
+            public void testManyParametersNoCapture_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/manyParametersNoCapture.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("manyParameters.kt")
             public void testManyParameters_1_3() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/manyParameters.kt", "kotlin.coroutines");

js/js.tests/test/org/jetbrains/kotlin/js/test/ir/semantics/IrJsCodegenBoxTestGenerated.java

@@ -6312,6 +6312,11 @@ public class IrJsCodegenBoxTestGenerated extends AbstractIrJsCodegenBoxTest {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/localVal.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("manyParametersNoCapture.kt")
+            public void testManyParametersNoCapture_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/manyParametersNoCapture.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("manyParameters.kt")
             public void testManyParameters_1_3() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/manyParameters.kt", "kotlin.coroutines");

js/js.tests/test/org/jetbrains/kotlin/js/test/semantics/JsCodegenBoxTestGenerated.java

@@ -7292,6 +7292,16 @@ public class JsCodegenBoxTestGenerated extends AbstractJsCodegenBoxTest {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/localVal.kt", "kotlin.coroutines");
             }
 
+            @TestMetadata("manyParametersNoCapture.kt")
+            public void testManyParametersNoCapture_1_2() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/manyParametersNoCapture.kt", "kotlin.coroutines.experimental");
+            }
+
+            @TestMetadata("manyParametersNoCapture.kt")
+            public void testManyParametersNoCapture_1_3() throws Exception {
+                runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/manyParametersNoCapture.kt", "kotlin.coroutines");
+            }
+
             @TestMetadata("manyParameters.kt")
             public void testManyParameters_1_2() throws Exception {
                 runTestWithPackageReplacement("compiler/testData/codegen/box/coroutines/suspendFunctionTypeCall/manyParameters.kt", "kotlin.coroutines.experimental");