FIR2IR: handle more vararg spreads for adapted callable references

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/generators/CallAndReferenceGenerator.kt

@@ -90,9 +90,14 @@ class CallAndReferenceGenerator(
                     )
                 }
                 is IrFunctionSymbol -> {
+                    assert(type.isFunctionTypeOrSubtype()) {
+                        "Callable reference whose symbol refers to a function should be of functional type."
+                    }
+                    type as IrSimpleType
                     val function = symbol.owner
                     // TODO: should refer to LanguageVersionSettings.SuspendConversion
-                    if (requiresCoercionToUnit(type, function) ||
+                    if (requiresVarargSpread(callableReferenceAccess, type, function) ||
+                        requiresCoercionToUnit(type, function) ||
                         requiresSuspendConversion(type, function)
                     ) {
                         val adaptedType = callableReferenceAccess.typeRef.coneType.kFunctionTypeToFunctionType()
@@ -115,15 +120,58 @@ class CallAndReferenceGenerator(
         }.applyTypeArguments(callableReferenceAccess).applyReceivers(callableReferenceAccess, explicitReceiverExpression)
     }
 
-    private fun requiresCoercionToUnit(type: IrType, function: IrFunction): Boolean {
+    /*
-        if (!type.isFunctionTypeOrSubtype()) {
+     * For example,
+     * fun referenceConsumer(f: (Char, Char) -> String): String = ... // e.g., f(char1, char2)
+     * fun referenced(vararg xs: Char) = ...
+     * fun useSite(...) = { ... referenceConsumer(::referenced) ... }
+     *
+     * At the use site, instead of referenced, we can put the adapter: { a, b -> referenced(a, b) }
+     */
+    private fun requiresVarargSpread(
+        callableReferenceAccess: FirCallableReferenceAccess,
+        type: IrSimpleType,
+        function: IrFunction
+    ): Boolean {
+        // Unbound callable reference 'A::foo'
+        val shift = if (callableReferenceAccess.explicitReceiver is FirResolvedQualifier) 1 else 0
+        val typeArguments = type.arguments
+        // Drop the return type from type arguments
+        val expectedParameterSize = typeArguments.size - 1 - shift
+        if (expectedParameterSize < function.valueParameters.size) {
             return false
         }
-        val expectedReturnType = (type as? IrSimpleType)?.arguments?.last()?.typeOrNull
+        var hasSpreadCase = false
+        function.valueParameters.forEachIndexed { index, irValueParameter ->
+            if (irValueParameter.isVararg && typeArguments[shift + index] == irValueParameter.varargElementType) {
+                hasSpreadCase = true
+            }
+        }
+        return hasSpreadCase
+    }
+
+    /*
+     * For example,
+     * fun referenceConsumer(f: () -> Unit) = f()
+     * fun referenced(...): Any { ... }
+     * fun useSite(...) = { ... referenceConsumer(::referenced) ... }
+     *
+     * At the use site, instead of referenced, we can put the adapter: { ... -> referenced(...) }
+     */
+    private fun requiresCoercionToUnit(type: IrSimpleType, function: IrFunction): Boolean {
+        val expectedReturnType = type.arguments.last().typeOrNull
         return expectedReturnType?.isUnit() == true && !function.returnType.isUnit()
     }
 
-    private fun requiresSuspendConversion(type: IrType, function: IrFunction): Boolean =
+    /*
+     * For example,
+     * fun referenceConsumer(f: suspend () -> Unit) = ...
+     * fun nonSuspendFunction(...) = ...
+     * fun useSite(...) = { ... referenceConsumer(::nonSuspendFunction) ... }
+     *
+     * At the use site, instead of referenced, we can put the suspend lambda as an adapter.
+     */
+    private fun requiresSuspendConversion(type: IrSimpleType, function: IrFunction): Boolean =
         type.isKSuspendFunction() && !function.isSuspend
 
     private fun ConeKotlinType.kFunctionTypeToFunctionType(): IrSimpleType {
@@ -301,7 +349,7 @@ class CallAndReferenceGenerator(
                 error("unknown callee kind: ${adapteeFunction.render()}")
         }
 
-        var shift = 0
+        var adapterParameterIndex = 0
         if (boundDispatchReceiver != null || boundExtensionReceiver != null) {
             val receiverValue = IrGetValueImpl(
                 startOffset, endOffset, adapterFunction.extensionReceiverParameter!!.symbol, IrStatementOrigin.ADAPTED_FUNCTION_REFERENCE
@@ -318,7 +366,7 @@ class CallAndReferenceGenerator(
             )
             if (adapteeFunction.extensionReceiverParameter != null) {
                 irCall.extensionReceiver = adaptedReceiverValue
-                shift = 1
+                adapterParameterIndex++
             } else {
                 irCall.dispatchReceiver = adaptedReceiverValue
             }
@@ -329,26 +377,39 @@ class CallAndReferenceGenerator(
 
         adapteeFunction.valueParameters.mapIndexed { index, valueParameter ->
             when {
-                valueParameter.hasDefaultValue() -> {
-                    irCall.putValueArgument(shift + index, null)
-                }
                 valueParameter.isVararg -> {
                     if (adapterFunction.valueParameters.size <= index) {
-                        irCall.putValueArgument(shift + index, null)
+                        irCall.putValueArgument(index, null)
                     } else {
                         val adaptedValueArgument =
                             IrVarargImpl(startOffset, endOffset, valueParameter.type, valueParameter.varargElementType!!)
-                        val irValueArgument = adapterFunction.valueParameters[index].toGetValue()
+                        var neitherArrayNorSpread = false
-                        if (irValueArgument.type == valueParameter.type) {
+                        while (adapterParameterIndex < adapterFunction.valueParameters.size) {
-                            adaptedValueArgument.addElement(IrSpreadElementImpl(startOffset, endOffset, irValueArgument))
+                            val irValueArgument = adapterFunction.valueParameters[adapterParameterIndex].toGetValue()
-                        } else {
+                            if (irValueArgument.type == valueParameter.type) {
-                            adaptedValueArgument.addElement(irValueArgument)
+                                adaptedValueArgument.addElement(IrSpreadElementImpl(startOffset, endOffset, irValueArgument))
+                                adapterParameterIndex++
+                                break
+                            } else if (irValueArgument.type == valueParameter.varargElementType) {
+                                adaptedValueArgument.addElement(irValueArgument)
+                                adapterParameterIndex++
+                            } else {
+                                neitherArrayNorSpread = true
+                                break
+                            }
+                        }
+                        if (neitherArrayNorSpread) {
+                            irCall.putValueArgument(index, null)
+                        } else {
+                            irCall.putValueArgument(index, adaptedValueArgument)
                         }
-                        irCall.putValueArgument(shift + index, adaptedValueArgument)
                     }
                 }
+                valueParameter.hasDefaultValue() -> {
+                    irCall.putValueArgument(index, null)
+                }
                 else -> {
-                    irCall.putValueArgument(shift + index, adapterFunction.valueParameters[index].toGetValue())
+                    irCall.putValueArgument(index, adapterFunction.valueParameters[adapterParameterIndex++].toGetValue())
                 }
             }
         }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/ResolutionStages.kt

@@ -332,7 +332,11 @@ private fun FirSession.createAdaptedKFunctionType(
 
     // If a function with vararg is passed to a place where a spread of elements is expected, we can adapt the function reference to
     // literally spread such vararg argument. E.g., foo(vararg xs: Char): String => bar(::foo) where bar(f: (Char, Char) -> String)
-    if (expectedParameterNumber != null && expectedParameterTypes != null && parameterTypes.size < expectedParameterNumber && lastVarargParameter != null) {
+    if (expectedParameterNumber != null &&
+        expectedParameterTypes != null &&
+        parameterTypes.size < expectedParameterNumber &&
+        lastVarargParameter != null
+    ) {
         val varargArrayType = lastVarargParameter.returnTypeRef.coneType
         val varargElementType = varargArrayType.varargElementType(session)
         val expectedParameterType = (expectedParameterTypes[parameterTypes.size + shift] as? ConeKotlinTypeProjection)?.type

compiler/testData/codegen/box/callableReference/adaptedReferences/inlineVararg.kt

@@ -1,5 +1,4 @@
 // !LANGUAGE: +NewInference
-// IGNORE_BACKEND_FIR: JVM_IR
 // WITH_RUNTIME
 // KJS_WITH_FULL_RUNTIME
 

compiler/testData/codegen/box/callableReference/adaptedReferences/inlineVarargInts.kt

@@ -1,5 +1,4 @@
 // !LANGUAGE: +NewInference
-// IGNORE_BACKEND_FIR: JVM_IR
 // WITH_RUNTIME
 // KJS_WITH_FULL_RUNTIME
 

compiler/testData/codegen/box/callableReference/equality/capturedVararg.kt

@@ -1,6 +1,5 @@
 // IGNORE_BACKEND: JS, JS_IR
 // IGNORE_BACKEND: JS_IR_ES6
-// IGNORE_BACKEND_FIR: JVM_IR
 // FILE: test.kt
 
 fun checkEqual(x: Any, y: Any) {

compiler/testData/codegen/box/callableReference/equality/varargAsArrayMemberOrExtension.kt

@@ -1,4 +1,3 @@
-// IGNORE_BACKEND_FIR: JVM_IR
 // FILE: test.kt
 
 fun checkNotEqual(x: Any, y: Any) {

compiler/testData/ir/irText/expressions/callableReferences/unboundMemberReferenceWithAdaptedArguments.fir.txt

@@ -67,8 +67,6 @@ FILE fqName:<root> fileName:/unboundMemberReferenceWithAdaptedArguments.kt
             BLOCK_BODY
               CALL 'public open fun foo (vararg xs: kotlin.Int): kotlin.Int declared in <root>.A' type=kotlin.Int origin=null
                 $this: GET_VAR 'p0: <root>.A declared in <root>.testUnbound.foo' type=<root>.A origin=null
-                xs: VARARG type=kotlin.IntArray varargElementType=kotlin.Int
-                  GET_VAR 'p0: <root>.A declared in <root>.testUnbound.foo' type=<root>.A origin=null
   FUN name:testBound visibility:public modality:FINAL <> (a:<root>.A) returnType:kotlin.Unit
     VALUE_PARAMETER name:a index:0 type:<root>.A
     BLOCK_BODY

compiler/testData/ir/irText/expressions/callableReferences/withAdaptedArguments.fir.txt

@@ -72,7 +72,14 @@ FILE fqName:<root> fileName:/withAdaptedArguments.kt
     BLOCK_BODY
       RETURN type=kotlin.Nothing from='public final fun testVararg (): kotlin.String declared in <root>'
         CALL 'public final fun use (fn: kotlin.Function1<kotlin.Int, kotlin.String>): kotlin.String declared in <root>' type=kotlin.String origin=null
-          fn: FUNCTION_REFERENCE 'public final fun fnWithVarargs (vararg xs: kotlin.Int): kotlin.String declared in <root>' type=kotlin.reflect.KFunction1<kotlin.Int, kotlin.String> origin=null reflectionTarget=<same>
+          fn: FUN_EXPR type=kotlin.Function1<kotlin.Int, kotlin.String> origin=ADAPTED_FUNCTION_REFERENCE
+            FUN ADAPTER_FOR_CALLABLE_REFERENCE name:fnWithVarargs visibility:local modality:FINAL <> (p0:kotlin.Int) returnType:kotlin.String
+              VALUE_PARAMETER ADAPTER_PARAMETER_FOR_CALLABLE_REFERENCE name:p0 index:0 type:kotlin.Int
+              BLOCK_BODY
+                RETURN type=kotlin.Nothing from='local final fun fnWithVarargs (p0: kotlin.Int): kotlin.String declared in <root>.testVararg'
+                  CALL 'public final fun fnWithVarargs (vararg xs: kotlin.Int): kotlin.String declared in <root>' type=kotlin.String origin=null
+                    xs: VARARG type=kotlin.IntArray varargElementType=kotlin.Int
+                      GET_VAR 'p0: kotlin.Int declared in <root>.testVararg.fnWithVarargs' type=kotlin.Int origin=null
   FUN name:testCoercionToUnit visibility:public modality:FINAL <> () returnType:kotlin.Unit
     BLOCK_BODY
       RETURN type=kotlin.Nothing from='public final fun testCoercionToUnit (): kotlin.Unit declared in <root>'
@@ -87,8 +94,17 @@ FILE fqName:<root> fileName:/withAdaptedArguments.kt
     BLOCK_BODY
       RETURN type=kotlin.Nothing from='public final fun testImportedObjectMember (): kotlin.String declared in <root>'
         CALL 'public final fun use (fn: kotlin.Function1<kotlin.Int, kotlin.String>): kotlin.String declared in <root>' type=kotlin.String origin=null
-          fn: FUNCTION_REFERENCE 'public final fun importedObjectMemberWithVarargs (vararg xs: kotlin.Int): kotlin.String declared in <root>.Host' type=kotlin.reflect.KFunction1<kotlin.Int, kotlin.String> origin=null reflectionTarget=<same>
+          fn: BLOCK type=kotlin.Function1<kotlin.Int, kotlin.String> origin=ADAPTED_FUNCTION_REFERENCE
-            $this: GET_OBJECT 'CLASS OBJECT name:Host modality:FINAL visibility:public superTypes:[kotlin.Any]' type=<root>.Host
+            FUN ADAPTER_FOR_CALLABLE_REFERENCE name:importedObjectMemberWithVarargs visibility:local modality:FINAL <> ($receiver:<root>.Host, p0:kotlin.Int) returnType:kotlin.String
+              $receiver: VALUE_PARAMETER ADAPTER_PARAMETER_FOR_CALLABLE_REFERENCE name:receiver type:<root>.Host
+              VALUE_PARAMETER ADAPTER_PARAMETER_FOR_CALLABLE_REFERENCE name:p0 index:0 type:kotlin.Int
+              BLOCK_BODY
+                RETURN type=kotlin.Nothing from='local final fun importedObjectMemberWithVarargs (p0: kotlin.Int): kotlin.String declared in <root>.testImportedObjectMember'
+                  CALL 'public final fun importedObjectMemberWithVarargs (vararg xs: kotlin.Int): kotlin.String declared in <root>.Host' type=kotlin.String origin=null
+                    $this: GET_VAR 'receiver: <root>.Host declared in <root>.testImportedObjectMember.importedObjectMemberWithVarargs' type=<root>.Host origin=ADAPTED_FUNCTION_REFERENCE
+                    xs: VARARG type=kotlin.IntArray varargElementType=kotlin.Int
+                      GET_VAR 'p0: kotlin.Int declared in <root>.testImportedObjectMember.importedObjectMemberWithVarargs' type=kotlin.Int origin=null
+            FUNCTION_REFERENCE 'local final fun importedObjectMemberWithVarargs (p0: kotlin.Int): kotlin.String declared in <root>.testImportedObjectMember' type=kotlin.Function1<kotlin.Int, kotlin.String> origin=ADAPTED_FUNCTION_REFERENCE reflectionTarget=null
   FUN name:testDefault0 visibility:public modality:FINAL <> () returnType:kotlin.String
     BLOCK_BODY
       RETURN type=kotlin.Nothing from='public final fun testDefault0 (): kotlin.String declared in <root>'

compiler/testData/ir/irText/expressions/callableReferences/withVarargViewedAsArray.fir.txt

@@ -61,7 +61,16 @@ FILE fqName:<root> fileName:/withVarargViewedAsArray.kt
   FUN name:testPlainArgs visibility:public modality:FINAL <> () returnType:kotlin.Unit
     BLOCK_BODY
       CALL 'public final fun usePlainArgs (fn: kotlin.Function2<kotlin.Int, kotlin.Int, kotlin.Int>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
-        fn: FUNCTION_REFERENCE 'public final fun sum (vararg args: kotlin.Int): kotlin.Int declared in <root>' type=kotlin.reflect.KFunction2<kotlin.Int, kotlin.Int, kotlin.Int> origin=null reflectionTarget=<same>
+        fn: FUN_EXPR type=kotlin.Function2<kotlin.Int, kotlin.Int, kotlin.Int> origin=ADAPTED_FUNCTION_REFERENCE
+          FUN ADAPTER_FOR_CALLABLE_REFERENCE name:sum visibility:local modality:FINAL <> (p0:kotlin.Int, p1:kotlin.Int) returnType:kotlin.Int
+            VALUE_PARAMETER ADAPTER_PARAMETER_FOR_CALLABLE_REFERENCE name:p0 index:0 type:kotlin.Int
+            VALUE_PARAMETER ADAPTER_PARAMETER_FOR_CALLABLE_REFERENCE name:p1 index:1 type:kotlin.Int
+            BLOCK_BODY
+              RETURN type=kotlin.Nothing from='local final fun sum (p0: kotlin.Int, p1: kotlin.Int): kotlin.Int declared in <root>.testPlainArgs'
+                CALL 'public final fun sum (vararg args: kotlin.Int): kotlin.Int declared in <root>' type=kotlin.Int origin=null
+                  args: VARARG type=kotlin.IntArray varargElementType=kotlin.Int
+                    GET_VAR 'p0: kotlin.Int declared in <root>.testPlainArgs.sum' type=kotlin.Int origin=null
+                    GET_VAR 'p1: kotlin.Int declared in <root>.testPlainArgs.sum' type=kotlin.Int origin=null
   FUN name:testPrimitiveArrayAsVararg visibility:public modality:FINAL <> () returnType:kotlin.Unit
     BLOCK_BODY
       CALL 'public final fun usePrimitiveArray (fn: kotlin.Function1<kotlin.IntArray, kotlin.Int>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null