From d921dd0eed46da9d53e7e32a55727b04d03e2bed Mon Sep 17 00:00:00 2001
From: Mikhail Zarechenskiy <mikhail.zarechenskiy@jetbrains.com>
Date: Mon, 17 Feb 2020 17:05:21 +0300
Subject: [PATCH] Provide candidate interceptor for NI infrastructure

---
 .../internal/CandidateInterceptor.kt          | 22 ++++++++++++++++++-
 .../internal/NonStableExtensionPoints.kt      | 11 ++++++++++
 .../resolve/calls/tower/PSICallResolver.kt    |  4 ++++
 3 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/compiler/frontend/src/org/jetbrains/kotlin/extensions/internal/CandidateInterceptor.kt b/compiler/frontend/src/org/jetbrains/kotlin/extensions/internal/CandidateInterceptor.kt
index c1c0a924686..a6c737565be 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/extensions/internal/CandidateInterceptor.kt
+++ b/compiler/frontend/src/org/jetbrains/kotlin/extensions/internal/CandidateInterceptor.kt
@@ -12,7 +12,9 @@ import org.jetbrains.kotlin.incremental.components.LookupLocation
 import org.jetbrains.kotlin.name.Name
 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.model.CallResolutionResult
 import org.jetbrains.kotlin.resolve.calls.tasks.TracingStrategy
 import org.jetbrains.kotlin.resolve.calls.tower.ImplicitScopeTower
 import org.jetbrains.kotlin.resolve.calls.tower.NewResolutionOldInference
@@ -30,10 +32,28 @@ class CandidateInterceptor(project: Project) {
         name: Name,
         kind: NewResolutionOldInference.ResolutionKind,
         tracing: TracingStrategy
-    ) = extensions.fold(candidates) { it, extension ->
+    ): Collection<NewResolutionOldInference.MyCandidate> = extensions.fold(candidates) { it, extension ->
         extension.interceptCandidates(it, context, candidateResolver, callResolver, name, kind, tracing)
     }
 
+    fun interceptResolvedCandidates(
+        callResolutionResult: CallResolutionResult,
+        context: BasicCallResolutionContext,
+        callResolver: KotlinCallResolver,
+        name: Name,
+        kind: NewResolutionOldInference.ResolutionKind,
+        tracing: TracingStrategy
+    ): CallResolutionResult {
+        var resolutionResult = callResolutionResult
+        for (extension in extensions) {
+            resolutionResult = extension.interceptCandidates(
+                resolutionResult, context, callResolver, name, kind, tracing
+            )
+        }
+
+        return resolutionResult
+    }
+
     fun interceptCandidates(
         candidates: Collection<FunctionDescriptor>,
         scopeTower: ImplicitScopeTower,
diff --git a/compiler/frontend/src/org/jetbrains/kotlin/extensions/internal/NonStableExtensionPoints.kt b/compiler/frontend/src/org/jetbrains/kotlin/extensions/internal/NonStableExtensionPoints.kt
index 9928eb9a7e0..3e3e83c6087 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/extensions/internal/NonStableExtensionPoints.kt
+++ b/compiler/frontend/src/org/jetbrains/kotlin/extensions/internal/NonStableExtensionPoints.kt
@@ -13,7 +13,9 @@ import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtLambdaExpression
 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.model.CallResolutionResult
 import org.jetbrains.kotlin.resolve.calls.tasks.TracingStrategy
 import org.jetbrains.kotlin.resolve.calls.tower.ImplicitScopeTower
 import org.jetbrains.kotlin.resolve.calls.tower.NewResolutionOldInference
@@ -58,6 +60,15 @@ interface CallResolutionInterceptorExtension {
         tracing: TracingStrategy
     ): Collection<NewResolutionOldInference.MyCandidate> = candidates
 
+    fun interceptCandidates(
+        callResolutionResult: CallResolutionResult,
+        context: BasicCallResolutionContext,
+        callResolver: KotlinCallResolver,
+        name: Name,
+        kind: NewResolutionOldInference.ResolutionKind,
+        tracing: TracingStrategy
+    ): CallResolutionResult = callResolutionResult
+
     fun interceptCandidates(
         candidates: Collection<FunctionDescriptor>,
         scopeTower: ImplicitScopeTower,
diff --git a/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/PSICallResolver.kt b/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/PSICallResolver.kt
index f76a1a1b78f..a1215385bd1 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/PSICallResolver.kt
+++ b/compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/tower/PSICallResolver.kt
@@ -109,6 +109,10 @@ class PSICallResolver(
             return OverloadResolutionResultsImpl.nameNotFound()
         }
 
+        result = candidateInterceptor.interceptResolvedCandidates(
+            result, context, kotlinCallResolver, name, resolutionKind, tracingStrategy
+        )
+
         val overloadResolutionResults = convertToOverloadResolutionResults<D>(context, result, tracingStrategy)
         return overloadResolutionResults.also {
             clearCacheForApproximationResults()