KT-51247 Fix resolution of functional types with context receivers

compiler/fir/fir2ir/tests-gen/org/jetbrains/kotlin/test/runners/codegen/FirBlackBoxCodegenTestGenerated.java

@@ -17468,6 +17468,12 @@ public class FirBlackBoxCodegenTestGenerated extends AbstractFirBlackBoxCodegenT
                 runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/iteratorOperator.kt");
             }
 
+            @Test
+            @TestMetadata("kt51247.kt")
+            public void testKt51247() throws Exception {
+                runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/kt51247.kt");
+            }
+
             @Test
             @TestMetadata("kt51284.kt")
             public void testKt51284() throws Exception {

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

@@ -137,7 +137,8 @@ private fun buildReflectionType(
                 receiverType = receiverType.takeIf { fir.receiverParameter != null },
                 rawReturnType = returnType,
                 isKFunctionType = true,
-                isSuspend = isSuspend
+                isSuspend = isSuspend,
+                contextReceivers = fir.contextReceivers.map { it.typeRef.coneType }
             ) to callableReferenceAdaptation
         }
         is FirVariable -> createKPropertyType(fir, receiverType, returnTypeRef, candidate) to null

compiler/ir/ir.tree/src/org/jetbrains/kotlin/ir/expressions/impl/IrFunctionReferenceImpl.kt

@@ -56,7 +56,7 @@ class IrFunctionReferenceImpl(
             type,
             symbol,
             typeArgumentsCount,
-            symbol.descriptor.valueParameters.size,
+            symbol.descriptor.valueParameters.size + symbol.descriptor.contextReceiverParameters.size,
             reflectionTarget,
             origin
         )

compiler/testData/codegen/box/extensionFunctions/contextReceivers/kt51247.kt

@@ -0,0 +1,55 @@
+//!LANGUAGE: +ContextReceivers
+//TARGET_BACKEND: JVM_IR
+
+interface Context
+interface Receiver
+interface Param
+
+fun foo(context: Context, receiver: Receiver, p: Param) {}
+
+context(Context)
+fun bar(receiver: Receiver, p: Param) {}
+
+context(Context)
+fun Receiver.baz(p: Param) {}
+
+
+fun box(): String {
+    var a: context(Context, Receiver, Param) () -> Unit // 3, 0, 0
+    a = ::foo            
+    a = ::bar         
+    a = Receiver::baz 
+
+    var b: context(Context, Receiver) Param.() -> Unit // 2, 1, 0
+    b = ::foo            
+    b = ::bar         
+    b = Receiver::baz 
+
+    var c: context(Context, Receiver) (Param) -> Unit // 2, 0, 1
+    c = ::foo            
+    c = ::bar         
+    c = Receiver::baz 
+
+    var d: context(Context) Receiver.(Param) -> Unit // 1, 1, 1. This is the same as in KEEP.
+    d = ::foo            
+    d = ::bar         
+    d = Receiver::baz 
+
+    var e: context(Context) (Receiver, Param) -> Unit // 1, 0, 2
+    e = ::foo            
+    e = ::bar         
+    e = Receiver::baz 
+
+    var f: Context.(Receiver, Param) -> Unit // 0, 1, 2
+    f = ::foo            
+    f = ::bar         
+    f = Receiver::baz 
+
+    var g: (Context, Receiver, Param) -> Unit // 0, 0, 3
+    g = ::foo            
+    g = ::bar         
+    g = Receiver::baz
+
+    return "OK"
+}
+

compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/functionalType.fir.kt

@@ -1,21 +0,0 @@
-// !LANGUAGE: +ContextReceivers
-// !DIAGNOSTICS: -ASSIGNED_BUT_NEVER_ACCESSED_VARIABLE
-
-class Context
-class Receiver
-class Param
-
-fun foo(context: Context, receiver: Receiver, p: Param) {}
-
-context(Context)
-fun bar(receiver: Receiver, p: Param) {}
-
-context(Context)
-fun Receiver.baz(p: Param) {}
-
-fun main() {
-    var g: context(Context) Receiver.(Param) -> Unit
-    g = ::foo         // OK
-    g = ::<!UNRESOLVED_REFERENCE!>bar<!>         // OK
-    g = Receiver::<!UNRESOLVED_REFERENCE!>baz<!> // OK
-}

compiler/testData/diagnostics/tests/extensions/contextReceivers/fromKEEP/functionalType.kt

@@ -1,3 +1,4 @@
+// FIR_IDENTICAL
 // !LANGUAGE: +ContextReceivers
 // !DIAGNOSTICS: -ASSIGNED_BUT_NEVER_ACCESSED_VARIABLE
 

compiler/tests-common-new/tests-gen/org/jetbrains/kotlin/test/runners/codegen/IrBlackBoxCodegenTestGenerated.java

@@ -17468,6 +17468,12 @@ public class IrBlackBoxCodegenTestGenerated extends AbstractIrBlackBoxCodegenTes
                 runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/iteratorOperator.kt");
             }
 
+            @Test
+            @TestMetadata("kt51247.kt")
+            public void testKt51247() throws Exception {
+                runTest("compiler/testData/codegen/box/extensionFunctions/contextReceivers/kt51247.kt");
+            }
+
             @Test
             @TestMetadata("kt51284.kt")
             public void testKt51284() throws Exception {