[FE] Support multiple candidates for extension receiver

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

@@ -35,10 +35,7 @@ import org.jetbrains.kotlin.resolve.calls.CallTransformer
 import org.jetbrains.kotlin.resolve.calls.CandidateResolver
 import org.jetbrains.kotlin.resolve.calls.context.*
 import org.jetbrains.kotlin.resolve.calls.inference.BuilderInferenceSupport
-import org.jetbrains.kotlin.resolve.calls.model.KotlinCallDiagnostic
+import org.jetbrains.kotlin.resolve.calls.model.*
-import org.jetbrains.kotlin.resolve.calls.model.MutableResolvedCall
-import org.jetbrains.kotlin.resolve.calls.model.ResolvedCallImpl
-import org.jetbrains.kotlin.resolve.calls.model.VariableAsFunctionResolvedCallImpl
 import org.jetbrains.kotlin.resolve.calls.results.OverloadResolutionResultsImpl
 import org.jetbrains.kotlin.resolve.calls.results.ResolutionResultsHandler
 import org.jetbrains.kotlin.resolve.calls.results.ResolutionStatus
@@ -492,6 +489,16 @@ class NewResolutionOldInference(
             }
         }
 
+        /**
+         * The function is called only inside [NoExplicitReceiverScopeTowerProcessor] with [TowerData.BothTowerLevelAndContextReceiversGroup].
+         * This case involves only [SimpleCandidateFactory].
+         */
+        override fun createCandidate(
+            towerCandidate: CandidateWithBoundDispatchReceiver,
+            explicitReceiverKind: ExplicitReceiverKind,
+            extensionReceiverCandidates: List<ReceiverValueWithSmartCastInfo>
+        ): MyCandidate = error("${this::class.simpleName} doesn't support candidates with multiple extension receiver candidates")
+
         override fun createErrorCandidate(): MyCandidate {
             throw IllegalStateException("Not supported creating error candidate for the old type inference candidate factory")
         }

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

@@ -22,6 +22,7 @@ import org.jetbrains.kotlin.resolve.calls.inference.model.*
 import org.jetbrains.kotlin.resolve.calls.model.*
 import org.jetbrains.kotlin.resolve.calls.smartcasts.getReceiverValueWithSmartCast
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind.*
+import org.jetbrains.kotlin.resolve.calls.tower.ContextReceiverAmbiguity
 import org.jetbrains.kotlin.resolve.calls.tower.InfixCallNoInfixModifier
 import org.jetbrains.kotlin.resolve.calls.tower.InvokeConventionCallNoOperatorModifier
 import org.jetbrains.kotlin.resolve.calls.tower.VisibilityError
@@ -657,7 +658,70 @@ private fun ResolutionCandidate.prepareExpectedType(expectedType: UnwrappedType)
     return resolvedCall.knownParametersSubstitutor.safeSubstitute(resultType)
 }
 
