FIR2IR: element-wise SAM conversion for vararg

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

@@ -16,6 +16,7 @@ import org.jetbrains.kotlin.fir.references.FirReference
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.references.FirSuperReference
 import org.jetbrains.kotlin.fir.render
+import org.jetbrains.kotlin.fir.resolve.calls.getExpectedTypeForSAMConversion
 import org.jetbrains.kotlin.fir.resolve.calls.isFunctional
 import org.jetbrains.kotlin.fir.resolve.inference.isBuiltinFunctionalType
 import org.jetbrains.kotlin.fir.resolve.toSymbol
@@ -529,12 +530,37 @@ class CallAndReferenceGenerator(
 
     private fun IrExpression.applySamConversionIfNeeded(
         argument: FirExpression,
-        parameter: FirValueParameter?
+        parameter: FirValueParameter?,
+        shouldUnwrapVarargType: Boolean = false
     ): IrExpression {
-        if (parameter == null || !needSamConversion(argument, parameter)) {
+        if (parameter == null) {
             return this
         }
-        val samType = parameter.returnTypeRef.toIrType()
+        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
+            }
+            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, shouldUnwrapVarargType = true)
+                } else
+                    irVarargElement
+            }
+            return this
+        }
+        if (!needSamConversion(argument, parameter)) {
+            return this
+        }
+        var samType = parameter.returnTypeRef.toIrType()
+        if (shouldUnwrapVarargType) {
+            samType = samType.getArrayElementType(irBuiltIns)
+        }
         // Make sure the converted IrType owner indeed has a single abstract method, since FunctionReferenceLowering relies on it.
         if (!samType.isSamType) return this
         return IrTypeOperatorCallImpl(this.startOffset, this.endOffset, samType, IrTypeOperator.SAM_CONVERSION, samType, this)
