FIR2IR: apply suspend conversions also to varargs

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

@@ -395,21 +395,9 @@ internal class AdapterGenerator(
         }
         if (this is IrVararg) {
             // element-wise SAM conversion if and only if we can build 1-to-1 mapping for elements.
-            if (argument !is FirVarargArgumentsExpression || argument.arguments.size != elements.size) {
-                return this
+            return applyConversionOnVararg(argument) { firVarargArgument ->
+                applySamConversionIfNeeded(firVarargArgument, parameter, substitutor, shouldUnwrapVarargType = true)
             }
-            val argumentMapping = this.elements.zip(argument.arguments).toMap()
-            // [IrElementTransformer] is not preferred, since it's hard to visit vararg elements only.
-            val irVarargElements = elements as MutableList<IrVarargElement>
-            irVarargElements.replaceAll { irVarargElement ->
-                if (irVarargElement is IrExpression) {
-                    val firVarargArgument =
-                        argumentMapping[irVarargElement] ?: error("Can't find the original FirExpression for ${irVarargElement.render()}")
-                    irVarargElement.applySamConversionIfNeeded(firVarargArgument, parameter, substitutor, shouldUnwrapVarargType = true)
-                } else
-                    irVarargElement
-            }
-            return this
         }
         if (!needSamConversion(argument, parameter)) {
             return this
@@ -426,6 +414,27 @@ internal class AdapterGenerator(
         return IrTypeOperatorCallImpl(this.startOffset, this.endOffset, samType, IrTypeOperator.SAM_CONVERSION, samType, this)
     }
 
+    private fun IrVararg.applyConversionOnVararg(
+        argument: FirExpression,
+        conversion: IrExpression.(FirExpression) -> IrExpression
+    ): IrExpression {
+        if (argument !is FirVarargArgumentsExpression || argument.arguments.size != elements.size) {
+            return this
+        }
+        val argumentMapping = this.elements.zip(argument.arguments).toMap()
+        // [IrElementTransformer] is not preferred, since it's hard to visit vararg elements only.
+        val irVarargElements = elements as MutableList<IrVarargElement>
+        irVarargElements.replaceAll { irVarargElement ->
+            if (irVarargElement is IrExpression) {
+                val firVarargArgument =
+                    argumentMapping[irVarargElement] ?: error("Can't find the original FirExpression for ${irVarargElement.render()}")
+                irVarargElement.conversion(firVarargArgument)
+            } else
+                irVarargElement
+        }
+        return this
+    }
+
     private fun needSamConversion(argument: FirExpression, parameter: FirValueParameter): Boolean {
         // If the type of the argument is already an explicitly subtype of the type of the parameter, we don't need SAM conversion.
         if (argument.typeRef !is FirResolvedTypeRef ||
@@ -470,22 +479,26 @@ internal class AdapterGenerator(
      */
     internal fun IrExpression.applySuspendConversionIfNeeded(
         argument: FirExpression,
-        parameterType: ConeKotlinType,
-        samFunctionType: ConeKotlinType?
+        parameterType: ConeKotlinType
     ): IrExpression {
         // TODO: should refer to LanguageVersionSettings.SuspendConversion
         if (this is IrBlock && origin == IrStatementOrigin.ADAPTED_FUNCTION_REFERENCE) {
             return this
         }
         // Expect the expected type to be a suspend functional type.
-        val potentialFunctionType = samFunctionType ?: parameterType
-        if (!potentialFunctionType.isSuspendFunctionType(session)) {
+        if (!parameterType.isSuspendFunctionType(session)) {
             return this
         }
-        val expectedFunctionalType = potentialFunctionType.suspendFunctionTypeToFunctionType(session)
+        val expectedFunctionalType = parameterType.suspendFunctionTypeToFunctionType(session)
+        if (this is IrVararg) {
+            // element-wise conversion if and only if we can build 1-to-1 mapping for elements.
+            return applyConversionOnVararg(argument) { firVarargArgument ->
+                applySuspendConversionIfNeeded(firVarargArgument, parameterType)
+            }
+        }
 
         val invokeSymbol = findInvokeSymbol(expectedFunctionalType, argument) ?: return this
-        val suspendConvertedType = potentialFunctionType.toIrType() as IrSimpleType
+        val suspendConvertedType = parameterType.toIrType() as IrSimpleType
         return argument.convertWithOffsets { startOffset, endOffset ->
             val irAdapterFunction = createAdapterFunctionForArgument(startOffset, endOffset, suspendConvertedType, type, invokeSymbol)
             // TODO add a bound receiver property to IrFunctionExpressionImpl?
@@ -500,7 +513,10 @@ internal class AdapterGenerator(
         }
     }
 
-    private fun findInvokeSymbol(expectedFunctionalType: ConeClassLikeType, argument: FirExpression): IrSimpleFunctionSymbol? {
+    private fun findInvokeSymbol(
+        expectedFunctionalType: ConeClassLikeType,
+        argument: FirExpression
+    ): IrSimpleFunctionSymbol? {
         val argumentType = argument.typeRef.coneType
         val argumentTypeWithInvoke = argumentType.findSubtypeOfNonSuspendFunctionalType(session, expectedFunctionalType) ?: return null
 

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

@@ -658,8 +658,9 @@ class CallAndReferenceGenerator(
             if (parameter?.returnTypeRef is FirResolvedTypeRef) {
                 // Java type case (from annotations)
                 val parameterType = parameter.returnTypeRef.coneType
-                val samFunctionType = getFunctionTypeForPossibleSamType(parameterType)
-                irArgument = irArgument.applySuspendConversionIfNeeded(argument, parameterType, samFunctionType)
+                val unwrappedParameterType = if (parameter.isVararg) parameterType.arrayElementType()!! else parameterType
+                val samFunctionType = getFunctionTypeForPossibleSamType(unwrappedParameterType)
+                irArgument = irArgument.applySuspendConversionIfNeeded(argument, samFunctionType ?: unwrappedParameterType)
                 irArgument = irArgument.applySamConversionIfNeeded(argument, parameter, substitutor)
             }
         }

compiler/testData/codegen/box/suspendConversion/suspendConversionOnVarargElements.kt

@@ -1,8 +1,6 @@
 // !LANGUAGE: +SuspendConversion
 // !DIAGNOSTICS: -UNUSED_PARAMETER
 // IGNORE_BACKEND: JVM
-// IGNORE_BACKEND_FIR: JVM_IR
-// FIR status: SuspendConversionOnVarargElementsKt$box$2 cannot be cast to kotlin.jvm.functions.Function1
 
 fun useSuspendVararg(vararg sfn: suspend () -> Unit) {}
 

compiler/testData/ir/irText/expressions/suspendConversionInVararg.fir.ir.txt

@@ -9,9 +9,23 @@ FILE fqName:<root> fileName:/suspendConversionInVararg.kt
     BLOCK_BODY
       CALL 'public final fun useSuspendVararg (vararg sfn: kotlin.coroutines.SuspendFunction0<kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
         sfn: VARARG type=kotlin.Array<out kotlin.coroutines.SuspendFunction0<kotlin.Unit>> varargElementType=kotlin.coroutines.SuspendFunction0<kotlin.Unit>
-          GET_VAR 'f1: kotlin.Function0<kotlin.Unit> declared in <root>.testSuspendConversionInVarargElementsSome' type=kotlin.Function0<kotlin.Unit> origin=null
+          BLOCK type=kotlin.coroutines.SuspendFunction0<kotlin.Unit> origin=SUSPEND_CONVERSION
+            FUN ADAPTER_FOR_SUSPEND_CONVERSION name:suspendConversion visibility:local modality:FINAL <> ($receiver:kotlin.Function0<kotlin.Unit>) returnType:kotlin.Unit [suspend]
+              $receiver: VALUE_PARAMETER ADAPTER_PARAMETER_FOR_SUSPEND_CONVERSION name:callee type:kotlin.Function0<kotlin.Unit>
+              BLOCK_BODY
+                CALL 'public abstract fun invoke (): R of kotlin.Function0 [operator] declared in kotlin.Function0' type=kotlin.Unit origin=null
+                  $this: GET_VAR 'callee: kotlin.Function0<kotlin.Unit> declared in <root>.testSuspendConversionInVarargElementsSome.suspendConversion' type=kotlin.Function0<kotlin.Unit> origin=null
+            FUNCTION_REFERENCE 'local final fun suspendConversion (): kotlin.Unit [suspend] declared in <root>.testSuspendConversionInVarargElementsSome' type=kotlin.coroutines.SuspendFunction0<kotlin.Unit> origin=SUSPEND_CONVERSION reflectionTarget=null
+              $receiver: GET_VAR 'f1: kotlin.Function0<kotlin.Unit> declared in <root>.testSuspendConversionInVarargElementsSome' type=kotlin.Function0<kotlin.Unit> origin=null
           GET_VAR 'sf2: kotlin.coroutines.SuspendFunction0<kotlin.Unit> declared in <root>.testSuspendConversionInVarargElementsSome' type=kotlin.coroutines.SuspendFunction0<kotlin.Unit> origin=null
-          GET_VAR 'f3: kotlin.Function0<kotlin.Unit> declared in <root>.testSuspendConversionInVarargElementsSome' type=kotlin.Function0<kotlin.Unit> origin=null
+          BLOCK type=kotlin.coroutines.SuspendFunction0<kotlin.Unit> origin=SUSPEND_CONVERSION
+            FUN ADAPTER_FOR_SUSPEND_CONVERSION name:suspendConversion visibility:local modality:FINAL <> ($receiver:kotlin.Function0<kotlin.Unit>) returnType:kotlin.Unit [suspend]
+              $receiver: VALUE_PARAMETER ADAPTER_PARAMETER_FOR_SUSPEND_CONVERSION name:callee type:kotlin.Function0<kotlin.Unit>
+              BLOCK_BODY
+                CALL 'public abstract fun invoke (): R of kotlin.Function0 [operator] declared in kotlin.Function0' type=kotlin.Unit origin=null
+                  $this: GET_VAR 'callee: kotlin.Function0<kotlin.Unit> declared in <root>.testSuspendConversionInVarargElementsSome.suspendConversion' type=kotlin.Function0<kotlin.Unit> origin=null
+            FUNCTION_REFERENCE 'local final fun suspendConversion (): kotlin.Unit [suspend] declared in <root>.testSuspendConversionInVarargElementsSome' type=kotlin.coroutines.SuspendFunction0<kotlin.Unit> origin=SUSPEND_CONVERSION reflectionTarget=null
+              $receiver: GET_VAR 'f3: kotlin.Function0<kotlin.Unit> declared in <root>.testSuspendConversionInVarargElementsSome' type=kotlin.Function0<kotlin.Unit> origin=null
   FUN name:testSuspendConversionInVarargElementsAll visibility:public modality:FINAL <> (f1:kotlin.Function0<kotlin.Unit>, f2:kotlin.Function0<kotlin.Unit>, f3:kotlin.Function0<kotlin.Unit>) returnType:kotlin.Unit
     VALUE_PARAMETER name:f1 index:0 type:kotlin.Function0<kotlin.Unit>
     VALUE_PARAMETER name:f2 index:1 type:kotlin.Function0<kotlin.Unit>
@@ -19,6 +33,27 @@ FILE fqName:<root> fileName:/suspendConversionInVararg.kt
     BLOCK_BODY
       CALL 'public final fun useSuspendVararg (vararg sfn: kotlin.coroutines.SuspendFunction0<kotlin.Unit>): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
         sfn: VARARG type=kotlin.Array<out kotlin.coroutines.SuspendFunction0<kotlin.Unit>> varargElementType=kotlin.coroutines.SuspendFunction0<kotlin.Unit>
-          GET_VAR 'f1: kotlin.Function0<kotlin.Unit> declared in <root>.testSuspendConversionInVarargElementsAll' type=kotlin.Function0<kotlin.Unit> origin=null
-          GET_VAR 'f2: kotlin.Function0<kotlin.Unit> declared in <root>.testSuspendConversionInVarargElementsAll' type=kotlin.Function0<kotlin.Unit> origin=null
-          GET_VAR 'f3: kotlin.Function0<kotlin.Unit> declared in <root>.testSuspendConversionInVarargElementsAll' type=kotlin.Function0<kotlin.Unit> origin=null
+          BLOCK type=kotlin.coroutines.SuspendFunction0<kotlin.Unit> origin=SUSPEND_CONVERSION
+            FUN ADAPTER_FOR_SUSPEND_CONVERSION name:suspendConversion visibility:local modality:FINAL <> ($receiver:kotlin.Function0<kotlin.Unit>) returnType:kotlin.Unit [suspend]
+              $receiver: VALUE_PARAMETER ADAPTER_PARAMETER_FOR_SUSPEND_CONVERSION name:callee type:kotlin.Function0<kotlin.Unit>
+              BLOCK_BODY
+                CALL 'public abstract fun invoke (): R of kotlin.Function0 [operator] declared in kotlin.Function0' type=kotlin.Unit origin=null
+                  $this: GET_VAR 'callee: kotlin.Function0<kotlin.Unit> declared in <root>.testSuspendConversionInVarargElementsAll.suspendConversion' type=kotlin.Function0<kotlin.Unit> origin=null
+            FUNCTION_REFERENCE 'local final fun suspendConversion (): kotlin.Unit [suspend] declared in <root>.testSuspendConversionInVarargElementsAll' type=kotlin.coroutines.SuspendFunction0<kotlin.Unit> origin=SUSPEND_CONVERSION reflectionTarget=null
+              $receiver: GET_VAR 'f1: kotlin.Function0<kotlin.Unit> declared in <root>.testSuspendConversionInVarargElementsAll' type=kotlin.Function0<kotlin.Unit> origin=null
+          BLOCK type=kotlin.coroutines.SuspendFunction0<kotlin.Unit> origin=SUSPEND_CONVERSION
+            FUN ADAPTER_FOR_SUSPEND_CONVERSION name:suspendConversion visibility:local modality:FINAL <> ($receiver:kotlin.Function0<kotlin.Unit>) returnType:kotlin.Unit [suspend]
+              $receiver: VALUE_PARAMETER ADAPTER_PARAMETER_FOR_SUSPEND_CONVERSION name:callee type:kotlin.Function0<kotlin.Unit>
+              BLOCK_BODY
+                CALL 'public abstract fun invoke (): R of kotlin.Function0 [operator] declared in kotlin.Function0' type=kotlin.Unit origin=null
+                  $this: GET_VAR 'callee: kotlin.Function0<kotlin.Unit> declared in <root>.testSuspendConversionInVarargElementsAll.suspendConversion' type=kotlin.Function0<kotlin.Unit> origin=null
+            FUNCTION_REFERENCE 'local final fun suspendConversion (): kotlin.Unit [suspend] declared in <root>.testSuspendConversionInVarargElementsAll' type=kotlin.coroutines.SuspendFunction0<kotlin.Unit> origin=SUSPEND_CONVERSION reflectionTarget=null
+              $receiver: GET_VAR 'f2: kotlin.Function0<kotlin.Unit> declared in <root>.testSuspendConversionInVarargElementsAll' type=kotlin.Function0<kotlin.Unit> origin=null
+          BLOCK type=kotlin.coroutines.SuspendFunction0<kotlin.Unit> origin=SUSPEND_CONVERSION
+            FUN ADAPTER_FOR_SUSPEND_CONVERSION name:suspendConversion visibility:local modality:FINAL <> ($receiver:kotlin.Function0<kotlin.Unit>) returnType:kotlin.Unit [suspend]
+              $receiver: VALUE_PARAMETER ADAPTER_PARAMETER_FOR_SUSPEND_CONVERSION name:callee type:kotlin.Function0<kotlin.Unit>
+              BLOCK_BODY
+                CALL 'public abstract fun invoke (): R of kotlin.Function0 [operator] declared in kotlin.Function0' type=kotlin.Unit origin=null
+                  $this: GET_VAR 'callee: kotlin.Function0<kotlin.Unit> declared in <root>.testSuspendConversionInVarargElementsAll.suspendConversion' type=kotlin.Function0<kotlin.Unit> origin=null
+            FUNCTION_REFERENCE 'local final fun suspendConversion (): kotlin.Unit [suspend] declared in <root>.testSuspendConversionInVarargElementsAll' type=kotlin.coroutines.SuspendFunction0<kotlin.Unit> origin=SUSPEND_CONVERSION reflectionTarget=null
+              $receiver: GET_VAR 'f3: kotlin.Function0<kotlin.Unit> declared in <root>.testSuspendConversionInVarargElementsAll' type=kotlin.Function0<kotlin.Unit> origin=null