Extend CandidateInterceptor to intercept ResolvedCallAtom candidate

compiler/frontend/src/org/jetbrains/kotlin/extensions/internal/CandidateInterceptor.kt

@@ -6,16 +6,21 @@
 package org.jetbrains.kotlin.extensions.internal
 
 import com.intellij.openapi.project.Project
+import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
 import org.jetbrains.kotlin.extensions.ProjectExtensionDescriptor
 import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.resolve.BindingTrace
 import org.jetbrains.kotlin.resolve.calls.CallResolver
 import org.jetbrains.kotlin.resolve.calls.CandidateResolver
 import org.jetbrains.kotlin.resolve.calls.KotlinCallResolver
 import org.jetbrains.kotlin.resolve.calls.context.BasicCallResolutionContext
+import org.jetbrains.kotlin.resolve.calls.inference.components.NewTypeSubstitutor
 import org.jetbrains.kotlin.resolve.calls.model.CallResolutionResult
+import org.jetbrains.kotlin.resolve.calls.model.KotlinCallDiagnostic
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCallAtom
 import org.jetbrains.kotlin.resolve.calls.tasks.TracingStrategy
 import org.jetbrains.kotlin.resolve.calls.tower.ImplicitScopeTower
 import org.jetbrains.kotlin.resolve.calls.tower.NewResolutionOldInference