+private data class ApplicableArgumentWithConstraint(
+    val argument: SimpleKotlinCallArgument,
+    val argumentType: UnwrappedType,
+    val expectedType: UnwrappedType,
+    val position: ConstraintPosition
+)
+
+private fun ResolutionCandidate.getArgumentWithConstraintIfCompatible(
+    argument: SimpleKotlinCallArgument,
+    parameter: ParameterDescriptor
+): ApplicableArgumentWithConstraint? {
+    val csBuilder = getSystem().getBuilder()
+    val expectedTypeUnprepared = argument.getExpectedType(parameter, callComponents.languageVersionSettings)
+    val expectedType = prepareExpectedType(expectedTypeUnprepared)
+    val argumentType = captureFromTypeParameterUpperBoundIfNeeded(argument.receiver.stableType, expectedType)
+    val position = ReceiverConstraintPositionImpl(argument)
+    return if (csBuilder.isSubtypeConstraintCompatible(argumentType, expectedType, position))
+        ApplicableArgumentWithConstraint(argument, argumentType, expectedType, position)
+    else null
+}
+
 internal object CheckReceivers : ResolutionPart() {
+    override fun ResolutionCandidate.process(workIndex: Int) {
+        when (workIndex) {
+            0 -> checkReceiver(
+                resolvedCall.dispatchReceiverArgument,
+                candidateDescriptor.dispatchReceiverParameter,
+                shouldCheckImplicitInvoke = true,
+            )
+            1 -> {
+                if (resolvedCall.extensionReceiverArgument == null) {
+                    resolvedCall.extensionReceiverArgument = chooseExtensionReceiverCandidate() ?: return
+                }
+                checkReceiver(
+                    resolvedCall.extensionReceiverArgument,
+                    candidateDescriptor.extensionReceiverParameter,
+                    shouldCheckImplicitInvoke = false, // reproduce old inference behaviour
+                )
+            }
+        }
+    }
+
+    override fun ResolutionCandidate.workCount() = 2
+
+    private fun ResolutionCandidate.chooseExtensionReceiverCandidate(): SimpleKotlinCallArgument? {
+        val receiverCandidates = resolvedCall.extensionReceiverArgumentCandidates
+        if (receiverCandidates.isNullOrEmpty()) {
+            return null
+        }
+        if (receiverCandidates.size == 1) {
+            return receiverCandidates.single()
+        }
+        val extensionReceiverParameter = candidateDescriptor.extensionReceiverParameter ?: return null
+        val compatible = receiverCandidates.mapNotNull { getArgumentWithConstraintIfCompatible(it, extensionReceiverParameter) }
+        return when (compatible.size) {
+            0 -> null
+            1 -> compatible.single().argument
+            else -> {
+                addDiagnostic(ContextReceiverAmbiguity())
+                null
+            }
+        }
+    }
+
     private fun ResolutionCandidate.checkReceiver(
         receiverArgument: SimpleKotlinCallArgument?,
         receiverParameter: ReceiverParameterDescriptor?,
@@ -680,25 +744,6 @@ internal object CheckReceivers : ResolutionPart() {
 
         resolveKotlinArgument(receiverArgument, receiverParameter, receiverInfo)
     }
-
-    override fun ResolutionCandidate.process(workIndex: Int) {
-        when (workIndex) {
-            0 -> checkReceiver(
-                resolvedCall.dispatchReceiverArgument,
-                candidateDescriptor.dispatchReceiverParameter,
-                shouldCheckImplicitInvoke = true,
-            )
-            1 -> {
-                checkReceiver(
-                    resolvedCall.extensionReceiverArgument,
-                    candidateDescriptor.extensionReceiverParameter,
-                    shouldCheckImplicitInvoke = false, // reproduce old inference behaviour
-                )
-            }
-        }
-    }
-
-    override fun ResolutionCandidate.workCount() = 2
 }
 
 internal object CheckArgumentsInParenthesis : ResolutionPart() {
@@ -822,44 +867,27 @@ internal object CheckContextReceiversResolutionPart : ResolutionPart() {
         resolvedCall.contextReceiversArguments = contextReceiversArguments
     }
 
-    private data class ApplicableArgumentWithConstraint(
-        val argument: SimpleKotlinCallArgument,
-        val argumentType: UnwrappedType,
-        val expectedType: UnwrappedType,
-        val position: ConstraintPosition
-    )
-
     private fun ResolutionCandidate.findContextReceiver(
         implicitReceiversGroups: List<List<ReceiverValueWithSmartCastInfo>>,
         candidateContextReceiverParameter: ReceiverParameterDescriptor
     ): SimpleKotlinCallArgument? {
-
+        val csBuilder = getSystem().getBuilder()
-        fun ReceiverValueWithSmartCastInfo.createArgumentIfCompatible(): ApplicableArgumentWithConstraint? {
-            val csBuilder = getSystem().getBuilder()
-            val argument = ReceiverExpressionKotlinCallArgument(this)
-            val expectedTypeUnprepared = argument.getExpectedType(candidateContextReceiverParameter, callComponents.languageVersionSettings)
-
-            val expectedType = prepareExpectedType(expectedTypeUnprepared)
-            val argumentType = captureFromTypeParameterUpperBoundIfNeeded(argument.receiver.stableType, expectedType)
-            val position = ReceiverConstraintPositionImpl(argument)
-            return if (csBuilder.isSubtypeConstraintCompatible(argumentType, expectedType, position))
-                ApplicableArgumentWithConstraint(argument, argumentType, expectedType, position)
-            else null
-        }
-
         for (implicitReceiverGroup in implicitReceiversGroups) {
-            val applicableArguments = implicitReceiverGroup.mapNotNull { it.createArgumentIfCompatible() }.toList()
+            val applicableArguments = implicitReceiverGroup.mapNotNull {
+                val argument = ReceiverExpressionKotlinCallArgument(it)
+                getArgumentWithConstraintIfCompatible(argument, candidateContextReceiverParameter)
+            }.toList()
             if (applicableArguments.size == 1) {
                 val (argument, argumentType, expectedType, position) = applicableArguments.single()
                 csBuilder.addSubtypeConstraint(argumentType, expectedType, position)
                 return argument
             }
             if (applicableArguments.size > 1) {
-                diagnosticsFromResolutionParts.add(MultipleArgumentsApplicableForContextReceiver(candidateContextReceiverParameter))
+                addDiagnostic(MultipleArgumentsApplicableForContextReceiver(candidateContextReceiverParameter))
                 return null
             }
         }
-        diagnosticsFromResolutionParts.add(NoContextReceiver(candidateContextReceiverParameter))
+        addDiagnostic(NoContextReceiver(candidateContextReceiverParameter))
         return null
     }
 }

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/model/CallableReferencesCandidateFactory.kt

@@ -119,6 +119,17 @@ class CallableReferencesCandidateFactory(
         return createCallableReferenceCallCandidate(diagnostics)
     }
 
+    /**
+     * The function is called only inside [NoExplicitReceiverScopeTowerProcessor] with [TowerData.BothTowerLevelAndContextReceiversGroup].
+     * This case involves only [SimpleCandidateFactory].
+     */
+    override fun createCandidate(
+        towerCandidate: CandidateWithBoundDispatchReceiver,
+        explicitReceiverKind: ExplicitReceiverKind,
+        extensionReceiverCandidates: List<ReceiverValueWithSmartCastInfo>
+    ): CallableReferenceResolutionCandidate =
+        error("${this::class.simpleName} doesn't support candidates with multiple extension receiver candidates")
+
     fun createCallableProcessor(explicitReceiver: DetailedReceiver?) =
         createCallableReferenceProcessor(scopeTower, kotlinCall.rhsName, this, explicitReceiver)
 
@@ -377,7 +388,7 @@ class CallableReferencesCandidateFactory(
                         (suspendConversionStrategy == SuspendConversionStrategy.SUSPEND_CONVERSION && buildTypeWithConversions)
 
                 callComponents.reflectionTypes.getKFunctionType(
-                    Annotations.EMPTY, null, argumentsAndReceivers, null,
+                    Annotations.EMPTY, null, emptyList(), argumentsAndReceivers, null,
                     returnType, descriptor.builtIns, isSuspend
                 ) to callableReferenceAdaptation
             }

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/model/ResolutionAtoms.kt

@@ -75,7 +75,8 @@ abstract class ResolvedCallAtom : ResolvedAtom() {
     abstract val candidateDescriptor: CallableDescriptor
     abstract val explicitReceiverKind: ExplicitReceiverKind
     abstract val dispatchReceiverArgument: SimpleKotlinCallArgument?
-    abstract val extensionReceiverArgument: SimpleKotlinCallArgument?
+    abstract var extensionReceiverArgument: SimpleKotlinCallArgument?
+    abstract val extensionReceiverArgumentCandidates: List<SimpleKotlinCallArgument>?
     abstract var contextReceiversArguments: List<SimpleKotlinCallArgument>
     abstract val typeArgumentMappingByOriginal: TypeArgumentsToParametersMapper.TypeArgumentsMapping
     abstract val argumentMappingByOriginal: Map<ValueParameterDescriptor, ResolvedCallArgument>

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/model/ResolvedCallAtoms.kt

@@ -73,7 +73,7 @@ class ResolvedCallableReferenceCallAtom(
     reflectionCandidateType: UnwrappedType? = null,
     candidate: CallableReferenceResolutionCandidate? = null
 ) : MutableResolvedCallAtom(
-    atom, candidateDescriptor, explicitReceiverKind, dispatchReceiverArgument, extensionReceiverArgument, reflectionCandidateType, candidate
+    atom, candidateDescriptor, explicitReceiverKind, dispatchReceiverArgument, extensionReceiverArgument, emptyList(), reflectionCandidateType, candidate
 ), ResolvedCallableReferenceAtom
 
 open class MutableResolvedCallAtom(
@@ -81,7 +81,8 @@ open class MutableResolvedCallAtom(
     originalCandidateDescriptor: CallableDescriptor, // original candidate descriptor
     override val explicitReceiverKind: ExplicitReceiverKind,
     override val dispatchReceiverArgument: SimpleKotlinCallArgument?,
-    override val extensionReceiverArgument: SimpleKotlinCallArgument?,
+    override var extensionReceiverArgument: SimpleKotlinCallArgument?,
+    override val extensionReceiverArgumentCandidates: List<SimpleKotlinCallArgument>?,
     open val reflectionCandidateType: UnwrappedType? = null,
     open val candidate: CallableReferenceResolutionCandidate? = null
 ) : ResolvedCallAtom() {

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/model/SimpleCandidateFactory.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.resolve.calls.model
 
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.descriptors.ClassConstructorDescriptor
 import org.jetbrains.kotlin.resolve.calls.components.InferenceSession
 import org.jetbrains.kotlin.resolve.calls.components.KotlinResolutionCallbacks
 import org.jetbrains.kotlin.resolve.calls.components.NewConstraintSystemImpl
@@ -77,7 +78,7 @@ class SimpleCandidateFactory(
             ReceiverExpressionKotlinCallArgument(it, isSafeCall)
         }
         return createCandidate(
-            givenCandidate.descriptor, explicitReceiverKind, dispatchArgumentReceiver, null,
+            givenCandidate.descriptor, explicitReceiverKind, dispatchArgumentReceiver, null, null,
             listOf(), givenCandidate.knownTypeParametersResultingSubstitutor
         )
     }
@@ -96,7 +97,26 @@ class SimpleCandidateFactory(
 
         return createCandidate(
             towerCandidate.descriptor, explicitReceiverKind, dispatchArgumentReceiver,
-            extensionArgumentReceiver, towerCandidate.diagnostics, knownSubstitutor = null
+            extensionArgumentReceiver, null, towerCandidate.diagnostics, knownSubstitutor = null
+        )
+    }
+
+    override fun createCandidate(
+        towerCandidate: CandidateWithBoundDispatchReceiver,
+        explicitReceiverKind: ExplicitReceiverKind,
+        extensionReceiverCandidates: List<ReceiverValueWithSmartCastInfo>
+    ): SimpleResolutionCandidate {
+        val dispatchArgumentReceiver = createReceiverArgument(
+            kotlinCall.getExplicitDispatchReceiver(explicitReceiverKind),
+            towerCandidate.dispatchReceiver
+        )
+        val extensionArgumentReceiverCandidates = extensionReceiverCandidates.mapNotNull {
+            createReceiverArgument(kotlinCall.getExplicitExtensionReceiver(explicitReceiverKind), it)
+        }
+
+        return createCandidate(
+            towerCandidate.descriptor, explicitReceiverKind, dispatchArgumentReceiver,
+            null, extensionArgumentReceiverCandidates, towerCandidate.diagnostics, knownSubstitutor = null
         )
     }
 
@@ -105,12 +125,13 @@ class SimpleCandidateFactory(
         explicitReceiverKind: ExplicitReceiverKind,
         dispatchArgumentReceiver: SimpleKotlinCallArgument?,
         extensionArgumentReceiver: SimpleKotlinCallArgument?,
+        extensionArgumentReceiverCandidates: List<SimpleKotlinCallArgument>?,
         initialDiagnostics: Collection<KotlinCallDiagnostic>,
         knownSubstitutor: TypeSubstitutor?
     ): SimpleResolutionCandidate {
         val resolvedKtCall = MutableResolvedCallAtom(
             kotlinCall, descriptor, explicitReceiverKind,
-            dispatchArgumentReceiver, extensionArgumentReceiver
+            dispatchArgumentReceiver, extensionArgumentReceiver, extensionArgumentReceiverCandidates
         )
 
         if (ErrorUtils.isError(descriptor)) {
@@ -126,7 +147,7 @@ class SimpleCandidateFactory(
             candidate.addDiagnostic(HiddenDescriptor)
         }
 
-        if (extensionArgumentReceiver != null) {
+        if (extensionArgumentReceiver != null && descriptor !is ClassConstructorDescriptor) {
             val parameterIsDynamic = descriptor.extensionReceiverParameter!!.value.type.isDynamic()
             val argumentIsDynamic = extensionArgumentReceiver.receiver.receiverValue.type.isDynamic()
 
@@ -153,7 +174,7 @@ class SimpleCandidateFactory(
 
         return createCandidate(
             errorDescriptor, explicitReceiverKind, dispatchReceiver, extensionArgumentReceiver = null,
-            initialDiagnostics = listOf(), knownSubstitutor = null
+            extensionArgumentReceiverCandidates = null, initialDiagnostics = listOf(), knownSubstitutor = null
         )
     }
 }

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/tower/ScopeTowerProcessors.kt

@@ -154,27 +154,22 @@ private class NoExplicitReceiverScopeTowerProcessor<C : Candidate>(
             data.implicitReceiver
         )
         is TowerData.BothTowerLevelAndContextReceiversGroup -> {
-            val collected = mutableListOf<CandidateWithBoundDispatchReceiver>()
+            val groupsOfDuplicateCandidates = data.contextReceiversGroup.flatMap { receiver ->
-            val receiversWithCandidates = mutableMapOf<ReceiverValueWithSmartCastInfo, List<CandidateWithBoundDispatchReceiver>>()
+                data.level.collectCandidates(receiver).map { it to receiver }
-            for (contextReceiver in data.contextReceiversGroup) {
+            }.filter { (candidate, _) ->
-                val collectedFromReceiver = data.level.collectCandidates(contextReceiver).toMutableList()
+                candidate.requiresExtensionReceiver
-                collectedFromReceiver.removeIf { collectedCandidate ->
+            }.groupBy { it.first.descriptor }.values
-                    val duplicate = collected.find { it.descriptor == collectedCandidate.descriptor }
+
-                    if (duplicate != null) {
+            val candidateToReceivers = groupsOfDuplicateCandidates.map { l ->
-                        duplicate.diagnostics.add(ContextReceiverAmbiguity())
+                val candidate = l.first().first
-                        true
+                val receivers = l.map { it.second }
-                    } else {
+                candidate to receivers
-                        false
-                    }
-                }
-                collected.addAll(collectedFromReceiver)
-                receiversWithCandidates[contextReceiver] = collectedFromReceiver
             }
-            receiversWithCandidates.flatMap {
+            candidateToReceivers.map {
-                createCandidates(
+                candidateFactory.createCandidate(
-                    it.value,
+                    it.first,
                     ExplicitReceiverKind.NO_EXPLICIT_RECEIVER,
-                    it.key
+                    it.second
                 )
             }
         }
@@ -186,6 +181,7 @@ private class NoExplicitReceiverScopeTowerProcessor<C : Candidate>(
             when (data) {
                 is TowerData.TowerLevel -> data.level.recordLookup(name)
                 is TowerData.BothTowerLevelAndImplicitReceiver -> data.level.recordLookup(name)
+                is TowerData.BothTowerLevelAndContextReceiversGroup -> data.level.recordLookup(name)
                 is TowerData.ForLookupForNoExplicitReceiver -> data.level.recordLookup(name)
                 else -> {}
             }

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/tower/TowerResolver.kt

@@ -48,6 +48,12 @@ interface CandidateFactory<out C : Candidate> {
     ): C
 
     fun createErrorCandidate(): C
+
+    fun createCandidate(
+        towerCandidate: CandidateWithBoundDispatchReceiver,
+        explicitReceiverKind: ExplicitReceiverKind,
+        extensionReceiverCandidates: List<ReceiverValueWithSmartCastInfo>
+    ): C
 }
 
 interface CandidateFactoryProviderForInvoke<C : Candidate> {

compiler/testData/diagnostics/tests/extensions/contextReceivers/ambiguityInGroup.kt

@@ -20,6 +20,6 @@ fun test() {
     <!OVERLOAD_RESOLUTION_AMBIGUITY!>supertypeMember<!>()
     <!OVERLOAD_RESOLUTION_AMBIGUITY!>member<!>()
     <!AMBIGUOUS_CALL_WITH_IMPLICIT_CONTEXT_RECEIVER!>supertypeExtension()<!>
-    <!AMBIGUOUS_CALL_WITH_IMPLICIT_CONTEXT_RECEIVER!>supertypeExtensionGeneric()<!>
+    <!AMBIGUOUS_CALL_WITH_IMPLICIT_CONTEXT_RECEIVER!><!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>supertypeExtensionGeneric<!>()<!>
     <!MULTIPLE_ARGUMENTS_APPLICABLE_FOR_CONTEXT_RECEIVER!>supertypeContextual()<!>
 }