FIR: Fix overload resolution with context receivers

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/FirDeclarationInspector.kt

@@ -161,6 +161,10 @@ interface FirDeclarationPresenter {
     }
 
     fun represent(it: FirSimpleFunction) = buildString {
+        it.contextReceivers.forEach {
+            appendRepresentation(it)
+            append(',')
+        }
         append('<')
         it.typeParameters.forEach {
             appendRepresentation(it)
@@ -195,6 +199,10 @@ interface FirDeclarationPresenter {
     }
 
     fun represent(it: FirConstructor, owner: FirRegularClass) = buildString {
+        it.contextReceivers.forEach {
+            appendRepresentation(it)
+            append(',')
+        }
         append('<')
         it.typeParameters.forEach {
             appendRepresentation(it)

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

@@ -41,6 +41,9 @@ abstract class AbstractConeCallConflictResolver(
         if (!call1.isExpect && call2.isExpect) return true
         if (call1.isExpect && !call2.isExpect) return false
 
+        if (call1.contextReceiverCount > call2.contextReceiverCount) return true
+        if (call1.contextReceiverCount < call2.contextReceiverCount) return false
+
         return createEmptyConstraintSystem().isSignatureNotLessSpecific(
             call1,
             call2,
@@ -114,7 +117,7 @@ abstract class AbstractConeCallConflictResolver(
             (variable as? FirProperty)?.typeParameters?.map { it.symbol.toLookupTag() }.orEmpty(),
             computeSignatureTypes(call, variable),
             variable.receiverTypeRef != null,
-            0, // TODO
+            variable.contextReceivers.size,
             false,
             0,
             (variable as? FirProperty)?.isExpect == true,
@@ -129,7 +132,7 @@ abstract class AbstractConeCallConflictResolver(
             computeSignatureTypes(call, constructor),
             //constructor.receiverTypeRef != null,
             false,
-            0, // TODO
+            constructor.contextReceivers.size,
             constructor.valueParameters.any { it.isVararg },
             call.numDefaults,
             constructor.isExpect,
@@ -143,7 +146,7 @@ abstract class AbstractConeCallConflictResolver(
             function.typeParameters.map { it.symbol.toLookupTag() },
             computeSignatureTypes(call, function),
             function.receiverTypeRef != null,
-            0, // TODO
+            function.contextReceivers.size,
             function.valueParameters.any { it.isVararg },
             call.numDefaults,
             function.isExpect,
@@ -171,6 +174,7 @@ abstract class AbstractConeCallConflictResolver(
                         (it as ConeKotlinType).removeTypeVariableTypes(inferenceComponents.session.typeContext)
                     }
             } else {
+                called.contextReceivers.mapTo(this) { it.typeRef.coneType }
                 call.argumentMapping?.mapTo(this) { it.value.argumentType() }
             }
         }

compiler/testData/diagnostics/tests/extensions/contextReceivers/conflictingWithDifferentOrder.fir.kt

@@ -10,9 +10,9 @@ fun f(): Unit<!> = TODO()
 fun f(): Unit<!> = TODO()
 
 fun test(a: A, b: B) {
-    with(a) {
+    <!NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER, NEW_INFERENCE_NO_INFORMATION_FOR_PARAMETER!>with<!>(a) {
         with(b) {
-            f()
+            <!ARGUMENT_TYPE_MISMATCH!><!OVERLOAD_RESOLUTION_AMBIGUITY!>f<!>()<!>
         }
     }
 }

compiler/testData/diagnostics/tests/extensions/contextReceivers/overloadPriority.fir.kt

@@ -1,14 +0,0 @@
-// !LANGUAGE: +ContextReceivers
-
-class Context
-
-<!CONFLICTING_OVERLOADS!>context(Context)
-fun f(): String<!> = TODO()
-
-<!CONFLICTING_OVERLOADS!>fun f(): Any<!> = TODO()
-
-fun test() {
-    with(Context()) {
-        f().length
-    }
-}

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

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !LANGUAGE: +ContextReceivers
 
 class Context

compiler/testData/diagnostics/tests/extensions/contextReceivers/overloading.fir.kt

@@ -1,17 +0,0 @@
-// !LANGUAGE: +ContextReceivers
-
-<!CONFLICTING_OVERLOADS!>context(Int, String)
-fun foo(): Int<!> {
-    return this@Int + 42
-}
-
-<!CONFLICTING_OVERLOADS!>context(Int)
-fun foo(): Int<!> {
-    return this@Int + 42
-}
-
-fun test() {
-    with(42) {
-        foo()
-    }
-}

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

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !LANGUAGE: +ContextReceivers
 
 context(Int, String)

compiler/testData/diagnostics/tests/extensions/contextReceivers/thisWithReceiverLabelsFunctions.fir.kt

@@ -5,16 +5,16 @@ class B(val b: Any)
 class C(val c: Any)
 
 context(A<String>) fun A<Int>.f() {
-    this@A.a.length
+    this@A.a.<!UNRESOLVED_REFERENCE!>length<!>
 }
 
-<!CONFLICTING_OVERLOADS!>context(A<String>, B) fun f()<!> {
+context(A<String>, B) fun f() {
     this@A.a.length
     this@B.b
     <!NO_THIS!>this<!>
 }
 
-<!CONFLICTING_OVERLOADS!>context(A<Int>, A<String>, B) fun f()<!> {
+context(A<Int>, A<String>, B) fun f() {
     this@A.a.length
     this@B.b
     <!NO_THIS!>this<!>

compiler/testData/diagnostics/tests/extensions/contextReceivers/unsupported.fir.kt

@@ -28,7 +28,7 @@ class A {
 fun useWithContextReceivers() {
     with(42) {
         with("") {
-            <!OVERLOAD_RESOLUTION_AMBIGUITY!>f<!>({}, 42)
+            f({}, 42)
             sameAsFWithoutNonContextualCounterpart({}, 42)
             p
             val a = A()