@@ -542,7 +568,8 @@ class CallAndReferenceGenerator(
 
     private fun needSamConversion(argument: FirExpression, parameter: FirValueParameter): Boolean {
         // If the expected type is a built-in functional type, we don't need SAM conversion.
-        if (parameter.returnTypeRef.coneType.isBuiltinFunctionalType(session)) {
+        val expectedType = argument.getExpectedTypeForSAMConversion(parameter)
+        if (expectedType is ConeTypeParameterType || expectedType.isBuiltinFunctionalType(session)) {
             return false
         }
         // On the other hand, the actual type should be a functional type.

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

@@ -360,7 +360,7 @@ private fun Candidate.prepareExpectedType(
     context: ResolutionContext
 ): ConeKotlinType? {
     if (parameter == null) return null
-    val basicExpectedType = argument.getExpectedType(session, parameter/*, LanguageVersionSettings*/)
+    val basicExpectedType = argument.getExpectedTypeForSAMConversion(parameter/*, LanguageVersionSettings*/)
     val expectedType = getExpectedTypeWithSAMConversion(session, argument, basicExpectedType, context) ?: basicExpectedType
     return this.substitutor.substituteOrSelf(expectedType)
 }
@@ -391,8 +391,7 @@ fun FirExpression.isFunctional(session: FirSession): Boolean =
         else -> typeRef.coneTypeSafe<ConeKotlinType>()?.isBuiltinFunctionalType(session) == true
     }
 
-internal fun FirExpression.getExpectedType(
-    session: FirSession,
+fun FirExpression.getExpectedTypeForSAMConversion(
     parameter: FirValueParameter/*, languageVersionSettings: LanguageVersionSettings*/
 ): ConeKotlinType {
     val shouldUnwrapVarargType = when (this) {

compiler/testData/codegen/box/funInterface/funConversionInVararg.kt

@@ -1,5 +1,4 @@
 // !LANGUAGE: +NewInference +FunctionalInterfaceConversion +SamConversionPerArgument
-// IGNORE_BACKEND_FIR: JVM_IR
 
 fun interface MyRunnable {
     fun run()

compiler/testData/codegen/box/sam/kt31908.kt

@@ -1,5 +1,4 @@
 // TARGET_BACKEND: JVM
-// IGNORE_BACKEND_FIR: JVM_IR
 // FILE: kt31908.kt
 fun box(): String {
     var result = "failed"

compiler/testData/ir/irText/expressions/funInterface/samConversionInVarargs.fir.txt

@@ -24,34 +24,38 @@ FILE fqName:<root> fileName:/samConversionInVarargs.kt
     BLOCK_BODY
       CALL 'public final fun useVararg (vararg foos: <root>.IFoo): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
         foos: VARARG type=kotlin.Array<out <root>.IFoo> varargElementType=<root>.IFoo
-          FUN_EXPR type=kotlin.Function1<kotlin.Int, kotlin.Unit> origin=LAMBDA
-            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (it:kotlin.Int) returnType:kotlin.Unit
-              VALUE_PARAMETER name:it index:0 type:kotlin.Int
-              BLOCK_BODY
-                RETURN type=kotlin.Nothing from='local final fun <anonymous> (it: kotlin.Int): kotlin.Unit declared in <root>.testLambda'
-                  GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
+          TYPE_OP type=<root>.IFoo origin=SAM_CONVERSION typeOperand=<root>.IFoo
+            FUN_EXPR type=kotlin.Function1<kotlin.Int, kotlin.Unit> origin=LAMBDA
+              FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (it:kotlin.Int) returnType:kotlin.Unit
+                VALUE_PARAMETER name:it index:0 type:kotlin.Int
+                BLOCK_BODY
+                  RETURN type=kotlin.Nothing from='local final fun <anonymous> (it: kotlin.Int): kotlin.Unit declared in <root>.testLambda'
+                    GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
   FUN name:testSeveralLambdas visibility:public modality:FINAL <> () returnType:kotlin.Unit
     BLOCK_BODY
       CALL 'public final fun useVararg (vararg foos: <root>.IFoo): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
         foos: VARARG type=kotlin.Array<out <root>.IFoo> varargElementType=<root>.IFoo
-          FUN_EXPR type=kotlin.Function1<kotlin.Int, kotlin.Unit> origin=LAMBDA
-            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (it:kotlin.Int) returnType:kotlin.Unit
-              VALUE_PARAMETER name:it index:0 type:kotlin.Int
-              BLOCK_BODY
-                RETURN type=kotlin.Nothing from='local final fun <anonymous> (it: kotlin.Int): kotlin.Unit declared in <root>.testSeveralLambdas'
-                  GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
-          FUN_EXPR type=kotlin.Function1<kotlin.Int, kotlin.Unit> origin=LAMBDA
-            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (it:kotlin.Int) returnType:kotlin.Unit
-              VALUE_PARAMETER name:it index:0 type:kotlin.Int
-              BLOCK_BODY
-                RETURN type=kotlin.Nothing from='local final fun <anonymous> (it: kotlin.Int): kotlin.Unit declared in <root>.testSeveralLambdas'
-                  GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
-          FUN_EXPR type=kotlin.Function1<kotlin.Int, kotlin.Unit> origin=LAMBDA
-            FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (it:kotlin.Int) returnType:kotlin.Unit
-              VALUE_PARAMETER name:it index:0 type:kotlin.Int
-              BLOCK_BODY
-                RETURN type=kotlin.Nothing from='local final fun <anonymous> (it: kotlin.Int): kotlin.Unit declared in <root>.testSeveralLambdas'
-                  GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
+          TYPE_OP type=<root>.IFoo origin=SAM_CONVERSION typeOperand=<root>.IFoo
+            FUN_EXPR type=kotlin.Function1<kotlin.Int, kotlin.Unit> origin=LAMBDA
+              FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (it:kotlin.Int) returnType:kotlin.Unit
+                VALUE_PARAMETER name:it index:0 type:kotlin.Int
+                BLOCK_BODY
+                  RETURN type=kotlin.Nothing from='local final fun <anonymous> (it: kotlin.Int): kotlin.Unit declared in <root>.testSeveralLambdas'
+                    GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
+          TYPE_OP type=<root>.IFoo origin=SAM_CONVERSION typeOperand=<root>.IFoo
+            FUN_EXPR type=kotlin.Function1<kotlin.Int, kotlin.Unit> origin=LAMBDA
+              FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (it:kotlin.Int) returnType:kotlin.Unit
+                VALUE_PARAMETER name:it index:0 type:kotlin.Int
+                BLOCK_BODY
+                  RETURN type=kotlin.Nothing from='local final fun <anonymous> (it: kotlin.Int): kotlin.Unit declared in <root>.testSeveralLambdas'
+                    GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
+          TYPE_OP type=<root>.IFoo origin=SAM_CONVERSION typeOperand=<root>.IFoo
+            FUN_EXPR type=kotlin.Function1<kotlin.Int, kotlin.Unit> origin=LAMBDA
+              FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> (it:kotlin.Int) returnType:kotlin.Unit
+                VALUE_PARAMETER name:it index:0 type:kotlin.Int
+                BLOCK_BODY
+                  RETURN type=kotlin.Nothing from='local final fun <anonymous> (it: kotlin.Int): kotlin.Unit declared in <root>.testSeveralLambdas'
+                    GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
   FUN name:withVarargOfInt visibility:public modality:FINAL <> (xs:kotlin.IntArray) returnType:kotlin.String
     VALUE_PARAMETER name:xs index:0 type:kotlin.IntArray varargElementType:kotlin.Int [vararg]
     BLOCK_BODY
@@ -61,10 +65,11 @@ FILE fqName:<root> fileName:/samConversionInVarargs.kt
     BLOCK_BODY
       CALL 'public final fun useVararg (vararg foos: <root>.IFoo): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
         foos: VARARG type=kotlin.Array<out <root>.IFoo> varargElementType=<root>.IFoo
-          FUN_EXPR type=kotlin.Function1<kotlin.Int, kotlin.Unit> origin=ADAPTED_FUNCTION_REFERENCE
-            FUN ADAPTER_FOR_CALLABLE_REFERENCE name:withVarargOfInt visibility:local modality:FINAL <> (p0:kotlin.Int) returnType:kotlin.Unit
-              VALUE_PARAMETER ADAPTER_PARAMETER_FOR_CALLABLE_REFERENCE name:p0 index:0 type:kotlin.Int
-              BLOCK_BODY
-                CALL 'public final fun withVarargOfInt (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>.testAdaptedCR.withVarargOfInt' type=kotlin.Int origin=null
+          TYPE_OP type=<root>.IFoo origin=SAM_CONVERSION typeOperand=<root>.IFoo
+            FUN_EXPR type=kotlin.Function1<kotlin.Int, kotlin.Unit> origin=ADAPTED_FUNCTION_REFERENCE
+              FUN ADAPTER_FOR_CALLABLE_REFERENCE name:withVarargOfInt visibility:local modality:FINAL <> (p0:kotlin.Int) returnType:kotlin.Unit
+                VALUE_PARAMETER ADAPTER_PARAMETER_FOR_CALLABLE_REFERENCE name:p0 index:0 type:kotlin.Int
+                BLOCK_BODY
+                  CALL 'public final fun withVarargOfInt (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>.testAdaptedCR.withVarargOfInt' type=kotlin.Int origin=null

compiler/testData/ir/irText/expressions/funInterface/samConversionInVarargsMixed.fir.txt

@@ -26,11 +26,12 @@ FILE fqName:<root> fileName:/samConversionInVarargsMixed.kt
             GET_VAR 'a: kotlin.Any declared in <root>.test' type=kotlin.Any origin=null
           then: CALL 'public final fun foo (vararg rs: <root>.MyRunnable): kotlin.Unit declared in <root>' type=kotlin.Unit origin=null
             rs: VARARG type=kotlin.Array<out <root>.MyRunnable> varargElementType=<root>.MyRunnable
-              FUN_EXPR type=kotlin.Function0<kotlin.Unit> origin=LAMBDA
-                FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> () returnType:kotlin.Unit
-                  BLOCK_BODY
-                    RETURN type=kotlin.Nothing from='local final fun <anonymous> (): kotlin.Unit declared in <root>.test'
-                      GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
+              TYPE_OP type=<root>.MyRunnable origin=SAM_CONVERSION typeOperand=<root>.MyRunnable
+                FUN_EXPR type=kotlin.Function0<kotlin.Unit> origin=LAMBDA
+                  FUN LOCAL_FUNCTION_FOR_LAMBDA name:<anonymous> visibility:local modality:FINAL <> () returnType:kotlin.Unit
+                    BLOCK_BODY
+                      RETURN type=kotlin.Nothing from='local final fun <anonymous> (): kotlin.Unit declared in <root>.test'
+                        GET_OBJECT 'CLASS IR_EXTERNAL_DECLARATION_STUB OBJECT name:Unit modality:FINAL visibility:public superTypes:[kotlin.Any]' type=kotlin.Unit
               GET_VAR 'r: <root>.MyRunnable declared in <root>.test' type=<root>.MyRunnable origin=null
               TYPE_OP type=<root>.MyRunnable origin=IMPLICIT_CAST typeOperand=<root>.MyRunnable
                 GET_VAR 'a: kotlin.Any declared in <root>.test' type=kotlin.Any origin=null