FIR: Fix overload resolution for some invoke cases

Namely, it's about the cases when there are multiple variable candidates
with the same priority and all of them should be collected and compared

compiler/fir/analysis-tests/testData/resolve/callResolution/invokeAmbiguity.kt

@@ -0,0 +1,17 @@
+open class A {
+    operator fun invoke(f: () -> Unit): Int = 1
+}
+
+class B {
+    operator fun invoke(f: () -> Unit): CharSequence = ""
+}
+
+open class C
+val C.attr: A get() = TODO()
+
+open class D: C()
+val D.attr: B get() = TODO()
+
+fun box(d: D) {
+    (d.attr {}).length
+}

compiler/fir/analysis-tests/testData/resolve/callResolution/invokeAmbiguity.txt

@@ -0,0 +1,47 @@
+FILE: invokeAmbiguity.kt
+    public open class A : R|kotlin/Any| {
+        public constructor(): R|A| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final operator fun invoke(f: R|() -> kotlin/Unit|): R|kotlin/Int| {
+            ^invoke Int(1)
+        }
+
+    }
+    public final class B : R|kotlin/Any| {
+        public constructor(): R|B| {
+            super<R|kotlin/Any|>()
+        }
+
+        public final operator fun invoke(f: R|() -> kotlin/Unit|): R|kotlin/CharSequence| {
+            ^invoke String()
+        }
+
+    }
+    public open class C : R|kotlin/Any| {
+        public constructor(): R|C| {
+            super<R|kotlin/Any|>()
+        }
+
+    }
+    public final val R|C|.attr: R|A|
+        public get(): R|A| {
+            ^ R|kotlin/TODO|()
+        }
+    public open class D : R|C| {
+        public constructor(): R|D| {
+            super<R|C|>()
+        }
+
+    }
+    public final val R|D|.attr: R|B|
+        public get(): R|B| {
+            ^ R|kotlin/TODO|()
+        }
+    public final fun box(d: R|D|): R|kotlin/Unit| {
+        R|<local>/d|.R|/attr|.R|/B.invoke|(<L> = attr@fun <anonymous>(): R|kotlin/Unit| {
+            Unit
+        }
+        ).R|kotlin/CharSequence.length|
+    }

compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/invokePriority.kt

@@ -29,6 +29,6 @@ class E {
 }
 
 fun main() {
-    E.f() // Resolves to (2) in old FE (to (1) in FIR with object implicit invoke support)
+    E.<!AMBIGUITY!>f<!>() // Resolves to (2) in old FE (ambiguity in FIR)
     E.f.invoke() // Resolves to (1)
-}
+}

compiler/fir/analysis-tests/testData/resolveWithStdlib/problems/invokePriority.txt

@@ -74,6 +74,6 @@ FILE: invokePriority.kt
 
     }
     public final fun main(): R|kotlin/Unit| {
-        Q|E|.R|/E.Companion.f|.R|/E.f.invoke|()
+        Q|E|.<Ambiguity: f, [/E.f.invoke, kotlin/Function0.invoke]>#()
         Q|E.f|.R|/E.f.invoke|()
     }

compiler/fir/analysis-tests/tests/org/jetbrains/kotlin/fir/FirDiagnosticsTestGenerated.java

@@ -584,6 +584,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
             runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/companionInvoke.kt");
         }
 
+        @TestMetadata("invokeAmbiguity.kt")
+        public void testInvokeAmbiguity() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/invokeAmbiguity.kt");
+        }
+
         @TestMetadata("objectInvoke.kt")
         public void testObjectInvoke() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/objectInvoke.kt");

compiler/fir/analysis-tests/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithLightTreeTestGenerated.java

@@ -584,6 +584,11 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
             runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/companionInvoke.kt");
         }
 
