[FIR] Introduce & use remapArgumentsWithVararg

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt

@@ -25,6 +25,7 @@ import org.jetbrains.kotlin.fir.resolve.inference.isSuspendFunctionType
 import org.jetbrains.kotlin.fir.resolve.inference.returnType
 import org.jetbrains.kotlin.fir.resolve.propagateTypeFromQualifiedAccessAfterNullCheck
 import org.jetbrains.kotlin.fir.resolve.substitution.ConeSubstitutor
+import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.remapArgumentsWithVararg
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.resultType
 import org.jetbrains.kotlin.fir.resolvedTypeFromPrototype
 import org.jetbrains.kotlin.fir.scopes.impl.FirIntegerOperatorCall
@@ -169,24 +170,7 @@ class FirCallCompletionResultsWriterTransformer(
             // Create a FirVarargArgumentExpression for the vararg arguments
             val varargParameterTypeRef = varargParameter.returnTypeRef
             val resolvedArrayType = varargParameterTypeRef.substitute(this)
-            val resolvedElementType = resolvedArrayType.arrayElementType()
+            this.argumentMapping = remapArgumentsWithVararg(varargParameter, resolvedArrayType, argumentList, argumentMapping)
-            var firstIndex = argumentList.arguments.size
-            val newArgumentMapping = mutableMapOf<FirExpression, FirValueParameter>()
-            val varargArgument = buildVarargArgumentsExpression {
-                varargElementType = varargParameterTypeRef.withReplacedConeType(resolvedElementType)
-                this.typeRef = varargParameterTypeRef.withReplacedConeType(resolvedArrayType)
-                for ((i, arg) in argumentList.arguments.withIndex()) {
-                    val valueParameter = argumentMapping[arg] ?: continue
-                    if (valueParameter.isVararg) {
-                        firstIndex = min(firstIndex, i)
-                        arguments += arg
-                    } else {
-                        newArgumentMapping[arg] = valueParameter
-                    }
-                }
-            }
-            newArgumentMapping[varargArgument] = varargParameter
-            this.argumentMapping = newArgumentMapping
         }
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveUtils.kt

@@ -7,14 +7,21 @@ package org.jetbrains.kotlin.fir.resolve.transformers.body.resolve
 
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirSession
+import org.jetbrains.kotlin.fir.declarations.FirValueParameter
+import org.jetbrains.kotlin.fir.expressions.FirArgumentList
 import org.jetbrains.kotlin.fir.expressions.FirExpression
+import org.jetbrains.kotlin.fir.expressions.builder.buildVarargArgumentsExpression
 import org.jetbrains.kotlin.fir.typeContext
 import org.jetbrains.kotlin.fir.renderWithType
 import org.jetbrains.kotlin.fir.resolve.ScopeSession
 import org.jetbrains.kotlin.fir.resolve.inference.InferenceComponents
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
+import org.jetbrains.kotlin.fir.scopes.impl.withReplacedConeType
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
+import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirTypeRef
+import org.jetbrains.kotlin.fir.types.arrayElementType
+import kotlin.math.min
 
 inline fun <reified T : FirElement> FirBasedSymbol<*>.firUnsafe(): T {
     val fir = this.fir
@@ -38,3 +45,31 @@ internal fun inferenceComponents(
     val inferenceContext = session.typeContext
     return InferenceComponents(inferenceContext, session, returnTypeCalculator, scopeSession)
 }
+
+internal fun remapArgumentsWithVararg(
+    varargParameter: FirValueParameter,
+    varargArrayType: ConeKotlinType,
+    argumentList: FirArgumentList,
+    argumentMapping: Map<FirExpression, FirValueParameter>
+): Map<FirExpression, FirValueParameter> {
+    // Create a FirVarargArgumentExpression for the vararg arguments
+    val varargParameterTypeRef = varargParameter.returnTypeRef
+    val varargElementType = varargArrayType.arrayElementType()
+    var firstIndex = argumentList.arguments.size
+    val newArgumentMapping = mutableMapOf<FirExpression, FirValueParameter>()
+    val varargArgument = buildVarargArgumentsExpression {
+        this.varargElementType = varargParameterTypeRef.withReplacedConeType(varargElementType)
+        this.typeRef = varargParameterTypeRef.withReplacedConeType(varargArrayType)
+        for ((i, arg) in argumentList.arguments.withIndex()) {
+            val valueParameter = argumentMapping[arg] ?: continue
+            if (valueParameter.isVararg) {
+                firstIndex = min(firstIndex, i)
+                arguments += arg
+            } else {
+                newArgumentMapping[arg] = valueParameter
+            }
+        }
+    }
+    newArgumentMapping[varargArgument] = varargParameter
+    return newArgumentMapping
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -700,26 +700,9 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
                     if (varargParameter == null) {
                         it.replaceArgumentList(buildResolvedArgumentList(argumentMapping))
                     } else {
-                        // TODO: refactor/reuse [Candidate#handleVarargs] in [FirCallCompletionResultsWriterTransformer]
                         val varargParameterTypeRef = varargParameter.returnTypeRef
                         val arrayType = varargParameterTypeRef.coneType
-                        val elementType = arrayType.arrayElementType()
+                        val newArgumentMapping = remapArgumentsWithVararg(varargParameter, arrayType, it.argumentList, argumentMapping)
-                        var firstIndex = it.argumentList.arguments.size
-                        val newArgumentMapping = mutableMapOf<FirExpression, FirValueParameter>()
-                        val varargArgument = buildVarargArgumentsExpression {
-                            varargElementType = varargParameterTypeRef.withReplacedConeType(elementType)
-                            this@buildVarargArgumentsExpression.typeRef = varargParameterTypeRef
-                            for ((i, arg) in it.argumentList.arguments.withIndex()) {
-                                val valueParameter = argumentMapping[arg] ?: continue
-                                if (valueParameter.isVararg) {
-                                    firstIndex = min(firstIndex, i)
-                                    arguments += arg
-                                } else {
-                                    newArgumentMapping[arg] = valueParameter
-                                }
-                            }
-                        }
-                        newArgumentMapping[varargArgument] = varargParameter
                         it.replaceArgumentList(buildResolvedArgumentList(newArgumentMapping))
                     }
                 }