From a97c107a2b11370bfff708d285a43cda87039d3c Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Tue, 10 Nov 2020 10:15:37 +0300
Subject: [PATCH] [FIR] Replace callee reference before checking completion
 from inference session

This is needed because in some cases inference session checks candidate
  and candidates of all sub calls and for that it uses candidate from
  callee reference of call. So we should replace reference with specific
  candidate even if we didn't choose proper candidate yet

#KT-43129
---
 .../jetbrains/kotlin/fir/FirCallResolver.kt   | 37 ++++++++++++++-----
 1 file changed, 28 insertions(+), 9 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirCallResolver.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirCallResolver.kt
index 47a4f0c21a4..02030483a62 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirCallResolver.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/FirCallResolver.kt
@@ -165,16 +165,35 @@ class FirCallResolver(
 
         if (
             reducedCandidates.size > 1 &&
-            session.languageVersionSettings.supportsFeature(LanguageFeature.OverloadResolutionByLambdaReturnType) &&
-            bestCandidates.all { components.context.inferenceSession.shouldRunCompletion(qualifiedAccess) }
+            session.languageVersionSettings.supportsFeature(LanguageFeature.OverloadResolutionByLambdaReturnType)
         ) {
-            val newCandidates = chooseCandidateRegardingOverloadResolutionByLambdaReturnType(
-                qualifiedAccess,
-                reducedCandidates,
-                bestCandidates
-            )
-            if (newCandidates != null) {
-                reducedCandidates = newCandidates
+            /*
+             * Inference session may look into candidate of call, and for that it uses callee reference.
+             * So we need replace reference with proper candidate before calling inference session
+             */
+            val shouldRunCompletion = if (components.context.inferenceSession != FirInferenceSession.DEFAULT) {
+                var shouldRunCompletion = true
+                val originalReference = qualifiedAccess.calleeReference
+                val inferenceSession = components.context.inferenceSession
+                for (candidate in bestCandidates) {
+                    qualifiedAccess.replaceCalleeReference(FirNamedReferenceWithCandidate(null, candidate.callInfo.name, candidate))
+                    shouldRunCompletion = shouldRunCompletion && inferenceSession.shouldRunCompletion(qualifiedAccess)
+                    if (!shouldRunCompletion) break
+                }
+                qualifiedAccess.replaceCalleeReference(originalReference)
+                shouldRunCompletion
+            } else {
+                true
+            }
+            if (shouldRunCompletion) {
+                val newCandidates = chooseCandidateRegardingOverloadResolutionByLambdaReturnType(
+                    qualifiedAccess,
+                    reducedCandidates,
+                    bestCandidates
+                )
+                if (newCandidates != null) {
+                    reducedCandidates = newCandidates
+                }
             }
         }