JVM_IR: Support suspend lambdas with many arguments.

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

@@ -15,9 +15,9 @@ import org.jetbrains.kotlin.backend.jvm.JvmLoweredDeclarationOrigin
 import org.jetbrains.kotlin.backend.jvm.codegen.isInvokeSuspendForInlineOfLambda
 import org.jetbrains.kotlin.backend.jvm.codegen.isInvokeSuspendOfContinuation
 import org.jetbrains.kotlin.backend.jvm.codegen.isInvokeSuspendOfLambda
-import org.jetbrains.kotlin.backend.jvm.ir.IrInlineReferenceLocator
+import org.jetbrains.kotlin.backend.jvm.ir.*
-import org.jetbrains.kotlin.backend.jvm.ir.defaultValue
 import org.jetbrains.kotlin.backend.jvm.localDeclarationsPhase
+import org.jetbrains.kotlin.builtins.functions.FunctionInvokeDescriptor.Factory.BIG_ARITY
 import org.jetbrains.kotlin.codegen.coroutines.*
 import org.jetbrains.kotlin.codegen.inline.coroutines.FOR_INLINE_SUFFIX
 import org.jetbrains.kotlin.descriptors.Modality
@@ -808,11 +808,34 @@ private fun IrCall.createSuspendFunctionCallViewIfNeeded(context: JvmBackendCont
     val view = (symbol.owner as IrSimpleFunction).suspendFunctionViewOrStub(context)
     if (view == symbol.owner) return this
 
-    return IrCallImpl(startOffset, endOffset, view.returnType, view.symbol, superQualifierSymbol = superQualifierSymbol).also {
+    val isBigAritySuspendFunctionN =
+        symbol.owner.parentAsClass?.defaultType?.let { it.isSuspendFunction() || it.isKSuspendFunction() } && valueArgumentsCount + 1 >= BIG_ARITY
+    val calleeSymbol =
+        if (isBigAritySuspendFunctionN) context.ir.symbols.functionN.functions.single { it.owner.name.toString() == "invoke" } else view.symbol
+    return IrCallImpl(startOffset, endOffset, view.returnType, calleeSymbol, superQualifierSymbol = superQualifierSymbol).also {
         it.copyTypeArgumentsFrom(this)
-        it.copyAttributes(this)
         it.dispatchReceiver = dispatchReceiver
         it.extensionReceiver = extensionReceiver
+        val valueArguments = computeSuspendFunctionCallViewValueArguments(caller, context, view)
+        if (isBigAritySuspendFunctionN) {
+            // If we are calling FunctionN.invoke, wrap arguments in an array.
+            it.putValueArgument(0, context.createJvmIrBuilder(caller.symbol).run {
+                irArray(irSymbols.array.typeWith(context.irBuiltIns.anyNType)) {
+                    valueArguments.forEach { argument -> +argument }
+                }
+            })
+        } else {
+            valueArguments.forEachIndexed { index, argument -> it.putValueArgument(index, argument) }
+        }
+    }
+}
+
+private fun IrCall.computeSuspendFunctionCallViewValueArguments(
+    caller: IrFunction,
+    context: JvmBackendContext,
+    view: IrFunction
+): List<IrExpression> {
+    // Locate the caller continuation parameter. The continuation parameter is before default argument mask(s) and handler params.
     val callerNumberOfMasks = caller.valueParameters.count { it.origin == IrDeclarationOrigin.MASK_FOR_DEFAULT_FUNCTION }
     val callerContinuationIndex = caller.valueParameters.size - 1 - (if (callerNumberOfMasks != 0) callerNumberOfMasks + 1 else 0)
     val continuationParameter =
@@ -827,14 +850,15 @@ private fun IrCall.createSuspendFunctionCallViewIfNeeded(context: JvmBackendCont
     // continuation before the default parameter mask(s) and handler.
     val numberOfMasks = view.valueParameters.count { it.origin == IrDeclarationOrigin.MASK_FOR_DEFAULT_FUNCTION }
     val continuationIndex = valueArgumentsCount - (if (numberOfMasks != 0) numberOfMasks + 1 else 0)
+    val valueArguments = mutableListOf<IrExpression>()
     for (i in 0 until continuationIndex) {
-            it.putValueArgument(i, getValueArgument(i))
+        valueArguments.add(getValueArgument(i)!!)
     }
-        it.putValueArgument(continuationIndex, continuationParameter)
+    valueArguments.add(continuationParameter)
     if (numberOfMasks != 0) {
         for (i in 0 until numberOfMasks + 1) {
-                it.putValueArgument(valueArgumentsCount + i - 1, getValueArgument(continuationIndex + i))
+            valueArguments.add(getValueArgument(continuationIndex + i)!!)
-            }
         }
     }
+    return valueArguments
 }

compiler/testData/codegen/box/coroutines/featureIntersection/callableReference/bigArity.kt

@@ -1,5 +1,5 @@
 // IGNORE_BACKEND_FIR: JVM_IR
-// IGNORE_BACKEND: JS_IR, JS, NATIVE, JVM_IR
+// IGNORE_BACKEND: JS_IR, JS, NATIVE
 // WITH_REFLECT
 // WITH_COROUTINES
 
@@ -27,6 +27,20 @@ suspend fun expectsLambdaWithBigArity(c: suspend (Long, Long, Long, Long, Long,
     return c.invoke(1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, 1L, "OK")
 }
 
+// Uses Function22.invoke with individual arguments. Function22 and not Function21 because of the added continuation parameter.
+suspend fun expectsLambdaWithArity21(c: suspend (Int, Int, Int, Int, Int, Int, Int, Int, Int, Int,
+                                                 Int, Int, Int, Int, Int, Int, Int, Int, Int, Int,
+                                                 String) -> String): String {
+    return c.invoke(1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, "OK")
+}
+
+// Uses FunctionN.invoke with varargs array of arguments.
+suspend fun expectsLambdaWithArity22(c: suspend (Int, Int, Int, Int, Int, Int, Int, Int, Int, Int,
+                                                 Int, Int, Int, Int, Int, Int, Int, Int, Int, Int,
+                                                 Int, String) -> String): String {
+    return c.invoke(1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, "OK")
+}
+
 fun box(): String {
     val a = A()
     var res = "FAIL 1"
@@ -43,5 +57,19 @@ fun box(): String {
     builder {
         res = expectsLambdaWithBigArity { _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, s -> s }
     }
+    if (res != "OK") return res
+    res = "FAIL 4"
+    builder {
+        res = expectsLambdaWithArity21 {
+            i1, i2, i3, i4, i5, i6, i7, i8, i9, i10, i11, i12, i13, i14, i15, i16, i17, i18, i19, i20, s -> s
+        }
+    }
+    if (res != "OK") return res
+    res = "FAIL 5"
+    builder {
+        res = expectsLambdaWithArity22 {
+            i1, i2, i3, i4, i5, i6, i7, i8, i9, i10, i11, i12, i13, i14, i15, i16, i17, i18, i19, i20, i21, s -> s
+        }
+    }
     return res
 }