+        @TestMetadata("invokeAmbiguity.kt")
+        public void testInvokeAmbiguity() throws Exception {
+            runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/invokeAmbiguity.kt");
+        }
+
         @TestMetadata("objectInvoke.kt")
         public void testObjectInvoke() throws Exception {
             runTest("compiler/fir/analysis-tests/testData/resolve/callResolution/objectInvoke.kt");

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirCallResolver.kt

@@ -354,6 +354,7 @@ class FirCallResolver(
             session,
             file,
             transformer.components.implicitReceiverStack,
+            candidateForCommonInvokeReceiver = null,
             // Additional things for callable reference resolve
             expectedType,
             outerConstraintSystemBuilder,

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

@@ -43,6 +43,8 @@ data class CallInfo(
     val containingFile: FirFile,
     val implicitReceiverStack: ImplicitReceiverStack,
 
+    val candidateForCommonInvokeReceiver: Candidate? = null,
+
     // Four properties for callable references only
     val expectedType: ConeKotlinType? = null,
     val outerCSBuilder: ConstraintSystemBuilder? = null,

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

@@ -44,6 +44,9 @@ open class CandidateCollector(
 
     fun bestCandidates() = candidates
 
+    fun shouldStopAtTheLevel(group: TowerGroup) =
+        isSuccess() && bestGroup < group
+
     fun isSuccess(): Boolean {
         return currentApplicability >= CandidateApplicability.SYNTHETIC_RESOLVED
     }

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

@@ -35,11 +35,31 @@ class ConeOverloadConflictResolver(
         discriminateGenerics: Boolean,
         discriminateAbstracts: Boolean
     ): Set<Candidate> {
+        if (candidates.size == 1) return candidates
+        val fixedCandidates =
+            if (candidates.first().callInfo.candidateForCommonInvokeReceiver != null)
+                chooseCandidatesWithMostSpecificInvokeReceiver(candidates)
+            else
+                candidates
+
         return chooseMaximallySpecificCandidates(
-            candidates, discriminateGenerics, discriminateAbstracts, discriminateSAMs = true
+            fixedCandidates, discriminateGenerics, discriminateAbstracts, discriminateSAMs = true
         )
     }
 
+    private fun chooseCandidatesWithMostSpecificInvokeReceiver(candidates: Set<Candidate>): Set<Candidate> {
+        val propertyReceiverCandidates = candidates.mapTo(mutableSetOf()) {
+            it.callInfo.candidateForCommonInvokeReceiver
+                ?: error("If one candidate within a group is property+invoke, other should be the same, but $it found")
+        }
+
+        val bestInvokeReceiver =
+            chooseMaximallySpecificCandidates(propertyReceiverCandidates, discriminateGenerics = false)
+                .singleOrNull() ?: return candidates
+
+        return candidates.filterTo(mutableSetOf()) { it.callInfo.candidateForCommonInvokeReceiver == bestInvokeReceiver }
+    }
+
     private fun chooseMaximallySpecificCandidates(
         candidates: Set<Candidate>,
         discriminateGenerics: Boolean,

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

@@ -101,8 +101,6 @@ class FirTowerResolverSession internal constructor(
             enqueueResolverTasksForInvokeReceiverCandidates(
                 invokeResolveMode, callInfo
             )
-        }.also {
-            manager.stopIfSuccess()
         }
     }
 
@@ -521,7 +519,11 @@ class FirTowerResolverSession internal constructor(
         group: TowerGroup,
         explicitReceiverKind: ExplicitReceiverKind,
         candidateFactory: CandidateFactory
-    ) = processLevel(towerLevel, callInfo, group, explicitReceiverKind, InvokeResolveMode.IMPLICIT_CALL_ON_GIVEN_RECEIVER, candidateFactory)
+    ) = processLevel(
+        towerLevel, callInfo,
+        group.InvokeResolvePriority(InvokeResolvePriority.COMMON_INVOKE),
+        explicitReceiverKind, InvokeResolveMode.IMPLICIT_CALL_ON_GIVEN_RECEIVER, candidateFactory
+    )
 
     // Here we already know extension receiver for invoke, and it's stated in info as first argument
     private suspend fun runResolverForBuiltinInvokeExtensionWithExplicitArgument(
@@ -531,7 +533,7 @@ class FirTowerResolverSession internal constructor(
     ) {
         processLevel(
             invokeReceiverValue.toMemberScopeTowerLevel(),
-            info, TowerGroup.Member,
+            info, TowerGroup.Member.InvokeResolvePriority(InvokeResolvePriority.INVOKE_EXTENSION),
             ExplicitReceiverKind.DISPATCH_RECEIVER,
             InvokeResolveMode.IMPLICIT_CALL_ON_GIVEN_RECEIVER,
             invokeOnGivenReceiverCandidateFactory
@@ -545,13 +547,13 @@ class FirTowerResolverSession internal constructor(
         invokeOnGivenReceiverCandidateFactory: CandidateFactory
     ) {
         for ((implicitReceiverValue, depth) in implicitReceiversUsableAsValues) {
-            val parentGroup = TowerGroup.Implicit(depth)
+            val towerGroup = TowerGroup.Implicit(depth).InvokeExtension.InvokeResolvePriority(InvokeResolvePriority.INVOKE_EXTENSION)
             processLevel(
                 invokeReceiverValue.toMemberScopeTowerLevel(
                     extensionReceiver = implicitReceiverValue,
                     implicitExtensionInvokeMode = true
                 ),
-                info, parentGroup.InvokeExtension,
+                info, towerGroup,
                 ExplicitReceiverKind.DISPATCH_RECEIVER,
                 InvokeResolveMode.IMPLICIT_CALL_ON_GIVEN_RECEIVER,
                 invokeOnGivenReceiverCandidateFactory
@@ -587,7 +589,10 @@ class FirTowerResolverSession internal constructor(
                 )
 
             val invokeFunctionInfo =
-                info.copy(explicitReceiver = invokeReceiverExpression, name = OperatorNameConventions.INVOKE).let {
+                info.copy(
+                    explicitReceiver = invokeReceiverExpression, name = OperatorNameConventions.INVOKE,
+                    candidateForCommonInvokeReceiver = invokeReceiverCandidate.takeUnless { invokeBuiltinExtensionMode }
+                ).let {
                     when {
                         invokeBuiltinExtensionMode -> it.withReceiverAsArgument(info.explicitReceiver!!)
                         else -> it

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

@@ -43,7 +43,11 @@ sealed class TowerGroupKind(private val index: Int) : Comparable<TowerGroupKind>
 }
 
 @Suppress("FunctionName", "unused", "PropertyName")
-class TowerGroup private constructor(private val kinds: Array<TowerGroupKind>) : Comparable<TowerGroup> {
+class TowerGroup
+private constructor(
+    private val kinds: Array<TowerGroupKind>,
+    private val invokeResolvePriority: InvokeResolvePriority = InvokeResolvePriority.NONE
+) : Comparable<TowerGroup> {
     companion object {
         private fun kindOf(kind: TowerGroupKind): TowerGroup = TowerGroup(arrayOf(kind))
 
@@ -84,6 +88,14 @@ class TowerGroup private constructor(private val kinds: Array<TowerGroupKind>) :
 
     fun Static(depth: Int) = kindOf(TowerGroupKind.Static(depth))
 
+    // Treating `a.foo()` common calls as more prioritized than `a.foo.invoke()`
+    // It's not the same as TowerGroupKind because it's not about tower levels, but rather a different dimension semantically.
+    // It could be implemented via another TowerGroupKind, but it's not clear what priority should be assigned to the new TowerGroupKind
+    fun InvokeResolvePriority(invokeResolvePriority: InvokeResolvePriority): TowerGroup {
+        if (invokeResolvePriority == InvokeResolvePriority.NONE) return this
+        return TowerGroup(kinds, invokeResolvePriority)
+    }
+
     override fun compareTo(other: TowerGroup): Int {
         var index = 0
         while (index < kinds.size) {
@@ -95,6 +107,29 @@ class TowerGroup private constructor(private val kinds: Array<TowerGroupKind>) :
             index++
         }
         if (index < other.kinds.size) return -1
-        return 0
+
+        return invokeResolvePriority.compareTo(other.invokeResolvePriority)
+    }
+
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (javaClass != other?.javaClass) return false
+
+        other as TowerGroup
+
+        if (!kinds.contentEquals(other.kinds)) return false
+        if (invokeResolvePriority != other.invokeResolvePriority) return false
+
+        return true
+    }
+
+    override fun hashCode(): Int {
+        var result = kinds.contentHashCode()
+        result = 31 * result + invokeResolvePriority.hashCode()
+        return result
     }
 }
+
+enum class InvokeResolvePriority {
+    NONE, COMMON_INVOKE, INVOKE_EXTENSION;
+}

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

@@ -8,9 +8,9 @@ package org.jetbrains.kotlin.fir.resolve.calls
 import java.util.*
 import kotlin.coroutines.*
 
-class TowerResolveManager private constructor(private val isSuccess: () -> Boolean) {
+class TowerResolveManager private constructor(private val shouldStopAtTheLevel: (TowerGroup) -> Boolean) {
 
-    constructor(collector: CandidateCollector) : this({ collector.isSuccess() })
+    constructor(collector: CandidateCollector) : this(collector::shouldStopAtTheLevel)
 
     private val queue = PriorityQueue<SuspendedResolverTask>()
 
@@ -21,6 +21,9 @@ class TowerResolveManager private constructor(private val isSuccess: () -> Boole
     private suspend fun suspendResolverTask(group: TowerGroup) = suspendCoroutine<Unit> { queue += SuspendedResolverTask(it, group) }
 
     suspend fun requestGroup(requested: TowerGroup) {
+        if (shouldStopAtTheLevel(requested)) {
+            stopResolverTask()
+        }
         val peeked = queue.peek()
 
         // Task ordering should be FIFO
@@ -32,10 +35,6 @@ class TowerResolveManager private constructor(private val isSuccess: () -> Boole
 
     private suspend fun stopResolverTask(): Nothing = suspendCoroutine { }
 
-    suspend fun stopIfSuccess() {
-        if (isSuccess()) stopResolverTask()
-    }
-
     private data class SuspendedResolverTask(
         val continuation: Continuation<Unit>,
         val group: TowerGroup
@@ -62,8 +61,11 @@ class TowerResolveManager private constructor(private val isSuccess: () -> Boole
 
     fun runTasks() {
         while (queue.isNotEmpty()) {
-            queue.poll().continuation.resume(Unit)
+            val current = queue.poll()
-            if (isSuccess()) return
+            if (shouldStopAtTheLevel(current.group)) {
+                return
+            }
+            current.continuation.resume(Unit)
         }
     }
 }