@@ -27,6 +32,16 @@ import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValueWithSmartCastI
 class CandidateInterceptor(project: Project) {
     private val extensions = getInstances(project)
 
+    fun interceptResolvedCallAtomCandidate(
+        candidateDescriptor: CallableDescriptor,
+        completedCallAtom: ResolvedCallAtom,
+        trace: BindingTrace?,
+        resultSubstitutor: NewTypeSubstitutor?,
+        diagnostics: Collection<KotlinCallDiagnostic>
+    ): CallableDescriptor = extensions.fold(candidateDescriptor) { it, extension ->
+        extension.interceptResolvedCallAtomCandidate(it, completedCallAtom, trace, resultSubstitutor, diagnostics)
+    }
+
     fun interceptResolvedCandidates(
         candidates: Collection<NewResolutionOldInference.MyCandidate>,
         context: BasicCallResolutionContext,

compiler/frontend/src/org/jetbrains/kotlin/extensions/internal/NonStableExtensionPoints.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.extensions.internal
 
+import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.VariableDescriptor
 import org.jetbrains.kotlin.descriptors.impl.AnonymousFunctionDescriptor
@@ -12,9 +13,13 @@ import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtLambdaExpression
+import org.jetbrains.kotlin.resolve.BindingTrace
 import org.jetbrains.kotlin.resolve.calls.CallResolver
 import org.jetbrains.kotlin.resolve.calls.CandidateResolver
 import org.jetbrains.kotlin.resolve.calls.context.BasicCallResolutionContext
+import org.jetbrains.kotlin.resolve.calls.inference.components.NewTypeSubstitutor
+import org.jetbrains.kotlin.resolve.calls.model.KotlinCallDiagnostic
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCallAtom
 import org.jetbrains.kotlin.resolve.calls.tasks.TracingStrategy
 import org.jetbrains.kotlin.resolve.calls.tower.ImplicitScopeTower
 import org.jetbrains.kotlin.resolve.calls.tower.NewResolutionOldInference
@@ -52,6 +57,16 @@ interface TypeResolutionInterceptorExtension {
 
 @InternalNonStableExtensionPoints
 interface CallResolutionInterceptorExtension {
+    @Suppress("DEPRECATION")
+    @JvmDefault
+    fun interceptResolvedCallAtomCandidate(
+        candidateDescriptor: CallableDescriptor,
+        completedCallAtom: ResolvedCallAtom,
+        trace: BindingTrace?,
+        resultSubstitutor: NewTypeSubstitutor?,
+        diagnostics: Collection<KotlinCallDiagnostic>
+    ): CallableDescriptor = candidateDescriptor
+
     @Suppress("DEPRECATION")
     @JvmDefault
     fun interceptCandidates(

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

@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
 import org.jetbrains.kotlin.diagnostics.Errors
+import org.jetbrains.kotlin.extensions.internal.CandidateInterceptor
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.resolve.*
 import org.jetbrains.kotlin.resolve.calls.ArgumentTypeResolver
@@ -60,6 +61,7 @@ class KotlinToResolvedCallTransformer(
     private val smartCastManager: SmartCastManager,
     private val typeApproximator: TypeApproximator,
     private val missingSupertypesResolver: MissingSupertypesResolver,
+    private val candidateInterceptor: CandidateInterceptor,
 ) {
     companion object {
         private val REPORT_MISSING_NEW_INFERENCE_DIAGNOSTIC
@@ -182,6 +184,16 @@ class KotlinToResolvedCallTransformer(
     ): NewAbstractResolvedCall<D> {
         val psiKotlinCall = completedCallAtom.atom.psiKotlinCall
 
+        completedCallAtom.setCandidateDescriptor(
+            candidateInterceptor.interceptResolvedCallAtomCandidate(
+                completedCallAtom.candidateDescriptor,
+                completedCallAtom,
+                trace,
+                resultSubstitutor,
+                diagnostics
+            )
+        )
+
         return if (psiKotlinCall is PSIKotlinCallForInvoke) {
             val diagnosticsForVariableCall = if (completedCallAtom.candidateDescriptor is FunctionDescriptor) emptyList() else diagnostics
             val diagnosticsForFunctionCall = if (completedCallAtom.candidateDescriptor is FunctionDescriptor) diagnostics else emptyList()

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

@@ -84,6 +84,7 @@ abstract class ResolvedCallAtom : ResolvedAtom() {
     abstract val argumentsWithSuspendConversion: Map<KotlinCallArgument, UnwrappedType>
     abstract val argumentsWithUnitConversion: Map<KotlinCallArgument, UnwrappedType>
     abstract val argumentsWithConstantConversion: Map<KotlinCallArgument, IntegerValueTypeConstant>
+    abstract fun setCandidateDescriptor(newCandidateDescriptor: CallableDescriptor)
 }
 
 class SamConversionDescription(

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

@@ -78,7 +78,7 @@ class ResolvedCallableReferenceCallAtom(
 
 open class MutableResolvedCallAtom(
     override val atom: KotlinCall,
-    override val candidateDescriptor: CallableDescriptor, // original candidate descriptor
+    originalCandidateDescriptor: CallableDescriptor, // original candidate descriptor
     override val explicitReceiverKind: ExplicitReceiverKind,
     override val dispatchReceiverArgument: SimpleKotlinCallArgument?,
     override val extensionReceiverArgument: SimpleKotlinCallArgument?,
@@ -94,6 +94,10 @@ open class MutableResolvedCallAtom(
     private var suspendAdapterMap: HashMap<KotlinCallArgument, UnwrappedType>? = null
     private var unitAdapterMap: HashMap<KotlinCallArgument, UnwrappedType>? = null
     private var signedUnsignedConstantConversions: HashMap<KotlinCallArgument, IntegerValueTypeConstant>? = null
+    private var _candidateDescriptor = originalCandidateDescriptor
+
+    override val candidateDescriptor: CallableDescriptor
+        get() = _candidateDescriptor
 
     val hasSamConversion: Boolean
         get() = samAdapterMap != null
@@ -146,6 +150,11 @@ open class MutableResolvedCallAtom(
         super.setAnalyzedResults(subResolvedAtoms)
     }
 
+    override fun setCandidateDescriptor(newCandidateDescriptor: CallableDescriptor) {
+        if (newCandidateDescriptor == candidateDescriptor) return
+        _candidateDescriptor = newCandidateDescriptor
+    }
+
     override fun toString(): String = "$atom, candidate = $candidateDescriptor"
 }