[NI] Introduce basic resolver for many candidates

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

@@ -32,6 +32,7 @@ import org.jetbrains.kotlin.resolve.calls.results.ResolutionStatus
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowInfo
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory
 import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
+import org.jetbrains.kotlin.resolve.calls.tasks.TracingStrategy
 import org.jetbrains.kotlin.resolve.constants.evaluate.ConstantExpressionEvaluator
 import org.jetbrains.kotlin.resolve.scopes.receivers.CastImplicitClassReceiver
 import org.jetbrains.kotlin.resolve.scopes.receivers.ImplicitClassReceiver
@@ -71,7 +72,8 @@ class KotlinToResolvedCallTransformer(
 
     fun <D : CallableDescriptor> transformAndReport(
         baseResolvedCall: CallResolutionResult,
-        context: BasicCallResolutionContext
+        context: BasicCallResolutionContext,
+        tracingStrategy: TracingStrategy
     ): ResolvedCall<D> {
         return when (baseResolvedCall) {
             is PartialCallResolutionResult -> {
@@ -84,7 +86,9 @@ class KotlinToResolvedCallTransformer(
                     psiKotlinCall.psiCall
 
                 context.trace.record(BindingContext.ONLY_RESOLVED_CALL, psiCall, baseResolvedCall)
-                context.inferenceSession.addPartiallyResolvedCall(baseResolvedCall)
+                context.inferenceSession.addPartialCallInfo(
+                    PSIPartialCallInfo(baseResolvedCall, context, tracingStrategy)
+                )
 
                 createStubResolvedCallAndWriteItToTrace(candidate, context.trace, baseResolvedCall.diagnostics)
             }

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

@@ -5,17 +5,110 @@
 
 package org.jetbrains.kotlin.resolve.calls.tower
 
-import org.jetbrains.kotlin.resolve.calls.components.InferenceSession
+import org.jetbrains.kotlin.builtins.KotlinBuiltIns
-import org.jetbrains.kotlin.resolve.calls.model.CallResolutionResult
+import org.jetbrains.kotlin.descriptors.CallableDescriptor
+import org.jetbrains.kotlin.resolve.BindingTrace
+import org.jetbrains.kotlin.resolve.calls.components.*
+import org.jetbrains.kotlin.resolve.calls.context.BasicCallResolutionContext
+import org.jetbrains.kotlin.resolve.calls.inference.NewConstraintSystem
+import org.jetbrains.kotlin.resolve.calls.inference.components.KotlinConstraintSystemCompleter
+import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
+import org.jetbrains.kotlin.resolve.calls.inference.model.NewConstraintSystemImpl
+import org.jetbrains.kotlin.resolve.calls.model.*
+import org.jetbrains.kotlin.resolve.calls.results.OverloadResolutionResults
+import org.jetbrains.kotlin.resolve.calls.tasks.TracingStrategy
 
-abstract class ManyCandidatesResolver : InferenceSession {
+abstract class ManyCandidatesResolver<D : CallableDescriptor>(
-    private val partiallyResolvedCalls = arrayListOf<CallResolutionResult>()
+    private val psiCallResolver: PSICallResolver,
+    private val postponedArgumentsAnalyzer: PostponedArgumentsAnalyzer,
+    private val kotlinConstraintSystemCompleter: KotlinConstraintSystemCompleter,
+    private val callComponents: KotlinCallComponents,
+    val builtIns: KotlinBuiltIns
+) : InferenceSession {
+    private val partiallyResolvedCallsInfo = arrayListOf<PSIPartialCallInfo>()
+    private val errorCallsInfo = arrayListOf<PSIErrorCallInfo<D>>()
 
-    override fun shouldFixTypeVariables(): Boolean {
+    abstract fun prepareForCompletion(commonSystem: NewConstraintSystem, resolvedCallsInfo: List<PSIPartialCallInfo>)
+
+    override fun shouldFixTypeVariables(candidate: KotlinResolutionCandidate): Boolean {
         return false
     }
 
-    override fun addPartiallyResolvedCall(call: CallResolutionResult) {
+    override fun addPartialCallInfo(callInfo: PartialCallInfo) {
-        partiallyResolvedCalls.add(call)
+        if (callInfo !is PSIPartialCallInfo) {
+            throw AssertionError("Call info for $callInfo should be instance of PSIPartialCallInfo")
+        }
+        partiallyResolvedCallsInfo.add(callInfo)
     }
-}
+
+    override fun addErrorCallInfo(callInfo: ErrorCallInfo) {
+        if (callInfo !is PSIErrorCallInfo<*>) {
+            throw AssertionError("Error call info for $callInfo should be instance of PSIErrorCallInfo")
+        }
+        errorCallsInfo.add(callInfo as PSIErrorCallInfo<D>)
+    }
+
+    override fun currentConstraintSystem(): ConstraintStorage {
+        return partiallyResolvedCallsInfo.lastOrNull()?.callResolutionResult?.constraintSystem ?: ConstraintStorage.Empty
+    }
+
+    fun resolveCandidates(resolutionCallbacks: KotlinResolutionCallbacks): List<ResolutionResultCallInfo<D>> {
+        val resolvedCallsInfo = partiallyResolvedCallsInfo.toList()
+
+        val commonSystem = NewConstraintSystemImpl(callComponents.constraintInjector, builtIns).apply {
+            addOtherSystem(currentConstraintSystem())
+        }
+
+        prepareForCompletion(commonSystem, resolvedCallsInfo)
+
+        val diagnosticHolder = KotlinDiagnosticsHolder.SimpleHolder()
+
+        kotlinConstraintSystemCompleter.runCompletion(
+            commonSystem.asConstraintSystemCompleterContext(),
+            KotlinConstraintSystemCompleter.ConstraintSystemCompletionMode.FULL,
+            resolvedCallsInfo.map { it.callResolutionResult },
+            builtIns.unitType
+        ) {
+            postponedArgumentsAnalyzer.analyze(
+                commonSystem.asPostponedArgumentsAnalyzerContext(), resolutionCallbacks, it, diagnosticHolder
+            )
+        }
+
+        val allCandidates = arrayListOf<ResolutionResultCallInfo<D>>()
+
+        resolvedCallsInfo.mapTo(allCandidates) {
+            val resolutionResult = it.asCallResolutionResult(diagnosticHolder, commonSystem)
+            ResolutionResultCallInfo(
+                resolutionResult, psiCallResolver.convertToOverloadResolutionResults(it.context, resolutionResult, it.tracingStrategy)
+            )
+        }
+
+        errorCallsInfo.mapTo(allCandidates) { ResolutionResultCallInfo(it.callResolutionResult, it.result) }
+
+        return allCandidates
+    }
+
+    private fun PartialCallInfo.asCallResolutionResult(
+        diagnosticsHolder: KotlinDiagnosticsHolder.SimpleHolder,
+        commonSystem: NewConstraintSystem
+    ): CallResolutionResult {
+        val diagnostics = diagnosticsHolder.getDiagnostics() + callResolutionResult.diagnostics + commonSystem.diagnostics
+        return CompletedCallResolutionResult(callResolutionResult.resultCallAtom, diagnostics, commonSystem.asReadOnlyStorage())
+    }
+}
+
+data class ResolutionResultCallInfo<D : CallableDescriptor>(
+    val resolutionResult: CallResolutionResult,
+    val overloadResolutionResults: OverloadResolutionResults<D>
+)
+
+class PSIPartialCallInfo(
+    override val callResolutionResult: PartialCallResolutionResult,
+    val context: BasicCallResolutionContext,
+    val tracingStrategy: TracingStrategy
+) : PartialCallInfo
+
+class PSIErrorCallInfo<D : CallableDescriptor>(
+    override val callResolutionResult: CallResolutionResult,
+    val result: OverloadResolutionResults<D>
+) : ErrorCallInfo

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

@@ -173,7 +173,7 @@ class PSICallResolver(
         }
     }
 
-    private fun <D : CallableDescriptor> convertToOverloadResolutionResults(
+    fun <D : CallableDescriptor> convertToOverloadResolutionResults(
         context: BasicCallResolutionContext,
         result: CallResolutionResult,
         tracingStrategy: TracingStrategy
@@ -189,26 +189,12 @@ class PSICallResolver(
 
         val trace = context.trace
 
-        result.diagnostics.firstIsInstanceOrNull<NoneCandidatesCallDiagnostic>()?.let {
+        handleErrorResolutionResult<D>(trace, result, tracingStrategy)?.let { errorResult ->
-            tracingStrategy.unresolvedReference(trace)
+            context.inferenceSession.addErrorCallInfo(PSIErrorCallInfo(result, errorResult))
-            return OverloadResolutionResultsImpl.nameNotFound()
+            return errorResult
         }
 
-        result.diagnostics.firstIsInstanceOrNull<ManyCandidatesCallDiagnostic>()?.let {
+        val resolvedCall = kotlinToResolvedCallTransformer.transformAndReport<D>(result, context, tracingStrategy)
-            return transformManyCandidatesAndRecordTrace(it, tracingStrategy, trace)
-        }
-
-        val isInapplicableReceiver =
-            getResultApplicability(result.diagnostics) == ResolutionCandidateApplicability.INAPPLICABLE_WRONG_RECEIVER
-
-        val resolvedCall = if (isInapplicableReceiver) {
-            val singleCandidate = result.resultCallAtom() ?: error("Should be not null for result: $result")
-            kotlinToResolvedCallTransformer.onlyTransform<D>(singleCandidate, result.diagnostics).also {
-                tracingStrategy.unresolvedReferenceWrongReceiver(trace, listOf(it))
-            }
-        } else {
-            kotlinToResolvedCallTransformer.transformAndReport<D>(result, context)
-        }
 
         // NB. Be careful with moving this invocation, as effect system expects resolution results to be written in trace
         // (see EffectSystem for details)
@@ -217,6 +203,34 @@ class PSICallResolver(
         return SingleOverloadResolutionResult(resolvedCall)
     }
 
+    private fun <D : CallableDescriptor> handleErrorResolutionResult(
+        trace: BindingTrace,
+        result: CallResolutionResult,
+        tracingStrategy: TracingStrategy
+    ): OverloadResolutionResults<D>? {
+        val diagnostics = result.diagnostics
+
+        diagnostics.firstIsInstanceOrNull<NoneCandidatesCallDiagnostic>()?.let {
+            tracingStrategy.unresolvedReference(trace)
+            return OverloadResolutionResultsImpl.nameNotFound()
+        }
+
+        diagnostics.firstIsInstanceOrNull<ManyCandidatesCallDiagnostic>()?.let {
+            return transformManyCandidatesAndRecordTrace(it, tracingStrategy, trace)
+        }
+
+        if (getResultApplicability(diagnostics) == ResolutionCandidateApplicability.INAPPLICABLE_WRONG_RECEIVER) {
+            val singleCandidate = result.resultCallAtom() ?: error("Should be not null for result: $result")
+            val resolvedCall = kotlinToResolvedCallTransformer.onlyTransform<D>(singleCandidate, diagnostics).also {
+                tracingStrategy.unresolvedReferenceWrongReceiver(trace, listOf(it))
+            }
+
+            return SingleOverloadResolutionResult(resolvedCall)
+        }
+
+        return null
+    }
+
     private fun <D : CallableDescriptor> transformManyCandidatesAndRecordTrace(
         diagnostic: ManyCandidatesCallDiagnostic,
         tracingStrategy: TracingStrategy,
@@ -360,7 +374,7 @@ class PSICallResolver(
         override fun factoryForVariable(stripExplicitReceiver: Boolean): CandidateFactory<KotlinResolutionCandidate> {
             val explicitReceiver = if (stripExplicitReceiver) null else kotlinCall.explicitReceiver
             val variableCall = PSIKotlinCallForVariable(kotlinCall, explicitReceiver, kotlinCall.name)
-            return SimpleCandidateFactory(callComponents, scopeTower, variableCall)
+            return SimpleCandidateFactory(callComponents, scopeTower, variableCall, context.inferenceSession)
         }
 
         override fun factoryForInvoke(variable: KotlinResolutionCandidate, useExplicitReceiver: Boolean):
@@ -380,7 +394,9 @@ class PSICallResolver(
                 PSIKotlinCallForInvoke(kotlinCall, variable, variableCallArgument, null)
             }
 
-            return variableCallArgument.receiver to SimpleCandidateFactory(callComponents, scopeTower, callForInvoke)
+            return variableCallArgument.receiver to SimpleCandidateFactory(
+                callComponents, scopeTower, callForInvoke, context.inferenceSession
+            )
         }
 
         // todo: create special check that there is no invoke on variable

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/KotlinCallResolver.kt

@@ -44,7 +44,7 @@ class KotlinCallResolver(
     ): CallResolutionResult {
         kotlinCall.checkCallInvariants()
 
-        val candidateFactory = SimpleCandidateFactory(callComponents, scopeTower, kotlinCall)
+        val candidateFactory = SimpleCandidateFactory(callComponents, scopeTower, kotlinCall, resolutionCallbacks.inferenceSession)
         val processor = when (kotlinCall.callKind) {
             KotlinCallKind.VARIABLE -> {
                 createVariableAndObjectProcessor(scopeTower, kotlinCall.name, candidateFactory, kotlinCall.explicitReceiver?.receiver)
@@ -85,7 +85,7 @@ class KotlinCallResolver(
         collectAllCandidates: Boolean
     ): CallResolutionResult {
         kotlinCall.checkCallInvariants()
-        val candidateFactory = SimpleCandidateFactory(callComponents, scopeTower, kotlinCall)
+        val candidateFactory = SimpleCandidateFactory(callComponents, scopeTower, kotlinCall, resolutionCallbacks.inferenceSession)
 
         val resolutionCandidates = givenCandidates.map { candidateFactory.createCandidate(it).forceResolution() }
 

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/InferenceSession.kt

@@ -5,19 +5,31 @@
 
 package org.jetbrains.kotlin.resolve.calls.components
 
-import org.jetbrains.kotlin.resolve.calls.inference.NewConstraintSystem
+import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage
 import org.jetbrains.kotlin.resolve.calls.model.CallResolutionResult
+import org.jetbrains.kotlin.resolve.calls.model.KotlinResolutionCandidate
+import org.jetbrains.kotlin.resolve.calls.model.PartialCallResolutionResult
 
 interface InferenceSession {
     companion object {
         val default = object : InferenceSession {
-            override fun prepareBeforeCompletion(commonSystem: NewConstraintSystem) {}
+            override fun shouldFixTypeVariables(candidate: KotlinResolutionCandidate): Boolean = true
-            override fun shouldFixTypeVariables(): Boolean = false
+            override fun addPartialCallInfo(callInfo: PartialCallInfo) {}
-            override fun addPartiallyResolvedCall(call: CallResolutionResult) {}
+            override fun addErrorCallInfo(callInfo: ErrorCallInfo) {}
+            override fun currentConstraintSystem(): ConstraintStorage = ConstraintStorage.Empty
         }
     }
 
-    fun prepareBeforeCompletion(commonSystem: NewConstraintSystem)
+    fun shouldFixTypeVariables(candidate: KotlinResolutionCandidate): Boolean
-    fun shouldFixTypeVariables(): Boolean
+    fun addPartialCallInfo(callInfo: PartialCallInfo)
-    fun addPartiallyResolvedCall(call: CallResolutionResult)
+    fun addErrorCallInfo(callInfo: ErrorCallInfo)
+    fun currentConstraintSystem(): ConstraintStorage
+}
+
+interface PartialCallInfo {
+    val callResolutionResult: PartialCallResolutionResult
+}
+
+interface ErrorCallInfo {
+    val callResolutionResult: CallResolutionResult
 }

compiler/resolution/src/org/jetbrains/kotlin/resolve/calls/components/KotlinCallCompleter.kt

@@ -8,6 +8,7 @@ package org.jetbrains.kotlin.resolve.calls.components
 import org.jetbrains.kotlin.resolve.calls.inference.NewConstraintSystem
 import org.jetbrains.kotlin.resolve.calls.inference.components.KotlinConstraintSystemCompleter
 import org.jetbrains.kotlin.resolve.calls.inference.components.KotlinConstraintSystemCompleter.ConstraintSystemCompletionMode
+import org.jetbrains.kotlin.resolve.calls.inference.model.ConstraintStorage.Empty.hasContradiction
 import org.jetbrains.kotlin.resolve.calls.inference.model.ExpectedTypeConstraintPosition
 import org.jetbrains.kotlin.resolve.calls.model.*
 import org.jetbrains.kotlin.resolve.calls.tower.forceResolution
@@ -37,7 +38,7 @@ class KotlinCallCompleter(
         val candidate = prepareCandidateForCompletion(factory, candidates, resolutionCallbacks)
         val completionType = candidate.prepareForCompletion(expectedType, resolutionCallbacks)
 
-        return if (resolutionCallbacks.inferenceSession.shouldFixTypeVariables())
+        return if (resolutionCallbacks.inferenceSession.shouldFixTypeVariables(candidate))
             candidate.runCompletion(completionType, diagnosticHolder, resolutionCallbacks)
         else
             candidate.asCallResolutionResult(ConstraintSystemCompletionMode.PARTIAL, diagnosticHolder)
@@ -68,9 +69,9 @@ class KotlinCallCompleter(
         diagnosticHolder: KotlinDiagnosticsHolder.SimpleHolder,
         resolutionCallbacks: KotlinResolutionCallbacks
     ): CallResolutionResult {
-        if (ErrorUtils.isError(resolvedCall.candidateDescriptor) || csBuilder.hasContradiction) {
+        if (isErrorCandidate()) {
             runCompletion(resolvedCall, ConstraintSystemCompletionMode.FULL, diagnosticHolder, getSystem(), resolutionCallbacks)
-            return asCallResolutionResult(completionType, diagnosticHolder, isError = true)
+            return asCallResolutionResult(completionType, diagnosticHolder)
         }
 
         runCompletion(resolvedCall, completionType, diagnosticHolder, getSystem(), resolutionCallbacks)
@@ -155,15 +156,14 @@ class KotlinCallCompleter(
         return resolutionCallbacks.createReceiverWithSmartCastInfo(resolvedCall)?.stableType ?: returnType
     }
 
-    private fun KotlinResolutionCandidate.asCallResolutionResult(
+    fun KotlinResolutionCandidate.asCallResolutionResult(
         type: ConstraintSystemCompletionMode,
-        diagnosticsHolder: KotlinDiagnosticsHolder.SimpleHolder,
+        diagnosticsHolder: KotlinDiagnosticsHolder.SimpleHolder
-        isError: Boolean = false
     ): CallResolutionResult {
         val systemStorage = getSystem().asReadOnlyStorage()
         val allDiagnostics = diagnosticsHolder.getDiagnostics() + this.diagnosticsFromResolutionParts
 
-        if (isError) {
+        if (isErrorCandidate()) {
             return ErrorCallResolutionResult(resolvedCall, allDiagnostics, systemStorage)
         }
 
@@ -173,4 +173,8 @@ class KotlinCallCompleter(
             PartialCallResolutionResult(resolvedCall, allDiagnostics, systemStorage)
         }
     }
+
+    private fun KotlinResolutionCandidate.isErrorCandidate(): Boolean {
+        return ErrorUtils.isError(resolvedCall.candidateDescriptor) || hasContradiction
+    }
 }

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

@@ -48,7 +48,8 @@ class KotlinCallComponents(
 class SimpleCandidateFactory(
     val callComponents: KotlinCallComponents,
     val scopeTower: ImplicitScopeTower,
-    val kotlinCall: KotlinCall
+    val kotlinCall: KotlinCall,
+    val inferenceSession: InferenceSession
 ) : CandidateFactory<KotlinResolutionCandidate> {
     val baseSystem: ConstraintStorage
 
@@ -61,6 +62,8 @@ class SimpleCandidateFactory(
         }
         baseSystem.addSubsystemFromArgument(kotlinCall.externalArgument)
 
+        baseSystem.addOtherSystem(inferenceSession.currentConstraintSystem())
+
         this.baseSystem = baseSystem.asReadOnlyStorage()
     }