[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

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
+                }
             }
         }