FE: provide argument mapping info for adapted callable reference

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/NewCallArguments.kt

@@ -181,6 +181,18 @@ class FakeValueArgumentForLeftCallableReference(val ktExpression: KtCallableRefe
     override fun isExternal(): Boolean = false
 }
 
+class FakePositionalValueArgumentForCallableReferenceImpl(
+    private val callElement: KtElement,
+    override val index: Int
+) : FakePositionalValueArgumentForCallableReference {
+    override fun getArgumentExpression(): KtExpression? = null
+    override fun getArgumentName(): ValueArgumentName? = null
+    override fun isNamed(): Boolean = false
+    override fun asElement(): KtElement = callElement
+    override fun getSpreadElement(): LeafPsiElement? = null
+    override fun isExternal(): Boolean = false
+}
+
 class EmptyLabeledReturn(
     val returnExpression: KtReturnExpression,
     builtIns: KotlinBuiltIns

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/ResolvedAtomCompleter.kt

@@ -13,6 +13,7 @@ import org.jetbrains.kotlin.descriptors.impl.ReceiverParameterDescriptorImpl
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtExpression
 import org.jetbrains.kotlin.psi.KtNamedFunction
+import org.jetbrains.kotlin.psi.ValueArgument
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.BindingTrace
 import org.jetbrains.kotlin.resolve.MissingSupertypesResolver
@@ -20,6 +21,8 @@ import org.jetbrains.kotlin.resolve.TemporaryBindingTrace
 import org.jetbrains.kotlin.resolve.calls.ArgumentTypeResolver
 import org.jetbrains.kotlin.resolve.calls.NewCommonSuperTypeCalculator
 import org.jetbrains.kotlin.resolve.calls.checkers.CallCheckerContext
+import org.jetbrains.kotlin.resolve.calls.components.CallableReferenceAdaptation
+import org.jetbrains.kotlin.resolve.calls.components.isVararg
 import org.jetbrains.kotlin.resolve.calls.context.BasicCallResolutionContext
 import org.jetbrains.kotlin.resolve.calls.context.ContextDependency
 import org.jetbrains.kotlin.resolve.calls.inference.components.NewTypeSubstitutor
@@ -216,7 +219,7 @@ class ResolvedAtomCompleter(
         }
 
         val functionDescriptor = trace.bindingContext.get(BindingContext.FUNCTION, ktFunction) as? FunctionDescriptorImpl
-                ?: throw AssertionError("No function descriptor for resolved lambda argument")
+            ?: throw AssertionError("No function descriptor for resolved lambda argument")
         functionDescriptor.setReturnType(returnType)
 
         val existingLambdaType = trace.getType(ktArgumentExpression)
@@ -277,7 +280,10 @@ class ResolvedAtomCompleter(
         val resultTypeParameters =
             callableCandidate.freshSubstitutor!!.freshVariables.map { resultSubstitutor.safeSubstitute(it.defaultType) }
 
-        val typeParametersSubstitutor = NewTypeSubstitutorByConstructorMap((callableCandidate.candidate.typeParameters.map { it.typeConstructor } zip resultTypeParameters).toMap())
+        val typeParametersSubstitutor =
+            NewTypeSubstitutorByConstructorMap(
+                (callableCandidate.candidate.typeParameters.map { it.typeConstructor } zip resultTypeParameters).toMap()
+            )
 
         val firstSubstitution = typeParametersSubstitutor.toOldSubstitution()
         val secondSubstitution = resultSubstitutor.toOldSubstitution()
@@ -317,6 +323,8 @@ class ResolvedAtomCompleter(
         )
         resolvedCall.setResultingSubstitutor(resultSubstitutor)
 
+        recordArgumentAdaptationForCallableReference(resolvedCall, callableCandidate.callableReferenceAdaptation)
+
         tracing.bindCall(topLevelTrace, psiCall)
         tracing.bindReference(topLevelTrace, resolvedCall)
         tracing.bindResolvedCall(topLevelTrace, resolvedCall)
@@ -352,6 +360,47 @@ class ResolvedAtomCompleter(
         kotlinToResolvedCallTransformer.runCallCheckers(resolvedCall, topLevelCallCheckerContext)
     }
 
+    private fun recordArgumentAdaptationForCallableReference(
+        resolvedCall: ResolvedCallImpl<CallableDescriptor>,
+        callableReferenceAdaptation: CallableReferenceAdaptation?
+    ) {
+        if (callableReferenceAdaptation == null) return
+
+        val callElement = resolvedCall.call.callElement
+
+        for ((valueParameter, resolvedCallArgument) in callableReferenceAdaptation.mappedArguments) {
+            resolvedCall.recordValueArgument(
+                valueParameter,
+                when (resolvedCallArgument) {
+                    ResolvedCallArgument.DefaultArgument ->
+                        DefaultValueArgument.DEFAULT
+                    is ResolvedCallArgument.SimpleArgument -> {
+                        val valueArgument = makeFakeValueArgument(resolvedCallArgument.callArgument, callElement)
+                        if (valueParameter.isVararg)
+                            VarargValueArgument(listOf(valueArgument))
+                        else
+                            ExpressionValueArgument(valueArgument)
+                    }
+                    is ResolvedCallArgument.VarargArgument ->
+                        VarargValueArgument(
+                            resolvedCallArgument.arguments.map {
+                                makeFakeValueArgument(it, callElement)
+                            }
+                        )
+                }
+            )
+        }
+    }
+
+    private fun makeFakeValueArgument(
+        callArgument: KotlinCallArgument,
+        callElement: KtElement
+    ): ValueArgument {
+        val fakeCallArgument = callArgument as? FakeKotlinCallArgumentForCallableReference
+            ?: throw AssertionError("FakeKotlinCallArgumentForCallableReference expected: $callArgument")
+        return FakePositionalValueArgumentForCallableReferenceImpl(callElement, fakeCallArgument.index)
+    }
+
     private fun completeCollectionLiteralCalls(collectionLiteralArgument: ResolvedCollectionLiteralAtom) {
         val psiCallArgument = collectionLiteralArgument.atom.psiCallArgument as CollectionLiteralKotlinCallArgumentImpl
         val context = psiCallArgument.outerCallContext

compiler/psi/src/org/jetbrains/kotlin/psi/ValueArgument.kt

@@ -36,6 +36,10 @@ interface ValueArgument {
     fun isExternal(): Boolean
 }
 
+interface FakePositionalValueArgumentForCallableReference : ValueArgument {
+    val index: Int
+}
+
 interface LambdaArgument : ValueArgument {
     fun getLambdaExpression(): KtLambdaExpression?
 }

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/CallableReferenceResolution.kt

@@ -59,7 +59,7 @@ class CallableReferenceCandidate(
     val extensionReceiver: CallableReceiver?,
     val explicitReceiverKind: ExplicitReceiverKind,
     val reflectionCandidateType: UnwrappedType,
-    val numDefaults: Int,
+    val callableReferenceAdaptation: CallableReferenceAdaptation?,
     val diagnostics: List<KotlinCallDiagnostic>
 ) : Candidate {
     override val resultingApplicability = getResultApplicability(diagnostics)
@@ -67,8 +67,17 @@ class CallableReferenceCandidate(
 
     var freshSubstitutor: FreshVariableNewTypeSubstitutor? = null
         internal set
+
+    val numDefaults get() = callableReferenceAdaptation?.defaults ?: 0
 }
 
+class CallableReferenceAdaptation(
+    val argumentTypes: Array<KotlinType>,
+    val coercionStrategy: CoercionStrategy,
+    val defaults: Int,
+    val mappedArguments: Map<ValueParameterDescriptor, ResolvedCallArgument>
+)
+
 /**
  * cases: class A {}, class B { companion object }, object C, enum class D { E }
  * A::foo <-> Type
@@ -179,7 +188,7 @@ class CallableReferencesCandidateFactory(
         val candidateDescriptor = towerCandidate.descriptor
         val diagnostics = SmartList<KotlinCallDiagnostic>()
 
-        val (reflectionCandidateType, defaults) = buildReflectionType(
+        val (reflectionCandidateType, callableReferenceAdaptation) = buildReflectionType(
             candidateDescriptor,
             dispatchCallableReceiver,
             extensionCallableReceiver,
@@ -187,16 +196,17 @@ class CallableReferencesCandidateFactory(
             callComponents.builtIns
         )
 
-        if (defaults != 0 &&
+        if (callableReferenceAdaptation != null &&
+            callableReferenceAdaptation.defaults != 0 &&
             !callComponents.languageVersionSettings.supportsFeature(LanguageFeature.FunctionReferenceWithDefaultValueAsOtherType)
         ) {
-            diagnostics.add(CallableReferencesDefaultArgumentUsed(argument, candidateDescriptor, defaults))
+            diagnostics.add(CallableReferencesDefaultArgumentUsed(argument, candidateDescriptor, callableReferenceAdaptation.defaults))
         }
 
         if (candidateDescriptor !is CallableMemberDescriptor) {
             return CallableReferenceCandidate(
                 candidateDescriptor, dispatchCallableReceiver, extensionCallableReceiver,
-                explicitReceiverKind, reflectionCandidateType, defaults,
+                explicitReceiverKind, reflectionCandidateType, callableReferenceAdaptation,
                 listOf(NotCallableMemberReference(argument, candidateDescriptor))
             )
         }
@@ -226,7 +236,7 @@ class CallableReferencesCandidateFactory(
 
         return CallableReferenceCandidate(
             candidateDescriptor, dispatchCallableReceiver, extensionCallableReceiver,
-            explicitReceiverKind, reflectionCandidateType, defaults, diagnostics
+            explicitReceiverKind, reflectionCandidateType, callableReferenceAdaptation, diagnostics
         )
     }
 
@@ -234,18 +244,18 @@ class CallableReferencesCandidateFactory(
         UNMAPPED, MAPPED_WITH_PLAIN_ARGS, MAPPED_WITH_ARRAY
     }
 
-    private fun getArgumentAndReturnTypeUseMappingByExpectedType(
+    private fun getCallableReferenceAdaptation(
         descriptor: FunctionDescriptor,
         expectedType: UnwrappedType?,
         unboundReceiverCount: Int,
         builtins: KotlinBuiltIns
-    ): Triple<Array<KotlinType>, CoercionStrategy, Int>? {
+    ): CallableReferenceAdaptation? {
         val inputOutputTypes = extractInputOutputTypesFromCallableReferenceExpectedType(expectedType) ?: return null
 
         val expectedArgumentCount = inputOutputTypes.inputTypes.size - unboundReceiverCount
         if (expectedArgumentCount < 0) return null
 
-        val fakeArguments = (0..(expectedArgumentCount - 1)).map { FakeKotlinCallArgumentForCallableReference(it) }
+        val fakeArguments = (0 until expectedArgumentCount).map { FakeKotlinCallArgumentForCallableReference(it) }
         val argumentMapping =
             callComponents.argumentsToParametersMapper.mapArguments(fakeArguments, externalArgument = null, descriptor = descriptor)
         if (argumentMapping.diagnostics.any { !it.candidateApplicability.isSuccess }) return null
@@ -288,8 +298,11 @@ class CallableReferencesCandidateFactory(
 
         val coercion = if (returnExpectedType.isUnit()) CoercionStrategy.COERCION_TO_UNIT else CoercionStrategy.NO_COERCION
 
-        @Suppress("UNCHECKED_CAST")
+
-        return Triple(mappedArguments as Array<KotlinType>, coercion, defaults)
+        return CallableReferenceAdaptation(
+            @Suppress("UNCHECKED_CAST") (mappedArguments as Array<KotlinType>),
+            coercion, defaults, argumentMapping.parameterToCallArgumentMap
+        )
     }
 
     private fun varargParameterTypeByExpectedParameter(
@@ -328,7 +341,7 @@ class CallableReferencesCandidateFactory(
         extensionReceiver: CallableReceiver?,
         expectedType: UnwrappedType?,
         builtins: KotlinBuiltIns
-    ): Pair<UnwrappedType, /*defaults*/ Int> {
+    ): Pair<UnwrappedType, CallableReferenceAdaptation?> {
         val argumentsAndReceivers = ArrayList<KotlinType>(descriptor.valueParameters.size + 2)
 
         if (dispatchReceiver is CallableReceiver.UnboundReference) {
@@ -356,35 +369,35 @@ class CallableReferencesCandidateFactory(
                     argumentsAndReceivers,
                     descriptorReturnType,
                     mutable
-                ) to 0
+                ) to null
             }
             is FunctionDescriptor -> {
-                val returnType: KotlinType
+                val callableReferenceAdaptation = getCallableReferenceAdaptation(
-                val defaults: Int
-                val argumentsAndExpectedTypeCoercion = getArgumentAndReturnTypeUseMappingByExpectedType(
                     descriptor, expectedType,
                     unboundReceiverCount = argumentsAndReceivers.size,
                     builtins = builtins
                 )
 
-                if (argumentsAndExpectedTypeCoercion == null) {
+                val returnType = if (callableReferenceAdaptation == null) {
                     descriptor.valueParameters.mapTo(argumentsAndReceivers) { it.type }
-                    returnType = descriptorReturnType
+                    descriptorReturnType
-                    defaults = 0
                 } else {
-                    val (arguments, coercion) = argumentsAndExpectedTypeCoercion
+                    val arguments = callableReferenceAdaptation.argumentTypes
-                    defaults = argumentsAndExpectedTypeCoercion.third
+                    val coercion = callableReferenceAdaptation.coercionStrategy
                     argumentsAndReceivers.addAll(arguments)
 
-                    returnType = if (coercion == CoercionStrategy.COERCION_TO_UNIT) descriptor.builtIns.unitType else descriptorReturnType
+                    if (coercion == CoercionStrategy.COERCION_TO_UNIT)
+                        descriptor.builtIns.unitType
+                    else
+                        descriptorReturnType
                 }
 
                 return callComponents.reflectionTypes.getKFunctionType(
                     Annotations.EMPTY, null, argumentsAndReceivers, null,
                     returnType, descriptor.builtIns, descriptor.isSuspend
-                ) to defaults
+                ) to callableReferenceAdaptation
             }
-            else -> return ErrorUtils.createErrorType("Unsupported descriptor type: $descriptor") to 0
+            else -> return ErrorUtils.createErrorType("Unsupported descriptor type: $descriptor") to null
         }
     }