FIR: Add constraints for extension receiver of callable reference

^KT-32725 In Progress

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

@@ -34,7 +34,6 @@ import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.fir.types.impl.FirResolvedTypeRefImpl
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
-import org.jetbrains.kotlin.resolve.calls.inference.model.SimpleConstraintSystemConstraintPosition
 import org.jetbrains.kotlin.resolve.calls.results.TypeSpecificityComparator
 
 class FirCallResolver(
@@ -257,10 +256,7 @@ class FirCallResolver(
 
         val chosenCandidate = reducedCandidates.single()
         constraintSystemBuilder.runTransaction {
-            addOtherSystem(chosenCandidate.system.asReadOnlyStorage())
-
-            val position = SimpleConstraintSystemConstraintPosition //TODO
-            addSubtypeConstraint(chosenCandidate.resultingTypeForCallableReference!!, expectedType, position)
+            chosenCandidate.outerConstraintBuilderEffect!!(this)
 
             true
         }

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

@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.fir.types.ConeKotlinType
 import org.jetbrains.kotlin.fir.types.FirTypeProjection
 import org.jetbrains.kotlin.fir.types.FirTypeRef
 import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemBuilder
+import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemOperation
 import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
 import org.jetbrains.kotlin.resolve.calls.model.PostponedResolvedAtomMarker
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
@@ -73,6 +74,7 @@ class Candidate(
 
     lateinit var substitutor: ConeSubstitutor
     var resultingTypeForCallableReference: ConeKotlinType? = null
+    var outerConstraintBuilderEffect: (ConstraintSystemOperation.() -> Unit)? = null
 
     var argumentMapping: Map<FirExpression, FirValueParameter>? = null
     val postponedAtoms = mutableListOf<PostponedResolvedAtomMarker>()

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

@@ -6,10 +6,7 @@
 package org.jetbrains.kotlin.fir.resolve.calls
 
 import org.jetbrains.kotlin.descriptors.Visibilities
-import org.jetbrains.kotlin.fir.declarations.FirMemberDeclaration
-import org.jetbrains.kotlin.fir.declarations.FirProperty
-import org.jetbrains.kotlin.fir.declarations.FirSimpleFunction
-import org.jetbrains.kotlin.fir.declarations.visibility
+import org.jetbrains.kotlin.fir.declarations.*
 import org.jetbrains.kotlin.fir.expressions.FirResolvedQualifier
 import org.jetbrains.kotlin.fir.resolve.DoubleColonLHS
 import org.jetbrains.kotlin.fir.resolve.createFunctionalType
@@ -25,6 +22,7 @@ import org.jetbrains.kotlin.fir.types.FirResolvedTypeRef
 import org.jetbrains.kotlin.fir.types.coneTypeSafe
 import org.jetbrains.kotlin.load.java.JavaVisibilities
 import org.jetbrains.kotlin.name.FqName
+import org.jetbrains.kotlin.resolve.calls.inference.ConstraintSystemOperation
 import org.jetbrains.kotlin.resolve.calls.inference.model.SimpleConstraintSystemConstraintPosition
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind.*
@@ -184,22 +182,35 @@ internal object CheckCallableReferenceExpectedType : CheckerStage() {
             else -> null
         }
 
-        val resultingType: ConeKotlinType = when (val fir = candidate.symbol.fir) {
+        val fir = candidate.symbol.fir
+        val resultingType: ConeKotlinType = when (fir) {
             is FirSimpleFunction -> createKFunctionType(fir, resultingReceiverType)
             is FirProperty -> createKPropertyType(fir, resultingReceiverType)
             else -> ConeKotlinErrorType("Unknown callable kind: ${fir::class}")
         }.let(candidate.substitutor::substituteOrSelf)
 
         candidate.resultingTypeForCallableReference = resultingType
-
-        var isApplicable = true
-
-        outerCsBuilder.runTransaction {
+        candidate.outerConstraintBuilderEffect = fun ConstraintSystemOperation.() {
             addOtherSystem(candidate.system.asReadOnlyStorage())
 
             val position = SimpleConstraintSystemConstraintPosition //TODO
 
             addSubtypeConstraint(resultingType, expectedType, position)
+
+            val declarationReceiverType: ConeKotlinType? =
+                (fir as? FirCallableMemberDeclaration<*>)?.receiverTypeRef?.coneTypeSafe<ConeKotlinType>()
+                    ?.let(candidate.substitutor::substituteOrSelf)
+
+            if (resultingReceiverType != null && declarationReceiverType != null) {
+                addSubtypeConstraint(resultingReceiverType, declarationReceiverType, position)
+            }
+        }
+
+        var isApplicable = true
+
+        outerCsBuilder.runTransaction {
+            candidate.outerConstraintBuilderEffect!!(this)
+
             isApplicable = !hasContradiction
 
             false

compiler/fir/resolve/testData/diagnostics/callableReferences/extensionReceiverInference.kt

@@ -0,0 +1,10 @@
+class A
+val <X> X.prop: Int get() = 1
+fun <X> X.baz(): Int = 1
+
+fun foo(x: (A) -> Int) {}
+
+fun main() {
+    foo(A::prop)
+    foo(A::baz)
+}

compiler/fir/resolve/testData/diagnostics/callableReferences/extensionReceiverInference.txt

@@ -0,0 +1,20 @@
+FILE: extensionReceiverInference.kt
+    public final class A : R|kotlin/Any| {
+        public constructor(): R|A| {
+            super<R|kotlin/Any|>()
+        }
+
+    }
+    public final val <X> R|X|.prop: R|kotlin/Int|
+        public get(): R|kotlin/Int| {
+            ^ Int(1)
+        }
+    public final fun <X> R|X|.baz(): R|kotlin/Int| {
+        ^baz Int(1)
+    }
+    public final fun foo(x: R|kotlin/Function1<A, kotlin/Int>|): R|kotlin/Unit| {
+    }
+    public final fun main(): R|kotlin/Unit| {
+        R|/foo|(Q|A|::R|/prop<A>|)
+        R|/foo|(Q|A|::R|/baz<A>|)
+    }

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsTestGenerated.java

@@ -50,6 +50,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
             runTest("compiler/fir/resolve/testData/diagnostics/callableReferences/differentLevels.kt");
         }
 
+        @TestMetadata("extensionReceiverInference.kt")
+        public void testExtensionReceiverInference() throws Exception {
+            runTest("compiler/fir/resolve/testData/diagnostics/callableReferences/extensionReceiverInference.kt");
+        }
+
         @TestMetadata("inferenceFromCallableReferenceType.kt")
         public void testInferenceFromCallableReferenceType() throws Exception {
             runTest("compiler/fir/resolve/testData/diagnostics/callableReferences/inferenceFromCallableReferenceType.kt");